import os import re import json import torch import gradio as gr from transformers import AutoModelForCausalLM, AutoTokenizer from peft import PeftModel from huggingface_hub import login from huggingface_hub.utils import RepositoryNotFoundError, GatedRepoError # Authenticate for gated Llama access (token supplied as a Space secret) hf_token = os.environ.get("HF_TOKEN") if hf_token: login(hf_token) GPU = torch.cuda.is_available() DTYPE = torch.float16 if GPU else torch.float32 # ---------------------------------------------------------------------------- # Evidence panel (UniProt / PubMed / AlphaFold) — same source as the project's # guarded report layer. Used to attach real citation source IDs to claims. # ---------------------------------------------------------------------------- EVIDENCE = {} _ev_path = os.path.join(os.path.dirname(__file__), "kaubiomed_gene_evidence.jsonl") if os.path.exists(_ev_path): with open(_ev_path) as fh: for line in fh: if line.strip(): row = json.loads(line) if row.get("gene"): EVIDENCE[row["gene"]] = row # ClinVar significance bucket -> evidence-label score (NOT calibrated model). SIG_TO_BUCKET = { "Unknown / Not provided": "other", "Pathogenic": "pathogenic", "Likely_pathogenic": "likely_pathogenic", "Uncertain_significance": "vus", "Likely_benign": "likely_benign", "Benign": "benign", "Conflicting_interpretations": "conflicting", } BUCKET_SCORE = { "pathogenic": 0.95, "likely_pathogenic": 0.80, "benign": 0.05, "likely_benign": 0.15, "conflicting": 0.50, "vus": 0.50, "other": 0.50, } EVIDENCE_TYPES = [ "uniprot_function", "pubmed_literature", "alphafold_structure", "vep_mane_transcript_mapping", "residue_level_plddt", "clinical_population_functional_evidence", ] QWEN_INSTRUCTION = ( "Create one compact JSON variant interpretation report for this exact human variant. " "Return ONLY valid JSON, no prose, no markdown. The JSON object must contain keys " "task_type, gene, variant, evidence_summary, reasoning_summary, final_interpretation, " "confidence, limitations. For evidence_summary, cite only the given ClinVar significance " "and molecular consequence; if either field is 'Unknown / Not provided', state that it was " "not supplied and infer what you can from gene and variant coordinates alone, clearly marking " "the inference as uncertain. Mention uncertainty when appropriate." ) LLAMA_INSTRUCTION = ( "Create one compact JSON variant interpretation report for this exact human variant. " "Return ONLY valid JSON. Include top-level fields task_type, gene, variant, kaubiomed_score, " "score_basis, genomic_context, evidence_summary, reasoning_summary, final_interpretation, " "confidence, limitations. kaubiomed_score must be a numeric value between 0 and 1, not a string. " "If clinvar_significance or molecular_consequence is 'Unknown / Not provided', state this clearly " "in evidence_summary and set confidence to 'low' with appropriate limitations." ) MODELS = { "Qwen2.5-1.5B (KAU-BioMedLLM v0.1 — runs on free CPU)": { "base": "Qwen/Qwen2.5-1.5B-Instruct", "adapter": "Babajaan/KAU-BioMedLLM", "subfolder": "qwen_v0_1", "instruction": QWEN_INSTRUCTION, "gpu_only": False, }, "Llama-3.1-8B (KAU-BioMedLLM v0.2 — requires GPU / PRO hardware)": { "base": "meta-llama/Llama-3.1-8B-Instruct", "adapter": "Babajaan/KAU-BioMedLLM", "subfolder": None, "instruction": LLAMA_INSTRUCTION, "gpu_only": True, }, } DEFAULT_MODEL = list(MODELS.keys())[0] _CACHE = {} def load_model(choice): # noqa: E302 if choice in _CACHE: return _CACHE[choice] cfg = MODELS[choice] tok = AutoTokenizer.from_pretrained(cfg["base"]) if tok.pad_token is None: tok.pad_token = tok.eos_token # device_map="auto" triggers a PEFT offload path that breaks on CPU; only use it on GPU. load_kwargs = {"dtype": DTYPE} if GPU: load_kwargs["device_map"] = "auto" base = AutoModelForCausalLM.from_pretrained(cfg["base"], **load_kwargs) peft_kwargs = {"subfolder": cfg["subfolder"]} if cfg.get("subfolder") else {} model = PeftModel.from_pretrained(base, cfg["adapter"], **peft_kwargs) if not GPU: model = model.to("cpu") model.eval() _CACHE[choice] = (tok, model) return _CACHE[choice] def run_llm(choice, gene, variant, significance, consequence, disease): """Return the parsed LLM narrative dict (and raw text fallback).""" cfg = MODELS[choice] tokenizer, model = load_model(choice) evidence = { "task_type": "variant_interpretation", "gene": gene, "variant": variant, "clinvar_significance": significance, "molecular_consequence": consequence, "disease_context": disease, } prompt = cfg["instruction"] + "\n\nInput evidence:\n" + json.dumps(evidence, ensure_ascii=False) inputs = tokenizer.apply_chat_template( [{"role": "user", "content": prompt}], tokenize=True, add_generation_prompt=True, return_tensors="pt", return_dict=True, ).to(model.device) with torch.no_grad(): outputs = model.generate(**inputs, max_new_tokens=320, do_sample=False, pad_token_id=tokenizer.eos_token_id) raw = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True) parsed, ok = {}, False s, e = raw.find("{"), raw.rfind("}") + 1 if s >= 0 and e > s: try: parsed = json.loads(raw[s:e]); ok = True except Exception: parsed = {} return parsed, raw, ok def parse_variant(variant): m = re.match(r"chr?([\dXYMT]+):(\d+):([ACGTN]+)>([ACGTN]+)", variant.strip(), re.I) if m: return {"assembly": "GRCh38/hg38", "chrom": m.group(1), "pos": int(m.group(2)), "ref": m.group(3).upper(), "alt": m.group(4).upper()} return {"assembly": "GRCh38/hg38", "raw": variant} def assemble_report(choice, gene, variant, significance, consequence, disease): """Build the full guarded report (all components) around the LLM narrative.""" cfg = MODELS[choice] if cfg["gpu_only"] and not GPU: return None, {"error": "gpu_required"} try: narrative, raw, schema_ok = run_llm(choice, gene, variant, significance, consequence, disease) except Exception as ex: # noqa: BLE001 msg = str(ex) if isinstance(ex, (RepositoryNotFoundError, GatedRepoError)) or "adapter_config.json" in msg: return None, {"error": "adapter_missing", "detail": cfg["adapter"]} return None, {"error": "load", "detail": f"{type(ex).__name__}: {ex}"} bucket = SIG_TO_BUCKET.get(significance, "other") evidence_label_score = BUCKET_SCORE.get(bucket, 0.5) ge = EVIDENCE.get(gene, {}) uni = ge.get("uniprot") or {} pub = ge.get("pubmed") or {} af = ge.get("alphafold") or {} accession = uni.get("accession") protein_name = uni.get("protein_name") function_text = uni.get("function") or "" articles = (pub.get("articles") or [])[:5] # PubMed IDs cited in the UniProt function text + retrieved articles pmids_in_function = sorted(set(re.findall(r"PubMed:(\d+)", function_text)))[:8] article_pmids = [a.get("pmid") for a in articles if a.get("pmid")] # ---- 1. score_model_prediction (calibrated model abstains for arbitrary variants) ---- score_model_prediction = { "status": "abstain", "model_version": "broad_no_review_status_first_appearance_v0 / missense_v1", "reason": "variant_not_in_precomputed_score_tables", "message": ("The leakage-free calibrated HistGradientBoosting score_model is evaluated " "on precomputed temporal test tables; arbitrary-variant scoring is deferred to v2. " "This demo therefore abstains from a calibrated probability."), "not_for_clinical_diagnosis": True, } # ---- 7. score_basis (what the evidence-label score means) ---- sig_supplied = significance != "Unknown / Not provided" score_basis = { "score_version": "KAUBioMED_score_schema_v2", "value": evidence_label_score, "source_label_bucket": bucket, "derivation": (f"Evidence-label score mapped from the supplied ClinVar significance bucket '{bucket}'." if sig_supplied else "No ClinVar significance supplied; evidence-label score defaults to 0.5 (uncertain)."), "calibration_note": ("Prototype evidence-limited score, NOT the leakage-free calibrated " "pathogenicity probability. Used only to organize the report."), } # ---- 2. citation source IDs ---- source_ids = [] if accession: source_ids.append(f"UniProt:{accession}") for p in (pmids_in_function or article_pmids): source_ids.append(f"PubMed:{p}") if af.get("status") == "ok" and af.get("uniprot"): source_ids.append(f"AlphaFold:{af['uniprot']}") # ---- 5. structure_context (residue mapping abstains without live VEP) ---- if af.get("status") == "ok": structure_context = { "status": "abstain_residue_level", "alphafold_model": af.get("entry_id"), "alphafold_uniprot": af.get("uniprot"), "model_url": af.get("cif_url"), "reason": "no_live_vep_mane_reconciliation_in_demo", "policy": ("Residue-level pLDDT/MANE structural claims require Ensembl VEP transcript + " "HGVSp + UniProt reconciliation, which is not run in this demo, so it abstains."), } else: structure_context = {"status": "abstain", "reason": "alphafold_structure_unavailable"} # ---- 6. evidence_completeness ---- available, missing = [], [] (available if accession else missing).append("uniprot_function") (available if articles else missing).append("pubmed_literature") (available if af.get("status") == "ok" else missing).append("alphafold_structure") missing += ["vep_mane_transcript_mapping", "residue_level_plddt", "clinical_population_functional_evidence"] evidence_completeness = { "available_evidence_types": available, "missing_evidence_types": missing, "available_count": len(available), "total_tracked": len(EVIDENCE_TYPES), "fraction_available": round(len(available) / len(EVIDENCE_TYPES), 3), } # ---- 4. abstentions ---- abstentions = [ {"area": "score_model", "reason": "calibrated_score_deferred_v2", "message": "Calibrated pathogenicity score abstains for arbitrary variants."}, {"area": "protein_structure_residue", "reason": "no_reconciled_vep_mapping", "message": "Residue-level structural interpretation abstains."}, {"area": "clinical_population_functional_evidence", "reason": "not_in_demo_panel", "message": "Population/functional clinical evidence not retrieved in this demo."}, ] if significance == "Unknown / Not provided": abstentions.append({"area": "clinvar_significance", "reason": "not_supplied_by_user", "message": ("ClinVar significance was not provided. Evidence-label score " "defaults to 0.5 (uncertain). Supply the significance label " "for a more informative report.")}) if consequence == "Unknown / Not provided": abstentions.append({"area": "molecular_consequence", "reason": "not_supplied_by_user", "message": ("Molecular consequence was not provided. " "LLM will infer from coordinates with lower confidence.")}) if not accession: abstentions.append({"area": "gene_function", "reason": "gene_not_in_evidence_panel", "message": f"No UniProt source for {gene} in the 20-gene demo panel."}) # ---- 3. guard_status (citation + schema + clinical use) ---- citation_pass = bool(accession) and bool(source_ids) guard_status = { "citation_check": "pass" if citation_pass else "fail", "schema_validation": "pass" if schema_ok else "fail", "abstention_allowed": True, "clinical_use": "not permitted", } # ---- 8. llm_interpretation (kept SEPARATE from the score) ---- llm_interpretation = { "source": f"KAU-BioMedLLM report generator ({cfg['base'].split('/')[-1]})", "final_interpretation": narrative.get("final_interpretation") or significance.replace("_", " "), "confidence": narrative.get("confidence", "limited"), "reasoning_summary": narrative.get("reasoning_summary") or (raw.strip() if not schema_ok else ""), "limitations": narrative.get("limitations", "Research record only. Not a clinical diagnosis. Requires expert review."), "schema_valid_json": schema_ok, } report = { "gene": gene, "variant": variant, "genomic_context": parse_variant(variant), "input_evidence": {"clinvar_significance": significance, "molecular_consequence": consequence, "disease_context": disease}, "score_model_prediction": score_model_prediction, "evidence_label_score": evidence_label_score, "score_basis": score_basis, "llm_interpretation": llm_interpretation, "evidence_examples": { "uniprot_accession": accession, "protein_name": protein_name, "gene_function": (function_text[:600] + "…") if function_text else None, "literature": [{"pmid": a.get("pmid"), "title": a.get("title"), "journal": a.get("journal"), "year": a.get("year")} for a in articles], "source_ids": source_ids, }, "structure_context": structure_context, "evidence_completeness": evidence_completeness, "abstentions": abstentions, "guard_status": guard_status, "score_semantics": { "field": "evidence_label_score", "current_meaning": "evidence_label_score_from_supplied_source_bucket", "not_yet": "leakage_free_calibrated_pathogenicity_probability", }, "disclaimer": "Research use only. Not for clinical diagnosis, treatment, or patient management.", } return report, {"ok": True} def badge(text, color): return (f"{text}") def bar(frac, color="#0072B2"): pct = int(round(frac * 100)) return (f"
" f"
" f"{pct}%") def render(report): g = report["guard_status"] li = report["llm_interpretation"] ev = report["evidence_examples"] sb = report["score_basis"] sc = report["structure_context"] ec = report["evidence_completeness"] ok = "#2e7d32"; warn = "#ed6c02"; bad = "#c62828"; blue = "#0072B2" sev_color = {"pathogenic": bad, "likely_pathogenic": warn, "vus": "#757575", "conflicting": warn, "likely_benign": "#558b2f", "benign": ok}.get( sb["source_label_bucket"], "#757575") h = [] h.append(f"

