training/scripts/kaggle_three_llms_text.py

#!/usr/bin/env python3
"""
Kaggle-ready: run AutoDataLab++ once with 3 different LLM CoS pickers and print
the *full agent text* for each run. No scores, just answers.

Three pickers:
  1) base LLM            (BASE_MODEL_ID, no adapter)
  2) base LLM + SFT LoRA (BASE_MODEL_ID + SFT_ADAPTER_ID)
  3) base LLM + GRPO LoRA (BASE_MODEL_ID + GRPO_ADAPTER_ID, optional subfolder)

Usage on Kaggle (Settings → Internet: ON, GPU: T4 or better):
    !pip install -q "transformers>=4.45,<4.49" "peft>=0.13,<0.16" \
        "accelerate>=0.33,<1.1" bitsandbytes huggingface_hub pydantic
    !python3 training/kaggle_three_llms_text.py --task expert_brief

Set TASK and adapter ids below or via CLI flags.

The CoS LLM only chooses the *first* action. The env then runs deterministic
continuation to make sure all required experts (analyst, finance, strategy, hr)
report. That guarantees you always get a full text readout to compare.
"""
from __future__ import annotations

import argparse
import json
import os
import re
import sys
from pathlib import Path
from typing import Any, Optional

# Kaggle can trip torch.compile/triton paths when transformers imports optional
# packages. Keep inference on the plain eager path.
os.environ.setdefault("TORCH_COMPILE_DISABLE", "1")
os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")

# Make the env package importable both as a script and from a Kaggle notebook.
REPO = Path(__file__).resolve().parents[2]
if str(REPO) not in sys.path:
    sys.path.insert(0, str(REPO))


# ---------- Config defaults (override via CLI) ----------
DEFAULTS = {
    "task": "expert_brief",
    "use_rag": False,
    "base_model_id": "Qwen/Qwen2.5-1.5B-Instruct",
    "sft_adapter_id": "uchihakamal/qwen2.5-1.5b-autodatalab-sft",   # change to your repo
    "sft_subfolder": "",
    "grpo_adapter_id": "uchihakamal/qwen2.5-7b-autodatalab-grpo",   # change to your repo
    "grpo_subfolder": "final",
    "use_4bit": True,
    "max_new_tokens": 96,
    "shaping": "strict",
}

# ---------- Env imports (deferred so --help works without deps) ----------


def _load_env():
    from ceo_brief_env.environment import (  # noqa: F401
        CEOBriefEnvironment,
        oracle_action_for_observation,
        required_experts_for_task,
    )
    from ceo_brief_env.models import CoSAction, CoSObservation  # noqa: F401

    return (
        CEOBriefEnvironment,
        oracle_action_for_observation,
        required_experts_for_task,
        CoSAction,
        CoSObservation,
    )


# ---------- LLM action policy (LLM picks first action; rest is deterministic) ----------

_VALID_ACTIONS = {"consult", "ask", "summarize", "submit", "noop"}
_VALID_EXPERTS = {"analyst", "finance", "hr", "strategy"}
_JSON_RE = re.compile(r"\{[^{}]*\}", re.S)

_SYSTEM_PROMPT = (
    "You are the Chief of Staff in AutoDataLab++. You orchestrate four "
    "specialists: analyst, finance, strategy, hr. Reply with STRICT JSON only.\n"
    'Schema: {"action_type": one of [consult, ask, summarize, submit, noop], '
    '"expert_id": one of [analyst, finance, hr, strategy] or null}.\n'
    "Rules: consult each required expert at most once -> summarize -> submit."
)


def _render_obs(obs) -> str:
    return (
        f"task={obs.task_name} step={obs.step_count}/{obs.max_steps} "
        f"rag={obs.rag_enabled} consulted={obs.consulted_experts} "
        f"brief_done={obs.current_brief is not None} available={obs.available_experts}"
    )


