Update download_models.py

download_models.py

@@ -1,40 +1,37 @@
 import os
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
-from faster_whisper import WhisperModel
+import sys
 
-MODEL_CACHE_DIR = "./models_cache"
-DISTILBERT_PATH = os.path.join(MODEL_CACHE_DIR, "distilbert")
-
-# Pastikan folder cache ada
+# Gunakan path absolut agar tidak salah lokasi di Docker
+MODEL_CACHE_DIR = "/home/user/app/models_cache"
 os.makedirs(MODEL_CACHE_DIR, exist_ok=True)
 
-print("=== MODEL INITIALIZER ===")
-
-# ================= DISTILBERT =================
-# Catatan: Hapus bagian ini jika app.py tidak menggunakan fitur klasifikasi teks
-DISTILBERT_HF = "distilbert-base-uncased-finetuned-sst-2-english"
-
-if not os.path.exists(DISTILBERT_PATH):
-    print("Downloading DistilBERT model...")
-    tokenizer = AutoTokenizer.from_pretrained(DISTILBERT_HF)
-    model = AutoModelForSequenceClassification.from_pretrained(DISTILBERT_HF)
-
-    tokenizer.save_pretrained(DISTILBERT_PATH)
-    model.save_pretrained(DISTILBERT_PATH)
-    print("DistilBERT downloaded successfully.")
-else:
-    print("DistilBERT already exists.")
-
-# ================= WHISPER =================
-print("Downloading Whisper model (tiny)...")
-
-# Kita paksa download ke folder cache agar tidak hilang saat restart
-whisper_model = WhisperModel(
-    "tiny", 
-    device="cpu", 
-    compute_type="int8", 
-    download_root=MODEL_CACHE_DIR  # PENTING: simpan di folder yang sama
-)
+print("=== STARTING MODEL DOWNLOAD ===")
+
+try:
+    from faster_whisper import WhisperModel
+    print("Downloading Whisper model (tiny)...")
+    # Simpan langsung ke folder cache
+    WhisperModel("tiny", device="cpu", compute_type="int8", download_root=MODEL_CACHE_DIR)
+    print("Whisper model downloaded successfully.")
+except Exception as e:
+    print(f"ERROR downloading Whisper: {e}")
+    sys.exit(1) # Beritahu Docker build ini gagal
+
+try:
+    from transformers import AutoTokenizer, AutoModelForSequenceClassification
+    # Jika Anda pakai DistilBERT
+    model_name = "distilbert-base-uncased-finetuned-sst-2-english"
+    save_path = os.path.join(MODEL_CACHE_DIR, "distilbert")
+    
+    print(f"Downloading {model_name}...")
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForSequenceClassification.from_pretrained(model_name)
+    
+    tokenizer.save_pretrained(save_path)
+    model.save_pretrained(save_path)
+    print("Transformers models saved.")
+except Exception as e:
+    print(f"ERROR downloading Transformers: {e}")
+    sys.exit(1)
 
-print("Whisper model downloaded & verified.")
 print("=== ALL MODELS READY ===")