run_allqs.py

import argparse
import json
import re
from pathlib import Path

import pandas as pd
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from tqdm import tqdm


FIELD_RE = re.compile(r"(?<!\{)\{([a-zA-Z0-9_]+)\}(?!\})")


def parse_args():
    p = argparse.ArgumentParser("Run CAARMS Q1–Q15 from /allqs/Q-folders using CSV mapping; save per-subject JSONL.")
    p.add_argument("--model_id", type=str, default="Llama-3.3-70B-Instruct")

    # Base and structure
    p.add_argument("--base_path", type=str, default="",
                   help="Base folder. Expects transcripts and csv under <base_path>/allqs/")
    p.add_argument("--allqs_dirname", type=str, default="allqs",
                   help="Folder name under base_path containing Q1..Q15 and the CSV.")

    # CSV + prompts + outputs root
    p.add_argument("--csv_name", type=str, required=True,
                   help="CSV filename under <base_path>/<allqs_dirname>/ (e.g. matching_1wk.csv)")
    p.add_argument("--run_root", type=str, default="",
                   help="Folder containing prompts_*/ raw_outputs/")
    p.add_argument("--prompt_name", type=str, required=True,
                   help="Prompt set suffix, e.g. CAARMS -> <run_root>/prompts_CAARMS/q1.txt..q15.txt")

    # Gen params
    p.add_argument("--max_new_tokens", type=int, default=512)
    p.add_argument("--temperature", type=float, default=0.2)
    p.add_argument("--top_p", type=float, default=0.9)
    p.add_argument("--top_k", type=int, default=40)
    p.add_argument("--do_sample", action="store_true", default=False)

    # Resume/overwrite
    p.add_argument("--overwrite", action="store_true", default=False,
                   help="Overwrite per-subject JSONL (otherwise resume by skipping existing entries).")

    # New skip controls
    p.add_argument("--skip_if_same_transcript", action="store_false", default=True,
                   help="Skip if JSONL already has the same (instance_id, question, transcript_file). (default: True)")
    p.add_argument("--skip_if_any_exists", action="store_true", default=False,
                   help="Skip if JSONL already has (instance_id, question) regardless of transcript_file.")

    return p.parse_args()


def process_prompt(case: dict, template: str) -> str:
    def repl(m):
        key = m.group(1)
        return str(case.get(key, f"{{{key}}}"))
    return FIELD_RE.sub(repl, template)


def load_prompts(prompts_dir: Path) -> dict[str, str]:
    prompts = {}
    for i in range(1, 16):
        qfile = prompts_dir / f"q{i}.txt"
        if not qfile.exists():
            raise FileNotFoundError(f"Missing prompt file: {qfile}")
        prompts[f"Q{i}"] = qfile.read_text(encoding="utf-8")
    return prompts


def build_chat_input(tokenizer, user_prompt: str) -> str:
    messages = [
        {"role": "system", "content": "You are a clinician specialised in the clinical high risk for psychosis state."},
        {"role": "user", "content": user_prompt},
    ]
    return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)


def generate_chat(model, tokenizer, prompt: str, args) -> str:
    chat_text = build_chat_input(tokenizer, prompt)
    inputs = tokenizer(
        chat_text,
        return_tensors="pt",
        truncation=True,
        max_length=min(getattr(tokenizer, "model_max_length", 32768), 32768),
    ).to(model.device)

    with torch.inference_mode():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=args.max_new_tokens,
            do_samplez=args.do_sample,
            temperature=args.temperature,
            top_p=args.top_p,
            top_k=args.top_k,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )

    return tokenizer.decode(
        output_ids[0][inputs["input_ids"].shape[-1]:],
        skip_special_tokens=True,
    ).strip()


