Update download_models.py

download_models.py

@@ -1,34 +1,61 @@
 import os
 import sys
 
-# Path harus sinkron dengan Dockerfile
-HF_CACHE = os.getenv("HF_HOME", "/home/user/app/models_cache/transformers")
-WHISPER_CACHE = "/home/user/app/models_cache/whisper"
+CACHE_DIR = "/home/user/app/models_cache"
+os.makedirs(CACHE_DIR, exist_ok=True)
 
-os.makedirs(HF_CACHE, exist_ok=True)
-os.makedirs(WHISPER_CACHE, exist_ok=True)
+print("================================")
+print("DOWNLOADING AI MODELS")
+print("================================")
 
-print("=== STARTING MODEL DOWNLOAD ===")
+# ==========================
+# WHISPER MODEL
+# ==========================
 
 try:
     from faster_whisper import WhisperModel
-    print(f"Downloading Whisper tiny to {WHISPER_CACHE}...")
-    WhisperModel("tiny", device="cpu", compute_type="int8", download_root=WHISPER_CACHE)
-    print("Whisper model OK.")
+
+    print("Downloading Whisper tiny model...")
+
+    WhisperModel(
+        "tiny",
+        device="cpu",
+        compute_type="int8",
+        download_root=CACHE_DIR
+    )
+
+    print("Whisper download complete")
+
 except Exception as e:
-    print(f"ERROR Whisper: {e}")
+    print("WHISPER ERROR:", e)
     sys.exit(1)
 
+# ==========================
+# DISTILBERT SENTIMENT
+# ==========================
+
 try:
-    from transformers import AutoTokenizer, AutoModelForSequenceClassification
+    from transformers import AutoTokenizer
+    from transformers import AutoModelForSequenceClassification
+
     model_name = "distilbert-base-uncased-finetuned-sst-2-english"
-    
-    print(f"Downloading {model_name} to {HF_CACHE}...")
-    AutoTokenizer.from_pretrained(model_name, cache_dir=HF_CACHE)
-    AutoModelForSequenceClassification.from_pretrained(model_name, cache_dir=HF_CACHE)
-    print("Transformers models OK.")
+
+    save_path = os.path.join(CACHE_DIR, "distilbert")
+
+    print("Downloading DistilBERT...")
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForSequenceClassification.from_pretrained(model_name)
+
+    tokenizer.save_pretrained(save_path)
+    model.save_pretrained(save_path)
+
+    print("DistilBERT saved")
+
 except Exception as e:
-    print(f"ERROR Transformers: {e}")
+    print("TRANSFORMER ERROR:", e)
     sys.exit(1)
 
-print("=== ALL MODELS READY ===")
+print("================================")
+print("ALL MODELS READY")
+print("================================")