🧬 KAU-BioMedLLM — Guarded Variant Report

") def _display(v): return "—" if not v or v == "Unknown / Not provided" else v h.append(f"
{report['gene']}  ·  {report['variant']} " f" ·  {_display(report['input_evidence']['molecular_consequence'])} " f" ·  {_display(report['input_evidence']['disease_context'])}
") h.append(f"
{badge('RESEARCH USE ONLY', bad)} " f"{badge('NOT FOR CLINICAL DIAGNOSIS', '#555')}
") h.append("
") # Score block (model abstain + evidence-label score, kept separate from narrative) h.append("

① Score

") h.append(f"

{badge('calibrated score_model: ABSTAIN', warn)} " f"{report['score_model_prediction']['reason']}

") h.append(f"

Evidence-label score: {badge(str(sb['value']) + ' (' + sb['source_label_bucket'] + ')', sev_color)}

") h.append(bar(sb["value"], sev_color)) h.append(f"

score_basis: {sb['derivation']} {sb['calibration_note']}

") # LLM interpretation (separate named block) h.append("

② LLM interpretation (narrative — separate from score)

") h.append(f"

Final interpretation: {li['final_interpretation']}   " f"Confidence: {li['confidence']}

") if li["reasoning_summary"]: h.append(f"

{li['reasoning_summary']}

