app.py

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"<span style='background:{color};color:#fff;padding:2px 9px;border-radius:10px;"
            f"font-size:0.82em;font-weight:600;'>{text}</span>")


def bar(frac, color="#0072B2"):
    pct = int(round(frac * 100))
    return (f"<div style='background:#e9ecef;border-radius:6px;height:14px;width:100%;max-width:320px;'>"
            f"<div style='background:{color};width:{pct}%;height:14px;border-radius:6px;'></div></div>"
            f"<small>{pct}%</small>")


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"<h2 style='margin-bottom:2px'>🧬 KAU-BioMedLLM — Guarded Variant Report</h2>")
    def _display(v):
        return "—" if not v or v == "Unknown / Not provided" else v
    h.append(f"<div style='color:#555'><b>{report['gene']}</b> &nbsp;·&nbsp; <code>{report['variant']}</code> "
             f"&nbsp;·&nbsp; {_display(report['input_evidence']['molecular_consequence'])} "
             f"&nbsp;·&nbsp; {_display(report['input_evidence']['disease_context'])}</div>")
    h.append(f"<div style='margin:8px 0'>{badge('RESEARCH USE ONLY', bad)} "
             f"{badge('NOT FOR CLINICAL DIAGNOSIS', '#555')}</div>")
    h.append("<hr>")

    # Score block (model abstain + evidence-label score, kept separate from narrative)
    h.append("<h3>① Score</h3>")
    h.append(f"<p>{badge('calibrated score_model: ABSTAIN', warn)} "
             f"<small>{report['score_model_prediction']['reason']}</small></p>")
    h.append(f"<p><b>Evidence-label score:</b> {badge(str(sb['value']) + '  (' + sb['source_label_bucket'] + ')', sev_color)}</p>")
    h.append(bar(sb["value"], sev_color))
    h.append(f"<p><small><b>score_basis:</b> {sb['derivation']} {sb['calibration_note']}</small></p>")

    # LLM interpretation (separate named block)
    h.append("<h3>② LLM interpretation <small style='color:#888'>(narrative — separate from score)</small></h3>")
    h.append(f"<p style='color:#000'><b>Final interpretation:</b> {li['final_interpretation']} &nbsp; "
             f"<b>Confidence:</b> {li['confidence']}</p>")
    if li["reasoning_summary"]:
        h.append(f"<p style='background:#f6f8fa;padding:10px;border-radius:6px'>{li['reasoning_summary']}</p>")
    h.append(f"<p><small><b>Limitations:</b> {li['limitations']}</small></p>")

    # Evidence + citations
    h.append("<h3>③ Evidence &amp; citation source IDs</h3>")
    if ev["uniprot_accession"]:
        h.append(f"<p><b>UniProt:</b> <code>{ev['uniprot_accession']}</code> — {ev['protein_name']}</p>")
    if ev["literature"]:
        h.append("<ul>")
        for a in ev["literature"]:
            h.append(f"<li><code>PMID:{a['pmid']}</code> — {a['title']} <i>({a['journal']}, {a['year']})</i></li>")
        h.append("</ul>")
    if ev["source_ids"]:
        chips = " ".join(badge(s, blue) for s in ev["source_ids"])
        h.append(f"<p>{chips}</p>")

    # Structure
    h.append("<h3>④ Structure context</h3>")
    if sc.get("alphafold_model"):
        h.append(f"<p>AlphaFold model <code>{sc['alphafold_model']}</code> available "
                 f"({badge('residue mapping: ABSTAIN', warn)})<br><small>{sc['policy']}</small></p>")
    else:
        h.append(f"<p>{badge('structure: ABSTAIN', warn)} <small>{sc.get('reason')}</small></p>")

    # Guard status
    h.append("<h3>⑤ Guard status</h3>")
    def mark(v):
        return badge("✓ " + v, ok) if v == "pass" else badge("✗ " + v, bad)
    h.append(f"<p>Citation check: {mark(g['citation_check'])} &nbsp; "
             f"Schema validation: {mark(g['schema_validation'])} &nbsp; "
             f"Clinical use: {badge('NOT PERMITTED', bad)}</p>")

    # Abstentions
    h.append("<h3>⑥ Abstentions</h3><ul>")
    for a in report["abstentions"]:
        h.append(f"<li>{badge(a['area'], warn)} {a['message']}</li>")
    h.append("</ul>")

    # Completeness
    h.append("<h3>⑦ Evidence completeness</h3>")
    h.append(bar(ec["fraction_available"], blue))
    h.append(f"<p><small><b>Available:</b> {', '.join(ec['available_evidence_types']) or 'none'}<br>"
             f"<b>Missing/abstained:</b> {', '.join(ec['missing_evidence_types'])}</small></p>")
    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("<small>⏳ On free CPU, Qwen takes roughly 1–2 minutes per report. "
                "Please wait for the spinner to finish.</small>")
    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)