def extract_json_obj(text: str) -> dict:
    """
    Robust JSON extractor:
    1) Try to parse the first {...} block as JSON.
    2) If JSON decoding fails, recover severity/frequency via regex and
       recover summary as raw text, then return a valid dict.
    """
    start = text.find("{")
    end = text.rfind("}")
    if start == -1 or end == -1 or end <= start:
        print(f"No JSON object found in model output:\n{text[:800]}")
        return {
            "severity": None,
            "frequency": None,
            "summary": f"No JSON object found in model output:{text}",
        }


    block = text[start:end + 1].strip()

    # 1) Normal path
    try:
        obj = json.loads(block)
        for k in ("severity", "frequency", "summary"):
            if k not in obj:
                return {
                    "severity": None,
                    "frequency": None,
                    "summary": f"No JSON object found in model output:{text}",
                }


        if not isinstance(obj["severity"], int) or not (0 <= obj["severity"] <= 6):
            raise ValueError(f"Invalid severity: {obj.get('severity')}")
        if not isinstance(obj["frequency"], int) or not (0 <= obj["frequency"] <= 6):
            raise ValueError(f"Invalid frequency: {obj.get('frequency')}")
        if not isinstance(obj["summary"], str) or not obj["summary"].strip():
            raise ValueError("Invalid summary (empty or non-string).")

        return obj

    except json.JSONDecodeError:
        # 2) Recovery path for "almost JSON" (e.g., unescaped quotes/newlines in summary)

        sev_m = re.search(r'"\s*severity\s*"\s*:\s*([0-6])', block)
        freq_m = re.search(r'"\s*frequency\s*"\s*:\s*([0-6])', block)

        if sev_m is None or freq_m is None:
            return {
                "severity": None,
                "frequency": None,
                "summary": f"No JSON object found in model output:{text}",
            }


        severity = int(sev_m.group(1))
        frequency = int(freq_m.group(1))

        # Try to capture summary value even if it contains quotes/newlines.
        # We take everything after "summary": until the last quote before the final }
        summary = ""
        sum_m = re.search(r'"\s*summary\s*"\s*:\s*"(.*)"\s*\}', block, flags=re.DOTALL)
        if sum_m is not None:
            summary = sum_m.group(1).strip()
        else:
            # fallback: allow summary without quotes (rare)
            sum2 = re.search(r'"\s*summary\s*"\s*:\s*(.*)\s*\}', block, flags=re.DOTALL)
            if sum2 is not None:
                summary = str(sum2.group(1)).strip().strip('"').strip()

        # Clean up common trailing junk if model appended extra fields/commas
        summary = summary.strip().rstrip(",")

        return {
            "severity": severity,
            "frequency": frequency,
            "summary": summary,
        }


def normalize_status(val) -> str:
    if val is None:
        return ""
    s = str(val).strip()
    if not s:
        return ""
    sl = s.lower()
    if sl.startswith("comp"):
        return "Complete"
    if sl.startswith("part"):
        return "Partial"
    if sl == "no":
        return "No"
    return s


def find_status_col(columns: list[str], qkey: str) -> str:
    n = int(qkey[1:])
    cols = list(columns)

    for c in cols:
        cl = c.lower().strip()
        if cl.startswith(f"q{n}") or cl.startswith(f"q{n}.") or (n == 9 and cl.startswith("9.")):
            return c

    token = f"q{n}"
    for c in cols:
        if token in c.lower():
            return c

    raise ValueError(f"Could not find status column for {qkey}.")


def as_int_day(v) -> int | None:
    if pd.isna(v):
        return None
    s = str(v).strip()
    if not s:
        return None
    try:
        return int(float(s))
    except Exception:
        return None


def build_instance_id(row: pd.Series) -> str:
    sid = str(row.get("src_subject_id", "")).strip()
    visit = str(row.get("visit", "")).strip()
    v5 = as_int_day(row.get("V5", None))
    v6 = str(row.get("V6", "")).strip()
    v5s = "" if v5 is None else str(v5)
    return f"{sid}__{visit}__day{v5s}__{v6}"


