Spaces:
Paused
Paused
| #!/usr/bin/env python3 | |
| from __future__ import annotations | |
| # ================== STD / CORE ================== | |
| import os, re, json, time, sys, csv, uuid, datetime | |
| from typing import List, Dict, Any, Optional | |
| from functools import lru_cache | |
| from xml.etree import ElementTree as ET | |
| # ================== TRANSFORMERS / TORCH ================== | |
| from transformers import AutoTokenizer, AutoModelForCausalLM | |
| try: | |
| from transformers import BitsAndBytesConfig # may exist even if bnb isn't installed | |
| except Exception: | |
| BitsAndBytesConfig = None | |
| import numpy as np | |
| import requests | |
| import gradio as gr | |
| # ================== CONFIG / PATHS ================== | |
| ASSETS_DIR = os.environ.get("ASSETS_DIR", "assets") | |
| FAISS_PATH = os.environ.get("FAISS_PATH", f"{ASSETS_DIR}/index.faiss") | |
| META_PATH = os.environ.get("META_PATH", f"{ASSETS_DIR}/index_meta.filtered.jsonl") | |
| REL_CONFIG_PATH = os.environ.get("REL_CONFIG_PATH", f"{ASSETS_DIR}/relevance_config.json") | |
| # Normalize QUANTIZE env (default: no quantization) | |
| QUANTIZE = os.environ.get("QUANTIZE", "none").strip().lower() # "none" | "8bit" | "4bit" | |
| # Turn logging | |
| TRANSCRIPT_PATH = os.environ.get("TRANSCRIPT_PATH", "transcripts.jsonl") | |
| PUSH_TRANSCRIPTS = os.environ.get("PUSH_TRANSCRIPTS", "1") == "1" # set "0" to disable | |
| # Models | |
| BASE_MODEL = os.environ.get("BASE_MODEL", "mistralai/Mistral-7B-Instruct-v0.2") | |
| ADAPTER_PATH = os.environ.get("ADAPTER_PATH", f"{ASSETS_DIR}/lora_adapter") # local folder if present | |
| ADAPTER_REPO = os.environ.get("ADAPTER_REPO", "") # optional HF repo id for adapter | |
| # Embeddings used to build FAISS | |
| EMBED_MODEL_NAME = os.environ.get("EMBED_MODEL_NAME", "pritamdeka/S-PubMedBERT-MS-MARCO") | |
| # PubMed / NCBI (set in Space Secrets if you have them) | |
| NCBI_EMAIL = os.environ.get("NCBI_EMAIL", "") | |
| NCBI_TOOL = os.environ.get("NCBI_TOOL", "askstein") | |
| NCBI_APIKEY = os.environ.get("NCBI_APIKEY", "") | |
| # Feedback logging | |
| FEEDBACK_PATH = os.environ.get("FEEDBACK_PATH", "feedback.csv") | |
| PUSH_FEEDBACK = os.environ.get("PUSH_FEEDBACK", "0") == "1" # set "1" to enable Hub upload | |
| HF_READ_TOKEN = os.environ.get("HF_READ_TOKEN", os.environ.get("HF_TOKEN", "")) | |
| HF_WRITE_TOKEN = os.environ.get("HF_WRITE_TOKEN", HF_READ_TOKEN) | |
| SPACE_REPO_ID = os.environ.get("SPACE_REPO_ID", "") | |
| # Generation / toggles | |
| ALLOW_WIKIPEDIA = False | |
| DEBUG = True | |
| MAX_NEW_TOKENS_GROUNDED = 512 | |
| MAX_NEW_TOKENS_FALLBACK = 256 | |
| MIN_USEFUL_CHARS = 260 | |
| # Auto-continue if we hit the cap without EOS | |
| AUTO_CONTINUE = True | |
| AUTO_CONT_MAX_STEPS = 2 # continue up to 2 extra chunks | |
| AUTO_CONT_NEW_TOKENS = 256 # tokens per continuation step | |
| def dlog(tag, msg): | |
| if DEBUG: print(f"[{tag}] {msg}") | |
| # --- Logging/upload config (OFF by default unless env says "1") | |
| PUSH_TRANSCRIPTS = os.environ.get("PUSH_TRANSCRIPTS", "0") == "1" | |
| PUSH_FEEDBACK = os.environ.get("PUSH_FEEDBACK", "0") == "1" | |
| HF_WRITE_TOKEN = os.environ.get("HF_WRITE_TOKEN", os.environ.get("HF_TOKEN", "")) | |
| LOGS_DATASET_ID = os.environ.get("LOGS_DATASET_ID", "") | |
| def _log_path(prefix: str, ext: str) -> str: | |
| # rotate daily, e.g., transcripts-20250929.jsonl | |
| today = datetime.datetime.utcnow().strftime("%Y%m%d") | |
| return f"{prefix}-{today}.{ext}" | |
| # rotate by default if user didn't override via env | |
| TRANSCRIPT_PATH = os.environ.get("TRANSCRIPT_PATH", _log_path("transcripts", "jsonl")) | |
| FEEDBACK_PATH = os.environ.get("FEEDBACK_PATH", _log_path("feedback", "csv")) | |
| # ================== HEAVY IMPORTS ================== | |
| import faiss | |
| from sentence_transformers import SentenceTransformer | |
| import torch | |
| from peft import PeftModel | |
| import wikipedia | |
| from wikipedia.exceptions import DisambiguationError, PageError | |
| from huggingface_hub import login, snapshot_download, HfApi | |
| # ================== LOW-VRAM RUNTIME KNOBS ================== | |
| os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "max_split_size_mb:128,expandable_segments:True") | |
| # ================== GPU CHECK ================== | |
| if not torch.cuda.is_available(): | |
| with gr.Blocks() as demo: | |
| gr.Markdown("### 🚦 GPU not detected.\nThis Space needs a GPU (A10G or better). Go to **Settings → Hardware** and select a GPU.") | |
| demo.launch() | |
| sys.exit(0) | |
| device = "cuda" | |
| dtype = torch.float16 | |
| torch.manual_seed(42) | |
| torch.backends.cuda.matmul.allow_tf32 = True | |
| torch.backends.cudnn.allow_tf32 = True | |
| # ================== RELEVANCE CONFIG ================== | |
| DEFAULT_REL_CONFIG = { | |
| "positive_terms": [ | |
| "ctra","rigidity","ct-based","qct","micro-ct","hounsfield", | |
| "femur","femoral","hip","proximal femur", | |
| "bending","torsional","axial","failure load","modulus", | |
| "nazarian","freedman","alboro" | |
| ], | |
| "negative_terms": [ | |
| "t cell","lymph","immunolog","synapse","receptor","egfr", | |
| "tumor","oncolog","immune","lymph node","cardio","myocard","neuro","skull","heart","brain" | |
| ], | |
| "weights": {"positive": 2, "negative": 1}, | |
| "author_whitelist": ["Nazarian","Freedman","Alboro"], | |
| "mesh_positive": [ | |
| "Femur", "Femoral Neck", "Hip", "Bone Density", | |
| "Tomography, X-Ray Computed", "Finite Element Analysis", | |
| "Bone and Bones", "Elastic Modulus", "Biomechanical Phenomena" | |
| ], | |
| "mesh_weight": 2, | |
| "author_weight": 3, | |
| "min_rel_to_use_faiss": 3, | |
| "ncbi_email": NCBI_EMAIL, | |
| "ncbi_tool": NCBI_TOOL, | |
| "ncbi_apikey": NCBI_APIKEY, | |
| } | |
| def load_rel_config(path: str) -> Dict[str, Any]: | |
| cfg = DEFAULT_REL_CONFIG.copy() | |
| try: | |
| if os.path.exists(path): | |
| with open(path, "r", encoding="utf-8") as f: | |
| user_cfg = json.load(f) | |
| cfg.update(user_cfg) | |
| except Exception as e: | |
| dlog("rel-config", f"using defaults ({e})") | |
| return cfg | |
| REL_CFG = load_rel_config(REL_CONFIG_PATH) | |
| dlog("rel-config", f"Loaded keys: {list(REL_CFG.keys())}") | |
| # ================== HTTP / NCBI HELPERS ================== | |
| class _Http: | |
| session = requests.Session() | |
| session.headers.update({"User-Agent": "Askstein/1.0 (https://huggingface.co)"}) | |
| def _ncbi_params(extra: Dict[str, Any] | None = None) -> Dict[str, Any]: | |
| p = {"retmode": "xml"} | |
| if REL_CFG.get("ncbi_email"): p["email"] = REL_CFG["ncbi_email"] | |
| if REL_CFG.get("ncbi_tool"): p["tool"] = REL_CFG["ncbi_tool"] | |
| if REL_CFG.get("ncbi_apikey"): p["api_key"] = REL_CFG["ncbi_apikey"] | |
| if extra: p.update(extra) | |
| return p | |
| def _get_with_backoff(url: str, params: Dict[str, Any], tries: int = 2, base_sleep: float = 0.5, timeout: int = 6) -> str: | |
| for i in range(tries): | |
| try: | |
| if not REL_CFG.get("ncbi_apikey"): | |
| time.sleep(0.34) | |
| r = _Http.session.get(url, params=params, timeout=timeout) | |
| r.raise_for_status() | |
| return r.text | |
| except Exception: | |
| if i == tries - 1: raise | |
| time.sleep(base_sleep * (2 ** i)) | |
| # ================== WIKIPEDIA ================== | |
| def wiki_summary_allow(query: str, sentences: int = 2) -> Optional[str]: | |
| prev = globals().get("ALLOW_WIKIPEDIA", False) | |
| globals()["ALLOW_WIKIPEDIA"] = True | |
| try: | |
| q = re.sub(r'^(what is|what are|define|where is|where are)\s+', '', query, flags=re.IGNORECASE) | |
| q = re.sub(r'\s+(located|location)\s*\?*$', '', q, flags=re.IGNORECASE).strip('?').strip() | |
| return wikipedia.summary(q, sentences=sentences) | |
| except (DisambiguationError, PageError, Exception): | |
| return None | |
| finally: | |
| globals()["ALLOW_WIKIPEDIA"] = prev | |
| # ================== LOAD FAISS + META + EMBEDDER ================== | |
| dlog("FAISS", "Loading index…") | |
| index = faiss.read_index(FAISS_PATH) | |
| dlog("FAISS", f"ntotal={index.ntotal}") | |
| dlog("FAISS", "Loading metadata…") | |
| all_chunks: List[Dict[str, Any]] = [] | |
| with open(META_PATH, "r", encoding="utf-8") as f: | |
| for line in f: | |
| try: | |
| all_chunks.append(json.loads(line)) | |
| except Exception: | |
| pass | |
| dlog("FAISS", f"Metadata records: {len(all_chunks)}") | |
| if len(all_chunks) != index.ntotal: | |
| raise RuntimeError(f"[ALIGNMENT] Metadata rows ({len(all_chunks)}) != FAISS ntotal ({index.ntotal}).") | |
| dlog("EMBED", f"Loading embedder: {EMBED_MODEL_NAME}") | |
| embed_model = SentenceTransformer(EMBED_MODEL_NAME) | |
| try: | |
| _probe = embed_model.encode(["__dimcheck__"], convert_to_numpy=True).astype("float32") | |
| _dim = _probe.shape[1] if _probe.ndim == 2 else len(_probe) | |
| assert index.d == _dim, f"FAISS dim {index.d} != embed dim {_dim} (model={EMBED_MODEL_NAME}). Rebuild index." | |
| except Exception as e: | |
| raise RuntimeError(f"[FAISS] Dimension check failed: {e}") | |
| _IS_IP = isinstance(index, faiss.IndexFlatIP) or "IndexFlatIP" in type(index).__name__ | |
| # ================== HUGGING FACE LOGIN & ADAPTER PATH ================== | |
| if HF_READ_TOKEN: | |
| try: | |
| login(token=HF_READ_TOKEN) | |
| dlog("HF", "Login (read) OK") | |
| except Exception as e: | |
| dlog("HF", f"Login issue: {e}") | |
| if ADAPTER_REPO: | |
| ADAPTER_PATH = snapshot_download(repo_id=ADAPTER_REPO, allow_patterns=["*"]) | |
| # ================== QUANTIZATION AVAILABILITY ================== | |
| try: | |
| import bitsandbytes as _bnb # noqa: F401 | |
| _BNB_AVAILABLE = True | |
| except Exception: | |
| _BNB_AVAILABLE = False | |
| # ================== LLM (BASE + LORA) ================== | |
| dlog("LLM", f"Loading base model: {BASE_MODEL}") | |
| tokenizer_lm = AutoTokenizer.from_pretrained(BASE_MODEL, use_fast=False) | |
| use_bnb = QUANTIZE in {"8bit", "4bit"} and BitsAndBytesConfig is not None and _BNB_AVAILABLE | |
| if use_bnb: | |
| # Quantized path (only if explicitly requested and bnb is installed) | |
| bnb_config = BitsAndBytesConfig( | |
| load_in_8bit=(QUANTIZE == "8bit"), | |
| load_in_4bit=(QUANTIZE == "4bit"), | |
| bnb_4bit_quant_type="nf4", | |
| bnb_4bit_use_double_quant=True, | |
| bnb_4bit_compute_dtype=torch.float16, | |
| ) | |
| base_model = AutoModelForCausalLM.from_pretrained( | |
| BASE_MODEL, | |
| device_map="auto", | |
| quantization_config=bnb_config, | |
| ) | |
| else: | |
| # fp16 path with SDPA attention (lower VRAM). Fallback if not supported. | |
| try: | |
| base_model = AutoModelForCausalLM.from_pretrained( | |
| BASE_MODEL, | |
| torch_dtype=dtype, | |
| device_map="auto", | |
| attn_implementation="sdpa", | |
| ) | |
| except TypeError: | |
| base_model = AutoModelForCausalLM.from_pretrained( | |
| BASE_MODEL, | |
| torch_dtype=dtype, | |
| device_map="auto", | |
| ) | |
| dlog("LLM", f"Loading LoRA adapter from: {ADAPTER_PATH}") | |
| model_lm = PeftModel.from_pretrained(base_model, ADAPTER_PATH) | |
| model_lm.eval() | |
| GEN_ARGS_GROUNDED = dict( | |
| max_new_tokens=MAX_NEW_TOKENS_GROUNDED, | |
| min_new_tokens=220, | |
| do_sample=False, | |
| num_beams=1, | |
| no_repeat_ngram_size=3, | |
| repetition_penalty=1.08, | |
| eos_token_id=tokenizer_lm.eos_token_id, | |
| ) | |
| GEN_ARGS_FALLBACK = dict( | |
| max_new_tokens=MAX_NEW_TOKENS_FALLBACK, | |
| min_new_tokens=120, | |
| do_sample=False, | |
| num_beams=1, | |
| no_repeat_ngram_size=3, | |
| repetition_penalty=1.05, | |
| eos_token_id=tokenizer_lm.eos_token_id, | |
| ) | |
| # ================== UTILITIES ================== | |
| _SANITIZE = re.compile(r"```.*?```|<\s*script[^>]*>.*?<\s*/\s*script\s*>", re.DOTALL|re.IGNORECASE) | |
| def _to_text(rec: Any) -> str: | |
| if isinstance(rec, str): return rec.strip() | |
| for k in ("text","chunk_text","content","body","passage","raw_text","section_text","abstract"): | |
| v = rec.get(k) | |
| if isinstance(v, str) and v.strip(): | |
| return _SANITIZE.sub("", v.strip()) | |
| segs = rec.get("segments") | |
| if isinstance(segs, list): | |
| return _SANITIZE.sub("", " ".join(s.get("text","").strip() for s in segs if isinstance(s, dict)).strip()) | |
| return "" | |
| def _hydrate_text(item: Dict[str, Any]) -> str: | |
| t = _to_text(item) | |
| if t: return t | |
| pmid = item.get("pmid") | |
| if pmid: | |
| abs_chunks = fetch_pubmed_chunks(str(pmid), max_papers=1) | |
| if abs_chunks: return abs_chunks[0].get("text", "").strip() | |
| title = item.get("title") | |
| if isinstance(title, str) and title.strip(): | |
| return title.strip() | |
| return "" | |
| def _dedup_by_pmid_or_title(records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: | |
| seen, out = set(), [] | |
| for r in records: | |
| key = str(r.get("pmid") or "").strip() | |
| if not key: key = (r.get("title") or "").strip().lower()[:120] | |
| if not key: key = (r.get("text") or "").strip().lower()[:200] | |
| if key in seen: continue | |
| seen.add(key); out.append(r) | |
| return out | |
| def _split_sentences(s: str) -> List[str]: | |
| s = s.replace("\r"," ").replace("\n"," ") | |
| parts = re.split(r"(?<=[\.\?\!])\s+", s) | |
| return [p.strip() for p in parts if p.strip()] | |
| _BAD_BULLETS = re.compile(r"^\s*(?:\d+\s*\)|[•\-\*])\s*$", re.M) | |
| _DANGLING = re.compile(r"[\[\(][^\]\)]$") | |
| def _post_clean(text: str) -> str: | |
| t = re.sub(r"[ \t]+\n", "\n", text) | |
| t = _BAD_BULLETS.sub("", t) | |
| t = re.sub(r"\n{3,}", "\n\n", t).strip() | |
| sents = _split_sentences(t) | |
| seen, kept = set(), [] | |
| for s in sents: | |
| key = s.lower() | |
| if key in seen: continue | |
| seen.add(key); kept.append(s) | |
| t = " ".join(kept) | |
| t = re.sub(_DANGLING, "", t).strip(" -,:;") | |
| return t | |
| def _ensure_min_answer(ans: str) -> str: | |
| if len(ans) >= MIN_USEFUL_CHARS: return ans | |
| tail = " If you want, I can add a short checklist of assumptions, units, and typical parameter ranges." | |
| return (ans + tail) if not ans.endswith(".") else (ans + tail) | |
| # ================== RELEVANCE / GATING ================== | |
| def _rel_score(text: str, title: str = "", cfg: Dict[str, Any] | None = None) -> int: | |
| cfg = cfg or REL_CFG | |
| blob = (title + " " + text).lower() | |
| pos = sum(1 for k in cfg.get("positive_terms", []) if k.lower() in blob) | |
| neg = sum(1 for k in cfg.get("negative_terms", []) if k.lower() in blob) | |
| w_pos = int(cfg.get("weights", {}).get("positive", 2)) | |
| w_neg = int(cfg.get("weights", {}).get("negative", 1)) | |
| return pos * w_pos - neg * w_neg | |
| def _mesh_by_pmid(pmid: str) -> List[str]: | |
| try: | |
| xml = _get_with_backoff( | |
| "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/efetch.fcgi", | |
| _ncbi_params({"db":"pubmed","id":str(pmid)}) | |
| ) | |
| root = ET.fromstring(xml) | |
| heads = [] | |
| for mh in root.findall(".//MeshHeading"): | |
| dn = mh.find("DescriptorName") | |
| if dn is not None and dn.text: | |
| heads.append(dn.text.strip()) | |
| return heads | |
| except Exception: | |
| return [] | |
| def _authors_by_pmid(pmid: str) -> List[str]: | |
| try: | |
| xml = _get_with_backoff( | |
| "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esummary.fcgi", | |
| _ncbi_params({"db":"pubmed","id":str(pmid)}) | |
| ) | |
| root = ET.fromstring(xml) | |
| names = [] | |
| for docsum in root.findall(".//DocSum"): | |
| for item in docsum.findall("Item"): | |
| if item.get("Name") == "AuthorList": | |
| for au in item.findall("Item"): | |
| if au.text: | |
| last = au.text.split(",")[0].split(" ")[-1] | |
| names.append(last) | |
| return names | |
| except Exception: | |
| return [] | |
| def _boost_by_author(pmid: str | int, rel_base: int, cfg: Dict[str, Any] | None = None) -> int: | |
| cfg = cfg or REL_CFG | |
| wl = set(cfg.get("author_whitelist", [])) | |
| if not pmid or not wl: return rel_base | |
| names = _authors_by_pmid(str(pmid)) | |
| return rel_base + int(cfg.get("author_weight", 3)) if any(n in wl for n in names) else rel_base | |
| def _mesh_boost(pmid: str | int, rel_base: int, cfg: Dict[str, Any] | None = None) -> int: | |
| cfg = cfg or REL_CFG | |
| if not pmid: return rel_base | |
| targets = set(x.lower() for x in cfg.get("mesh_positive", [])) | |
| weight = int(cfg.get("mesh_weight", 2)) | |
| heads = [h.lower() for h in _mesh_by_pmid(str(pmid))] | |
| hit = sum(1 for h in heads if h in targets) | |
| return rel_base + hit * weight | |
| _MSK_MUST = re.compile( | |
| r"\b(femur|femoral|hip|proximal\s+femur|ctra|qct|ct-?based|rigidity|bending|torsional|axial|failure\s+load)\b", | |
| re.I | |
| ) | |
| _CT_RIGIDITY_TOKENS = re.compile(r"\b(qct|ct[-\s]?based|ctra|rigidity|bending|torsion|hounsfield|finite\s+element|fe[am])\b", re.I) | |
| _FE_TOKENS = re.compile(r"\b(fe|fea|finite\s+element|micromotion|boundary\s+conditions|nonlinear|yield|ultimate|fracture\s+load)\b", re.I) | |
| _ANATOMY_OR_HISTORY = re.compile( | |
| r"(?:\bhistory\b.*\b(femur|hip|bone)\b|\bwhat\s+is\s+the\s+(femur|hip)\b|\banatomy\b.*\b(hip|femur)\b)", | |
| re.I | |
| ) | |
| _PAPERS_INTENT = re.compile(r"\b(key\s+papers|suggest\s+papers|landmark|seminal|important|top\s+papers)\b", re.I) | |
| CITE_TRIGGER = re.compile( | |
| r"\b(?:and\s+(?:cite|citations?|references?|studies?|papers?))\.?\s*$", | |
| re.IGNORECASE, | |
| ) | |
| # ================== PUBMED & RETRIEVAL ================== | |
| def fetch_pubmed_chunks(query_or_pmid: str, max_papers: int = 3) -> List[Dict[str, Any]]: | |
| retries = 1 | |
| chunks: List[Dict[str, Any]] = [] | |
| def _efetch(pmid: str): | |
| try: | |
| xml = _get_with_backoff( | |
| "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/efetch.fcgi", | |
| _ncbi_params({"db":"pubmed","id":pmid}) | |
| ) | |
| tree = ET.fromstring(xml) | |
| paras = [a.text for a in tree.findall(".//AbstractText") if a is not None and a.text] | |
| if paras: | |
| text = "\n".join(paras) | |
| chunks.append({"text": text, "source": "pubmed", "pmid": pmid}) | |
| except Exception: | |
| pass | |
| if query_or_pmid.isdigit(): | |
| _efetch(query_or_pmid); return chunks | |
| pmids: List[str] = [] | |
| for attempt in range(retries + 1): | |
| try: | |
| xml = _get_with_backoff( | |
| "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi", | |
| _ncbi_params({"db":"pubmed","term":query_or_pmid, "retmax":max_papers}) | |
| ) | |
| root = ET.fromstring(xml) | |
| pmids = [e.text for e in root.findall(".//Id") if e is not None and e.text] | |
| break | |
| except Exception: | |
| if attempt == retries: return [] | |
| time.sleep(0.5 * (2 ** attempt)) | |
| for pmid in pmids[:max_papers]: | |
| _efetch(pmid) | |
| return chunks | |
| STOPWORDS = set("the a an of and for with without to on in by from into over under how what why where when is are was were be been being this that these those it its as about".split()) | |
| def _parse_year(y: str) -> int: | |
| try: | |
| return int(re.findall(r"\d{4}", y or "")[0]) | |
| except Exception: | |
| return 0 | |
| def _is_msk_paper(title: str, journal: str, year: str = "") -> bool: | |
| t = f"{title or ''} {journal or ''}".lower() | |
| must_body = any(k in t for k in ["femur","femoral","hip","proximal femur","long bone","tibia","humerus"]) | |
| must_method = any(k in t for k in ["ct","qct","finite element","structural rigidity","ctra","rigidity","bending","torsion"]) | |
| if not (must_body and must_method): return False | |
| year_i = _parse_year(year) | |
| if year_i and not (2000 <= year_i <= 2025): return False | |
| return True | |
| def fetch_pubmed_citations(query: str, max_results: int = 5) -> List[str]: | |
| try: | |
| xml = _get_with_backoff( | |
| "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi", | |
| _ncbi_params({"db":"pubmed","term":query, "retmax":max_results}) | |
| ) | |
| root = ET.fromstring(xml) | |
| pmids = [elem.text for elem in root.findall(".//Id") if elem is not None and elem.text] | |
| if not pmids: return [] | |
| except Exception: | |
| return [] | |
| try: | |
| xml = _get_with_backoff( | |
| "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esummary.fcgi", | |
| _ncbi_params({"db":"pubmed","id":",".join(pmids)}) | |
| ) | |
| summary_root = ET.fromstring(xml) | |
| except Exception: | |
| return [] | |
| citations: List[str] = [] | |
| for docsum in summary_root.findall(".//DocSum"): | |
| pmid = docsum.findtext("Id", default="") | |
| title = journal = year = doi = "" | |
| authors: List[str] = [] | |
| for item in docsum.findall("Item"): | |
| name = item.get("Name", "") | |
| if name == "Title": title = item.text or "" | |
| elif name == "FullJournalName": journal = item.text or "" | |
| elif name == "PubDate": year = (item.text or "").split()[0] | |
| elif name == "AuthorList": | |
| for au in item.findall("Item"): | |
| if au.text: authors.append(au.text) | |
| elif name == "ArticleIds": | |
| for sub in item.findall("Item"): | |
| if sub.get("Name") == "doi": | |
| doi = sub.text or "" | |
| if not _is_msk_paper(title, journal, year): continue | |
| first_author = authors[0] if authors else "" | |
| auth_str = f"{first_author} et al." if first_author else "" | |
| parts = [p for p in [auth_str, title, journal, year] if p] | |
| cit = ", ".join(parts).strip().rstrip(",") | |
| if pmid: cit += f"; PMID:{pmid}" | |
| if doi: cit += f" DOI:{doi}" | |
| if cit: citations.append(cit) | |
| return citations[:max_results] | |
| def _compact_terms(q: str) -> str: | |
| words = re.findall(r"[A-Za-z0-9\-]+", q.lower()) | |
| keep = [w for w in words if w not in STOPWORDS and len(w) > 2] | |
| return " ".join(keep)[:200] | |
| def _biomechish(q: str) -> bool: | |
| return bool(re.search(r"\b(femur|femoral|hip|bone|qct|ctra|rigidity|bending|torsion|elastic modulus|finite\s+element|fea)\b", q, re.I)) | |
| def _is_fe_override(q: str) -> bool: | |
| return bool(_FE_TOKENS.search(q)) | |
| _CONTRA_NO_EFFECT = re.compile(r"\b(no\s+significant\s+difference|no\s+effect|not\s+significant)\b", re.I) | |
| _CONTRA_CHANGE = re.compile(r"\b(increase[ds]?|decrease[ds]?|higher|lower|greater|reduced?)\b", re.I) | |
| def _has_conflict(text: str) -> bool: | |
| return bool(_CONTRA_NO_EFFECT.search(text) and _CONTRA_CHANGE.search(text)) | |
| HARDCODED_CITS = { | |
| "EA": [ | |
| "Morgan EF et al., Mechanical properties of cortical bone…, J Biomech, 2003; PMID:12547357", | |
| "Turner CH, Burr DB., Experimental techniques for bone mechanics, Bone, 1993; PMID:8252072" | |
| ], | |
| "EI": [ | |
| "Courtney AC et al., Age-related reductions in the strength of the femur…, J Bone Miner Res, 1995; PMID:7584933", | |
| "Bell KL et al., Regional Heterogeneity of the Proximal Femur…, Bone, 1999; PMID:10574202" | |
| ], | |
| "GJ": [ | |
| "Cowin SC., Bone Mechanics Handbook (torsion of bone cross-sections), CRC Press, 2001.", | |
| "Vollmer M et al., Long bone torsion testing methods, J Biomech, 1987; PMID:3670157" | |
| ] | |
| } | |
| def _fallback_cits_for(term: str) -> List[str]: | |
| return HARDCODED_CITS.get(term.upper(), []) | |
| def detect_lab(q: str) -> str: | |
| ql = q.lower() | |
| if "freedman" in ql: return "freedman" | |
| if "alboro" in ql or "alborno" in ql: return "alboro" | |
| return "nazarian" | |
| def build_lab_query(core_q: str, lab: str = "nazarian") -> str: | |
| topics = [ | |
| "femur","femoral neck","hip","proximal femur", | |
| "CT","QCT","micro-CT","rigidity","CTRA","structural rigidity", | |
| "bending","torsional","axial","failure load","modulus","Hounsfield" | |
| ] | |
| ta = " OR ".join(f'"{t}"[Title/Abstract]' for t in topics) | |
| if lab == "freedman": | |
| author = '("Freedman BA"[Author] OR "Freedman"[Author])' | |
| elif lab == "alboro": | |
| author = '("Alboro"[Author] OR "Alborno"[Author])' | |
| else: | |
| author = '("Nazarian A"[Author] OR "Ara Nazarian"[Full Author Name])' | |
| date = '("2000"[Date - Publication] : "3000"[Date - Publication])' | |
| return f"{author} AND ({ta}) AND {date}" | |
| def retrieve_context(query: str, top_k: int = 10) -> List[Dict[str, Any]]: | |
| q = query.strip() | |
| if _ANATOMY_OR_HISTORY.search(q) and not _biomechish(q): | |
| wiki = wiki_summary_allow(q, sentences=4) | |
| if wiki: | |
| dlog("WIKI", "Wikipedia biomechanics fallback hit") | |
| return [{"text": wiki, "source": "wikipedia"}] | |
| pm = re.search(r"pmid[:\s]*(\d+)", q, re.IGNORECASE) | |
| if pm: | |
| dlog("PMID", f"PMID inline {pm.group(1)}") | |
| return fetch_pubmed_chunks(pm.group(1), max_papers=1) | |
| if not (_CT_RIGIDITY_TOKENS.search(q) or _is_fe_override(q)): | |
| dlog("FALLBACK", "No CT/rigidity/FE tokens → try PubMed/Wiki") | |
| results = fetch_pubmed_chunks(q) | |
| if results: | |
| dlog("PUBMED", "PubMed search hit") | |
| return results | |
| wiki = wiki_summary_allow(q, sentences=3) | |
| if wiki: | |
| dlog("WIKI", "Wikipedia fallback hit") | |
| return [{"text": wiki, "source": "wikipedia"}] | |
| dlog("RETRIEVAL", "No results found") | |
| return [] | |
| # FAISS path | |
| q_emb = embed_model.encode([q], convert_to_numpy=True).astype("float32") | |
| if _IS_IP: | |
| faiss.normalize_L2(q_emb) | |
| D, I = index.search(q_emb, top_k) | |
| results: List[Dict[str, Any]] = [] | |
| for dist, idx_ in zip(D[0], I[0]): | |
| if idx_ < 0: continue | |
| item = all_chunks[idx_].copy() | |
| item["score"] = float(dist) | |
| item["text"] = _hydrate_text(item) | |
| if not item["text"]: continue | |
| results.append(item) | |
| if results: | |
| scored: List[Dict[str, Any]] = [] | |
| for it in results: | |
| rel = _rel_score(it.get("text", ""), it.get("title", ""), REL_CFG) | |
| rel = _boost_by_author(it.get("pmid"), rel, REL_CFG) | |
| rel = _mesh_boost(it.get("pmid"), rel, REL_CFG) | |
| it["_rel"] = rel | |
| scored.append(it) | |
| results = sorted(scored, key=lambda x: (x.get("_rel", 0), x.get("score", 0)), reverse=True) | |
| min_rel = int(REL_CFG.get("min_rel_to_use_faiss", 3)) | |
| positives = [ | |
| r for r in results | |
| if r.get("_rel", 0) >= min_rel and _MSK_MUST.search((r.get("title","")+" "+r.get("text",""))) | |
| ] | |
| positives = _dedup_by_pmid_or_title(positives) | |
| if positives: | |
| dlog("FAISS", f"hit={len(positives)} (top rel={positives[0].get('_rel')} score={positives[0].get('score'):.3f})") | |
| return positives[:top_k] | |
| else: | |
| dlog("FALLBACK", "FAISS results off-topic → PubMed fallback") | |
| results = fetch_pubmed_chunks(q) | |
| if results: | |
| dlog("PUBMED", "PubMed search hit") | |
| return results | |
| wiki = wiki_summary_allow(q, sentences=3) | |
| if wiki: | |
| dlog("WIKI", "Wikipedia fallback hit") | |
| return [{"text": wiki, "source": "wikipedia"}] | |
| dlog("RETRIEVAL", "No results at all") | |
| return [] | |
| def build_prompt(chunks: List[Dict[str, Any]], question: str) -> str: | |
| header = ( | |
| "You are Askstein (orthopedic biomechanics). Use ONLY the [Context] to answer. " | |
| "If the context is insufficient, say 'I don’t know based on the provided context.' " | |
| "Stay within musculoskeletal CT-based rigidity (EA, EI, GJ), femur/hip, CTRA/QCT, or FE modeling of these. " | |
| "Do not discuss cardiology, neurology, or unrelated domains." | |
| ) | |
| cleaned = [] | |
| per_chunk_chars = 900 # lower prompt length = lower KV memory | |
| for c in chunks: | |
| t = _to_text(c) | |
| if t: cleaned.append(t[:per_chunk_chars]) | |
| context = "\n\n".join(cleaned) | |
| return f"{header}\n\n[Context]:\n{context}\n\n[Question]:\n{question}\n" | |
| def _decode_generated(out_ids, in_len: int) -> str: | |
| gen = out_ids[0][in_len:] | |
| return tokenizer_lm.decode(gen, skip_special_tokens=True).lstrip(". \n").strip() | |
| def _gen_once(inputs, args) -> Any: | |
| with torch.inference_mode(): | |
| return model_lm.generate(**inputs, **args, use_cache=True) | |
| def _generate(inputs, grounded: bool): | |
| args = GEN_ARGS_GROUNDED if grounded else GEN_ARGS_FALLBACK | |
| in_len = inputs["input_ids"].shape[-1] | |
| # First attempt | |
| try: | |
| out = _gen_once(inputs, args) | |
| except torch.cuda.OutOfMemoryError: | |
| try: | |
| torch.cuda.empty_cache() | |
| except Exception: | |
| pass | |
| small_args = dict(args) | |
| small_args["max_new_tokens"] = min(256, args.get("max_new_tokens", 256)) | |
| # disable cache to save VRAM | |
| with torch.inference_mode(): | |
| out = model_lm.generate(**inputs, **small_args, use_cache=False) | |
| if not AUTO_CONTINUE: | |
| return out | |
| # Auto-continue if we hit cap without EOS | |
| steps = 0 | |
| while steps < AUTO_CONT_MAX_STEPS: | |
| seq = out[0] | |
| ended_with_eos = (seq[-1].item() == tokenizer_lm.eos_token_id) | |
| hit_cap = (seq.shape[0] - in_len) >= args["max_new_tokens"] | |
| if ended_with_eos or not hit_cap: | |
| break | |
| cont_inputs = { | |
| "input_ids": seq.unsqueeze(0), | |
| "attention_mask": torch.ones_like(seq).unsqueeze(0), | |
| } | |
| cont_inputs = {k: v.to(device) for k, v in cont_inputs.items()} | |
| cont_args = dict(args) | |
| cont_args["max_new_tokens"] = AUTO_CONT_NEW_TOKENS | |
| try: | |
| with torch.inference_mode(): | |
| out = model_lm.generate(**cont_inputs, **cont_args, use_cache=True) | |
| except torch.cuda.OutOfMemoryError: | |
| try: | |
| torch.cuda.empty_cache() | |
| except Exception: | |
| pass | |
| with torch.inference_mode(): | |
| out = model_lm.generate(**cont_inputs, **cont_args, use_cache=False) | |
| steps += 1 | |
| return out | |
| # ================== ANSWER SYNTHESIS ================== | |
| def _synthesize_answer(chunks: List[Dict[str, Any]], question: str) -> str: | |
| prompt = build_prompt(chunks, question) | |
| inputs = tokenizer_lm(prompt, return_tensors="pt").to(device) | |
| in_len = inputs["input_ids"].shape[-1] | |
| out = _generate(inputs, grounded=True) | |
| answer = _decode_generated(out, in_len) | |
| return _post_clean(answer) | |
| def _answer_from_chunks(chunks: List[Dict[str, Any]], question: str) -> str: | |
| joined = " ".join(_to_text(c) for c in chunks if _to_text(c)) | |
| if _has_conflict(joined): | |
| dlog("SYNTH", "Conflict detected → summarize") | |
| return _synthesize_answer(chunks, question) | |
| return _synthesize_answer(chunks, question) | |
| def direct_llm_fallback(question: str) -> str: | |
| sys_prompt = ( | |
| "You are Askstein (orthopedic biomechanics). If you lack enough domain context, say you don’t know. " | |
| "Avoid discussing non-musculoskeletal systems (cardiology, neurology). Do NOT invent references." | |
| ) | |
| llm_prompt = f"{sys_prompt}\n\nQuestion: {question}\nAnswer:" | |
| inputs = tokenizer_lm(llm_prompt, return_tensors="pt").to(device) | |
| out = _generate(inputs, grounded=False) | |
| in_len = inputs["input_ids"].shape[-1] | |
| ans = _post_clean(_decode_generated(out, in_len)) | |
| ans = re.sub(r"(?is)(^|\n)\s*references?:.*$", "", ans).strip() | |
| return "[LLM fallback — ungrounded]\n\n" + ans | |
| # ================== PUBLIC API ================== | |
| def ask(question: str) -> str: | |
| q = (question or "").strip() | |
| # --- PMID short-circuit --------------------------------------------------- | |
| m = re.search(r"pmid[:\s]*(\d+)", q, re.IGNORECASE) | |
| if m: | |
| pmid = m.group(1) | |
| chunks = fetch_pubmed_chunks(pmid, max_papers=1) | |
| return "\n".join(c.get("text", "") for c in chunks) or "Sorry, no abstract found." | |
| # --- Normalize query & detect 'and cite' / 'key papers' intent ------------ | |
| wants_cite = bool(CITE_TRIGGER.search(q) or _PAPERS_INTENT.search(q)) | |
| core_q = CITE_TRIGGER.sub("", q).strip().rstrip(".") | |
| core_q = _PAPERS_INTENT.sub("", core_q).strip() | |
| # --- Deterministic definitions/workflows (used in both paths) ------------- | |
| det_text = deterministic_definitions_text(core_q) | |
| # ========================================================================== | |
| # Citation path | |
| # ========================================================================== | |
| if wants_cite: | |
| if det_text: | |
| explanation = det_text | |
| lq = core_q.lower() | |
| used_term = None | |
| if ("torsion" in lq) or ("gj" in lq): | |
| used_term = "GJ" | |
| pm_query = ( | |
| '(torsion[TiAb] OR "polar moment"[TiAb] OR GJ[TiAb]) AND ' | |
| '("Bone and Bones"[MeSH] OR Femur[TiAb]) AND ' | |
| '("Finite Element Analysis"[MeSH] OR QCT[TiAb] OR CT[TiAb]) AND ' | |
| '("2000"[DP] : "2025"[DP])' | |
| ) | |
| elif ("bending" in lq) or ("ei" in lq): | |
| used_term = "EI" | |
| pm_query = ( | |
| '(bending[TiAb] OR "second moment"[TiAb] OR EI[TiAb]) AND ' | |
| '("Bone and Bones"[MeSH] OR Femur[TiAb]) AND ' | |
| '("Finite Element Analysis"[MeSH] OR QCT[TiAb] OR CT[TiAb]) AND ' | |
| '("2000"[DP] : "2025"[DP])' | |
| ) | |
| else: | |
| used_term = "EA" | |
| pm_query = ( | |
| '("axial rigidity"[TiAb] OR EA[TiAb] OR "axial stiffness"[TiAb]) AND ' | |
| '("Bone and Bones"[MeSH] OR Femur[TiAb]) AND ' | |
| '("Finite Element Analysis"[MeSH] OR QCT[TiAb] OR CT[TiAb]) AND ' | |
| '("2000"[DP] : "2025"[DP])' | |
| ) | |
| citations = fetch_pubmed_citations(pm_query, max_results=5) | |
| if not citations and used_term: | |
| citations = _fallback_cits_for(used_term) | |
| else: | |
| # Grounded explanation + progressively broader citation search | |
| explanation = _answer_from_chunks(retrieve_context(core_q, top_k=5), core_q) | |
| # Broad, de-biased search first | |
| compact = _compact_terms(core_q) | |
| pm_query = ( | |
| f'({compact}) AND ("Bone and Bones"[MeSH] OR Femur[TiAb] OR Hip[TiAb] ' | |
| f'OR Rigidity[TiAb] OR "Tomography, X-Ray Computed"[MeSH] OR "Finite Element Analysis"[MeSH]) ' | |
| f'NOT (heart[TiAb] OR cardiac[TiAb] OR brain[TiAb] OR skull[TiAb] OR EGFR[TiAb]) ' | |
| f'AND ("2000"[DP] : "2025"[DP])' | |
| ) | |
| citations = fetch_pubmed_citations(pm_query, max_results=5) | |
| # If empty, bias toward lab-authored queries that are likely relevant | |
| if not citations: | |
| lab = detect_lab(core_q) | |
| pm_query = build_lab_query(core_q, lab=lab) | |
| citations = fetch_pubmed_citations(pm_query, max_results=5) | |
| resp = _post_clean(explanation) | |
| if citations: | |
| resp += "\n\nCitations:\n" + "\n".join(citations) | |
| else: | |
| resp += f'\n\nSorry, no relevant citations found for “{core_q}.”' | |
| return _ensure_min_answer(resp) | |
| # ========================================================================== | |
| # Non-citation path | |
| # ========================================================================== | |
| if det_text: | |
| dlog("ASK", "Deterministic definition/workflow fired") | |
| return det_text | |
| # If the query doesn't clearly hit MSK/FE tokens, try retrieval then fallback | |
| if not (_MSK_MUST.search(q) or _is_fe_override(q)): | |
| chunks = retrieve_context(q, top_k=5) | |
| if chunks: | |
| answer = _answer_from_chunks(chunks, q) | |
| return _ensure_min_answer(_post_clean(answer)) or direct_llm_fallback(q) | |
| return direct_llm_fallback(q) | |
| # Clear MSK/FE intent → do the normal grounded path | |
| chunks = retrieve_context(q, top_k=5) | |
| if not chunks: | |
| return direct_llm_fallback(q) | |
| answer = _answer_from_chunks(chunks, q) | |
| return _ensure_min_answer(_post_clean(answer)) or direct_llm_fallback(q) | |
| def deterministic_definitions_text(core_q: str) -> Optional[str]: | |
| q_lower = core_q.lower() | |
| if "define axial rigidity" in q_lower or "what is axial rigidity" in q_lower: | |
| return ("Axial rigidity (EA) is Σ(Eᵢ·dAᵢ) across a CT slice; units: N. " | |
| "Modulus E per voxel comes from a density–modulus calibration; areas dAᵢ are voxel areas.") | |
| if "define bending rigidity" in q_lower or "what is bending rigidity" in q_lower: | |
| return ("Bending rigidity (EI) is Σ(Eᵢ·dAᵢ·yᵢ²) about a given axis; units: N·mm². " | |
| "yᵢ is distance to the neutral axis; computed slice-by-slice from QCT.") | |
| if ("define torsional rigidity" in q_lower) or ("what is torsional rigidity" in q_lower) or ("define gj" in q_lower): | |
| return ("Torsional rigidity (GJ) = shear modulus G times polar moment J. " | |
| "In QCT, J ≈ Σ(dAᵢ·rᵢ²) about the centroid; G ≈ E/(2(1+ν)).") | |
| if "qct" in q_lower and ("torsional" in q_lower or "gj" in q_lower): | |
| return ("From QCT, torsional rigidity is estimated as GJ, where J ≈ Σ(dAᵢ·rᵢ²) about the slice centroid and " | |
| "G = E/(2(1+ν)) from the voxel E map (ν≈0.3). Compute per-slice and report the minimum.") | |
| if re.search(r"\b(outline|steps|workflow|protocol)\b.*\b(ct|qct).*(rigidity|ea|ei|gj)", q_lower): | |
| return ( | |
| "CT-based structural rigidity (CTRA/QCT) workflow:\n" | |
| "1) Acquire QCT (≤1 mm; density phantom).\n" | |
| "2) Preprocess & segment bone.\n" | |
| "3) HU→ρ; ρ→E calibration.\n" | |
| "4) Cross-sections along neck axis.\n" | |
| "5) EA, EI_x/EI_y, GJ (G≈E/(2(1+ν))).\n" | |
| "6) Extract minima & validate vs FEA/mech tests." | |
| ) | |
| if re.search(r"\b(modulus)\b.*\brigidity\b|\bdefine\s+modulus\b", q_lower): | |
| return ("Elastic modulus (E) is a material property (Pa). " | |
| "Rigidity is structural (EA, EI, GJ). Modulus ≠ rigidity.") | |
| return None | |
| # ================== UI: NAME GATE + PER-ANSWER FEEDBACK ================== | |
| def _now_iso(): | |
| return datetime.datetime.utcnow().replace(tzinfo=datetime.timezone.utc).isoformat() | |
| def init_session(): | |
| return { | |
| "session_id": str(uuid.uuid4()), | |
| "first_name": "", | |
| "last_name": "", | |
| "last_q": "", | |
| "last_a": "", | |
| } | |
| def enter_app(first_name, last_name, state): | |
| first_name = (first_name or "").strip() | |
| last_name = (last_name or "").strip() | |
| if not first_name or not last_name: | |
| return gr.update(visible=True), gr.update(visible=False), state, "Please enter both first and last name." | |
| state["first_name"] = first_name | |
| state["last_name"] = last_name | |
| return gr.update(visible=False), gr.update(visible=True), state, f"Welcome, {first_name}! You can start chatting." | |
| from huggingface_hub import HfApi | |
| def _push_file_to_dataset(local_path: str, repo_path: str) -> None: | |
| """ | |
| Upload a local file to a *dataset* repo so the Space doesn't rebuild. | |
| Requires: HF_WRITE_TOKEN, LOGS_DATASET_ID | |
| """ | |
| if not os.path.exists(local_path): | |
| dlog("UPLOAD", f"Skip: {local_path} does not exist"); return | |
| if not LOGS_DATASET_ID: | |
| dlog("UPLOAD", "Skip: LOGS_DATASET_ID not set"); return | |
| if not HF_WRITE_TOKEN: | |
| dlog("UPLOAD", "Skip: HF_WRITE_TOKEN not set"); return | |
| try: | |
| api = HfApi(token=HF_WRITE_TOKEN) | |
| api.upload_file( | |
| path_or_fileobj=local_path, | |
| path_in_repo=repo_path, # keep short names at the root | |
| repo_id=LOGS_DATASET_ID, | |
| repo_type="dataset", # <- key change: dataset, not space | |
| commit_message=f"Update {repo_path}", | |
| ) | |
| dlog("UPLOAD", f"Uploaded {repo_path} to dataset {LOGS_DATASET_ID}") | |
| except Exception as e: | |
| dlog("UPLOAD", f"Upload failed: {e}") | |
| def _push_feedback_to_hub() -> None: | |
| """ | |
| Back-compat name, but uploads to a *dataset* repo so the Space won't rebuild. | |
| Requires: PUSH_FEEDBACK=1, HF_WRITE_TOKEN, LOGS_DATASET_ID, and _push_file_to_dataset(). | |
| """ | |
| if not PUSH_FEEDBACK: | |
| return | |
| try: | |
| # Use the rotated filename if you enabled daily rotation | |
| remote_name = os.path.basename(FEEDBACK_PATH) | |
| _push_file_to_dataset(FEEDBACK_PATH, remote_name) | |
| dlog("UPLOAD", f"Feedback uploaded: {remote_name}") | |
| except Exception as e: | |
| dlog("UPLOAD", f"Feedback upload failed: {e}") | |
| def _log_turn(state: Dict[str, Any], question: str, answer: str): | |
| """ | |
| Append locally first, then (optionally) upload to a *dataset* repo. | |
| This avoids commits to the Space repo and prevents rebuilds after each turn. | |
| """ | |
| rec = { | |
| "timestamp_utc": _now_iso(), | |
| "session_id": state.get("session_id", ""), | |
| "first_name": state.get("first_name", ""), | |
| "last_name": state.get("last_name", ""), | |
| "question": question, | |
| "answer": answer, | |
| } | |
| # Local append (ensure dir exists) | |
| try: | |
| os.makedirs(os.path.dirname(TRANSCRIPT_PATH) or ".", exist_ok=True) | |
| with open(TRANSCRIPT_PATH, "a", encoding="utf-8") as f: | |
| f.write(json.dumps(rec, ensure_ascii=False) + "\n") | |
| except Exception as e: | |
| dlog("LOG", f"Failed to write transcript locally: {e}") | |
| # Optional remote upload (dataset repo -> no Space rebuild) | |
| if PUSH_TRANSCRIPTS: | |
| try: | |
| remote_name = os.path.basename(TRANSCRIPT_PATH) | |
| _push_file_to_dataset(TRANSCRIPT_PATH, remote_name) | |
| dlog("UPLOAD", f"Transcript uploaded: {remote_name}") | |
| except Exception as e: | |
| dlog("UPLOAD", f"Transcript upload failed: {e}") | |
| def save_feedback(rating, comment, state): | |
| if rating is None: | |
| return "Please select a rating (1–5).", gr.update(visible=True) | |
| row = { | |
| "timestamp_utc": _now_iso(), | |
| "session_id": state["session_id"], | |
| "first_name": state["first_name"], | |
| "last_name": state["last_name"], | |
| "question": state["last_q"], | |
| "answer": state["last_a"], | |
| "rating": rating, | |
| "comment": (comment or "").strip(), | |
| } | |
| header = ["timestamp_utc","session_id","first_name","last_name","question","answer","rating","comment"] | |
| try: | |
| file_exists = os.path.exists(FEEDBACK_PATH) | |
| with open(FEEDBACK_PATH, "a", newline="", encoding="utf-8") as f: | |
| w = csv.DictWriter(f, fieldnames=header) | |
| if not file_exists: | |
| w.writeheader() | |
| w.writerow(row) | |
| _push_feedback_to_hub() | |
| return "Thanks for the feedback! ✅", gr.update(visible=False) | |
| except Exception as e: | |
| return f"Failed to save feedback: {e}", gr.update(visible=True) | |
| def predict(message, chat_history, state): | |
| msg = (message or "").strip() | |
| if not msg: | |
| return chat_history, "", gr.update(visible=False), None, "", state | |
| try: | |
| answer = ask(msg) | |
| except Exception as e: | |
| answer = f"Sorry — something went wrong: {e!r}" | |
| chat_history = (chat_history or []) + [(msg, answer)] | |
| state["last_q"] = msg | |
| state["last_a"] = answer | |
| try: | |
| _log_turn(state, msg, answer) | |
| except Exception: | |
| pass | |
| # free a bit of VRAM between turns | |
| try: | |
| torch.cuda.empty_cache() | |
| except Exception: | |
| pass | |
| return ( | |
| chat_history, | |
| "", # clear input | |
| gr.update(visible=True), # show feedback pane | |
| gr.update(value=None), # reset rating | |
| gr.update(value=""), # reset comment | |
| state | |
| ) | |
| # ================== UI ================== | |
| with gr.Blocks(theme="soft") as demo: | |
| gr.Markdown("# Askstein — Orthopedic Biomechanics Chat (CT/QCT Rigidity, FE)") | |
| gr.Markdown("Grounded answers (FAISS + PubMed). Please enter your name to continue.") | |
| state = gr.State(init_session()) | |
| # ---- Gate ---- | |
| gate = gr.Group(visible=True) | |
| with gate: | |
| with gr.Row(): | |
| first_tb = gr.Textbox(label="First name", placeholder="e.g., Shubh", scale=1) | |
| last_tb = gr.Textbox(label="Last name", placeholder="e.g., Laiwala", scale=1) | |
| enter_btn = gr.Button("Enter", variant="primary") | |
| gate_msg = gr.Markdown("", elem_classes=["text-sm"]) | |
| # ---- App (hidden until gate passes) ---- | |
| app = gr.Group(visible=False) | |
| with app: | |
| chat = gr.Chatbot(label="Askstein", height=430, type="tuples") | |
| with gr.Row(): | |
| user_in = gr.Textbox(placeholder="Ask a biomechanics question…", scale=5) | |
| send_btn = gr.Button("Send", variant="primary", scale=1) | |
| clear_btn = gr.Button("Clear chat") | |
| feedback_grp = gr.Group(visible=False) | |
| with feedback_grp: | |
| gr.Markdown("### How helpful was this answer?") | |
| rating = gr.Radio(choices=[1,2,3,4,5], label="Rating (1=poor, 5=great)") | |
| comment = gr.Textbox(label="Optional comment", placeholder="What was good or missing?") | |
| submit_fb = gr.Button("Submit feedback") | |
| fb_status = gr.Markdown("") | |
| # ---- Wiring (must stay inside Blocks) ---- | |
| enter_btn.click( | |
| fn=enter_app, | |
| inputs=[first_tb, last_tb, state], | |
| outputs=[gate, app, state, gate_msg], | |
| ) | |
| send_btn.click( | |
| fn=predict, | |
| inputs=[user_in, chat, state], | |
| outputs=[chat, user_in, feedback_grp, rating, comment, state], | |
| concurrency_limit=1, # serialize LLM calls | |
| ) | |
| user_in.submit( | |
| fn=predict, | |
| inputs=[user_in, chat, state], | |
| outputs=[chat, user_in, feedback_grp, rating, comment, state], | |
| concurrency_limit=1, # serialize LLM calls | |
| ) | |
| clear_btn.click( | |
| fn=lambda: ([], "", gr.update(visible=False), None, "", init_session()), | |
| inputs=None, | |
| outputs=[chat, user_in, feedback_grp, rating, comment, state], | |
| concurrency_limit=4, | |
| ) | |
| submit_fb.click( | |
| fn=save_feedback, | |
| inputs=[rating, comment, state], | |
| outputs=[fb_status, feedback_grp], | |
| concurrency_limit=4, | |
| ) | |
| # ================== QUEUE & LAUNCH ================== | |
| demo.queue(max_size=64) # no deprecated concurrency_count | |
| demo.launch(max_threads=8) | |