topoteretes--cognee/cognee/eval_framework/benchmark_adapters/beam_adapter.py

"""BEAM benchmark adapter — loads synthetic long-context conversations.

Dataset: https://huggingface.co/datasets/Mohammadta/BEAM
Paper: "Beyond a Million Tokens: Benchmarking Long-Term Memory in LLMs"

Each conversation contains multi-session batches with 20 probing questions
across 10 skill categories (information_extraction, temporal_reasoning,
multi_session_reasoning, contradiction_resolution, event_ordering,
knowledge_update, summarization, abstention, preference_following,
instruction_following).
"""

import ast
import json
from typing import Any, Dict, List, Optional, Tuple, Union

from cognee.eval_framework.benchmark_adapters.base_benchmark_adapter import BaseBenchmarkAdapter
from cognee.shared.logging_utils import get_logger

logger = get_logger()

# Map BEAM question types to a normalized key for the answer field
# (the dataset uses different field names for the ground-truth answer)
_ANSWER_FIELD_NAMES = [
    "answer",
    "ideal_response",
    "ideal_answer",
    "ideal_summary",
]


def _extract_answer(question_dict: dict) -> str:
    """Extract the ground-truth answer from a BEAM probing question."""
    for field in _ANSWER_FIELD_NAMES:
        if field in question_dict and question_dict[field]:
            val = question_dict[field]
            return val if isinstance(val, str) else str(val)
    return ""


def _flatten_chat(chat_batches: list) -> str:
    """Flatten the nested chat structure into a single text corpus.

    The chat field is a list of batches, each batch is a list of messages.
    We produce a readable conversation transcript.
    """
    lines = []
    for batch_idx, batch in enumerate(chat_batches):
        lines.append(f"--- Session {batch_idx + 1} ---")
        for msg in batch:
            role = msg.get("role", "unknown").capitalize()
            content = msg.get("content", "")
            time_anchor = msg.get("time_anchor")
            prefix = f"[{time_anchor}] " if time_anchor else ""
            lines.append(f"{prefix}{role}: {content}")
        lines.append("")
    return "\n".join(lines)


class BEAMAdapter(BaseBenchmarkAdapter):
    """Adapter for the BEAM long-context conversation benchmark.

    Loads conversations from the HuggingFace dataset and extracts
    probing questions with rubric-based ground truth.

    Args:
        split: Dataset split to use ("100K", "500K", "1M"). Default "100K".
        conversation_index: Which conversation to load (0-indexed). Default 0.
    """

    def __init__(
        self,
        split: str = "100K",
        conversation_index: int = 0,
        max_batches: Optional[int] = None,
    ):
        self.split = split
        self.conversation_index = conversation_index
        self.max_batches = max_batches

    def load_corpus(
        self,
        limit: Optional[int] = None,
        seed: int = 42,
        load_golden_context: bool = False,
        instance_filter: Optional[Union[str, List[str], List[int]]] = None,
    ) -> Tuple[List[str], List[Dict[str, Any]]]:
        """Load a single BEAM conversation as corpus + probing questions.

        Returns:
            corpus_list: List with one entry — the full conversation transcript.
            question_answer_pairs: List of dicts with question, answer, rubric,
                question_type, and optional source_chat_ids.
        """
        try:
            import datasets as _datasets_lib

            load_dataset = _datasets_lib.load_dataset
        except ImportError:
            raise ImportError(
                "The 'datasets' package is required for BEAM. Install it with: pip install datasets"
            )

        logger.info(
            f"Loading BEAM dataset split={self.split}, conversation_index={self.conversation_index}"
        )

        ds = load_dataset("Mohammadta/BEAM", split=self.split)
        if self.conversation_index >= len(ds):
            raise IndexError(
                f"conversation_index={self.conversation_index} out of range "
                f"(split '{self.split}' has {len(ds)} conversations)"
            )

        row = ds[self.conversation_index]

        # Build corpus — optionally truncate to first N batches for faster local runs
        chat_batches = row["chat"]
        if self.max_batches is not None and len(chat_batches) > self.max_batches:
            logger.info(
                f"Truncating conversation from {len(chat_batches)} batches "
                f"to {self.max_batches} (max_batches)"
            )
            chat_batches = chat_batches[: self.max_batches]
        corpus_text = _flatten_chat(chat_batches)
        corpus_list = [corpus_text]

        # Parse probing questions
        probing_raw = row.get("probing_questions", "")
        if isinstance(probing_raw, str):
            try:
                probing_data = ast.literal_eval(probing_raw)
            except (ValueError, SyntaxError):
                try:
                    probing_data = json.loads(probing_raw)
                except json.JSONDecodeError:
                    logger.error("Failed to parse probing_questions field")
                    probing_data = {}
        else:
            probing_data = probing_raw if isinstance(probing_raw, dict) else {}

        question_answer_pairs = []

        for question_type, questions in probing_data.items():
            if not isinstance(questions, list):
                continue
            for q in questions:
                if not isinstance(q, dict) or "question" not in q:
                    continue

                answer_text = _extract_answer(q)
                rubric = q.get("rubric", [])
                if isinstance(rubric, str):
                    rubric = [rubric]

                qa_pair: Dict[str, Any] = {
                    "question": q["question"],
                    "answer": answer_text,
                    "question_type": question_type,
                    "rubric": rubric,
                    "difficulty": q.get("difficulty", "unknown"),
                    "conversation_id": row.get("conversation_id", str(self.conversation_index)),
                }

                source_ids = q.get("source_chat_ids")
                if source_ids and load_golden_context:
                    golden = self._extract_golden_context(chat_batches, source_ids)
                    if golden:
                        qa_pair["golden_context"] = golden

                question_answer_pairs.append(qa_pair)

        # Apply instance filter if provided
        if instance_filter is not None:
            question_answer_pairs = self._filter_instances(
                question_answer_pairs, instance_filter, id_key="question"
            )

        # Apply limit
        if limit is not None and limit < len(question_answer_pairs):
            question_answer_pairs = question_answer_pairs[:limit]

        logger.info(
            f"Loaded BEAM conversation {self.conversation_index}: "
            f"{len(corpus_text)} chars corpus, "
            f"{len(question_answer_pairs)} probing questions"
        )

        return corpus_list, question_answer_pairs

    @staticmethod
    def _extract_golden_context(chat_batches: list, source_ids: Any) -> str:
        """Extract messages referenced by source_chat_ids as golden context."""
        # source_chat_ids can be a list of ints, or a dict with nested lists
        ids_to_find: set = set()

        if isinstance(source_ids, list):
            for item in source_ids:
                if isinstance(item, int):
                    ids_to_find.add(item)
        elif isinstance(source_ids, dict):
            for v in source_ids.values():
                if isinstance(v, list):
                    for item in v:
                        if isinstance(item, int):
                            ids_to_find.add(item)
                elif isinstance(v, int):
                    ids_to_find.add(v)

        if not ids_to_find:
            return ""

        # Build id → message content map
        context_parts = []
        for batch in chat_batches:
            for msg in batch:
                msg_id = msg.get("id")
                if msg_id in ids_to_find:
                    role = msg.get("role", "unknown").capitalize()
                    content = msg.get("content", "")
                    context_parts.append(f"{role}: {content}")

        return "\n".join(context_parts)