def find_transcript_file(q_dir: Path, src_subject_id: str, v5, v6: str, qkey: str) -> Path | None:
    sid = str(src_subject_id).strip()
    sess = str(v6).strip()
    day_int = as_int_day(v5)
    day_tok = None if day_int is None else f"day{day_int:04d}"

    files = list(q_dir.glob("*.txt"))
    if not files:
        return None

    strict = []
    for f in files:
        name = f.name
        if f"_{sid}_" not in name:
            continue
        if sess and sess not in name:
            continue
        if day_tok and day_tok not in name:
            continue
        if qkey not in name:
            continue
        strict.append(f)

    if len(strict) == 1:
        return strict[0]
    if len(strict) > 1:
        return sorted(strict, key=lambda x: len(x.name))[0]

    fallback = [f for f in files if f"_{sid}_" in f.name and (sess in f.name if sess else True)]
    if len(fallback) == 1:
        return fallback[0]
    if len(fallback) > 1:
        return sorted(fallback, key=lambda x: len(x.name))[0]

    return None


def read_existing_keys(jsonl_path: Path) -> set[tuple[str, str, str | None]]:
    """
    Return set of (instance_id, question, transcript_file) already written.
    transcript_file can be None for status-only entries or missing transcripts.
    """
    if not jsonl_path.exists():
        return set()

    done = set()
    for line in jsonl_path.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line:
            continue
        obj = json.loads(line)
        iid = obj.get("instance_id")
        q = obj.get("question")
        tf = obj.get("transcript_file")
        if isinstance(iid, str) and isinstance(q, str):
            done.add((iid, q, tf if isinstance(tf, str) else None))
    return done


def has_any_instance_question(done: set[tuple[str, str, str | None]], instance_id: str, qkey: str) -> bool:
    # skip if ANY transcript_file already exists for this instance/question
    return any((iid == instance_id and q == qkey) for (iid, q, _) in done)


def read_existing_state(jsonl_path: Path) -> dict[tuple[str, str, str | None], dict]:
    """
    Return LAST record per (instance_id, question, transcript_file).
    We key on transcript_file too (same as your done set) so resume is consistent.
    """
    if not jsonl_path.exists():
        return {}

    state = {}
    for line in jsonl_path.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line:
            continue
        try:
            obj = json.loads(line)
        except Exception:
            continue

        iid = obj.get("instance_id")
        q = obj.get("question")
        tf = obj.get("transcript_file")
        tf = tf if isinstance(tf, str) else None

        if isinstance(iid, str) and isinstance(q, str):
            state[(iid, q, tf)] = obj  # keep last occurrence

    return state

def should_skip_existing(existing: dict | None) -> bool:
    """
    Skip ONLY if we already have a successful run with scores.
    Re-run if:
      - ran == False (status-only cached row)
      - severity/frequency missing
      - missing_transcript == True (transcripts may have been added later)
    """
    if not existing:
        return False

    if bool(existing.get("missing_transcript", False)):
        return False

    if not bool(existing.get("ran", False)):
        return False

    sev = existing.get("severity", None)
    freq = existing.get("frequency", None)
    if sev is None or freq is None:
        return False

    return True

