Load model artifacts from archive at startup

predictor.py

@@ -1,5 +1,8 @@
 import json
 import os
+import urllib.request
+import tempfile
+import zipfile
 from functools import lru_cache
 from pathlib import Path
 from typing import Any, Dict, Mapping, Union
@@ -8,7 +11,6 @@ import joblib
 import numpy as np
 import torch
 import torch.nn as nn
-from huggingface_hub import snapshot_download
 
 
 MODEL_REPO_ID = os.getenv("MODEL_REPO_ID", "Dharunkumar9/battery-capacity-predictor")
@@ -24,6 +26,7 @@ class PositionalEncoding(nn.Module):
 		pe[:, 0::2] = torch.sin(position * div_term)
 		pe[:, 1::2] = torch.cos(position * div_term)
 		self.register_buffer("pe", pe.unsqueeze(0))
+		self.pe: torch.Tensor
 
 	def forward(self, x: torch.Tensor) -> torch.Tensor:
 		return x + self.pe[:, : x.size(1), :]
@@ -99,13 +102,28 @@ def _normalize_window(window: Any, expected_rows: int, expected_cols: int) -> np
 
 
 def _download_artifacts() -> Path:
-	return Path(
-		snapshot_download(
-			repo_id=MODEL_REPO_ID,
-			repo_type="model",
-			allow_patterns=["config.json", "pytorch_model.bin", "x_scalers.pkl", "y_scalers.pkl"],
-		)
-	)
+	archive_name_candidates = ["artifacts_v1.zip", "artifacts-v1.zip"]
+	archive_path = None
+
+	for archive_name in archive_name_candidates:
+		try:
+			archive_url = f"https://huggingface.co/{MODEL_REPO_ID}/resolve/main/{archive_name}?download=true"
+			archive_file = Path(tempfile.mkdtemp(prefix="battery-archive-")) / archive_name
+			with urllib.request.urlopen(archive_url) as response, archive_file.open("wb") as output:
+				output.write(response.read())
+			archive_path = str(archive_file)
+			break
+		except Exception:
+			continue
+
+	if archive_path is None:
+		raise FileNotFoundError("Could not download the model artifact zip from the Hugging Face model repo")
+
+	extract_dir = Path(tempfile.mkdtemp(prefix="battery-model-"))
+	with zipfile.ZipFile(archive_path) as archive:
+		archive.extractall(extract_dir)
+
+	return extract_dir
 
 
 class BatteryPredictor:

requirements.txt

@@ -3,5 +3,4 @@ uvicorn[standard]>=0.29.0
 torch>=2.2.0
 numpy>=1.24.0
 joblib>=1.3.0
-scikit-learn==1.7.0
-huggingface_hub>=0.23.0
+scikit-learn==1.7.0