models.py

# models.py
import whisper
from sentence_transformers import SentenceTransformer
import spacy
import os
import sys
import subprocess
from utils import setup_logger
import config

logger = setup_logger(__name__)

class AIModelManager:
    """Singleton-like manager to load heavy AI models."""
    
    def __init__(self):
        self.whisper = None
        self.embedding_model = None
        self.nlp = None

    def load_models(self):
        """Loads all required models with error handling."""
        try:
            logger.info("Loading Whisper model...")
            self.whisper = whisper.load_model(config.WHISPER_MODEL_SIZE)
            
            logger.info("Loading Sentence Transformer...")
            self.embedding_model = SentenceTransformer(config.TRANSFORMER_MODEL)
            
            logger.info("Loading spaCy model...")
            try:
                self.nlp = spacy.load(config.SPACY_MODEL)
            except OSError:
                logger.warning(f"SpaCy model '{config.SPACY_MODEL}' not found. Downloading...")
                subprocess.check_call([sys.executable, "-m", "spacy", "download", config.SPACY_MODEL])
                self.nlp = spacy.load(config.SPACY_MODEL)
                
            logger.info("All models loaded successfully.")
            
        except Exception as e:
            logger.critical(f"Failed to load models: {str(e)}")
            raise RuntimeError("Model loading failed.") from e

    def get_whisper(self):
        return self.whisper

    def get_embedding_model(self):
        return self.embedding_model

    def get_nlp(self):
        return self.nlp