def main():
    args = parse_args()

    base_path = Path(args.base_path)
    allqs_root = base_path / args.allqs_dirname
    csv_path = allqs_root / args.csv_name

    run_root = Path(args.run_root)
    prompts_dir = run_root / f"prompts_{args.prompt_name}"
    out_dir = run_root / f"raw_outputs/{args.model_id}"
    out_dir.mkdir(parents=True, exist_ok=True)

    prompts = load_prompts(prompts_dir)

    if not csv_path.exists():
        raise FileNotFoundError(f"CSV not found: {csv_path}")

    df = pd.read_csv(csv_path)

    if "src_subject_id" not in df.columns:
        raise ValueError("CSV must contain 'src_subject_id'")
    if "V5" not in df.columns or "V6" not in df.columns:
        raise ValueError("CSV must contain 'V5' and 'V6'")

    status_cols = {qkey: find_status_col(list(df.columns), qkey) for qkey in prompts.keys()}

    model_path = base_path / "models" / args.model_id
    print(f"🔍 Loading model: {model_path}")
    tokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True, use_fast=True)
    dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        device_map="auto",
        torch_dtype=dtype,
        # trust_remote_code=True,
        local_files_only=True,
    )
    model.eval()

    print(f"✅ Model loaded: {args.model_id}")
    print(f"✅ CSV: {csv_path}")
    print(f"✅ Prompts: {prompts_dir}")
    print(f"🧾 Transcripts root: {allqs_root}/Q1..Q15")
    print(f"💾 Output: {out_dir} (one JSONL per src_subject_id)")

    df["__sid__"] = df["src_subject_id"].astype(str).str.strip()
    subjects = sorted(df["__sid__"].unique())

    for sid in tqdm(subjects, desc="Subjects", unit="subj"):
        subj_rows = df[df["__sid__"] == sid].copy()
        subj_jsonl = out_dir / f"{sid}.jsonl"

        if args.overwrite and subj_jsonl.exists():
            subj_jsonl.unlink()

        done = read_existing_keys(subj_jsonl)
        state = read_existing_state(subj_jsonl)          # NEW: actual last payload per key

        with subj_jsonl.open("a", encoding="utf-8") as fout:
            for _, row in subj_rows.iterrows():
                instance_id = build_instance_id(row)
                visit = str(row.get("visit", "")).strip()
                v5 = row.get("V5", None)
                v6 = str(row.get("V6", "")).strip()

                for qkey, template in prompts.items():
                    q_dir = allqs_root / qkey
                    if not q_dir.exists():
                        raise FileNotFoundError(f"Missing question folder: {q_dir}")

                    tpath = find_transcript_file(q_dir, sid, v5, v6, qkey)
                    tname = None if tpath is None else tpath.name

                    # --- NEW SKIP LOGIC (patched) ---
                    key = (instance_id, qkey, tname)

                    if args.skip_if_any_exists:
                        # keep existing behavior, BUT only skip if existing entry is truly done
                        if has_any_instance_question(done, instance_id, qkey):
                            # find any matching existing record and decide if it is skippable
                            ex = None
                            for (iid, qq, tf) in state.keys():
                                if iid == instance_id and qq == qkey:
                                    ex = state.get((iid, qq, tf))
                                    if should_skip_existing(ex):
                                        break
                            if ex is not None and should_skip_existing(ex):
                                continue
                    else:
                        # default behaviour: skip only if same transcript already processed,
                        # BUT only if that existing record is truly done
                        if args.skip_if_same_transcript and key in done:
                            ex = state.get(key, None)
                            if should_skip_existing(ex):
                                continue
                    # --------------------------------

                    status_val = normalize_status(row.get(status_cols[qkey], ""))
                    should_run = status_val in {"Complete", "Partial"}

                    payload = {
                        "src_subject_id": sid,
                        "instance_id": instance_id,
                        "visit": visit,
                        "V5": as_int_day(v5),
                        "V6": v6,
                        "question": qkey,
                        "status": status_val,
                        "ran": False,
                        "missing_transcript": False,
                        "model_id": args.model_id,
                        "prompt_name": args.prompt_name,
                        "csv_name": args.csv_name,
                        "transcript_file": tname,
                    }

                    if not should_run:
                        fout.write(json.dumps(payload, ensure_ascii=False) + "\n")
                        done.add((instance_id, qkey, tname))
                        continue

                    if tpath is None:
                        payload["missing_transcript"] = True
                        fout.write(json.dumps(payload, ensure_ascii=False) + "\n")
                        done.add((instance_id, qkey, tname))
                        continue

                    transcript = tpath.read_text(encoding="utf-8")
                    prompt = process_prompt({"transcript": transcript}, template)

                    raw = generate_chat(model, tokenizer, prompt, args)
                    parsed = extract_json_obj(raw)

                    payload["ran"] = True
                    payload["severity"] = parsed["severity"]
                    payload["frequency"] = parsed["frequency"]
                    payload["summary"] = parsed["summary"]

                    fout.write(json.dumps(payload, ensure_ascii=False) + "\n")
                    done.add((instance_id, qkey, tname))

    print("🏁 Done.")


if __name__ == "__main__":
    main()