data.py

"""
HotpotQA loader for the PACT split-evidence interaction setting (Setting A).

Each HotpotQA item provides 10 context paragraphs: 2 gold supporting
paragraphs and 8 distractors. We split them evenly between two agents
(5 paragraphs each: 1 gold + 4 distractors), so neither agent can answer
the question alone and the two agents must exchange evidence.
"""

import json
import random
from typing import Dict, Iterable, List, Optional, Tuple

from utils import normalize_answer


def _format_context(context: List[Tuple[str, List[str]]]) -> str:
    """
    Convert a HotpotQA-style context (list of [title, sentences]) into a
    readable string to include in the prompt.
    """
    parts = []
    for title, sentences in context:
        body = " ".join(sentences).strip()
        parts.append(f"[{title}]\n{body}")
    return "\n\n".join(parts)


def _split_5_5_contexts(
    item: Dict, rng: random.Random
) -> Tuple[str, str, str, str]:
    """
    Split all 10 context paragraphs evenly: 5 to Agent A, 5 to Agent B.
    Each agent receives exactly 1 gold article + 4 distractor articles.

    The 5 paragraphs given to each agent are shuffled (deterministically,
    via `rng`) so the gold article does not appear in a fixed position.
    Neither agent knows which of its 5 paragraphs is gold, and the
    information needed to answer is partitioned across the two agents.

    Returns:
        gold_title_a, context_a_str, gold_title_b, context_b_str
    """
    seen = {}
    for title, _ in item["supporting_facts"]:
        if title not in seen:
            seen[title] = True
    gold_titles = list(seen.keys())

    gold_a = gold_titles[0] if len(gold_titles) >= 1 else None
    gold_b = gold_titles[1] if len(gold_titles) >= 2 else None

    # Separate distractors (maintain original order)
    distractors = [ctx for ctx in item["context"]
                   if ctx[0] not in (gold_a, gold_b)]

    # Assign gold + 4 distractors each (8 distractors total -> 4 each)
    dist_a = distractors[:4]
    dist_b = distractors[4:8]

    ctx_lookup = dict(item["context"])
    group_a = ([(gold_a, ctx_lookup[gold_a])] if gold_a else []) + dist_a
    group_b = ([(gold_b, ctx_lookup[gold_b])] if gold_b else []) + dist_b

    # Shuffle within each agent's group so gold is not always first.
    rng.shuffle(group_a)
    rng.shuffle(group_b)

    ctx_a_str = _format_context(group_a)
    ctx_b_str = _format_context(group_b)

    return (gold_a or ""), ctx_a_str, (gold_b or ""), ctx_b_str


def load_hotpotqa(
    path: str,
    max_samples: int = -1,
    shuffle_seed: int = 42,
) -> Iterable[Dict]:
    """
    Load a HotpotQA JSON file (distractor setting) and yield items prepared
    for the split-evidence interaction setting.

    Each yielded dict contains the two agents' partial contexts
    (`context_a55`, `context_b55`), the question, and the gold answer.
    """
    with open(path, "r", encoding="utf-8") as f:
        data = json.load(f)

    rng = random.Random(shuffle_seed)

    for i, item in enumerate(data):
        if max_samples != -1 and i >= max_samples:
            break

        _, ctx_a55, _, ctx_b55 = _split_5_5_contexts(item, rng)

        yield {
            "raw_question": item["question"].strip(),
            # 5-5 split (1 gold + 4 distractors each, paragraphs shuffled)
            "context_a55": ctx_a55,
            "context_b55": ctx_b55,
            "gold": normalize_answer(item["answer"]),
            "raw_answer": item["answer"],
            "type": item.get("type", ""),
            "level": item.get("level", ""),
            "supporting_facts": item.get("supporting_facts", []),
            "id": item.get("_id", ""),
            "solution": item["answer"],
        }