Update app.py

app.py

@@ -1,120 +1,33 @@
-# FastAPI BioBERT embeddings (Hub-first, no TF1 ckpt dependency)
-# Works free on Hugging Face Spaces (CPU). Auto-converts PyTorch -> TF.
+# app.py — FastAPI embeddings service using PyTorch BioBERT
+# Works on Hugging Face Spaces (CPU Basic, free)
 
-import os, tarfile, glob, json, shutil, urllib.request
+import os
 from typing import List, Optional
 
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 
-# Load TF before transformers' TF models
-import tensorflow as tf  # noqa: F401
-from transformers import AutoTokenizer, TFAutoModel, BertConfig
+import torch
+from transformers import AutoTokenizer, AutoModel
 
-# ------------------- Config -------------------
 HF_MODEL_ID = os.environ.get("HF_MODEL_ID", "monologg/biobert_v1.1_pubmed").strip()
-MODEL_ROOT  = os.environ.get("MODEL_ROOT", "/app/bert_tf").rstrip("/")
-WEIGHTS_URL = os.environ.get("WEIGHTS_URL_TAR_GZ", "").strip()  # optional direct .tar.gz (Dropbox must end with dl=1)
-MAX_LEN     = int(os.environ.get("MAX_LEN", "128"))
-os.makedirs(MODEL_ROOT, exist_ok=True)
-
-# ------------------- Utils --------------------
-def _safe_extract_tar_gz(src: str, dest: str) -> None:
-    with tarfile.open(src, "r:gz") as tar:
-        def _is_within(directory, target):
-            ad = os.path.abspath(directory); at = os.path.abspath(target)
-            return os.path.commonpath([ad]) == os.path.commonpath([ad, at])
-        for m in tar.getmembers():
-            tp = os.path.join(dest, m.name)
-            if not _is_within(dest, tp):
-                raise RuntimeError("Blocked path traversal in tar")
-        tar.extractall(dest)
-
-def _maybe_download_tar_into_model_root() -> Optional[str]:
-    """If WEIGHTS_URL is set, download + extract it into MODEL_ROOT. Return extracted dir if any."""
-    if not WEIGHTS_URL:
-        return None
-    print("[app] downloading weights:", WEIGHTS_URL)
-    local_tar = "/tmp/model.tar.gz"
-    urllib.request.urlretrieve(WEIGHTS_URL, local_tar)
-    print("[app] extracting:", local_tar, "->", MODEL_ROOT)
-    _safe_extract_tar_gz(local_tar, MODEL_ROOT)
-    # return shallowest dir inside MODEL_ROOT
-    candidates = [d for d in glob.glob(os.path.join(MODEL_ROOT, "*")) if os.path.isdir(d)]
-    if not candidates:
-        return MODEL_ROOT
-    candidates.sort(key=lambda p: len(os.path.relpath(p, MODEL_ROOT).split(os.sep)))
-    return candidates[0]
-
-def _detect_local_hf_dir(root: str) -> Optional[str]:
-    """
-    Return a directory under root that looks like a modern HF model folder:
-      - pytorch_model.bin / model.safetensors  (for from_pt=True)
-      - OR tf_model.h5 (native TF)
-    """
-    # search at depth 0/1/2
-    for depth in range(3):
-        pattern = os.path.join(root, *(["**"] if depth else []))
-        # prefer TF weights first if present
-        tf_h5 = glob.glob(os.path.join(pattern, "tf_model.h5"), recursive=True)
-        if tf_h5:
-            tf_h5.sort(key=lambda p: len(os.path.relpath(p, root).split(os.sep)))
-            return os.path.dirname(tf_h5[0])
-
-        # else look for PT/safetensors
-        pt = glob.glob(os.path.join(pattern, "pytorch_model.bin"), recursive=True)
-        st = glob.glob(os.path.join(pattern, "model.safetensors"), recursive=True)
-        have = (pt or st)
-        if have:
-            have.sort(key=lambda p: len(os.path.relpath(p, root).split(os.sep)))
-            return os.path.dirname(have[0])
-    return None
-
-def _looks_like_tf1_ckpt_dir(path: str) -> bool:
-    return bool(glob.glob(os.path.join(path, "model.ckpt-*.index")))
-
-# ------------------- Load strategy -------------------
-# 1) If a tar URL is provided, unpack it (optional convenience)
-extracted = _maybe_download_tar_into_model_root()
-
-# 2) If after extraction we have a local HF-style folder, use it
-LOCAL_DIR = _detect_local_hf_dir(MODEL_ROOT)
-
-# 3) If only TF1 ckpt found, refuse with a clear message (no fragile loaders)
-if not LOCAL_DIR:
-    # If there is any directory in MODEL_ROOT with TF1 ckpts, warn
-    for d in [MODEL_ROOT] + [p for p in glob.glob(os.path.join(MODEL_ROOT, "*")) if os.path.isdir(p)]:
-        if _looks_like_tf1_ckpt_dir(d):
-            raise RuntimeError(
-                "Found TF-1 checkpoint files (model.ckpt-*) but this app purposely avoids "
-                "runtime TF-1 → TF-2 weight mapping. Either:\n"
-                "  • Set HF_MODEL_ID to a BioBERT model on the Hub (recommended), e.g. 'monologg/biobert_v1.1_pubmed'\n"
-                "  • Or package modern HF weights (pytorch_model.bin/model.safetensors or tf_model.h5) in your tar."
-            )
-
-# 4) Tokenizer+Model
-if LOCAL_DIR:
-    print(f"[app] Using LOCAL_DIR: {LOCAL_DIR}")
-    # Prefer native TF if available, else convert from PT
-    if os.path.isfile(os.path.join(LOCAL_DIR, "tf_model.h5")):
-        tokenizer = AutoTokenizer.from_pretrained(LOCAL_DIR)
-        model     = TFAutoModel.from_pretrained(LOCAL_DIR)
-        USED = {"source": "local", "format": "tf_h5", "path": LOCAL_DIR}
-    else:
-        tokenizer = AutoTokenizer.from_pretrained(LOCAL_DIR)
-        model     = TFAutoModel.from_pretrained(LOCAL_DIR, from_pt=True)
-        USED = {"source": "local", "format": "pt/safetensors->tf", "path": LOCAL_DIR}
-else:
-    print(f"[app] Using HF_MODEL_ID: {HF_MODEL_ID}")
-    tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_ID)
-    # Most BioBERT repos are PyTorch; allow auto-conversion
-    model     = TFAutoModel.from_pretrained(HF_MODEL_ID, from_pt=True)
-    USED = {"source": "hub", "model_id": HF_MODEL_ID, "format": "pt->tf"}
-
-# ------------------- API -------------------
-app = FastAPI(title="BioBERT Embeddings API (Hub-first)", version="2.0")
+MAX_LEN = int(os.environ.get("MAX_LEN", "128"))
+TORCH_THREADS = int(os.environ.get("TORCH_THREADS", "1"))
+
+torch.set_num_threads(TORCH_THREADS)
+
+# --------- Load model & tokenizer (PyTorch) ----------
+tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_ID)
+model = AutoModel.from_pretrained(HF_MODEL_ID)
+model.eval()  # inference mode
+DEVICE = "cpu"
+model.to(DEVICE)
+
+# --------- FastAPI ----------
+app = FastAPI(title="BioBERT (PyTorch) Embeddings API", version="1.0")
 
