Update product_return_prediction/api.py

product_return_prediction/api.py

@@ -8,6 +8,7 @@ import pickle
 from pathlib import Path
 from product_return_prediction.dataset import prepare_inventory, scale_data_with_trained_scaler
 from product_return_prediction.config import MODELS_DIR, EXTERNAL_DATA_DIR
+from huggingface_hub import hf_hub_download
 
 app = FastAPI(
     title="Product Return Prediction API",
@@ -48,6 +49,18 @@ def load_json(file_path: Path) -> dict:
         raise HTTPException(status_code=500, detail=f"Error reading JSON file {file_path}: {e}")
 
 
+def download_file(url: str, local_path: Path, headers: dict = None):
+    """Download a file from a URL and save it locally."""
+    try:
+        response = requests.get(url, headers=headers)
+        response.raise_for_status()  # Raise an exception for HTTP errors
+        with open(local_path, 'wb') as f:
+            f.write(response.content)
+        return local_path
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error downloading file from {url}: {e}")
+
+        
 def filter_inventory_by_combinations(inventory: pd.DataFrame, models: list, fabrics: list, colours: list) -> pd.DataFrame:
     """Filter inventory based on the product combinations."""
     filtered_inventory = pd.DataFrame()
@@ -113,8 +126,8 @@ async def root():
 @app.post("/predict/")
 async def predict(products: ProductRequest):
     inventory_path: Path = EXTERNAL_DATA_DIR / "inventory.tsv"
-    model_path: Path = MODELS_DIR / "svm.pkl"
-    scaler_file: Path = MODELS_DIR / "scaler.pkl"
+    model_name: str = "svm.pkl"
+    scaler_name: str = "scaler.pkl"
 
     hf_token = os.getenv("inventory_data")
     dataset = load_dataset("molinari135/armani-inventory", token=hf_token, data_files="inventory.tsv")
@@ -131,6 +144,15 @@ async def predict(products: ProductRequest):
         filtered_inventory, products.total_customer_purchases, products.total_customer_returns
     )
 
+    models_uri = "https://huggingface.co/molinari135/se4ai-models/resolve/main/"
+    model_path = MODELS_DIR / model_name
+    scaler_path = MODELS_DIR / scaler_name
+
+    headers = f"Authorization: Bearer {hf_token}"
+
+    download_file(f{models_uri}{model_name}, model_path, headers)
+    download_file(f{models_uri}{scaler_name}, scaler_path, headers)
+
     model = load_model(model_path)
 
     scaled_inventory = apply_scaling(prepared_inventory, scaler_file)