def _parse_action(text: str, CoSAction):
    m = _JSON_RE.search(text or "")
    if not m:
        return CoSAction(action_type="noop")
    try:
        a = json.loads(m.group(0))
    except Exception:
        return CoSAction(action_type="noop")
    at = a.get("action_type")
    if at not in _VALID_ACTIONS:
        return CoSAction(action_type="noop")
    eid = a.get("expert_id")
    if eid is not None and eid not in _VALID_EXPERTS:
        eid = None
    return CoSAction(action_type=at, expert_id=eid)


# ---------- Model loading ----------


def _load_llm(base_model_id: str, adapter_id: Optional[str], adapter_subfolder: Optional[str], use_4bit: bool, hf_token: Optional[str]):
    """Load base model and (optionally) attach a PEFT/LoRA adapter.

    When an adapter is present we:
      * load the *adapter's* tokenizer (training-time chat template + special
        tokens; avoids vocab drift),
      * resize base model embeddings to that tokenizer **before** PEFT attach,
      * load adapter and ``merge_and_unload`` for a clean inference graph
        (also fixes ``set_input_embeddings``-style layer mismatch errors).
    """
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    tokenizer_source = adapter_id or base_model_id
    tok_kw: dict[str, Any] = {"token": hf_token}
    if adapter_id and adapter_subfolder:
        tok_kw["subfolder"] = adapter_subfolder
    try:
        tok = AutoTokenizer.from_pretrained(tokenizer_source, **tok_kw)
    except Exception as e:
        if not adapter_id:
            raise
        print(
            f"[load] warn: adapter tokenizer failed ({e}); falling back to base tokenizer.",
            flush=True,
        )
        tok = AutoTokenizer.from_pretrained(base_model_id, token=hf_token)
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token

    bnb_cfg = None
    if use_4bit:
        try:
            from transformers import BitsAndBytesConfig

            bnb_cfg = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_quant_type="nf4",
                bnb_4bit_compute_dtype=torch.float16,
                bnb_4bit_use_double_quant=True,
            )
        except Exception:
            bnb_cfg = None

    try:
        model = AutoModelForCausalLM.from_pretrained(
            base_model_id,
            token=hf_token,
            device_map="auto",
            quantization_config=bnb_cfg,
            torch_dtype=torch.float16,
        )
    except Exception:
        model = AutoModelForCausalLM.from_pretrained(
            base_model_id,
            token=hf_token,
            device_map="auto",
            torch_dtype=torch.float16,
        )

    if adapter_id:
        try:
            model.resize_token_embeddings(len(tok))
        except Exception as e:
            print(f"[load] warn: resize_token_embeddings failed: {e}", flush=True)

    model.eval()
    # Greedy generation; clear sampling defaults from model generation_config to
    # avoid noisy warnings (`temperature/top_p/top_k` are ignored when
    # do_sample=False).
    try:
        model.generation_config.do_sample = False
        model.generation_config.temperature = None
        model.generation_config.top_p = None
        model.generation_config.top_k = None
    except Exception:
        pass

    if adapter_id:
        from peft import PeftModel

        kw: dict[str, Any] = {"token": hf_token}
        if adapter_subfolder:
            kw["subfolder"] = adapter_subfolder
        model = PeftModel.from_pretrained(model, adapter_id, **kw)
        try:
            model = model.merge_and_unload()
        except Exception as e:
            print(
                f"[load] warn: merge_and_unload failed ({e}); keeping PEFT wrapper.",
                flush=True,
            )
        model.eval()

    return tok, model


def _action_label(action) -> str:
    if action.action_type in {"consult", "ask"}:
        return f"{action.action_type}:{action.expert_id or 'null'}"
    return action.action_type


