preload_models.py

from faster_whisper import WhisperModel
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
import os

print("--- Preloading models for Docker image ---")

# 1. Preload Whisper (Medium)
# This downloads the model files to the cache directory
print("Downloading Whisper model...")
try:
    model = WhisperModel("medium", device="cpu", compute_type="int8")
except Exception as e:
    print(f"Error downloading Whisper: {e}")

# 2. Preload Summarization Model (BART)
print("Downloading Summarization model...")
try:
    model_name = "sshleifer/distilbart-cnn-12-6"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
except Exception as e:
    print(f"Error downloading BART: {e}")

# 3. Preload Sentiment Model
print("Downloading Sentiment model...")
try:
    from transformers import pipeline
    sentiment_pipeline = pipeline(
        "sentiment-analysis",
        model="distilbert-base-uncased-finetuned-sst-2-english"
    )
except Exception as e:
    print(f"Error downloading Sentiment model: {e}")


print("--- Finished preloading ---")