chat.py

"""
Two-mode interactive chat
=========================
An interactive console where a clinician picks a patient and talks about them.

TWO MODES, because measurement showed one configuration cannot do both well:

  REPORT   (LoRA on, thinking off, soft tokens injected)
     The trained closed-set tasks. Faithfulness 39/40, fast, fixed format.
     Reads the diagnosis from the soft tokens — the channel is verified by
     ablation, not assumed.

  DISCUSS  (LoRA off, thinking on, NO soft tokens)
     Open-ended clinical questions. The LoRA model got an atypical case wrong
     ("precuneus/parietal atrophy is consistent with amnestic AD" — it is not);
     the base model answered the same case correctly and with reasoning.
     Soft tokens are unnecessary here: every measurement is already in the text.

CONTEXT PRESERVATION — the core of this design:
  Patient data, regional measurements, the ATN profile and Vbai-2.6AD's verdict are
  written into the SYSTEM PROMPT, not into a user message. The system prompt
  stays visible every turn and does not scroll out of context. A question asked
  on turn 10 is still answered against the patient introduced on turn 1.

FAITHFULNESS GUARD:
  After every generation the class stated in the text is compared against the
  head's verdict; a mismatch prints a warning. This catches CONTRADICTION WITH
  THE HEAD only — it does not catch general medical errors. That limit is
  permanent.

Note on language: the assistant answers in Turkish. The system rules and the
canonical questions are Turkish by design — that is what the adapter was
trained on. Only this console's interface is English.

Commands:
  /report     switch to report mode     /discuss   switch to discuss mode
  /patient N  jump to patient N         /data      show the system prompt
  /reset      clear the conversation    /questions list trained questions
  /help       command list              /quit      exit

Run:
    python chat.py --features features.pt --text dataset.json \
        --projector projector.pt
"""
from __future__ import annotations
import argparse
import inspect
import os
import re
import sys

import torch

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from hfx_runtime import Projector, build_examples, SENTINEL
from eval_faithfulness import find_class
from eval_thinking import strip_body, split_thought

# Kept in Turkish on purpose: this is the prompt the adapter was trained under.
# Translating it would move the input off the training distribution.
SYS_RULES = (
    "Sen bir klinik karar destek asistanısın. Kurallar:\n"
    "1. Yukarıdaki hasta verilerinin DIŞINA çıkma; verilmeyen bir bulgu uydurma.\n"
    "2. 'ölçülmedi' yazan bir değeri normal kabul etme; eksik olduğunu söyle.\n"
    "3. Vbai-2.6AD sınıflandırmasını DEĞİŞTİRME. Katılmadığın noktayı belirtebilirsin "
    "ama sınıflandırma bu değerlendirmenin çıpasıdır.\n"
    "4. Doktor itiraz ettiğinde fikrini kanıtsız değiştirme; gerekçeni göster.\n"
    "5. Bu çıktı karar desteğidir, tanı yerine geçmez.\n"
    "Yanıtlarını TÜRKÇE ver."
)


def build_system(ex, head: str, class_names, probs=None, will_progress=None) -> str:
    """
    Patient context + anchor. Lives in the system prompt, visible every turn.

    THE HEAD'S OUTPUTS ARE WRITTEN HERE AS NUMBERS.
    Why: every digit was stripped from the training targets (the model had
    memorized "0.76" and was writing it for every patient). So the model CANNOT
    produce numbers — and should not. But Vbai-2.6AD genuinely computes them, so
    they are placed in the text for the model to read out when asked. This is
    the "numbers from code, prose from the LLM" principle in practice.
    """
    lines = [f"Vbai-2.6AD değerlendirmesi: {head}"]
    if probs is not None:
        lines.append("  Sınıf olasılıkları: " +
                     ", ".join(f"{c} %{100*p:.1f}" for c, p in zip(class_names, probs)))
    if will_progress is not None:
        if head == "MCI":
            lines.append(f"  MCI→AD ilerleme riski (5 yıl): {float(will_progress):.2f}")
        else:
            # Vbai-2.6AD README: the progression head is only meaningful for MCI;
            # a score is produced for CN/AD but must not be interpreted.
            lines.append("  İlerleme riski: yalnızca MCI vakalarında anlamlıdır, "
                         "bu vakada raporlanmaz.")
    return (strip_body(ex["prompt"]).replace(SENTINEL, "").strip() +
            "\n\n" + "\n".join(lines) + "\n\n" + SYS_RULES)