def _generate_action_and_text(tok, model, CoSAction, obs, max_new_tokens: int = 96):
    """Ask the LLM for one CoS action and keep the raw completion for evidence."""
    import torch

    with torch.no_grad():
        msgs = [
            {"role": "system", "content": _SYSTEM_PROMPT},
            {"role": "user", "content": _render_obs(obs)},
        ]
        if hasattr(tok, "apply_chat_template"):
            prompt = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
        else:
            prompt = _SYSTEM_PROMPT + "\n\n" + _render_obs(obs) + "\n"
        ids = tok(prompt, return_tensors="pt").to(model.device)
        out = model.generate(
            **ids,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            pad_token_id=tok.pad_token_id or tok.eos_token_id,
        )
        text = tok.decode(out[0][ids["input_ids"].shape[1]:], skip_special_tokens=True)
        action = _parse_action(text, CoSAction)
        return action, text.strip()


def _make_llm_picker(tok, model, CoSAction, max_new_tokens: int = 96):
    """Returns a CoS picker that asks the LLM, parses JSON, falls back to noop."""

    def picker(obs):
        action, _completion = _generate_action_and_text(
            tok, model, CoSAction, obs, max_new_tokens=max_new_tokens
        )
        return action

    return picker


# ---------- Run one episode and pretty-print agent text ----------


def _format_episode(data: dict) -> str:
    """Mirror training/kaggle_agent_answers.format_episode_answers but stand-alone."""
    width = 88
    sep = "=" * width
    out: list[str] = []
    out.append(
        f"Task: {data.get('task')}\nPolicy: {data.get('policy_label')}\nRAG: {data.get('use_rag')}\n"
    )
    inst = data.get("final_instruction") or ""
    if inst:
        out += [sep, "INSTRUCTION (from metadata)", sep, inst, ""]

    policy_trace = data.get("policy_trace") or []
    fallback_trace = data.get("fallback_trace") or []
    if policy_trace:
        out += [sep, "TRAINING EVIDENCE - MODEL-CONTROLLED POLICY TRAJECTORY", sep]
        out.append(
            "These are the actions chosen directly by the LLM/adapter before any "
            "deterministic completion logic. This is the part to compare across "
            "base vs SFT vs GRPO."
        )
        for row in policy_trace:
            out.append(
                f"  step {row['step']:02d} | action={row['action_label']:<18} "
                f"reward={row['reward']:+.4f} done={row['done']} "
                f"model_routed={row['model_routed_required']} "
                f"env_consulted={row['consulted_after']}"
            )
            preview = str(row.get("completion_preview") or "").replace("\n", " ")
            if preview:
                out.append(f"       raw: {preview}")
        if fallback_trace:
            out.append(
                "\nFallback used after model-controlled steps to complete the report "
                "for qualitative comparison:"
            )
            out.append(
                "  " + " -> ".join(row["action_label"] for row in fallback_trace)
            )
        out.append("")

    reports: dict = data.get("expert_reports") or {}
    order = ("analyst", "finance", "strategy", "hr")
    labels = {
        "analyst": "DATA ANALYST",
        "finance": "FINANCE",
        "strategy": "STRATEGIST",
        "hr": "HR / COMMS",
    }
    for eid in order:
        r = reports.get(eid)
        if not r:
            continue
        out += [sep, f"{labels.get(eid, eid).upper()} — {r.get('title', eid)}", sep]
        out.append(r.get("summary", "").strip() or "(no summary)")
        bps = r.get("bullet_points") or []
        if bps:
            out.append("\nBullets:")
            out += [f"  • {b}" for b in bps]
        m_c = r.get("memory_citations") or []
        m_s = r.get("memory_snippets") or []
        n = min(len(m_c), len(m_s), 3)
        if n:
            out.append(f"\nTape & citations — first {n}:")
            for i in range(n):
                snip = m_s[i]
                snip = (snip[:400] + "…") if len(snip) > 400 else snip
                out.append(f"  [{m_c[i]}] {snip}")
        if eid == "hr" and r.get("memo"):
            out.append("\nHR memo:\n" + str(r["memo"]))
        out.append("")

    brief = data.get("current_brief")
    if brief:
        out += [sep, "COMPOSED BRIEF (to CEO — merged from reports)", sep]
        out.append(brief.get("summary", "") or "")
        recs = brief.get("recommendations") or []
        if recs:
            out.append("\nRecommendations:")
            out += [f"  • {x}" for x in recs]
        if brief.get("hr_memo"):
            out.append("\nHR memo (in brief object):\n" + str(brief["hr_memo"]))
        out.append("")
    return "\n".join(out)


