Update app.py

app.py

@@ -1,91 +1,884 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+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
+
+import numpy as np
+import requests
 import gradio as gr
-from askstein import ask
 
-def main():
+# ================== 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")
+
+# 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 to "1" to enable Hub upload
+HF_TOKEN      = os.environ.get("HF_TOKEN", "")               # add in Space Secrets
+SPACE_REPO_ID = os.environ.get("SPACE_REPO_ID", "")          # auto-provided in Spaces
+
+# Generation / toggles
+ALLOW_WIKIPEDIA = False
+DEBUG = True
+MAX_NEW_TOKENS_GROUNDED = 512
+MAX_NEW_TOKENS_FALLBACK = 256
+MIN_USEFUL_CHARS = 260
+
+def dlog(tag, msg):
+    if DEBUG: print(f"[{tag}] {msg}")
+
+# ================== HEAVY IMPORTS ==================
+import faiss
+from sentence_transformers import SentenceTransformer
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from peft import PeftModel
+import wikipedia
+from wikipedia.exceptions import DisambiguationError, PageError
+from huggingface_hub import login, snapshot_download, HfApi
+
+# ================== GPU CHECK ==================
+if not torch.cuda.is_available():
     with gr.Blocks() as demo:
-        # Store user info and chat history in state
-        user_info = gr.State({"first_name": None, "last_name": None})
-        chat_history = gr.State([])  # list of (user_msg, bot_response, feedback)
+        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)
 
-        # --- User info input ---
-        with gr.Row():
-            first_name = gr.Textbox(label="First Name", placeholder="Enter your first name", interactive=True)
-            last_name = gr.Textbox(label="Last Name", placeholder="Enter your last name", interactive=True)
-        submit_names = gr.Button("Submit")
-
-        # --- Chatbot interface (hidden initially) ---
-        chatbot = gr.Chatbot(visible=False)
-        user_msg = gr.Textbox(placeholder="Ask me anything...", visible=False)
-        send_btn = gr.Button("Send", visible=False)
-        feedback = gr.Radio(choices=["👍", "👎"], label="Your feedback", visible=False)
-        submit_feedback = gr.Button("Submit Feedback", visible=False)
-        clear_btn = gr.Button("Clear Chat", visible=False)
-
-        # --- Functions ---
-
-        def submit_user_info(fn, ln):
-            fn = fn.strip()
-            ln = ln.strip()
-            if not fn or not ln:
-                return gr.update(visible=True), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), "Please enter both first and last names."
-            user_info.value = {"first_name": fn, "last_name": ln}
-            return (
-                gr.update(visible=False),  # hide first_name
-                gr.update(visible=False),  # hide last_name
-                gr.update(visible=True),   # show chatbot
-                gr.update(visible=True),   # show user_msg
-                gr.update(visible=True),   # show send_btn
-                gr.update(visible=False),  # hide feedback (initially)
-                gr.update(visible=False),  # hide submit_feedback
-                gr.update(visible=True),   # show clear_btn
-                ""
+device = "cuda"
+dtype = torch.float16
+torch.manual_seed(42)
+
+# ================== 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__
+
+# ================== LOAD LLM (BASE + LORA) ==================
+if HF_TOKEN:
+    try:
+        login(token=HF_TOKEN)
+        dlog("HF", "Login successful via HF_TOKEN")
+    except Exception as e:
+        dlog("HF", f"Login issue: {e}")
+
+if ADAPTER_REPO:
+    ADAPTER_PATH = snapshot_download(repo_id=ADAPTER_REPO, allow_patterns=["*"])
+
+dlog("LLM", f"Loading base model: {BASE_MODEL}")
+tokenizer_lm = AutoTokenizer.from_pretrained(BASE_MODEL, use_fast=False)
+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.to(device)
+model_lm.eval()
+
+GEN_ARGS_GROUNDED = dict(
+    max_new_tokens=MAX_NEW_TOKENS_GROUNDED,
+    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,
+    do_sample=False,
+    num_beams=1,
+    no_repeat_ngram_size=3,
+    repetition_penalty=1.05,
+    eos_token_id=tokenizer_lm.eos_token_id,
+)
+
+def _generate(inputs, grounded: bool):
+    args = GEN_ARGS_GROUNDED if grounded else GEN_ARGS_FALLBACK
+    with torch.inference_mode():
+        return model_lm.generate(**inputs, **args)
+
+# ================== 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
+
+@lru_cache(maxsize=4096)
+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 []
+
+@lru_cache(maxsize=4096)
+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)
+
+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
 
