Update app.py

app.py

@@ -1,22 +1,95 @@
-import os
+# app.py — self-bootstrapping TF BioBERT embeddings API (HF Spaces-friendly)
+
+import os, tarfile, glob, json, shutil, urllib.request
 from fastapi import FastAPI
 from pydantic import BaseModel
+from typing import List
 from transformers import BertTokenizer, BertConfig, TFBertModel
-import tensorflow as tf
+import tensorflow as tf  # noqa
 
 app = FastAPI()
 
-# start.sh exports this after extraction; keep a fallback for local/dev
-MODEL_DIR = os.environ.get("MODEL_DIR", "/app/bert_tf")
-os.makedirs(MODEL_DIR, exist_ok=True)
+# --- Config
+MODEL_ROOT = os.environ.get("MODEL_ROOT", "/app/bert_tf")
+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"
+
+os.makedirs(MODEL_ROOT, exist_ok=True)
+
+def _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):
+                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:
+        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)
+
+    # 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])
+
+    # Ensure checkpoint file points at the basename
+    basename = os.path.basename(idx_files[0]).replace(".index", "")
+    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")
+    bcfg = os.path.join(model_dir, "bert_config.json")
+    if not os.path.isfile(cfg):
+        if os.path.isfile(bcfg):
+            shutil.copy(bcfg, cfg)
+        else:
+            with open(cfg, "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)
+
+    # Ensure 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)
 
-# extra safety: if no vocab here, look 2 levels deep
-if not os.path.isfile(os.path.join(MODEL_DIR, "vocab.txt")):
-    for root, dirs, files in os.walk(MODEL_DIR):
-        if "vocab.txt" in files and "config.json" in files:
-            MODEL_DIR = root
-            break
+    return model_dir
 
+# Prepare weights (download/extract if needed), then load model
+MODEL_DIR = ensure_weights_and_get_model_dir()
 print("[app] Using MODEL_DIR:", MODEL_DIR)
 
 tok  = BertTokenizer(vocab_file=os.path.join(MODEL_DIR, "vocab.txt"), do_lower_case=True)
@@ -31,9 +104,8 @@ def health():
     return {"ok": True}
 
 @app.post("/v1/embeddings")
-def emb(req: EmbReq):
-    ids = tok(req.input, return_tensors="tf", truncation=True, max_length=128)
-    out = model(**ids)
-    # [CLS] pooled output
+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()
     return {"embedding": vec, "dim": len(vec)}