def _run(picker, label: str, task: str, use_rag: bool) -> dict:
    """Run one episode with the given picker. Uses inference.run_episode_collect for parity."""
    from inference import run_episode_collect

    return run_episode_collect(task, picker, label, use_rag=use_rag, quiet=True)


def _deterministic_next_action(obs, task: str, required_experts_for_task, CoSAction):
    missing = [e for e in required_experts_for_task(task) if e not in obs.consulted_experts]
    if missing:
        return CoSAction(action_type="consult", expert_id=missing[0])
    if obs.current_brief is None:
        return CoSAction(action_type="summarize")
    return CoSAction(action_type="submit")


def _run_with_policy_trace(
    tok,
    model,
    CoSAction,
    CEOBriefEnvironment,
    required_experts_for_task,
    *,
    label: str,
    task: str,
    use_rag: bool,
    max_new_tokens: int,
    policy_steps: int,
    shaping: str,
) -> dict:
    """Run N model-controlled steps, then fallback only to finish the report.

    This exposes the actual learned policy behavior. The final expert reports are
    still included so you can quote agent answers in the presentation.
    """
    env = CEOBriefEnvironment(shaping=shaping, auto_fill_required=False)
    obs = env.reset(task=task, use_rag=use_rag)
    rewards: list[float] = []
    policy_trace: list[dict[str, Any]] = []
    fallback_trace: list[dict[str, Any]] = []
    required = required_experts_for_task(task)
    model_routed_required: list[str] = []

    for _ in range(max(1, policy_steps)):
        if obs.done:
            break
        before_consulted = list(obs.consulted_experts)
        action, completion = _generate_action_and_text(
            tok, model, CoSAction, obs, max_new_tokens=max_new_tokens
        )
        obs = env.step(action)
        if action.expert_id in required and action.expert_id not in model_routed_required:
            model_routed_required.append(action.expert_id)
        rewards.append(float(obs.reward))
        policy_trace.append(
            {
                "step": obs.step_count,
                "obs_before_consulted": before_consulted,
                "action": action.model_dump(exclude_none=True),
                "action_label": _action_label(action),
                "completion_preview": completion[:400],
                "reward": round(float(obs.reward), 4),
                "done": bool(obs.done),
                "consulted_after": list(obs.consulted_experts),
                "model_routed_required": list(model_routed_required),
            }
        )

    while not obs.done and obs.step_count < obs.max_steps:
        action = _deterministic_next_action(obs, task, required_experts_for_task, CoSAction)
        obs = env.step(action)
        rewards.append(float(obs.reward))
        fallback_trace.append(
            {
                "step": obs.step_count,
                "action": action.model_dump(exclude_none=True),
                "action_label": _action_label(action),
                "reward": round(float(obs.reward), 4),
                "done": bool(obs.done),
                "consulted_after": list(obs.consulted_experts),
            }
        )

    action_sequence = [row["action_label"] for row in policy_trace]
    policy_covered_required = list(model_routed_required)
    data = {
        "task": task,
        "policy_label": label,
        "use_rag": use_rag,
        "shaping": shaping,
        "success": bool((obs.terminal_grader_score or 0.0) >= 0.5),
        "steps": int(obs.step_count),
        "terminal_score": round(float(obs.terminal_grader_score or 0.0), 4),
        "cumulative_reward": round(sum(rewards), 4),
        "step_rewards": [round(x, 4) for x in rewards],
        "final_instruction": obs.instruction,
        "task_difficulty": obs.task_difficulty,
        "max_steps": obs.max_steps,
        "consulted_experts": list(obs.consulted_experts),
        "current_brief": obs.current_brief.model_dump() if obs.current_brief is not None else None,
        "expert_reports": {k: v.model_dump() for k, v in obs.expert_reports.items()},
        "policy_trace": policy_trace,
        "fallback_trace": fallback_trace,
        "policy_evidence": {
            "model_controlled_steps": len(policy_trace),
            "action_sequence": action_sequence,
            "unique_actions": len(set(action_sequence)),
            "covered_required_experts_in_policy_steps": policy_covered_required,
            "required_experts": required,
            "needed_fallback": bool(fallback_trace),
            "fallback_steps": len(fallback_trace),
            "policy_reward_sum": round(sum(row["reward"] for row in policy_trace), 4),
            "total_reward_sum": round(sum(rewards), 4),
        },
    }
    return data


