run.py

"""
Entry point for the PACT split-evidence interaction demo on HotpotQA.

PACT (Protocolized Action-state Communication and Transmission) runs a
two-agent interaction in which each agent holds half of the evidence and the
agents exchange compact action-state messages until the final turn produces
the answer.

Example:

  python run.py \\
      --model_name Qwen/Qwen3-14B \\
      --data_path data/hotpot_dev_distractor_v1.json \\
      --use_vllm --tensor_parallel_size 1 \\
      --seed 42 --generate_bs 64 --max_new_tokens 4096 --max_turns 4 \\
      --max_samples 50 \\
      --output_path results/pact_hotpot_demo.jsonl
"""

import argparse
import json
import time
from typing import Dict, List, Tuple

from tqdm import tqdm

from data import load_hotpotqa
from methods.pact import PACTMethod
from models import ModelWrapper
from utils import auto_device, set_seed


# -- Evaluation --------------------------------------------------------------

def evaluate(preds: List[Dict]) -> Tuple[float, float, int, float, float, float, float]:
    total = len(preds)
    if total == 0:
        return 0.0, 0.0, 0, 0.0, 0.0, 0.0, 0.0
    correct = sum(1 for p in preds if p.get("correct", False))
    avg_f1 = sum(p.get("f1", 0.0) for p in preds) / total
    em = correct / total
    avg_comm = sum(p.get("communication_tokens", 0) for p in preds) / total

    def _total(p):
        if "total_tokens" in p:
            return p["total_tokens"]
        return sum(
            len(a.get("input_tokens", [])) + a.get("output_tokens", 0)
            for a in p.get("agents", [])
        )

    avg_total = sum(_total(p) for p in preds) / total
    avg_turns = sum(len(p.get("agents", [])) for p in preds) / total
    avg_output_per_turn = sum(
        p.get("communication_tokens", 0) / max(len(p.get("agents", [])), 1)
        for p in preds
    ) / total
    return em, avg_f1, correct, avg_comm, avg_total, avg_turns, avg_output_per_turn


# -- Batch processing --------------------------------------------------------

def process_batch(
    method,
    batch: List[Dict],
    processed: int,
    preds: List[Dict],
    progress,
    max_samples: int,
) -> Tuple[int, List[Dict]]:
    remaining = max_samples - processed
    if remaining <= 0:
        return processed, preds

    current_batch = batch[:remaining]
    results = method.run_batch(current_batch)
    if len(results) > remaining:
        results = results[:remaining]

    batch_start = processed
    for offset, res in enumerate(results):
        preds.append(res)
        problem_idx = batch_start + offset + 1
        print(f"\n{'='*60}")
        print(f"Problem #{problem_idx}")
        print(f"Question: {res.get('question', '').strip()}")
        for a in res.get("agents", []):
            print(f"\n  ----- {a['name']} ({a['role']}) -----")
            print(f"  Output: {a.get('output', '').strip()[:300]}")
        print(f"\n  Pred : {res.get('prediction')}")
        print(f"  Gold : {res.get('gold')}")
        print(f"  EM   : {res.get('correct')}  |  F1: {res.get('f1', 0.0):.3f}")

    processed += len(results)
    if progress is not None:
        progress.update(len(results))
    return processed, preds


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

def main():
    parser = argparse.ArgumentParser(description="PACT split-evidence interaction on HotpotQA")

    # Core
    parser.add_argument("--model_name", type=str, required=True,
                        help="HuggingFace model id, e.g. Qwen/Qwen3-14B")
    parser.add_argument("--data_path", type=str,
                        default="data/hotpot_dev_distractor_v1.json",
                        help="Path to the HotpotQA dev-distractor JSON file")
    parser.add_argument("--max_samples", type=int, default=-1,
                        help="Number of samples to evaluate (-1 = all)")

    # Generation
    parser.add_argument("--max_new_tokens", type=int, default=4096)
    parser.add_argument("--temperature", type=float, default=0.6)
    parser.add_argument("--top_p", type=float, default=0.95)
    parser.add_argument("--generate_bs", type=int, default=1,
                        help="Number of samples processed in parallel per forward pass")
    parser.add_argument("--max_turns", type=int, default=4,
                        help="Number of dialogue turns (default 4 = 2 rounds A/B)")

    # Hardware
    parser.add_argument("--device", type=str, default="cuda")
    parser.add_argument("--seed", type=int, default=42)

    # vLLM (optional)
    parser.add_argument("--use_vllm", action="store_true")
    parser.add_argument("--tensor_parallel_size", type=int, default=1)
    parser.add_argument("--gpu_memory_utilization", type=float, default=0.9)
    parser.add_argument("--max_model_len", type=int, default=None)

    # Output
    parser.add_argument("--output_path", type=str, default=None,
                        help="If set, save per-sample results as JSON lines to this path")

    args = parser.parse_args()

    set_seed(args.seed)
    device = auto_device(args.device)

    # -- Load model --
    model = ModelWrapper(args.model_name, device, use_vllm=args.use_vllm, args=args)

    # -- Build method --
    method = PACTMethod(
        model,
        max_new_tokens_each=args.max_new_tokens,
        temperature=args.temperature,
        top_p=args.top_p,
        generate_bs=args.generate_bs,
        max_turns=args.max_turns,
        args=args,
    )

    # -- Load data --
    dataset = list(load_hotpotqa(args.data_path, shuffle_seed=args.seed))
    if args.max_samples != -1:
        dataset = dataset[: args.max_samples]
    total = len(dataset)

    preds: List[Dict] = []
    processed = 0
    batch: List[Dict] = []
    start_time = time.time()

    progress = tqdm(total=total, desc="pact")

    for item in dataset:
        if processed >= total:
            break
        batch.append(item)
        if len(batch) == args.generate_bs or processed + len(batch) == total:
            processed, preds = process_batch(method, batch, processed, preds, progress, total)
            batch = []

    if batch and processed < total:
        processed, preds = process_batch(method, batch, processed, preds, progress, total)

    progress.close()

    total_time = time.time() - start_time
    em, avg_f1, correct, avg_comm, avg_total, avg_turns, avg_output_per_turn = evaluate(preds)

    summary = {
        "method": "pact",
        "model": args.model_name,
        "dataset": "hotpotqa",
        "data_path": args.data_path,
        "seed": args.seed,
        "max_samples": total,
        "exact_match": round(em, 4),
        "avg_f1": round(avg_f1, 4),
        "correct": correct,
        "avg_turns": round(avg_turns, 2),
        "avg_output_tokens_per_turn": round(avg_output_per_turn, 1),
        "avg_communication_tokens": round(avg_comm, 1),
        "avg_total_tokens": round(avg_total, 1),
        "total_time_sec": round(total_time, 2),
        "time_per_sample_sec": round(total_time / total, 4) if total else 0,
    }
    print("\n" + "=" * 60)
    print("RESULTS:")
    print(json.dumps(summary, indent=2, ensure_ascii=False))

    if args.output_path:
        with open(args.output_path, "w", encoding="utf-8") as f:
            for p in preds:
                f.write(json.dumps(p, ensure_ascii=False) + "\n")
        print(f"\nPer-sample results saved to: {args.output_path}")


if __name__ == "__main__":
    main()