") h.append(f"

Limitations: {li['limitations']}

") # Evidence + citations h.append("

③ Evidence & citation source IDs

") if ev["uniprot_accession"]: h.append(f"

UniProt: {ev['uniprot_accession']} — {ev['protein_name']}

") if ev["literature"]: h.append("") if ev["source_ids"]: chips = " ".join(badge(s, blue) for s in ev["source_ids"]) h.append(f"

{chips}

") # Structure h.append("

④ Structure context

") if sc.get("alphafold_model"): h.append(f"

AlphaFold model {sc['alphafold_model']} available " f"({badge('residue mapping: ABSTAIN', warn)})
{sc['policy']}

") else: h.append(f"

{badge('structure: ABSTAIN', warn)} {sc.get('reason')}

") # Guard status h.append("

⑤ Guard status

") def mark(v): return badge("✓ " + v, ok) if v == "pass" else badge("✗ " + v, bad) h.append(f"

Citation check: {mark(g['citation_check'])}   " f"Schema validation: {mark(g['schema_validation'])}   " f"Clinical use: {badge('NOT PERMITTED', bad)}

") # Abstentions h.append("

⑥ Abstentions

") # Completeness h.append("

⑦ Evidence completeness

") h.append(bar(ec["fraction_available"], blue)) h.append(f"

