Load model from HF Hub

Dockerfile

@@ -1,16 +1,10 @@
-# Read the doc: https://huggingface.co/docs/hub/spaces-sdks-docker
-# you will also find guides on how best to write your Dockerfile
-
-FROM python:3.9
-
-RUN useradd -m -u 1000 user
-USER user
-ENV PATH="/home/user/.local/bin:$PATH"
+FROM python:3.10-slim
 
 WORKDIR /app
 
-COPY --chown=user ./requirements.txt requirements.txt
-RUN pip install --no-cache-dir --upgrade -r requirements.txt
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY . .
 
-COPY --chown=user . /app
 CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -1,7 +1,24 @@
 from fastapi import FastAPI
+import torch
+import pickle
+from huggingface_hub import hf_hub_download
 
 app = FastAPI()
 
-@app.get("/")
-def greet_json():
-    return {"Hello": "World!"}
+# Download files from model repo
+tag_vocab_path = hf_hub_download(
+    repo_id="SinaLab/Nested",
+    filename="tag_vocab.pkl"
+)
+
+checkpoint_path = hf_hub_download(
+    repo_id="SinaLab/Nested",
+    filename="checkpoints/checkpoint_2.pt"
+)
+
+# Load model
+with open(tag_vocab_path, "rb") as f:
+    id2label = pickle.load(f)
+
+model = torch.load(checkpoint_path, map_location="cpu")
+model.eval()

requirements.txt

@@ -1,2 +1,5 @@
+torch
 fastapi
-uvicorn[standard]
+uvicorn
+numpy
+huggingface_hub