Spaces:
Sleeping
Sleeping
| """ | |
| Model Loader - Loads models from local directory or downloads from HuggingFace | |
| """ | |
| import os | |
| import pickle | |
| import joblib | |
| from pathlib import Path | |
| from huggingface_hub import hf_hub_download | |
| from sentence_transformers import SentenceTransformer | |
| from config.settings import HUGGINGFACE_REPO, MODEL_FILES, EMBEDDING_MODEL | |
| class ModelLoader: | |
| def __init__(self): | |
| self.models_dir = Path("models") | |
| self.models_dir.mkdir(exist_ok=True) | |
| self.models = {} | |
| # Priority order: Check local models/ directory (for Docker), then parent directories | |
| if self.models_dir.exists() and any(self.models_dir.glob("*.pkl")): | |
| self.local_models_path = self.models_dir | |
| print(f"✅ Found local models in ./models directory") | |
| else: | |
| # Check for local models in parent directory (for local development) | |
| parent_models_improved = Path(__file__).parent.parent.parent / "models_improved" | |
| parent_models = Path(__file__).parent.parent.parent / "models" | |
| if parent_models_improved.exists(): | |
| self.local_models_path = parent_models_improved | |
| print(f"✅ Found local models at: {self.local_models_path}") | |
| elif parent_models.exists(): | |
| self.local_models_path = parent_models | |
| print(f"✅ Found local models at: {self.local_models_path}") | |
| else: | |
| self.local_models_path = None | |
| print(f"ℹ️ No local models found, will download from HuggingFace") | |
| def download_models(self): | |
| """Load models from local directory or download from HuggingFace""" | |
| # Try loading from local directory first | |
| if self.local_models_path and self.local_models_path.exists(): | |
| print(f"📂 Loading models from local directory: {self.local_models_path}") | |
| for model_file in MODEL_FILES: | |
| try: | |
| local_path = self.local_models_path / model_file | |
| if local_path.exists(): | |
| # Try joblib first (better for scikit-learn), then pickle | |
| try: | |
| self.models[model_file.replace('.pkl', '')] = joblib.load(local_path) | |
| print(f"✅ Loaded (joblib): {model_file}") | |
| except: | |
| with open(local_path, 'rb') as f: | |
| self.models[model_file.replace('.pkl', '')] = pickle.load(f) | |
| print(f"✅ Loaded (pickle): {model_file}") | |
| else: | |
| print(f"⚠️ File not found locally: {model_file}") | |
| except Exception as e: | |
| print(f"❌ Error loading {model_file}: {e}") | |
| raise | |
| else: | |
| # Download from HuggingFace | |
| print(f"📥 Downloading models from {HUGGINGFACE_REPO}...") | |
| for model_file in MODEL_FILES: | |
| try: | |
| local_path = hf_hub_download( | |
| repo_id=HUGGINGFACE_REPO, | |
| filename=model_file, | |
| cache_dir=str(self.models_dir) | |
| ) | |
| print(f"✅ Downloaded: {model_file}") | |
| # Load the model | |
| with open(local_path, 'rb') as f: | |
| self.models[model_file.replace('.pkl', '')] = pickle.load(f) | |
| except Exception as e: | |
| print(f"❌ Error downloading {model_file}: {e}") | |
| raise | |
| # Load Sentence-BERT for duplicate detection | |
| print(f"📥 Loading Sentence-BERT model: {EMBEDDING_MODEL}...") | |
| self.models['sentence_bert'] = SentenceTransformer(EMBEDDING_MODEL) | |
| print("✅ Sentence-BERT loaded") | |
| print(f"✅ All models loaded successfully!") | |
| return self.models | |
| def get_models(self): | |
| """Get loaded models""" | |
| if not self.models: | |
| self.download_models() | |
| return self.models | |