Upload folder using huggingface_hub

Dockerfile

@@ -10,4 +10,4 @@ RUN pip install -r requirements.txt
 
 EXPOSE 7860
 
-CMD ["uvicorn", "main:model_inference:inference_api", "--host", "0.0.0.0", "--port", "7860"]
+CMD ["uvicorn", "main.model_inference:inference_api", "--host", "0.0.0.0", "--port", "7860"]

main/helper.py

@@ -1,7 +1,17 @@
 # Helper functions for the model inference api
+import os
 import yaml
 import joblib
 from pathlib import Path
+from huggingface_hub import snapshot_download
+from dotenv import load_dotenv
+load_dotenv()
+
+# get environment variable
+hf_repo_id = os.getenv("HUGGINGFACE_REPO_ID")
+if not hf_repo_id:
+    raise EnvironmentError("HUGGINGFACE_REPO_ID is not set")
+
 
 # load yaml files to get model meta data.
 try:
@@ -14,12 +24,21 @@ except:
 def load_model():
     """ Loads ML model from location path and returns the model. """
     try:
-        with open(Path("serve_api/model/python_model.pkl"), "rb") as f:
+        model_dir = snapshot_download(
+            repo_id = str(hf_repo_id),
+            repo_type = "model",
+            local_dir = "/tmp/model",
+            local_dir_use_symlinks = False
+        )
+
+        model_path = Path(model_dir)/"python_model.pkl"
+        with open(model_path, "rb") as f:
             model = joblib.load(f)
+        
         return model
     
-    except:
-        raise FileNotFoundError("Failed to load model")
+    except Exception as e:
+        raise RuntimeError(f"Failed to load model from hub: {e}")
 
 
 def get_model_registry():

requirements.txt

@@ -1,4 +1,5 @@
 fastapi==0.116.1
 uvicorn==0.35.0
 joblib==1.5.1
-PyYAML==6.0.2
+PyYAML==6.0.2
+python-dotenv==1.1.1