Update app.py

app.py

@@ -1,72 +1,82 @@
-# app.py — self-bootstrapping TF BioBERT embeddings API (HF Spaces-friendly)
+# app.py — FastAPI TF-BioBERT embeddings service (handles TF1 checkpoints)
+# Requires: transformers==4.43.4, tensorflow-cpu==2.16.1, tf-keras, fastapi, uvicorn[standard]
 
 import os, tarfile, glob, json, shutil, urllib.request
+from typing import List, Optional
+
 from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
-from typing import List
+
+# Import TensorFlow before Transformers TF models to avoid odd init order issues
+import tensorflow as tf  # noqa: F401
 from transformers import BertTokenizer, BertConfig, TFBertModel
-import tensorflow as tf  # noqa
 
-app = FastAPI()
+# ---------------------------- Config ----------------------------
 
-# --- Config
-MODEL_ROOT = os.environ.get("MODEL_ROOT", "/app/bert_tf")
+MODEL_ROOT = os.environ.get("MODEL_ROOT", "/app/bert_tf").rstrip("/")
 WEIGHTS_URL = os.environ.get("WEIGHTS_URL_TAR_GZ", "").strip()  # direct .tar.gz link (Dropbox must end with dl=1)
 FALLBACK_VOCAB_URL = "https://huggingface.co/bert-base-uncased/resolve/main/vocab.txt"
+MAX_LEN = int(os.environ.get("MAX_LEN", "128"))
 
 os.makedirs(MODEL_ROOT, exist_ok=True)
 
-def _extract_tar_gz(src: str, dest: str) -> None:
+# ---------------------- Utils: safe extract ---------------------
+
+def _safe_extract_tar_gz(src: str, dest: str) -> None:
     with tarfile.open(src, "r:gz") as tar:
-        def is_within(directory, target):
+        def _is_within(directory, target):
             abs_directory = os.path.abspath(directory)
             abs_target = os.path.abspath(target)
             return os.path.commonpath([abs_directory]) == os.path.commonpath([abs_directory, abs_target])
         for member in tar.getmembers():
             target_path = os.path.join(dest, member.name)
-            if not is_within(dest, target_path):
+            if not _is_within(dest, target_path):
                 raise RuntimeError("Blocked path traversal in tar")
         tar.extractall(dest)
 
-def ensure_weights_and_get_model_dir() -> str:
-    # If already prepared (vocab + any ckpt index) → reuse
-    maybe_vocab = glob.glob(os.path.join(MODEL_ROOT, "**", "vocab.txt"), recursive=True)
-    maybe_idx   = glob.glob(os.path.join(MODEL_ROOT, "**", "model.ckpt-*.index"), recursive=True)
-    if maybe_vocab and maybe_idx:
-        # choose dir containing the first ckpt index
-        return os.path.dirname(maybe_idx[0])
-
-    # Otherwise download and extract the archive
-    if not WEIGHTS_URL:
-        print("[app] WEIGHTS_URL_TAR_GZ not set; will still try to run with fallback vocab if files exist.")
-    else:
+# ---------------------- Bootstrap weights ----------------------
+
+def ensure_weights_and_locate() -> (str, str):
+    """
+    Returns:
+      model_dir: directory containing vocab.txt/config.json/checkpoint + ckpt files
+      ckpt_prefix: full path WITHOUT extension, e.g. /app/bert_tf/bert_min/model.ckpt-150000
+    """
+    # Already present?
+    maybe_idx = glob.glob(os.path.join(MODEL_ROOT, "**", "model.ckpt-*.index"), recursive=True)
+    if not maybe_idx and WEIGHTS_URL:
         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)
-        _extract_tar_gz(local_tar, MODEL_ROOT)
+        _safe_extract_tar_gz(local_tar, MODEL_ROOT)
+        maybe_idx = glob.glob(os.path.join(MODEL_ROOT, "**", "model.ckpt-*.index"), recursive=True)
 
-    # Pick the folder that has a ckpt index
-    idx_files = glob.glob(os.path.join(MODEL_ROOT, "**", "model.ckpt-*.index"), recursive=True)
-    if not idx_files:
-        raise RuntimeError("No TensorFlow checkpoint index found under " + MODEL_ROOT)
-    model_dir = os.path.dirname(idx_files[0])
+    if not maybe_idx:
+        raise RuntimeError(f"No TensorFlow checkpoint *.index found under {MODEL_ROOT}")
 
