Add HF Hub model loading support

ml_models.py

@@ -1,6 +1,7 @@
 """
 ML Model Loader and Utilities
 Handles loading and using the conflict prediction model and package embeddings.
+Loads from local files if available, otherwise downloads from Hugging Face Hub.
 """
 
 import json
@@ -10,27 +11,70 @@ from typing import Dict, List, Tuple, Optional
 import numpy as np
 from packaging.requirements import Requirement
 
+# Try to import huggingface_hub for model downloading
+try:
+    from huggingface_hub import hf_hub_download
+    HF_HUB_AVAILABLE = True
+except ImportError:
+    HF_HUB_AVAILABLE = False
+    print("Warning: huggingface_hub not available. Models must be loaded locally.")
+
 
 class ConflictPredictor:
     """Load and use the conflict prediction model."""
     
-    def __init__(self, model_path: Optional[Path] = None):
-        """Initialize the conflict predictor."""
+    def __init__(self, model_path: Optional[Path] = None, repo_id: str = "ysakhale/dependency-conflict-models"):
+        """Initialize the conflict predictor.
+        
+        Args:
+            model_path: Local path to model file (optional)
+            repo_id: Hugging Face repository ID to download from if local file not found
+        """
+        self.repo_id = repo_id
+        self.model = None
+        self.model_path = model_path
+        
+        # Try local path first
         if model_path is None:
             model_path = Path(__file__).parent / "models" / "conflict_predictor.pkl"
         
-        self.model = None
         self.model_path = model_path
         
+        # Try loading from local file
         if model_path.exists():
             try:
                 with open(model_path, 'rb') as f:
                     self.model = pickle.load(f)
-                print(f"✅ Loaded conflict prediction model from {model_path}")
+                print(f"Loaded conflict prediction model from {model_path}")
+                return
             except Exception as e:
-                print(f"⚠️ Could not load conflict prediction model: {e}")
-        else:
-            print(f"⚠️ Conflict prediction model not found at {model_path}")
+                print(f"Could not load conflict prediction model from local: {e}")
+        
+        # If local file doesn't exist, try downloading from HF Hub
+        if HF_HUB_AVAILABLE:
+            try:
+                print(f"Model not found locally. Downloading from Hugging Face Hub: {repo_id}")
+                downloaded_path = hf_hub_download(
+                    repo_id=repo_id,
+                    filename="conflict_predictor.pkl",
+                    repo_type="model"
+                )
+                with open(downloaded_path, 'rb') as f:
+                    self.model = pickle.load(f)
+                print(f"Loaded conflict prediction model from Hugging Face Hub")
+                # Optionally cache it locally
+                try:
+                    model_path.parent.mkdir(parents=True, exist_ok=True)
+                    import shutil
+                    shutil.copy(downloaded_path, model_path)
+                    print(f"Cached model locally at {model_path}")
+                except:
+                    pass
+                return
+            except Exception as e:
+                print(f"Could not download model from Hugging Face Hub: {e}")
+        
+        print(f"Warning: Conflict prediction model not available")
     
     def extract_features(self, requirements_text: str) -> np.ndarray:
         """Extract features from requirements text (same as training)."""
@@ -130,24 +174,58 @@ class ConflictPredictor:
 class PackageEmbeddings:
     """Load and use package embeddings for similarity matching."""
     
-    def __init__(self, embeddings_path: Optional[Path] = None):
-        """Initialize package embeddings."""
-        if embeddings_path is None:
-            embeddings_path = Path(__file__).parent / "models" / "package_embeddings.json"
+    def __init__(self, embeddings_path: Optional[Path] = None, repo_id: str = "ysakhale/dependency-conflict-models"):
+        """Initialize package embeddings.
         
+        Args:
+            embeddings_path: Local path to embeddings file (optional)
+            repo_id: Hugging Face repository ID to download from if local file not found
+        """
+        self.repo_id = repo_id
         self.embeddings = {}
         self.embeddings_path = embeddings_path
         self.model = None
         
+        if embeddings_path is None:
+            embeddings_path = Path(__file__).parent / "models" / "package_embeddings.json"
+        
+        self.embeddings_path = embeddings_path
+        
+        # Try loading from local file
         if embeddings_path.exists():
             try:
                 with open(embeddings_path, 'r') as f:
                     self.embeddings = json.load(f)
