Update app.py

app.py

@@ -1,5 +1,10 @@
-# 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]
+# 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
 
 import os, tarfile, glob, json, shutil, urllib.request
 from typing import List, Optional
@@ -8,42 +13,42 @@ from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 
-# Import TensorFlow before Transformers TF models to avoid odd init order issues
+# Import TF first
 import tensorflow as tf  # noqa: F401
 from transformers import BertTokenizer, BertConfig, TFBertModel
 
-# ---------------------------- Config ----------------------------
+# 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 link (Dropbox must end with dl=1)
+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"))
 
 os.makedirs(MODEL_ROOT, exist_ok=True)
 
-# ---------------------- 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):
             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):
+        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)
 
-# ---------------------- Bootstrap weights ----------------------
-
 def ensure_weights_and_locate() -> (str, str):
     """
     Returns:
-      model_dir: directory containing vocab.txt/config.json/checkpoint + ckpt files
+      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
     """
-    # 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)
@@ -56,20 +61,20 @@ def ensure_weights_and_locate() -> (str, str):
     if not maybe_idx:
         raise RuntimeError(f"No TensorFlow checkpoint *.index found under {MODEL_ROOT}")
 
-    # Prefer shortest path depth (avoids weird nested dirs)
+    # 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", "")
 
-    # Ensure checkpoint meta file points to the basename
+    # 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')
 
-    # Ensure config.json (copy from bert_config.json if present, else write default BERT base config)
+    # 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):
@@ -90,13 +95,13 @@ def ensure_weights_and_locate() -> (str, str):
                     "vocab_size": 30522
                 }, f)
 
-    # Ensure vocab.txt (BioBERT uses BERT base uncased vocab)
+    # 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)
 
-    # Sanity: ensure data shard exists
+    # 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)")
@@ -107,27 +112,45 @@ def ensure_weights_and_locate() -> (str, str):
 
 MODEL_DIR, CKPT_PREFIX = ensure_weights_and_locate()
 
-# ---------------------- Load tokenizer & model ------------------
-
+# 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"))
 
-# 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")
+# 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")
 
-# Optional: allow your website to call this API directly
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # tighten in production
+    allow_origins=["*"],
     allow_credentials=False,
     allow_methods=["GET", "POST", "OPTIONS"],
     allow_headers=["*"],
@@ -148,7 +171,6 @@ def health():
 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: