use local cache

Dockerfile

@@ -1,18 +1,19 @@
-# Dockerfile
 FROM python:3.10-slim
-
-# Create working dir
 WORKDIR /app
 
-# Install Python dependencies
+# 1) Install HF tools + your deps
+RUN pip install --no-cache-dir huggingface_hub
 COPY requirements.txt .
 RUN pip install --no-cache-dir -r requirements.txt
 
-# Copy your code
-COPY . .
+# 2) Snapshot the model locally
+RUN python - <<EOF
+from huggingface_hub import snapshot_download
+snapshot_download("numind/NuExtract-1.5-tiny", cache_dir="/app/model_cache")
+EOF
 
-# Expose the port HF will route to
-ENV PORT 7860
+# 3) Copy your code
+COPY . .
 
-# Launch Uvicorn on HF’s $PORT
-CMD ["bash", "-lc", "uvicorn app:app --host 0.0.0.0 --port $PORT"]
+ENV PORT=7860
+CMD ["bash","-lc","uvicorn app:app --host 0.0.0.0 --port $PORT"]

app.py

@@ -14,9 +14,10 @@ from supabase import create_client
 load_dotenv()
 app = FastAPI()
 
-model_name = "numind/NuExtract-tiny-v1.5"
-device     = "mps" if torch.backends.mps.is_available() else "cpu"
-dtype      = torch.float16 if device=="mps" else torch.float32
+
+model_name = "numind/NuExtract-1.5-tiny"
+device     = "gpu" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
+dtype = torch.float16 if device in ("mps", "gpu") else torch.float32
 
 @app.on_event("startup")
 def startup_supabase():
@@ -30,10 +31,19 @@ def startup_supabase():
 def load_model():
     print("Loading model and tokenizer...", flush=True)
     global model, tokenizer
+    # model = AutoModelForCausalLM.from_pretrained(
+    # model_name, torch_dtype=dtype, trust_remote_code=True
+    # )
     model = AutoModelForCausalLM.from_pretrained(
-    model_name, torch_dtype=dtype, trust_remote_code=True
+    "/app/model_cache/numind/NuExtract-1.5-tiny",
+    torch_dtype=dtype,
+    trust_remote_code=True
     ).to(device).eval()
-    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+    # tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+    tokenizer = AutoTokenizer.from_pretrained(
+    "/app/model_cache/numind/NuExtract-1.5-tiny",
+    trust_remote_code=True
+)
     print("Model and tokenizer loaded.", flush=True)
 
 def predict_NuExtract(texts, template, batch_size=10, max_length=5096, max_new_tokens=1024):