Update app.py

app.py

@@ -4,11 +4,24 @@ from typing import List
 from transformers import AutoTokenizer, AutoModel
 import torch, os
 
-MODEL_ID = "dmis-lab/biobert-base-cased-v1"
-
-# Load once at startup
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-model = AutoModel.from_pretrained(MODEL_ID)
+MODEL_ID = os.getenv("MODEL_ID", "dmis-lab/biobert-base-cased-v1.2").strip()
+HF_TOKEN = (os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN") or "").strip() or None
+
+def load_model(model_id: str):
+    # Try public/anonymous first (works for public models)
+    try:
+        tok = AutoTokenizer.from_pretrained(model_id, token=None, trust_remote_code=False)
+        mdl = AutoModel.from_pretrained(model_id, token=None, trust_remote_code=False)
+        return tok, mdl
+    except Exception:
+        # If you actually use a private/gated model, fall back to an explicit token
+        if HF_TOKEN:
+            tok = AutoTokenizer.from_pretrained(model_id, token=HF_TOKEN, trust_remote_code=False)
+            mdl = AutoModel.from_pretrained(model_id, token=HF_TOKEN, trust_remote_code=False)
+            return tok, mdl
+        raise  # bubble up the original error
+
+tokenizer, model = load_model(MODEL_ID)
 model.eval()
 
 def mean_pooling(model_output, attention_mask):
@@ -29,7 +42,7 @@ app = FastAPI(title="BioBERT Embeddings", version="1.0")
 
 @app.get("/healthz")
 def health():
-    return {"ok": True}
+    return {"ok": True, "model_id": MODEL_ID}
 
 @app.post("/embed", response_model=EmbedResponse)
 def embed(req: EmbedRequest):
@@ -46,5 +59,5 @@ def embed(req: EmbedRequest):
     return {"embeddings": pooled.cpu().tolist()}
 
 if __name__ == "__main__":
-    import uvicorn
+    import uvicorn, os
     uvicorn.run("app:app", host="0.0.0.0", port=int(os.getenv("PORT", "7860")), workers=1)