Upload app.py with huggingface_hub

app.py

@@ -6,6 +6,7 @@ import numpy as np
 import pickle
 from huggingface_hub import hf_hub_download
 import os
+import sys
 from typing import List, Union
 
 app = FastAPI(title="Headache Predictor API")
@@ -30,8 +31,22 @@ async def load_model():
             cache_dir=cache_dir,
             token=hf_token  # Use token for private repo access
         )
+
+        # Load with custom unpickler to handle missing modules
+        class CustomUnpickler(pickle.Unpickler):
+            def find_class(self, module, name):
+                # Redirect src.* imports to current directory
+                if module.startswith('src.'):
+                    module = module.replace('src.', '')
+                try:
+                    return super().find_class(module, name)
+                except (ModuleNotFoundError, AttributeError):
+                    # If module not found, try without the prefix
+                    return super().find_class(name, name)
+
         with open(model_path, 'rb') as f:
-            model = pickle.load(f)
+            model = CustomUnpickler(f).load()
+
         print("✅ Model loaded successfully")
     except Exception as e:
         print(f"❌ Error loading model: {e}")