@torch.no_grad()
def generate(model, tok, embed_layer, projector, ex, device, supported,
             messages, mode: str, max_new_tokens: int):
    use_lora = (mode == "report")
    thinking = (mode == "discuss")

    kw = dict(tokenize=False, add_generation_prompt=True)
    try:
        full = tok.apply_chat_template(messages, enable_thinking=thinking, **kw)
    except TypeError:
        full = tok.apply_chat_template(messages, **kw)

    ids = tok(full, add_special_tokens=False, return_tensors="pt").input_ids.to(device)
    embeds = embed_layer(ids)
    off = n_soft = 0
    if use_lora:
        # In report mode the soft tokens are prepended (the position the LoRA
        # learned to read them from). In discuss mode the base model cannot read
        # them, so they are omitted.
        soft = projector(ex["feat"].unsqueeze(0).to(device)).to(embeds.dtype)
        embeds = torch.cat([soft, embeds], dim=1)
        n_soft = soft.size(1)

    attn = torch.ones(embeds.shape[:2], dtype=torch.long, device=device)
    gk = {"inputs_embeds": embeds, "attention_mask": attn,
          "max_new_tokens": max_new_tokens, "do_sample": False,
          "repetition_penalty": 1.1}
    if n_soft and "mm_token_type_ids" in supported:
        mm = torch.zeros(embeds.shape[:2], dtype=torch.long, device=device)
        mm[0, off:off + n_soft] = 1
        gk["mm_token_type_ids"] = mm

    ctx = model.disable_adapter() if (not use_lora and hasattr(model, "disable_adapter")) \
        else _null()
    with ctx:
        out = model.generate(**gk)
    # generate() with inputs_embeds returns ONLY the new tokens; hitting the
    # ceiling exactly means the output was cut off, not naturally finished.
    truncated = out.shape[1] >= max_new_tokens
    return tok.decode(out[0], skip_special_tokens=False), truncated


def trained_questions() -> list:
    """
    The canonical questions report mode was trained on — read from the dataset
    generator rather than copied by hand, so the two cannot drift apart.

    Why a menu: LoRA learns PHRASE PATTERNS, not task boundaries. A free-typed
    question landing between trained patterns gets answered with a blend of
    them — one run produced both "atrophy is marked" and "all values are within
    the normal range" in a single answer, and invented an anchor-undermining
    line ("the Vbai-2.6AD assessment can be changed"). Keep the input inside the
    distribution and blending cannot physically occur.

    The questions themselves stay Turkish: they are the training strings.
    """
    try:
        from build_multitask_dataset import TASKS, ABSENT_PROBES
    except Exception:
        return []
    out, dummy = [], (None, None, None, None, [])
    for name, fn in TASKS.items():
        if name == "absent":
            out.append((name, ABSENT_PROBES[0][0]))
            continue
        try:
            q, _ = fn(*dummy) if name != "probs" else fn(
                "MCI", "yüksek", None, None, [], probs=[0.1, 0.8, 0.1],
                class_names=["CN", "MCI", "AD"], will_progress=0.5)
            out.append((name, q))
        except Exception:
            continue
    return out


def find_verdict(text: str, class_names) -> str | None:
    """
    Return the class only when it appears in an EXPLICIT VERDICT pattern, not
    merely the first class name occurring in the text. Long discussion answers
    enumerate differentials, and those mentions must not be mistaken for a
    verdict.
    """
    lab = (r"(?:Değerlendirme|Sınıflandırma|Sonuç|Tanı|Vbai-2.6AD[^:\n]*|"
           r"Classification|Assessment|Diagnosis)")
    m = re.search(lab + r"[^\n:]*:\s*\**\s*(" +
                  "|".join(map(re.escape, class_names)) + r")\b",
                  text, flags=re.IGNORECASE)
    if not m:
        return None
    for c in class_names:
        if c.lower() == m.group(1).lower():
            return c
    return None


