import html
import os
import tempfile
import threading
from functools import lru_cache
from pathlib import Path
from urllib.parse import quote, urlparse
import gradio as gr
import h5py
import numpy as np
import pandas as pd
import tensorflow as tf
from huggingface_hub import hf_hub_download
from tensorflow import keras
from dl_model_def import build_two_tower_model
from runtime_data import build_result_annotations
# ============================================
# CONFIG
# ============================================
APP_ROOT = Path(__file__).resolve().parent
DATA_DIR = APP_ROOT / "data" / "proc"
MODEL_REPO_ID = os.environ.get("OTREC_MODEL_REPO_ID", "GrimSqueaker/OTRec")
MODEL_FILENAME = os.environ.get("OTREC_MODEL_FILENAME", "model.weights.h5")
MODEL_DOWNLOAD_ETAG_TIMEOUT = int(os.environ.get("OTREC_HF_ETAG_TIMEOUT", "30"))
FILTER_HIDE_PACKAGED = "Novel predictions only"
FILTER_ALL = "Known + novel"
FILTER_ONLY_PACKAGED = "Known clinical targets only"
DISPLAY_COLUMNS = [
"Rank",
"Gene",
"Gene Name",
"OTRec score",
"Open Targets score",
"OTTree score",
"Tractability",
"Known/novel",
"Function",
]
DISEASE_DISPLAY_COLUMNS = [
"Rank",
"Disease",
"Disease ID",
"OTRec score",
"Description",
]
SORT_OPTIONS = {
"OTRec score": ("otrec_score", True),
"Open Targets association score": ("ot_score", True),
"OTTree score": ("ottree_pred", True),
"Gene symbol": ("approvedSymbol", False),
}
RUNTIME_LOCK = threading.RLock()
# ============================================
# LOAD TRAINING DATA
# ============================================
df_learn = pd.read_parquet(DATA_DIR / "df_learn_sub.parquet")
disease_df = pd.read_parquet(DATA_DIR / "disease_df.parquet")
target_df = pd.read_parquet(DATA_DIR / "target_df.parquet")
# Ensure column names match training
df_learn = df_learn.rename(
columns={"disease_text_embed": "disease_text", "target_text_embed": "target_text"},
errors="ignore",
)
disease_df.rename(
columns={"disease_text_embed": "disease_text"}, errors="ignore", inplace=True
)
disease_df = disease_df.copy()
disease_df["diseaseId"] = disease_df["diseaseId"].astype(str)
disease_df["name"] = disease_df["name"].fillna("").astype(str)
target_df.rename(
columns={"target_text_embed": "target_text"}, errors="ignore", inplace=True
)
BATCH_SIZE = 1024
# Precomputed lowercase search blobs: one plain-substring scan per keystroke
# instead of five case-insensitive regex scans. The newline separator prevents
# accidental cross-field matches (queries never contain newlines).
_DISEASE_SEARCH_BLOB = (
disease_df["name"].astype(str)
+ "\n" + disease_df["diseaseId"].astype(str)
+ "\n" + disease_df["synonyms"].astype(str)
+ "\n" + disease_df["ExactSynonyms"].astype(str)
+ "\n" + disease_df["description"].astype(str)
).str.lower()
_TARGET_SEARCH_BLOB = (
target_df["approvedSymbol"].astype(str)
+ "\n" + target_df["targetId"].astype(str)
+ "\n" + target_df["approvedName"].astype(str)
).str.lower()
def _load_weights_file() -> str:
last_error = None
for attempt in range(2):
try:
print("Downloading model weights from Hugging Face Hub...")
weights_file = hf_hub_download(
repo_id=MODEL_REPO_ID,
filename=MODEL_FILENAME,
etag_timeout=MODEL_DOWNLOAD_ETAG_TIMEOUT,
)
print(f"Weights downloaded to: {weights_file}")
return weights_file
except Exception as error:
last_error = error
if attempt == 0:
print(f"Model download failed ({error}); retrying once...")
continue
raise RuntimeError(
"Unable to download OTRec model weights from Hugging Face. "
f"Check network access and that {MODEL_REPO_ID}/{MODEL_FILENAME} "
"is reachable, or set OTREC_MODEL_REPO_ID / OTREC_MODEL_FILENAME. "
f"Last error: {error}"
) from error
raise RuntimeError(f"Unable to download OTRec model weights: {last_error}")
def _load_model_with_weights():
weights_file = _load_weights_file()
print("Building TwoTowerDual...")
keras.backend.clear_session()
loaded_model = build_two_tower_model(df_learn)
# Restore the training-time vocabularies. Without this, TextVectorization
# is re-adapted from df_learn_sub, whose token FREQUENCIES differ from the
# full training frame; the same-sized vocabulary comes out in a different
# ORDER, weights load without error against permuted features, and
# predictions silently degrade to chance (measured ROC-AUC 0.52 vs 0.95).
vocab_path = DATA_DIR / "vocabs.json.gz"
if vocab_path.exists():
from vocab_io import apply_vocabularies, load_vocabularies
print("Applying saved training-time vocabularies...")
apply_vocabularies(loaded_model, load_vocabularies(vocab_path))
else:
print(
"WARNING: data/proc/vocabs.json.gz not found -- falling back to "
"re-adapted vocabularies, which are only correct if df_learn_sub "
"reproduces the training frame's exact token frequencies."
)
print("Loading weights...")
try:
loaded_model.load_weights(weights_file)
except ValueError as error:
print(f"Standard load failed ({error}). Attempting name-mismatch fix...")
with h5py.File(weights_file, "r") as handle:
h5_keys = list(handle.keys())
print(f"Weights file contains layers: {h5_keys}")
def match_layer_name(target_attr, prefix):
match = next((key for key in h5_keys if key.startswith(prefix)), None)
if match and hasattr(loaded_model, target_attr):
layer = getattr(loaded_model, target_attr)
print(
f"Renaming model layer '{layer.name}' to '{match}' to match file."
)
layer._name = match
match_layer_name("dise_emb", "dise_emb")
match_layer_name("q_tower", "tower")
loaded_model.load_weights(weights_file)
print("Weights loaded successfully.")
return loaded_model
def _weights_md5() -> str:
import hashlib
try:
return hashlib.md5(Path(_load_weights_file()).read_bytes()).hexdigest()
except Exception:
return ""
def _load_packaged_embeddings(loaded_model):
"""Load precomputed candidate/disease embeddings if present AND produced
from the exact weights file we just loaded (md5 guard). Returns
(cand_embs, dis_embs) or None -> caller computes live. Cuts the cold-start
embedding precompute (~1 min on the free CPU tier) to a file read."""
path = DATA_DIR / "embeddings.npz"
if not path.exists():
return None
try:
data = np.load(path, allow_pickle=False)
if str(data["weights_md5"]) != _weights_md5():
print("Packaged embeddings are for different weights; recomputing live.")
return None
cand, dis = data["candidate_embs"], data["disease_embs"]
if len(cand) != len(target_df) or len(dis) != len(disease_df):
print("Packaged embeddings shape mismatch; recomputing live.")
return None
print(f"Loaded packaged embeddings: candidates {cand.shape}, diseases {dis.shape}")
return cand.astype(np.float32), dis.astype(np.float32)
except Exception as error:
print(f"Packaged embeddings unreadable ({error}); recomputing live.")
return None
def _precompute_candidate_embeddings(loaded_model):
print("Precomputing candidate embeddings (batched)...")
target_texts = target_df["target_text"].astype(str).to_numpy()
target_ids = target_df["targetId"].astype(str).to_numpy()
cand_embs_list = []
total = len(target_texts)
for index in range(0, total, BATCH_SIZE):
end = min(index + BATCH_SIZE, total)
batch_txt = target_texts[index:end]
batch_ids = target_ids[index:end]
emb_batch = loaded_model.encode_k(batch_txt, batch_ids)
cand_embs_list.append(emb_batch)
if index % 5000 == 0:
print(f" Processed {index}/{total} candidates...")
cand_embs = tf.concat(cand_embs_list, axis=0)
cand_embs = tf.nn.l2_normalize(cand_embs, axis=1).numpy()
print(f"Candidate embeddings ready. Shape: {cand_embs.shape}")
return cand_embs
def _precompute_disease_embeddings(loaded_model):
print("Precomputing disease embeddings (batched)...")
disease_texts = disease_df["disease_text"].astype(str).to_numpy()
disease_ids = disease_df["diseaseId"].astype(str).to_numpy()
embs_list = []
total = len(disease_texts)
for index in range(0, total, BATCH_SIZE):
end = min(index + BATCH_SIZE, total)
batch_txt = disease_texts[index:end]
batch_ids = disease_ids[index:end]
emb_batch = loaded_model.encode_q(batch_txt, batch_ids)
embs_list.append(emb_batch)
if index % 5000 == 0:
print(f" Processed {index}/{total} diseases...")
dis_embs = tf.concat(embs_list, axis=0)
dis_embs = tf.nn.l2_normalize(dis_embs, axis=1).numpy()
print(f"Disease embeddings ready. Shape: {dis_embs.shape}")
return dis_embs
@lru_cache(maxsize=1)
def _get_runtime_cached():
loaded_model = _load_model_with_weights()
packaged = _load_packaged_embeddings(loaded_model)
if packaged is not None:
cand_embs, dis_embs = packaged
_PACKAGED_DISEASE_EMBS.append(dis_embs)
return loaded_model, cand_embs
cand_embs = _precompute_candidate_embeddings(loaded_model)
return loaded_model, cand_embs
_PACKAGED_DISEASE_EMBS: list = []
def get_runtime():
with RUNTIME_LOCK:
return _get_runtime_cached()
@lru_cache(maxsize=1)
def _get_disease_runtime_cached():
loaded_model, _ = get_runtime()
if _PACKAGED_DISEASE_EMBS:
return loaded_model, _PACKAGED_DISEASE_EMBS[0]
dis_embs = _precompute_disease_embeddings(loaded_model)
return loaded_model, dis_embs
def get_disease_runtime():
with RUNTIME_LOCK:
return _get_disease_runtime_cached()
target_df = target_df.copy()
target_df["targetId"] = target_df["targetId"].astype(str)
target_df["approvedSymbol"] = (
target_df["approvedSymbol"].fillna(target_df.get("sym", "")).astype(str)
)
target_df["approvedName"] = target_df["approvedName"].fillna("").astype(str)
FUNCTION_COLUMN = (
"functionDescriptions"
if "functionDescriptions" in target_df.columns
else "functionDescription"
)
def _empty_results() -> pd.DataFrame:
return pd.DataFrame(columns=DISPLAY_COLUMNS)
def _is_safe_url(value: str) -> bool:
parsed = urlparse(value)
return parsed.scheme in {"http", "https"} and bool(parsed.netloc)
def _render_results_table(
results: pd.DataFrame,
columns: list[str],
link_columns: dict[str, str] | None = None,
) -> str:
"""Render the ranked results as an HTML table.
link_columns maps a *display* column (e.g. "Gene") to the column holding
its URL (e.g. "Open Targets link"), turning that cell's text into a link to
the Open Targets platform page. The URL column does not need to be in
`columns` -- it is read from `results` directly.
"""
link_columns = link_columns or {}
header_cells = "".join(
f'| {html.escape(column)} | '
for column in columns
)
if results.empty:
body_rows = (
f'| '
"No rows to display for the current selection."
" |
"
)
else:
rendered_rows: list[str] = []
display_rows = results.reindex(columns=columns, fill_value="")
# URL sources are read positionally from `results`, so they survive the
# reindex above even when the URL column is not displayed.
url_sources = {
display_col: results[url_col].to_numpy()
for display_col, url_col in link_columns.items()
if url_col in results.columns
}
for position, (_, row) in enumerate(display_rows.iterrows()):
cells: list[str] = []
for column in columns:
value = row.get(column, "")
if pd.isna(value):
text = ""
else:
text = " ".join(str(value).split())
link_url = ""
if column in url_sources:
raw_url = url_sources[column][position]
if not pd.isna(raw_url):
link_url = str(raw_url).strip()
if column == "Open Targets link" and text and _is_safe_url(text):
cell_value = (
f'Open Targets'
)
elif text and link_url and _is_safe_url(link_url):
cell_value = (
f'{html.escape(text)}'
)
else:
cell_value = html.escape(text) if text else "—"
cells.append(
f'{cell_value} | '
)
rendered_rows.append(f"{''.join(cells)}
")
body_rows = "".join(rendered_rows)
return (
''
'
'
f"{header_cells}
"
f"{body_rows}"
"
"
"