def _write_evidence_artifacts(out_dir: Path, rows: list[dict[str, Any]], task: str) -> None:
    """Save compact JSON/CSV/Markdown and a reward plot for presentation use."""
    out_dir.mkdir(parents=True, exist_ok=True)
    summary = []
    for row in rows:
        ev = row.get("policy_evidence") or {}
        summary.append(
            {
                "policy": row.get("policy_label"),
                "task": row.get("task"),
                "action_sequence": " -> ".join(ev.get("action_sequence") or []),
                "covered_required": ",".join(ev.get("covered_required_experts_in_policy_steps") or []),
                "needed_fallback": ev.get("needed_fallback"),
                "policy_reward_sum": ev.get("policy_reward_sum"),
                "total_reward_sum": ev.get("total_reward_sum"),
                "terminal_score": row.get("terminal_score"),
                "steps": row.get("steps"),
            }
        )
    (out_dir / f"training_evidence__{task}.json").write_text(
        json.dumps(summary, indent=2, default=str), encoding="utf-8"
    )
    csv_lines = [
        "policy,task,action_sequence,covered_required,needed_fallback,policy_reward_sum,total_reward_sum,terminal_score,steps"
    ]
    for item in summary:
        csv_lines.append(
            ",".join(
                json.dumps(str(item[k])) if k in {"action_sequence", "covered_required"} else str(item[k])
                for k in [
                    "policy",
                    "task",
                    "action_sequence",
                    "covered_required",
                    "needed_fallback",
                    "policy_reward_sum",
                    "total_reward_sum",
                    "terminal_score",
                    "steps",
                ]
            )
        )
    (out_dir / f"training_evidence__{task}.csv").write_text("\n".join(csv_lines), encoding="utf-8")
    md = [
        "# Training Evidence",
        "",
        "| Policy | Model-controlled action sequence | Required experts covered | Needed fallback | Policy reward | Final score |",
        "|---|---|---|---:|---:|---:|",
    ]
    for item in summary:
        md.append(
            f"| {item['policy']} | `{item['action_sequence']}` | "
            f"{item['covered_required'] or '-'} | {item['needed_fallback']} | "
            f"{item['policy_reward_sum']} | {item['terminal_score']} |"
        )
    (out_dir / f"training_evidence__{task}.md").write_text("\n".join(md), encoding="utf-8")

    try:
        import matplotlib.pyplot as plt

        plt.figure(figsize=(8, 4))
        for row in rows:
            rewards = row.get("step_rewards") or []
            if not rewards:
                continue
            xs = list(range(1, len(rewards) + 1))
            cum = []
            total = 0.0
            for r in rewards:
                total += float(r)
                cum.append(total)
            plt.plot(xs, cum, marker="o", label=str(row.get("policy_label")))
        plt.xlabel("environment step")
        plt.ylabel("cumulative reward")
        plt.title(f"Training evidence: policy trajectories on {task}")
        plt.grid(True, alpha=0.3)
        plt.legend()
        plt.tight_layout()
        plt.savefig(out_dir / f"training_evidence_curve__{task}.png", dpi=160)
        plt.close()
    except Exception as e:
        print(f"[plot] skipped evidence curve: {e}", flush=True)


