Update main.py

main.py

@@ -33,6 +33,12 @@ from typing import Dict, List, Optional
 import aiohttp
 from urllib.parse import quote
 from fastapi.staticfiles import StaticFiles
+from huggingface_hub import hf_hub_download
+import joblib
+import logging
+
+logger = logging.getLogger(__name__)
+loaded_models = {}
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -135,18 +141,24 @@ class InteractionResult(BaseModel):
 loaded_models = {}
 
 def load_models():
-    """Load the pre-trained models"""
     try:
-        loaded_models['xgb_model'] = joblib.load(os.path.join("models", "xgb_model.joblib"))
-        loaded_models['label_encoder'] = joblib.load(os.path.join("models", "label_encoder.joblib"))
-        loaded_models['feature_order'] = joblib.load(os.path.join("models", "feature_order.joblib"))
-        logger.info("Models loaded successfully")
+        # Download from Hugging Face Hub dynamically
+        loaded_models['xgb_model'] = joblib.load(
+            hf_hub_download(repo_id="asritha22bce/FoodDrugInteraction", filename="models/xgb_model.joblib")
+        )
+        loaded_models['label_encoder'] = joblib.load(
+            hf_hub_download(repo_id="asritha22bce/FoodDrugInteraction", filename="models/label_encoder.joblib")
+        )
+        loaded_models['feature_order'] = joblib.load(
+            hf_hub_download(repo_id="asritha22bce/FoodDrugInteraction", filename="models/feature_order.joblib")
+        )
+
+        logger.info("Models loaded successfully from Hugging Face Hub")
         return True
     except Exception as e:
         logger.error(f"Error loading models: {e}")
         return False
 
-
 # Load models on startup
 @app.on_event("startup")
 async def startup_event():