Update app.py

app.py

@@ -1,10 +1,5 @@
-# app.py — FastAPI TF-BioBERT embeddings service (TF1 checkpoint loader)
-# Pin these (requirements.txt):
-#   fastapi
-#   uvicorn[standard]
-#   transformers==4.43.4
-#   tensorflow-cpu==2.16.1
-#   tf-keras
+# FastAPI BioBERT embeddings (Hub-first, no TF1 ckpt dependency)
+# Works free on Hugging Face Spaces (CPU). Auto-converts PyTorch -> TF.
 
 import os, tarfile, glob, json, shutil, urllib.request
 from typing import List, Optional
@@ -13,140 +8,112 @@ from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 
-# Import TF first
+# Load TF before transformers' TF models
 import tensorflow as tf  # noqa: F401
-from transformers import BertTokenizer, BertConfig, TFBertModel
-
-# For TF1 checkpoint loading
-try:
-    # Present in transformers TF BERT module
-    from transformers.models.bert.modeling_tf_bert import load_tf_weights_in_bert as _hf_load_tf_ckpt
-except Exception:
-    _hf_load_tf_ckpt = None
-
-MODEL_ROOT = os.environ.get("MODEL_ROOT", "/app/bert_tf").rstrip("/")
-WEIGHTS_URL = os.environ.get("WEIGHTS_URL_TAR_GZ", "").strip()  # direct .tar.gz (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"))
+from transformers import AutoTokenizer, TFAutoModel, BertConfig
 
+# ------------------- 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):
-            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])
+            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 ensure_weights_and_locate() -> (str, str):
+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]:
     """
-    Returns:
-      model_dir: folder containing vocab/config/checkpoint + ckpt files
-      ckpt_prefix: full path WITHOUT extension, e.g. /app/bert_tf/bert_min/model.ckpt-150000
+    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)
     """
-    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)
-        _safe_extract_tar_gz(local_tar, MODEL_ROOT)
-        maybe_idx = glob.glob(os.path.join(MODEL_ROOT, "**", "model.ckpt-*.index"), recursive=True)
-
-    if not maybe_idx:
-        raise RuntimeError(f"No TensorFlow checkpoint *.index found under {MODEL_ROOT}")
-
-    # Prefer shallowest
-    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", "")
-
-    # checkpoint meta
-    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')
-
-    # config.json (copy bert_config.json if present)
-    cfg_json = os.path.join(model_dir, "config.json")
-    bcfg = os.path.join(model_dir, "bert_config.json")
-    if not os.path.isfile(cfg_json):
-        if os.path.isfile(bcfg):
-            shutil.copy(bcfg, cfg_json)
-        else:
-            with open(cfg_json, "w") as f:
-                json.dump({
-                    "hidden_size": 768,
-                    "num_attention_heads": 12,
-                    "num_hidden_layers": 12,
-                    "intermediate_size": 3072,
-                    "hidden_act": "gelu",
-                    "hidden_dropout_prob": 0.1,
-                    "attention_probs_dropout_prob": 0.1,
-                    "max_position_embeddings": 512,
-                    "type_vocab_size": 2,
-                    "vocab_size": 30522
-                }, f)
-
-    # vocab.txt (BioBERT uses BERT base uncased vocab)
-    vocab = os.path.join(model_dir, "vocab.txt")
-    if not os.path.isfile(vocab):
-        print("[app] vocab.txt missing; fetching BERT base uncased vocab…")
-        urllib.request.urlretrieve(FALLBACK_VOCAB_URL, vocab)
-
-    # data shard sanity
-    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()
-
-# Tokenizer + config
-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"))
-
-# Build model skeleton
-model = TFBertModel(cfg)
-
-# Load TF1 checkpoint (no from_tf kwarg!)
-loaded = False
-err_stack = []
-
-if _hf_load_tf_ckpt is not None:
-    try:
-        # Some transformer versions: (model, ckpt_prefix)
-        _hf_load_tf_ckpt(model, CKPT_PREFIX)
-        loaded = True
-        print("[app] Loaded TF1 checkpoint via load_tf_weights_in_bert(model, ckpt_prefix)")
-    except TypeError as e1:
-        err_stack.append(str(e1))
-        try:
-            # Other versions: (model, config, ckpt_prefix)
-            _hf_load_tf_ckpt(model, cfg, CKPT_PREFIX)
-            loaded = True
-            print("[app] Loaded TF1 checkpoint via load_tf_weights_in_bert(model, config, ckpt_prefix)")
-        except Exception as e2:
-            err_stack.append(str(e2))
-
-if not loaded:
-    raise RuntimeError(
-        "Could not load TF1 checkpoint with transformers' loader. "
-        f"ckpt={CKPT_PREFIX}\nErrors: {err_stack or 'no loader available'}"
-    )
-
-# ---------- API ----------
-app = FastAPI(title="BioBERT-TF Embeddings API", version="1.1")
+    # 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")
 
 app.add_middleware(
     CORSMiddleware,
@@ -166,10 +133,10 @@ class BatchEmbReq(BaseModel):
 
 @app.get("/health")
 def health():
-    return {"ok": True, "model_dir": MODEL_DIR, "ckpt_prefix": CKPT_PREFIX}
+    return {"ok": True, "strategy": USED, "max_len_default": MAX_LEN}
 
 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)
+    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()
@@ -199,7 +166,8 @@ def embeddings_batch(req: BatchEmbReq):
 @app.get("/")
 def root():
     return {
-        "name": "BioBERT-TF Embeddings",
+        "name": "BioBERT Embeddings (Hub-first)",
         "endpoints": ["/health", "/v1/embeddings", "/v1/embeddings/batch"],
-        "hint": "POST to /v1/embeddings with {'input': 'your text'}"
+        "hint": "POST /v1/embeddings with {'input': 'your text'}",
+        "strategy": USED
     }