-        def send_message(message, history):
-            if not message.strip():
-                return history, "", gr.update(visible=False), gr.update(visible=True), gr.update(value=None)
-            response = ask(message)
-            history = history + [(message, response, None)]
-            return history, "", gr.update(visible=True), gr.update(visible=False), gr.update(value=None)
-
-        def submit_user_feedback(feedback_value, history):
-            if not history:
-                return history, gr.update(visible=False), "No conversation yet."
-            if feedback_value is None:
-                return history, gr.update(visible=True), "Please provide feedback."
-            # Update last message feedback
-            history[-1] = (history[-1][0], history[-1][1], feedback_value)
-            return history, gr.update(visible=False), "Thank you for your feedback!"
-
-        def clear_chat():
-            return [], ""
-
-        # --- Event bindings ---
-        submit_names.click(
-            submit_user_info,
-            inputs=[first_name, last_name],
-            outputs=[first_name, last_name, chatbot, user_msg, send_btn, feedback, submit_feedback, clear_btn, gr.Textbox(visible=True, interactive=False, label="Error Message")]
+    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
+
+@lru_cache(maxsize=4096)
+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 []
 
-        send_btn.click(
-            send_message,
-            inputs=[user_msg, chat_history],
-            outputs=[chatbot, user_msg, feedback, submit_feedback, feedback]
+    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
+        if _biomechish(q):
+            wiki = wiki_summary_allow(q, sentences=3)
+            if wiki:
+                dlog("WIKI", "Wikipedia biomechanics 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
+
+    if _biomechish(q):
+        wiki = wiki_summary_allow(q, sentences=3)
+        if wiki:
+            dlog("WIKI", "Wikipedia biomechanics 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 = 1600
+    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 _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)
 
-        submit_feedback.click(
-            submit_user_feedback,
-            inputs=[feedback, chat_history],
-            outputs=[chatbot, feedback, gr.Textbox(visible=True, interactive=False, label="Feedback Status")]
+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 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
+
+def ask(question: str) -> str:
+    q = question.strip()
+    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."
+
+    if _PAPERS_INTENT.search(q):
+        core_q = re.sub(_PAPERS_INTENT, "", q, flags=re.I).strip() or "CT/QCT structural rigidity femur hip finite element"
+        compact = _compact_terms(core_q)
+        pm_query = (
+            f'(({compact}) AND (hip[TiAb] OR femur[TiAb] OR femoral[TiAb])) AND '
+            '("Finite Element Analysis"[MeSH Terms] OR finite element[TiAb] OR QCT[TiAb] OR CT[TiAb] OR rigidity[TiAb]) '
+            'AND ("2000"[DP] : "2025"[DP])'
         )
+        cits = fetch_pubmed_citations(pm_query, max_results=5)
+        return "Recommended papers:\n" + "\n".join(f"- {c}" for c in cits) if cits else "Sorry, no good matches."
+
+    comp = re.match(r"(.+?)\s+and\s+(?:cite|references?|studies?|papers?)", q, flags=re.IGNORECASE)
+    if comp:
+        core_q = comp.group(1).strip()
+        det_text = deterministic_definitions_text(core_q)
+        used_term = None
+        if det_text:
+            explanation = det_text
+            lq = core_q.lower()
+            if ("torsional" 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:
+                dlog("CITE", f"PubMed empty → fallback {used_term}")
+                citations = _fallback_cits_for(used_term)
+        else:
+            explanation = _answer_from_chunks(retrieve_context(core_q, top_k=5), core_q)
+            pm_query  = f'"{core_q}"[Title/Abstract]'
+            citations = fetch_pubmed_citations(pm_query, max_results=5)
+            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)
+            if not citations:
+                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)
+        resp = 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(_post_clean(resp))
+
+    det_answer = deterministic_definitions_text(q)
+    if det_answer:
+        dlog("ASK", "Deterministic definition/workflow fired")
+        return det_answer
+
+    if not (_MSK_MUST.search(q) or _is_fe_override(q)):
+        chunks = retrieve_context(q, top_k=5)
+        if chunks:
+            dlog("CLEAN", "Post-clean applied")
+            answer = _answer_from_chunks(chunks, q)
+            return _ensure_min_answer(_post_clean(answer)) or "I don’t know."
+        return "I don’t know based on the provided context."
+
+    chunks = retrieve_context(q, top_k=5)
+    if not chunks:
+        return "I don’t know based on the provided context."
+    dlog("CLEAN", "Post-clean applied")
+    answer = _answer_from_chunks(chunks, q)
+    return _ensure_min_answer(_post_clean(answer))
+
+# ================== 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."
+
+def predict(message, chat_history, state):
+    answer = ask(message)
+    chat_history = chat_history + [(message, answer)]
+    state["last_q"] = message
+    state["last_a"] = answer
+    return (
+        chat_history,
+        "",  # clear input
+        gr.update(visible=True),  # show feedback pane
+        gr.update(value=None),    # reset rating
+        gr.update(value=""),      # reset comment
+        state
+    )
 
-        clear_btn.click(
-            clear_chat,
-            outputs=[chatbot, user_msg]
+def _push_feedback_to_hub():
+    """Optional: upload feedback.csv to Space repo path analytics/feedback.csv"""
+    if not (PUSH_FEEDBACK and HF_TOKEN and SPACE_REPO_ID and os.path.exists(FEEDBACK_PATH)):
+        return
+    try:
+        api = HfApi()
+        api.upload_file(
+            path_or_fileobj=FEEDBACK_PATH,
+            path_in_repo="analytics/feedback.csv",
+            repo_id=SPACE_REPO_ID,
+            repo_type="space",
+            token=HF_TOKEN
         )
+        dlog("FEEDBACK", "Uploaded analytics/feedback.csv to Hub")
+    except Exception as e:
+        dlog("FEEDBACK", f"Hub 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)
+
+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
+    enter_btn.click(
+        fn=enter_app,
+        inputs=[first_tb, last_tb, state],
+        outputs=[gate, app, state, gate_msg],
+    )
+
+    send_btn.click(predict, inputs=[user_in, chat, state],
+                   outputs=[chat, user_in, feedback_grp, rating, comment, state])
+    user_in.submit(predict, inputs=[user_in, chat, state],
+                   outputs=[chat, user_in, feedback_grp, rating, comment, state])
+
+    clear_btn.click(lambda: ([], "", gr.update(visible=False), None, "", init_session()),
+                    inputs=None,
+                    outputs=[chat, user_in, feedback_grp, rating, comment, state])
 
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    submit_fb.click(
+        fn=save_feedback,
+        inputs=[rating, comment, state],
+        outputs=[fb_status, feedback_grp],
+    )
 
-if __name__ == "__main__":
-    main()
+demo.queue(max_size=32).launch()