def clean(t: str) -> str:
    """Hide control tokens — they arrive because skip_special_tokens=False."""
    for m in ("<eos>", "<turn|>", "<|turn>", "<bos>", "<end_of_turn>"):
        t = t.replace(m, "")
    return t.strip()


class _null:
    def __enter__(self): return None
    def __exit__(self, *a): return False


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--features", required=True)
    ap.add_argument("--text", required=True)
    ap.add_argument("--projector", default="projector.pt",
                    help="projector checkpoint; the LoRA adapter is expected "
                         "alongside it at <projector>_lora/")
    ap.add_argument("--model", default=None)
    ap.add_argument("--split", default="test")
    ap.add_argument("--patient", type=int, default=0,
                    help="index of the patient within the split")
    ap.add_argument("--mode", default="report", choices=["report", "discuss"])
    # Budget per mode: in discuss mode the chain of thought alone eats 1000+
    # tokens and left no room for the answer. Report answers are short.
    ap.add_argument("--max-new-report", type=int, default=500)
    ap.add_argument("--max-new-discuss", type=int, default=2500)
    ap.add_argument("--no-4bit", dest="four_bit", action="store_false", default=True)
    args = ap.parse_args()

    import transformers
    from transformers import AutoTokenizer
    AutoCls = next(getattr(transformers, n) for n in
                   ("AutoModelForConditionalGeneration", "AutoModelForImageTextToText",
                    "AutoModelForCausalLM") if hasattr(transformers, n))

    device = "cuda" if torch.cuda.is_available() else "cpu"
    ck = torch.load(args.projector, map_location="cpu", weights_only=False)
    model_id = args.model or ck["model_id"]
    data, d = build_examples(args.features, args.text)
    class_names = list(d["class_names"])

    tok = AutoTokenizer.from_pretrained(model_id)
    load_kw = dict(device_map={"": 0} if device == "cuda" else None)
    if args.four_bit:
        from transformers import BitsAndBytesConfig
        load_kw["quantization_config"] = BitsAndBytesConfig(
            load_in_4bit=True, bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True)
    try:
        model = AutoCls.from_pretrained(model_id, dtype=torch.bfloat16, **load_kw)
    except TypeError:
        model = AutoCls.from_pretrained(model_id, torch_dtype=torch.bfloat16, **load_kw)
    supported = set(inspect.signature(model.forward).parameters)
    embed_layer = model.get_input_embeddings()

    if ck.get("lora") and os.path.isdir(args.projector + "_lora"):
        from peft import PeftModel
        model = PeftModel.from_pretrained(model, args.projector + "_lora")
    model.eval()

    projector = Projector(ck["in_dim"], ck["hidden"], ck["n_tokens"],
                          target_norm=ck.get("target_norm")).to(device)
    projector.load_state_dict(ck["projector"]); projector.eval()

    # De-duplicate by patient (the multitask file holds several records each)
    seen, patients = set(), []
    for e in data[args.split]:
        if e["ptid"] not in seen:
            seen.add(e["ptid"]); patients.append(e)

    probs_all = d["class_probs"]
    wp_all = d["will_progress"].reshape(-1)

    def mk_system(e):
        i = e.get("index")
        return build_system(e, e["head"], class_names,
                            probs_all[i].tolist() if i is not None else None,
                            float(wp_all[i]) if i is not None else None)

    idx, mode = args.patient, args.mode
    ex = patients[idx]
    system = mk_system(ex)
    messages = [{"role": "system", "content": system}]

    tq = trained_questions()

    def show_questions():
        if not tq:
            print("  (could not read the trained question list)"); return
        print("\n  Questions report mode was trained on — ask by number:")
        for k, (name, qq) in enumerate(tq, 1):
            print(f"   {k}. [{name}] {qq}")
        print("  Free text is allowed but falls outside the distribution in "
              "report mode.")

    print(f"\n{'='*70}\nPatient {ex['ptid']}  |  Vbai-2.6AD: {ex['head']}  |  "
          f"true: {class_names[ex['label']]}  |  mode: {mode}")
    print(f"{len(patients)} patients loaded. /help for commands.\n{'='*70}")
    show_questions()

    while True:
        try:
            q = input(f"\n[{mode}] clinician> ").strip()
        except (EOFError, KeyboardInterrupt):
            break
        if not q:
            continue
        if q in ("/quit", "/exit"):
            break
        if q == "/help":
            print("  /report  /discuss  /patient N  /data  /reset  /questions  /quit")
            print("  in report mode you can ask a question BY NUMBER (1-8)")
            continue
        if q == "/questions":
            show_questions(); continue
        if q in ("/report", "/discuss"):
            mode = q[1:]
            print(f"  → mode: {mode}" +
                  ("  (LoRA on, thinking off, soft tokens injected)" if mode == "report"
                   else "  (LoRA off, thinking on, reasoning from the base model)"))
            continue
        if q.startswith("/patient"):
            try:
                idx = int(q.split()[1]) % len(patients)
            except (IndexError, ValueError):
                print("  usage: /patient 3"); continue
            ex = patients[idx]
            system = mk_system(ex)
            messages = [{"role": "system", "content": system}]
            print(f"  → patient {ex['ptid']}  Vbai-2.6AD: {ex['head']}  "
                  f"true: {class_names[ex['label']]}  (conversation reset)")
            continue
        if q == "/data":
            print("\n" + system); continue
        if q == "/reset":
            messages = [{"role": "system", "content": system}]
            print("  → conversation reset (patient context preserved)"); continue

        # Report mode: a number selects a canonical question. The input stays
        # inside the distribution, so pattern blending cannot occur. Free text
        # is permitted but warned about.
        if mode == "report" and tq:
            if q.isdigit() and 1 <= int(q) <= len(tq):
                task_name, q = tq[int(q) - 1]
                print(f"  → [{task_name}] {q}")
            elif not q.startswith("/"):
                print("  ⚠ This question is outside the trained patterns; report "
                      "mode answers may blend. Use /questions for the list, or "
                      "switch to /discuss.")

        messages.append({"role": "user", "content": q})
        budget = args.max_new_report if mode == "report" else args.max_new_discuss
        raw, truncated = generate(model, tok, embed_layer, projector, ex, device,
                                  supported, messages, mode, budget)
        thought, answer = split_thought(raw)
        answer = clean(answer)

        if thought:
            print(f"\n  [reasoning, {len(thought)} chars — hidden]")
        print("\nassistant> " + answer)
        if truncated:
            print(f"\n  ⚠ Answer was CUT OFF at the {budget}-token limit — this is "
                  f"not the model's natural stop. Ask something shorter or raise "
                  f"--max-new-{mode}.")

        # --- faithfulness guard: does the stated class contradict the head? ---
        #
        # In discuss mode ONLY the labelled verdict pattern is searched
        # ("Değerlendirme: X", "Classification: X"). The first-class-mentioned
        # fallback produced false alarms here: while walking through
        # differentials the model would write "Early-Onset AD" and the guard
        # read AD as a verdict. A guard that cries wolf gets ignored when a real
        # contradiction appears.
        said = find_verdict(answer, class_names) if mode == "discuss" \
            else find_class(answer, class_names)
        if said and said != ex["head"]:
            print(f"\n  ⚠ WARNING: the model said '{said}' while Vbai-2.6AD says "
                  f"'{ex['head']}'. This contradiction needs review.")

        messages.append({"role": "assistant", "content": answer})
        # keep context bounded: system + last 8 messages
        if len(messages) > 9:
            messages = [messages[0]] + messages[-8:]

    print("\nexited.")


if __name__ == "__main__":
    main()