[BIOBERT-FTM] Add FastAPI Docker Space

Dockerfile

@@ -0,0 +1,15 @@
+FROM ghcr.io/huggingface/transformers-pytorch-gpu:latest
+
+ENV PIP_DISABLE_PIP_VERSION_CHECK=1         PYTHONDONTWRITEBYTECODE=1         PYTHONUNBUFFERED=1
+
+WORKDIR /app
+COPY requirements.txt /app/requirements.txt
+RUN pip install --no-cache-dir -r /app/requirements.txt
+
+COPY app.py /app/app.py
+
+EXPOSE 7860
+ENV HOST=0.0.0.0
+ENV PORT=7860
+
+CMD ["python", "-u", "app.py"]

README.md

@@ -1,10 +1,19 @@
 ---
-title: BIOBERT FTM Service
-emoji: 😻
-colorFrom: pink
-colorTo: gray
+title: BIOBERT-FTM FastAPI (Docker)
+emoji: 🧬
+colorFrom: blue
+colorTo: indigo
 sdk: docker
 pinned: false
+app_port: 7860
 ---
+# BIOBERT-FTM — FastAPI (Docker)
+REST endpoints:
+- `GET /health`
+- `POST /ner` { "text": "…", "score_threshold": 0.0 }
+- `POST /ner_batch` { "texts": ["…","…"], "score_threshold": 0.0 }
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Environment variables (override in Space settings):
+- `TOKENIZER_REPO_ID` (default: Milad96/BIOBERT-FTM-mlm)
+- `TASK_MODEL_REPO_ID` (default: Milad96/BIOBERT-FTM-tasks)
+- `NER_SUBFOLDER` (default: ner-spyysalo_bc2gm_corpus)

app.py

@@ -0,0 +1,77 @@
+import os
+from typing import List, Dict, Any
+from fastapi import FastAPI
+from pydantic import BaseModel
+from transformers import AutoTokenizer, AutoConfig, AutoModelForTokenClassification, pipeline
+import torch
+
+# ---- ENV with sensible defaults ----
+TOKENIZER_REPO_ID = os.getenv("TOKENIZER_REPO_ID", "Milad96/BIOBERT-FTM-mlm")
+TASK_MODEL_REPO_ID = os.getenv("TASK_MODEL_REPO_ID", "Milad96/BIOBERT-FTM-tasks")
+NER_SUBFOLDER = os.getenv("NER_SUBFOLDER", "ner-spyysalo_bc2gm_corpus")
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+# ---- Load tokenizer/model ----
+tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_REPO_ID, use_fast=True)
+config = AutoConfig.from_pretrained(TASK_MODEL_REPO_ID, subfolder=NER_SUBFOLDER)
+model = AutoModelForTokenClassification.from_pretrained(TASK_MODEL_REPO_ID, subfolder=NER_SUBFOLDER)
+model.to(DEVICE)
+model.eval()
+
+ner_pipe = pipeline(
+    "token-classification",
+    model=model,
+    tokenizer=tokenizer,
+    aggregation_strategy="simple",
+    device=0 if DEVICE == "cuda" else -1,
+    truncation=True
+)
+
+LABELS = sorted(set(l.replace("B-","").replace("I-","") for l in config.id2label.values()) | {"O"})
+
+class NerRequest(BaseModel):
+    text: str
+    score_threshold: float = 0.0
+
+class NerBatchRequest(BaseModel):
+    texts: List[str]
+    score_threshold: float = 0.0
+
+app = FastAPI(title="BIOBERT-FTM NER API", version="1.0")
+
+@app.get("/health")
+def health() -> Dict[str, Any]:
+    return {"status": "ok", "device": DEVICE, "labels": sorted(list(LABELS))}
+
+@app.post("/ner")
+def ner(req: NerRequest) -> Dict[str, Any]:
+    out = ner_pipe(req.text)
+    spans = []
+    for s in out:
+        sc = float(s.get("score", 0.0))
+        if sc < float(req.score_threshold):
+            continue
+        st = int(s.get("start", 0))
+        ed = int(s.get("end", 0))
+        spans.append({
+            "entity": str(s.get("entity_group", "")),
+            "start": st,
+            "end": ed,
+            "score": sc,
+            "text": req.text[st:ed]
+        })
+    return {"spans": spans, "count": len(spans)}
+
+@app.post("/ner_batch")
+def ner_batch(req: NerBatchRequest) -> List[Dict[str, Any]]:
+    results = []
+    for t in req.texts:
+        single = ner({"text": t, "score_threshold": req.score_threshold})  # reuse
+        results.append(single)
+    return results
+
+if __name__ == "__main__":
+    import uvicorn, os
+    host = os.getenv("HOST", "0.0.0.0")
+    port = int(os.getenv("PORT", "7860"))
+    uvicorn.run(app, host=host, port=port, log_level="info")

requirements.txt

@@ -0,0 +1,5 @@
+fastapi==0.115.5
+uvicorn[standard]==0.31.0
+torch
+transformers==4.57.1
+huggingface_hub>=0.25.0