Spaces:
Sleeping
Sleeping
File size: 2,341 Bytes
842daa2 d213edf 1c0323e d213edf 9bd55b2 5d7a5a6 d213edf 1c0323e d213edf 9b01576 d213edf 25bfd3b d213edf 9b01576 d213edf 1c0323e d213edf 25bfd3b 5d7a5a6 1c0323e d213edf 1c0323e d213edf 1c0323e d213edf 5d7a5a6 d213edf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 |
from fastapi import FastAPI
from pydantic import BaseModel, Field
from typing import List
from transformers import AutoTokenizer, AutoModel
import torch, os
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):
token_embeddings = model_output[0] # [batch, seq, hidden]
mask = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
summed = (token_embeddings * mask).sum(1)
counts = mask.sum(1).clamp(min=1e-9)
return summed / counts
class EmbedRequest(BaseModel):
texts: List[str] = Field(default_factory=list)
max_length: int = 256
class EmbedResponse(BaseModel):
embeddings: List[List[float]]
app = FastAPI(title="BioBERT Embeddings", version="1.0")
@app.get("/healthz")
def health():
return {"ok": True, "model_id": MODEL_ID}
@app.post("/embed", response_model=EmbedResponse)
def embed(req: EmbedRequest):
if not req.texts:
return {"embeddings": []}
enc = tokenizer(
req.texts, padding=True, truncation=True,
max_length=req.max_length, return_tensors="pt"
)
with torch.no_grad():
out = model(**enc)
pooled = mean_pooling(out, enc["attention_mask"])
pooled = torch.nn.functional.normalize(pooled, p=2, dim=1)
return {"embeddings": pooled.cpu().tolist()}
if __name__ == "__main__":
import uvicorn, os
uvicorn.run("app:app", host="0.0.0.0", port=int(os.getenv("PORT", "7860")), workers=1)
|