+# CORS (relax; tighten in production)
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -126,48 +39,73 @@ app.add_middleware(
 class EmbReq(BaseModel):
     input: str
     max_len: Optional[int] = None
+    pooling: Optional[str] = "cls"  # "cls" or "mean"
 
 class BatchEmbReq(BaseModel):
     inputs: List[str]
     max_len: Optional[int] = None
+    pooling: Optional[str] = "cls"  # "cls" or "mean"
+
+@app.get("/")
+def root():
+    return {
+        "name": "BioBERT Embeddings (PyTorch)",
+        "model": HF_MODEL_ID,
+        "device": DEVICE,
+        "endpoints": ["/health", "/v1/embeddings", "/v1/embeddings/batch"],
+        "hint": "POST to /v1/embeddings with {'input': 'your text'}",
+    }
 
 @app.get("/health")
 def health():
-    return {"ok": True, "strategy": USED, "max_len_default": MAX_LEN}
-
-def _embed(texts: List[str], max_len: int) -> List[List[float]]:
-    enc = tokenizer(texts, return_tensors="tf", truncation=True, padding=True, max_length=max_len)
-    out = model(**enc, training=False)
-    if hasattr(out, "pooler_output") and out.pooler_output is not None:
-        vecs = out.pooler_output.numpy()
-    else:
-        last = out.last_hidden_state.numpy()
-        vecs = last.mean(axis=1)
-    return [v.tolist() for v in vecs]
+    return {"ok": True, "model": HF_MODEL_ID, "device": DEVICE}
+
+def _pool(outputs, inputs, pooling: str):
+    """
+    pooling="cls": use CLS (pooler_output if present, else hidden_state[:,0])
+    pooling="mean": mean of token embeddings (mask-aware)
+    """
+    if pooling == "mean":
+        last = outputs.last_hidden_state  # [B,T,H]
+        mask = inputs["attention_mask"].unsqueeze(-1).type_as(last)  # [B,T,1]
+        summed = (last * mask).sum(dim=1)
+        counts = mask.sum(dim=1).clamp(min=1e-9)
+        return summed / counts
+    # cls
+    if hasattr(outputs, "pooler_output") and outputs.pooler_output is not None:
+        return outputs.pooler_output
+    return outputs.last_hidden_state[:, 0, :]  # CLS token
+
+def _embed(texts: List[str], max_len: int, pooling: str) -> List[List[float]]:
+    enc = tokenizer(
+        texts,
+        return_tensors="pt",
+        padding=True,
+        truncation=True,
+        max_length=max_len,
+    )
+    enc = {k: v.to(DEVICE) for k, v in enc.items()}
+    with torch.no_grad():
+        outputs = model(**enc)
+        vecs = _pool(outputs, enc, pooling=pooling)
+    return vecs.cpu().numpy().tolist()
 
 @app.post("/v1/embeddings")
 def embeddings(req: EmbReq):
-    text = req.input.strip()
+    text = (req.input or "").strip()
     if not text:
         return {"embedding": [], "dim": 0}
     L = int(req.max_len or MAX_LEN)
-    vec = _embed([text], L)[0]
-    return {"embedding": vec, "dim": len(vec)}
+    pooling = (req.pooling or "cls").lower()
+    vec = _embed([text], L, pooling)[0]
+    return {"embedding": vec, "dim": len(vec), "pooling": pooling}
 
 @app.post("/v1/embeddings/batch")
 def embeddings_batch(req: BatchEmbReq):
-    items = [t.strip() for t in req.inputs if str(t).strip()]
+    items = [str(t).strip() for t in (req.inputs or []) if str(t).strip()]
     if not items:
         return {"embeddings": [], "dim": 0}
     L = int(req.max_len or MAX_LEN)
-    vecs = _embed(items, L)
-    return {"embeddings": vecs, "dim": len(vecs[0])}
-
-@app.get("/")
-def root():
-    return {
-        "name": "BioBERT Embeddings (Hub-first)",
-        "endpoints": ["/health", "/v1/embeddings", "/v1/embeddings/batch"],
-        "hint": "POST /v1/embeddings with {'input': 'your text'}",
-        "strategy": USED
-    }
+    pooling = (req.pooling or "cls").lower()
+    vecs = _embed(items, L, pooling)
+    return {"embeddings": vecs, "dim": len(vecs[0]), "pooling": pooling}