Available: {', '.join(ec['available_evidence_types']) or 'none'}
" f"Missing/abstained: {', '.join(ec['missing_evidence_types'])}

") return "\n".join(h) def interpret(choice, gene, variant, significance, consequence, disease): if not gene.strip() or not variant.strip(): return "Please provide both a gene symbol and a variant.", "{}" report, status = assemble_report(choice, gene, variant, significance, consequence, disease) if report is None: if status.get("error") == "gpu_required": md = ("### ⚠️ Llama-3.1-8B requires GPU hardware\n" "This Space is on free CPU, where the 8B model is too slow to serve. " "Upgrade the Space to a T4 GPU (Settings → Hardware) or run on a GPU/PRO environment.\n\n" "For a working demo on free CPU, choose **Qwen2.5-1.5B (v0.1)** above.") return md, "{}" if status.get("error") == "adapter_missing": return f"### ⚠️ Adapter not available\n`{status.get('detail')}` was not found on the Hub.", "{}" return f"### Model error\n```\n{status.get('detail')}\n```", "{}" return render(report), json.dumps(report, indent=2, ensure_ascii=False) disclaimer = """ ## ⚠️ Research Use Only This system is **not clinically validated** and must not be used for diagnosis, treatment, or patient-management decisions. All outputs require expert review. Model: [KAU-BioMedLLM](https://huggingface.co/Babajaan/KAU-BioMedLLM) — King Abdulaziz University **Model versions:** Qwen2.5-1.5B (v0.1) runs on this **free CPU** Space · Llama-3.1-8B (v0.2) **requires GPU (paid T4 / PRO)**. """ with gr.Blocks(title="KAU-BioMedLLM Demo") as demo: gr.Markdown("# KAU-BioMedLLM — Guarded Biomedical Variant Interpretation") gr.Markdown(disclaimer) model_choice = gr.Radio(label="Model version", choices=list(MODELS.keys()), value=DEFAULT_MODEL, info="Qwen runs on free CPU. Llama needs a GPU (paid T4 / PRO).") with gr.Row(): gene = gr.Textbox(label="Gene Symbol", value="TP53") variant = gr.Textbox(label="Variant (chrN:POS:REF>ALT)", value="chr17:7674220:G>A") with gr.Row(): significance = gr.Dropdown( label="ClinVar Significance (optional — select if known)", choices=list(SIG_TO_BUCKET.keys()), value="Unknown / Not provided", info="Leave as 'Unknown / Not provided' if you only have coordinates.") consequence = gr.Dropdown( label="Molecular Consequence (optional — select if known)", choices=["Unknown / Not provided", "missense_variant", "stop_gained", "frameshift_variant", "synonymous_variant", "splice_donor_variant", "splice_acceptor_variant", "intron_variant"], value="Unknown / Not provided", info="Leave as 'Unknown / Not provided' if you only have coordinates.") disease = gr.Textbox(label="Disease Context (optional)", value="Li-Fraumeni syndrome") btn = gr.Button("Generate Guarded Report", variant="primary") report_md = gr.Markdown(label="Guarded Report") with gr.Accordion("Raw guarded report (JSON)", open=False): report_json = gr.Code(label="JSON", language="json", lines=24) gr.Markdown("⏳ On free CPU, Qwen takes roughly 1–2 minutes per report. " "Please wait for the spinner to finish.") btn.click(fn=interpret, inputs=[model_choice, gene, variant, significance, consequence, disease], outputs=[report_md, report_json], show_progress="full") gr.Examples( label="Quick Examples", examples=[ [DEFAULT_MODEL, "TP53", "chr17:7674220:G>A", "Pathogenic", "missense_variant", "Li-Fraumeni syndrome"], [DEFAULT_MODEL, "BRCA1", "chr17:43106455:G>A", "Likely_pathogenic", "missense_variant", "Hereditary breast cancer"], [DEFAULT_MODEL, "MC4R", "chr18:60371172:C>T", "Uncertain_significance", "missense_variant", "Obesity"], [DEFAULT_MODEL, "LDLR", "chr19:11089548:C>T", "Likely_benign", "missense_variant", "Familial hypercholesterolemia"], [DEFAULT_MODEL, "CFTR", "chr7:117559590:A>G", "Pathogenic", "missense_variant", "Cystic fibrosis"], ], inputs=[model_choice, gene, variant, significance, consequence, disease]) if __name__ == "__main__": demo.queue(default_concurrency_limit=1, max_size=8) demo.launch(ssr_mode=False)