scripts/run_mtrag.py

"""MTRAG benchmark (govt domain), with and without retrieval. Metrics:
retrieval recall@k, decline rate on answerable vs unanswerable items, and
ROUGE-L against the reference answers. Expects a clone of the IBM
mt-rag-benchmark repo at $MTRAG_REPO (see example.env).

Usage: python run_mtrag.py --mode rag --k 5   (--inspect to check schemas)
"""
import argparse, glob, json, os
from bench_common import (Retriever, generate, context_block, norm_ws,
                          save_outputs, SYSTEM_PREAMBLE, NORAG_PREAMBLE)

REPO = os.environ.get("MTRAG_REPO", "mt-rag-benchmark")
SCRATCH = os.environ.get("SCRATCH", ".")  # results/ and emb/ land here
DOMAIN = "govt"
DECLINE_MARKERS = ["do not answer", "does not answer", "cannot answer",
                   "don't know", "do not know", "not contain", "unable to",
                   "no information", "i'm sorry", "cannot find"]


def read_jsonl(path):
    with open(path, encoding="utf-8") as fh:
        return [json.loads(line) for line in fh if line.strip()]


def find_task_files():
    # reference.jsonl ONLY: RAG.jsonl and reference+RAG.jsonl are the same
    # tasks with retrieved (not gold) contexts; globbing all three would
    # triple-count. reference.jsonl carries the gold contexts we score
    # recall against.
    path = f"{REPO}/mtrag-human/generation_tasks/reference.jsonl"
    if not os.path.exists(path):
        raise FileNotFoundError(f"expected task file missing: {path}")
    return [path]


def load_corpus():
    path = f"{REPO}/corpora/passage_level/{DOMAIN}.jsonl"
    rows = read_jsonl(path)
    passages = []
    for r in rows:
        pid = r.get("_id") or r.get("id") or r.get("document_id")
        text = r.get("text") or r.get("passage") or ""
        title = r.get("title") or ""
        assert pid is not None and text, f"unrecognized corpus schema: {list(r)}"
        passages.append({"text": f"{title}\n{text}".strip(), "id": str(pid)})
    print(f"{len(passages)} {DOMAIN} passages")
    return passages


def extract_tasks(files):
    """Normalize task records: conversation turns, final question, reference
    answer, answerability label, gold passage ids, domain filter."""
    tasks = []
    for f in files:
        for r in read_jsonl(f):
            # domain lives in the Collection field, e.g.
            # "mt-rag-govt-elser-512-100-20240503"
            if DOMAIN not in str(r.get("Collection", "")).lower():
                continue
            turns = (r.get("input") or r.get("conversation")
                     or r.get("messages") or r.get("turns") or [])
            targets = r.get("targets") or r.get("reference_answers") or []
            if isinstance(targets, list) and targets and isinstance(targets[0], dict):
                ref = targets[0].get("text", "")
            elif isinstance(targets, list) and targets:
                ref = str(targets[0])
            else:
                ref = str(targets) if targets else ""
            answerability = (r.get("Answerability") or r.get("answerability")
                             or r.get("answerability_label") or "")
            if isinstance(answerability, list):
                answerability = answerability[0] if answerability else ""
            gold_ids = [str(c.get("document_id") or c.get("_id") or "")
                        for c in (r.get("contexts") or []) if isinstance(c, dict)]
            if not turns:
                continue
            last_user = None
            history = []
            for t in turns:
                role = t.get("speaker") or t.get("role") or ""
                text = t.get("text") or t.get("content") or ""
                if role.lower() in ("user", "human"):
                    last_user = text
                history.append(f"{role}: {text}")
            if not last_user:
                continue
            tasks.append({"task_id": r.get("task_id") or r.get("id") or "",
                          "history": history[:-1], "question": last_user,
                          "reference": ref,
                          "answerability": str(answerability).upper(),
                          "gold_ids": gold_ids})
    print(f"{len(tasks)} {DOMAIN} tasks from {len(files)} files")
    return tasks