-                print(f"✅ Loaded {len(self.embeddings)} package embeddings from {embeddings_path}")
+                print(f"Loaded {len(self.embeddings)} package embeddings from {embeddings_path}")
+                return
+            except Exception as e:
+                print(f"Could not load embeddings from local: {e}")
+        
+        # If local file doesn't exist, try downloading from HF Hub
+        if HF_HUB_AVAILABLE:
+            try:
+                print(f"Embeddings not found locally. Downloading from Hugging Face Hub: {repo_id}")
+                downloaded_path = hf_hub_download(
+                    repo_id=repo_id,
+                    filename="package_embeddings.json",
+                    repo_type="model"
+                )
+                with open(downloaded_path, 'r') as f:
+                    self.embeddings = json.load(f)
+                print(f"Loaded {len(self.embeddings)} package embeddings from Hugging Face Hub")
+                # Optionally cache it locally
+                try:
+                    embeddings_path.parent.mkdir(parents=True, exist_ok=True)
+                    import shutil
+                    shutil.copy(downloaded_path, embeddings_path)
+                    print(f"Cached embeddings locally at {embeddings_path}")
+                except:
+                    pass
+                return
             except Exception as e:
-                print(f"⚠️ Could not load embeddings: {e}")
-        else:
-            print(f"⚠️ Embeddings not found at {embeddings_path}")
+                print(f"Could not download embeddings from Hugging Face Hub: {e}")
+        
+        print(f"Warning: Package embeddings not available")
     
     def _load_model(self):
         """Lazy load the sentence transformer model."""

requirements.txt

@@ -5,4 +5,5 @@ requests>=2.31.0
 scikit-learn>=1.3.0
 sentence-transformers>=2.2.0
 numpy>=1.24.0
+huggingface-hub>=0.20.0
 

upload_models_to_hf.py

@@ -0,0 +1,86 @@
+"""
+Upload ML models to Hugging Face Hub
+This allows the models to be loaded in Hugging Face Spaces
+"""
+
+from pathlib import Path
+from huggingface_hub import HfApi, login
+import os
+
+def upload_models():
+    """Upload models to Hugging Face Hub."""
+    
+    # Check if models exist
+    models_dir = Path("models")
+    if not models_dir.exists():
+        print("Error: models/ directory not found!")
+        print("Please train the models first:")
+        print("  python train_conflict_model.py")
+        print("  python generate_embeddings.py")
+        return
+    
+    # Check for model files
+    model_files = {
+        "conflict_predictor.pkl": models_dir / "conflict_predictor.pkl",
+        "package_embeddings.json": models_dir / "package_embeddings.json",
+        "embedding_info.json": models_dir / "embedding_info.json"
+    }
+    
+    missing = [name for name, path in model_files.items() if not path.exists()]
+    if missing:
+        print(f"Error: Missing model files: {missing}")
+        print("Please train the models first:")
+        print("  python train_conflict_model.py")
+        print("  python generate_embeddings.py")
+        return
+    
+    # Login to Hugging Face
+    print("Logging in to Hugging Face...")
+    print("(You'll need to enter your HF token - get it from https://huggingface.co/settings/tokens)")
+    try:
+        login()
+    except Exception as e:
+        print(f"Login error: {e}")
+        print("\nYou can also set HF_TOKEN environment variable:")
+        print("  $env:HF_TOKEN='your_token_here'  # PowerShell")
+        return
+    
+    # Initialize API
+    api = HfApi()
+    
+    # Repository name for models
+    repo_id = "ysakhale/dependency-conflict-models"
+    
+    # Create repository if it doesn't exist
+    try:
+        api.create_repo(
+            repo_id=repo_id,
+            repo_type="model",
+            exist_ok=True,
+            private=False
+        )
+        print(f"Repository {repo_id} is ready!")
+    except Exception as e:
+        print(f"Note: {e}")
+    
+    # Upload each model file
+    print("\nUploading models...")
+    for filename, filepath in model_files.items():
+        print(f"Uploading {filename}...")
+        try:
+            api.upload_file(
+                path_or_fileobj=str(filepath),
+                path_in_repo=filename,
+                repo_id=repo_id,
+                repo_type="model"
+            )
+            print(f"  ✓ {filename} uploaded successfully!")
+        except Exception as e:
+            print(f"  ✗ Error uploading {filename}: {e}")
+    
+    print(f"\n✅ Models uploaded to: https://huggingface.co/{repo_id}")
+    print("\nNext step: Update ml_models.py to load from this repository")
+
+if __name__ == "__main__":
+    upload_models()
+