# ---------- Main ----------


def main() -> int:
    p = argparse.ArgumentParser()
    p.add_argument("--task", default=DEFAULTS["task"])
    p.add_argument("--use-rag", action="store_true", default=DEFAULTS["use_rag"])
    p.add_argument("--base-model-id", default=DEFAULTS["base_model_id"])
    p.add_argument("--sft-adapter-id", default=DEFAULTS["sft_adapter_id"])
    p.add_argument("--sft-subfolder", default=DEFAULTS["sft_subfolder"])
    p.add_argument("--grpo-adapter-id", default=DEFAULTS["grpo_adapter_id"])
    p.add_argument("--grpo-subfolder", default=DEFAULTS["grpo_subfolder"])
    p.add_argument("--no-4bit", action="store_true", help="disable 4-bit (fallback to bf16)")
    p.add_argument("--max-new-tokens", type=int, default=DEFAULTS["max_new_tokens"])
    p.add_argument(
        "--policy-steps",
        type=int,
        default=6,
        help="number of model-controlled CoS steps before deterministic fallback",
    )
    p.add_argument(
        "--shaping",
        choices=("default", "strict"),
        default=DEFAULTS["shaping"],
        help="env reward shaping for evidence run; strict makes lazy policies obvious",
    )
    p.add_argument("--hf-token", default=os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") or "")
    p.add_argument("--out-dir", type=Path, default=Path("/kaggle/working/three_llms_text") if Path("/kaggle/working").is_dir() else Path("./three_llms_text"))
    args = p.parse_args()

    use_4bit = not args.no_4bit
    args.out_dir.mkdir(parents=True, exist_ok=True)

    (
        CEOBriefEnvironment,
        oracle_action_for_observation,
        required_experts_for_task,
        CoSAction,
        CoSObservation,
    ) = _load_env()

    runs = [
        ("base_llm", args.base_model_id, None, None),
        ("base_llm+sft", args.base_model_id, (args.sft_adapter_id or None), (args.sft_subfolder or None)),
        ("base_llm+grpo", args.base_model_id, (args.grpo_adapter_id or None), (args.grpo_subfolder or None)),
    ]

    big_sep = "#" * 88
    evidence_rows: list[dict[str, Any]] = []
    for label, base_id, adapter_id, sub in runs:
        print(f"\n{big_sep}\n# RUN: {label}\n# base={base_id}\n# adapter={adapter_id or '(none)'}{(' subfolder='+sub) if sub else ''}\n{big_sep}", flush=True)
        try:
            tok, model = _load_llm(
                base_id,
                adapter_id,
                sub,
                use_4bit=use_4bit,
                hf_token=args.hf_token or None,
            )
        except Exception as e:
            print(f"[load failed for {label}] {e}", flush=True)
            continue

        try:
            data = _run_with_policy_trace(
                tok,
                model,
                CoSAction,
                CEOBriefEnvironment,
                required_experts_for_task,
                label=label,
                task=args.task,
                use_rag=args.use_rag,
                max_new_tokens=args.max_new_tokens,
                policy_steps=args.policy_steps,
                shaping=args.shaping,
            )
        except Exception as e:
            print(f"[episode failed for {label}] {e}", flush=True)
            continue
        finally:
            # Free GPU memory between runs.
            try:
                import torch

                del model
                del tok
                torch.cuda.empty_cache()
            except Exception:
                pass

        print(_format_episode(data))
        out_path = args.out_dir / f"{label}__{args.task}.json"
        out_path.write_text(json.dumps(data, indent=2, default=str), encoding="utf-8")
        print(f"[saved] {out_path}", flush=True)
        evidence_rows.append(data)

    if evidence_rows:
        _write_evidence_artifacts(args.out_dir, evidence_rows, args.task)
        print(
            f"[saved] training evidence artifacts under {args.out_dir} "
            f"(json/csv/md/png)",
            flush=True,
        )

    return 0


if __name__ == "__main__":
    raise SystemExit(main())