def rouge_l_f(a, b):
    """Plain LCS-based ROUGE-L F1, no external deps."""
    ta, tb = norm_ws(a).split(), norm_ws(b).split()
    if not ta or not tb:
        return 0.0
    dp = [[0] * (len(tb) + 1) for _ in range(len(ta) + 1)]
    for i in range(len(ta)):
        for j in range(len(tb)):
            dp[i + 1][j + 1] = (dp[i][j] + 1 if ta[i] == tb[j]
                                else max(dp[i][j + 1], dp[i + 1][j]))
    lcs = dp[-1][-1]
    p, r = lcs / len(tb), lcs / len(ta)
    return 2 * p * r / (p + r) if p + r else 0.0


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--mode", choices=["rag", "norag"])
    ap.add_argument("--k", type=int, default=5)
    ap.add_argument("--inspect", action="store_true")
    ap.add_argument("--limit", type=int, default=0)
    args = ap.parse_args()

    files = find_task_files()
    if args.inspect:
        print("task files found:")
        for f in files:
            print(" ", f)
            rows = read_jsonl(f)
            if rows:
                print("   first record keys:", sorted(rows[0].keys()))
        return
    assert args.mode, "--mode required unless --inspect"

    tasks = extract_tasks(files)
    if args.limit:
        tasks = tasks[:args.limit]

    retriever = None
    if args.mode == "rag":
        retriever = Retriever(load_corpus(), f"{SCRATCH}/emb/mtrag_{DOMAIN}.npy")

    samples = []
    for t in tasks:
        hist = "\n".join(t["history"][-6:])
        instruction = (f"Conversation so far:\n{hist}\n\n"
                       f"Current question: {t['question']}")
        if args.mode == "rag":
            hits = retriever.search(t["question"], k=args.k)
            block = f"Context:\n{context_block(hits)}\n\n{instruction}"
            resp = generate(block, SYSTEM_PREAMBLE)
            retrieved_ids = [h["id"] for h, _ in hits]
            recall = (len(set(retrieved_ids) & set(t["gold_ids"]))
                      / len(set(t["gold_ids"]))) if t["gold_ids"] else None
        else:
            resp = generate(instruction, NORAG_PREAMBLE)
            retrieved_ids, recall = [], None
        declined = any(m in resp.lower() for m in DECLINE_MARKERS)
        unanswerable = t["answerability"] in ("UNANSWERABLE", "NO")
        rouge = rouge_l_f(resp, t["reference"]) if t["reference"] else None
        samples.append({**{k: t[k] for k in
                           ("task_id", "question", "reference", "answerability")},
                        "response": resp, "declined": declined,
                        "unanswerable": unanswerable,
                        "retrieved": retrieved_ids, "recall_at_k": recall,
                        "rouge_l": rouge})

    unans = [s for s in samples if s["unanswerable"]]
    ans = [s for s in samples if not s["unanswerable"]]
    recs = [s["recall_at_k"] for s in samples if s["recall_at_k"] is not None]
    rls = [s["rouge_l"] for s in ans if s["rouge_l"] is not None]
    metrics = {
        "n_answerable": len(ans), "n_unanswerable": len(unans),
        "decline_rate_on_unanswerable":
            (sum(s["declined"] for s in unans) / len(unans)) if unans else None,
        "decline_rate_on_answerable":
            (sum(s["declined"] for s in ans) / len(ans)) if ans else None,
        "rouge_l_answerable": (sum(rls) / len(rls)) if rls else None,
        "mean_recall_at_k": (sum(recs) / len(recs)) if recs else None,
    }
    save_outputs(f"{SCRATCH}/results", "mtrag", args.mode, metrics, samples,
                 extra={"domain": DOMAIN, "k": args.k})


if __name__ == "__main__":
    main()