-    # Ensure checkpoint file points at the basename
-    basename = os.path.basename(idx_files[0]).replace(".index", "")
+    # Prefer shortest path depth (avoids weird nested dirs)
+    maybe_idx.sort(key=lambda p: len(os.path.relpath(p, MODEL_ROOT).split(os.sep)))
+    ckpt_index = maybe_idx[0]
+    model_dir = os.path.dirname(ckpt_index)
+    ckpt_prefix = ckpt_index.replace(".index", "")
+
+    # Ensure checkpoint meta file points to the basename
+    basename = os.path.basename(ckpt_prefix)
     ckpt_meta = os.path.join(model_dir, "checkpoint")
     if not os.path.isfile(ckpt_meta):
         with open(ckpt_meta, "w") as f:
             f.write(f'model_checkpoint_path: "{basename}"\n')
 
-    # Ensure config.json
-    cfg = os.path.join(model_dir, "config.json")
+    # Ensure config.json (copy from bert_config.json if present, else write default BERT base config)
+    cfg_json = os.path.join(model_dir, "config.json")
     bcfg = os.path.join(model_dir, "bert_config.json")
-    if not os.path.isfile(cfg):
+    if not os.path.isfile(cfg_json):
         if os.path.isfile(bcfg):
-            shutil.copy(bcfg, cfg)
+            shutil.copy(bcfg, cfg_json)
         else:
-            with open(cfg, "w") as f:
+            with open(cfg_json, "w") as f:
                 json.dump({
                     "hidden_size": 768,
                     "num_attention_heads": 12,
@@ -86,26 +96,88 @@ def ensure_weights_and_get_model_dir() -> str:
         print("[app] vocab.txt missing; fetching BERT base uncased vocab…")
         urllib.request.urlretrieve(FALLBACK_VOCAB_URL, vocab)
 
-    return model_dir
+    # Sanity: ensure data shard exists
+    data_glob = glob.glob(os.path.join(model_dir, "model.ckpt-*.data-00000-of-00001"))
+    if not data_glob:
+        raise RuntimeError(f"Checkpoint data file missing in {model_dir} (model.ckpt-*.data-00000-of-00001)")
+
+    print("[app] Using MODEL_DIR:", model_dir)
+    print("[app] Using CKPT_PREFIX:", ckpt_prefix)
+    return model_dir, ckpt_prefix
+
+MODEL_DIR, CKPT_PREFIX = ensure_weights_and_locate()
 
-# Prepare weights (download/extract if needed), then load model
-MODEL_DIR = ensure_weights_and_get_model_dir()
-print("[app] Using MODEL_DIR:", MODEL_DIR)
+# ---------------------- Load tokenizer & model ------------------
 
-tok  = BertTokenizer(vocab_file=os.path.join(MODEL_DIR, "vocab.txt"), do_lower_case=True)
-cfg  = BertConfig.from_json_file(os.path.join(MODEL_DIR, "config.json"))
-model= TFBertModel.from_pretrained(MODEL_DIR, from_tf=True, config=cfg)
+tok = BertTokenizer(vocab_file=os.path.join(MODEL_DIR, "vocab.txt"), do_lower_case=True)
+cfg = BertConfig.from_json_file(os.path.join(MODEL_DIR, "config.json"))
+
+# IMPORTANT: load from TF1 checkpoint using the PREFIX (not folder)
+model = TFBertModel.from_pretrained(
+    CKPT_PREFIX,
+    from_tf=True,   # TF1 .ckpt import
+    from_pt=False,
+    config=cfg
+)
+
+# ---------------------------- API ------------------------------
+
+app = FastAPI(title="BioBERT-TF Embeddings API", version="1.0")
+
+# Optional: allow your website to call this API directly
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # tighten in production
+    allow_credentials=False,
+    allow_methods=["GET", "POST", "OPTIONS"],
+    allow_headers=["*"],
+)
 
 class EmbReq(BaseModel):
     input: str
+    max_len: Optional[int] = None
+
+class BatchEmbReq(BaseModel):
+    inputs: List[str]
+    max_len: Optional[int] = None
 
 @app.get("/health")
 def health():
-    return {"ok": True}
+    return {"ok": True, "model_dir": MODEL_DIR, "ckpt_prefix": CKPT_PREFIX}
+
+def _embed(texts: List[str], max_len: int) -> List[List[float]]:
+    enc = tok(texts, return_tensors="tf", truncation=True, padding=True, max_length=max_len)
+    out = model(**enc, training=False)
+    # Prefer pooled output if available; fallback to mean of last_hidden_state
+    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]
 
 @app.post("/v1/embeddings")
 def embeddings(req: EmbReq):
-    enc = tok(req.input, return_tensors="tf", truncation=True, max_length=128)
-    out = model(**enc)
-    vec = out.pooler_output[0].numpy().tolist()
+    text = req.input.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)}
+
+@app.post("/v1/embeddings/batch")
+def embeddings_batch(req: BatchEmbReq):
+    items = [t.strip() for t in req.inputs 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-TF Embeddings",
+        "endpoints": ["/health", "/v1/embeddings", "/v1/embeddings/batch"],
+        "hint": "POST to /v1/embeddings with {'input': 'your text'}"
+    }