mcpbr 0.4.15__py3-none-any.whl → 0.5.0__py3-none-any.whl

mcpbr/benchmarks/longbench.py

@@ -0,0 +1,623 @@
+"""LongBench benchmark implementation for long-context understanding."""
+
+import re
+import string
+from collections import Counter
+from typing import Any
+
+from datasets import load_dataset
+
+from ..docker_env import DockerEnvironmentManager, TaskEnvironment
+from .base import BenchmarkTask
+
+# Mapping of LongBench subsets to their high-level categories
+SUBSET_TO_CATEGORY: dict[str, str] = {
+    # Single-Document QA
+    "narrativeqa": "single_doc_qa",
+    "qasper": "single_doc_qa",
+    "multifieldqa_en": "single_doc_qa",
+    "multifieldqa_zh": "single_doc_qa",
+    # Multi-Document QA
+    "hotpotqa": "multi_doc_qa",
+    "2wikimqa": "multi_doc_qa",
+    "musique": "multi_doc_qa",
+    "dureader": "multi_doc_qa",
+    # Summarization
+    "gov_report": "summarization",
+    "qmsum": "summarization",
+    "multi_news": "summarization",
+    "vcsum": "summarization",
+    "samsum": "summarization",
+    # Few-Shot Learning
+    "triviaqa": "few_shot",
+    "trec": "few_shot",
+    "lsht": "few_shot",
+    # Synthetic Tasks
+    "passage_count": "synthetic",
+    "passage_retrieval_en": "synthetic",
+    "passage_retrieval_zh": "synthetic",
+    # Code Completion
+    "lcc": "code",
+    "repobench-p": "code",
+}
+
+# Subsets that use F1 scoring
+F1_SUBSETS = {
+    "narrativeqa",
+    "qasper",
+    "multifieldqa_en",
+    "multifieldqa_zh",
+    "hotpotqa",
+    "2wikimqa",
+    "musique",
+    "triviaqa",
+}
+
+# Subsets that use ROUGE-L scoring
+ROUGE_SUBSETS = {
+    "dureader",
+    "gov_report",
+    "qmsum",
+    "multi_news",
+    "vcsum",
+    "samsum",
+}
+
+# Subsets that use classification accuracy
+ACCURACY_SUBSETS = {
+    "trec",
+    "lsht",
+    "passage_count",
+    "passage_retrieval_en",
+    "passage_retrieval_zh",
+}
+
+# Subsets that use edit similarity
+EDIT_SIM_SUBSETS = {
+    "lcc",
+    "repobench-p",
+}
+
+# All available subsets
+ALL_SUBSETS = list(SUBSET_TO_CATEGORY.keys())
+
+
+class LongBenchBenchmark:
+    """LongBench benchmark implementation.
+
+    LongBench is a bilingual, multitask benchmark for long context understanding.
+    It covers 6 categories across 21 tasks with contexts ranging from thousands
+    to tens of thousands of tokens:
+
+    - Single-Document QA: narrativeqa, qasper, multifieldqa_en, multifieldqa_zh
+    - Multi-Document QA: hotpotqa, 2wikimqa, musique, dureader
+    - Summarization: gov_report, qmsum, multi_news, vcsum, samsum
+    - Few-Shot Learning: triviaqa, trec, lsht
+    - Synthetic Tasks: passage_count, passage_retrieval_en, passage_retrieval_zh
+    - Code Completion: lcc, repobench-p
+
+    Evaluation uses task-appropriate metrics: F1 for QA, ROUGE-L for summarization,
+    accuracy for classification, and edit similarity for code tasks.
+    """
+
+    name = "longbench"
+
+    def __init__(
+        self,
+        dataset: str = "THUDM/LongBench",
+        subset: str = "hotpotqa",
+    ):
+        """Initialize LongBench benchmark.
+
+        Args:
+            dataset: HuggingFace dataset identifier.
+            subset: Dataset subset/config name (e.g., 'hotpotqa', 'narrativeqa').
+        """
+        self.dataset = dataset
+        self.subset = subset
+
+    def load_tasks(
+        self,
+        sample_size: int | None = None,
+        task_ids: list[str] | None = None,
+        level: int | None = None,
+        filter_difficulty: list[str] | None = None,
+        filter_category: list[str] | None = None,
+        filter_tags: list[str] | None = None,
+    ) -> list[dict[str, Any]]:
+        """Load tasks from LongBench dataset.
+
+        Args:
+            sample_size: Maximum number of tasks to load (None for all).
+            task_ids: Specific task IDs to load (None for all).
+            level: Unused for LongBench.
+            filter_difficulty: Unused for LongBench.
+            filter_category: Filter by task category. Valid categories:
+                single_doc_qa, multi_doc_qa, summarization, few_shot,
+                synthetic, code. When specified, loads tasks from all
+                subsets matching the given categories.
+            filter_tags: Unused for LongBench.
+
+        Returns:
+            List of LongBench task dictionaries.
+        """
+        _ = level
+        _ = filter_difficulty
+        _ = filter_tags
+
+        # Determine which subsets to load based on filter_category
+        subsets_to_load = self._resolve_subsets(filter_category)
+
+        all_tasks: list[dict[str, Any]] = []
+        for subset_name in subsets_to_load:
+            try:
+                ds = load_dataset(self.dataset, subset_name, split="test")
+                for idx, item in enumerate(ds):
+                    task = dict(item)
+                    task["_subset"] = subset_name
+                    task["_original_index"] = idx
+                    all_tasks.append(task)
+            except Exception:
+                # Skip subsets that fail to load (e.g., unavailable configs)
+                continue
+
+        # Filter by task_ids if specified
+        if task_ids:
+            task_id_set = set(task_ids)
+            all_tasks = [
+                t
+                for t in all_tasks
+                if f"longbench_{t['_subset']}_{t['_original_index']}" in task_id_set
+            ]
+
+        # Apply sample size limit
+        if sample_size is not None and len(all_tasks) > sample_size:
+            all_tasks = all_tasks[:sample_size]
+
+        # Augment tasks with instance_id and problem_statement
+        augmented_tasks = []
+        for task in all_tasks:
+            augmented = dict(task)
+            subset_name = task["_subset"]
+            orig_idx = task["_original_index"]
+            augmented["instance_id"] = f"longbench_{subset_name}_{orig_idx}"
+            augmented["problem_statement"] = self._generate_problem_statement(augmented)
+            augmented["ground_truth_answers"] = task.get("answers", [])
+            augmented["task_category"] = SUBSET_TO_CATEGORY.get(subset_name, "unknown")
+            augmented_tasks.append(augmented)
+
+        return augmented_tasks
+
+    def _resolve_subsets(self, filter_category: list[str] | None) -> list[str]:
+        """Resolve which subsets to load based on category filter.
+
+        Args:
+            filter_category: List of category names to include.
+
+        Returns:
+            List of subset names to load.
+        """
+        if not filter_category:
+            # If no category filter, use the configured subset
+            return [self.subset]
+
+        category_set = {c.lower() for c in filter_category}
+        resolved = []
+        for subset_name, category in SUBSET_TO_CATEGORY.items():
+            if category in category_set:
+                resolved.append(subset_name)
+
+        return resolved if resolved else [self.subset]
+
+    def normalize_task(self, task: dict[str, Any]) -> BenchmarkTask:
+        """Convert LongBench task to normalized format.
+
+        Args:
+            task: LongBench task dictionary.
+
+        Returns:
+            Normalized BenchmarkTask.
+
+        Raises:
+            ValueError: If required fields are missing.
+        """
+        instance_id = task.get("instance_id")
+        if not instance_id:
+            msg = f"Task missing required 'instance_id' field: {list(task.keys())}"
+            raise ValueError(msg)
+
+        input_text = task.get("input", "")
+        if not input_text and not task.get("context", ""):
+            msg = f"Task missing required 'input' or 'context' field: {list(task.keys())}"
+            raise ValueError(msg)
+
+        problem_statement = self._generate_problem_statement(task)
+        subset_name = task.get("_subset", self.subset)
+
+        return BenchmarkTask(
+            task_id=instance_id,
+            problem_statement=problem_statement,
+            repo="longbench/context",
+            commit="HEAD",
+            metadata={
+                "input": input_text,
+                "context_length": task.get("length", 0),
+                "dataset": task.get("dataset", subset_name),
+                "language": task.get("language", "en"),
+                "answers": task.get("answers", []),
+                "all_classes": task.get("all_classes"),
+                "subset": subset_name,
+                "category": SUBSET_TO_CATEGORY.get(subset_name, "unknown"),
+            },
+        )
+
+    def _generate_problem_statement(self, task: dict[str, Any]) -> str:
+        """Generate problem statement from task.
+
+        Includes the full long context in the problem statement so agents
+        must process and understand it.
+
+        Args:
+            task: LongBench task dictionary.
+
+        Returns:
+            Problem statement for the agent.
+        """
+        input_text = task.get("input", "No input provided")
+        context = task.get("context", "")
+        subset_name = task.get("_subset", task.get("dataset", self.subset))
+        category = SUBSET_TO_CATEGORY.get(subset_name, "unknown")
+
+        # Build category-specific instructions
+        if category in ("single_doc_qa", "multi_doc_qa"):
+            task_instruction = (
+                "Read the following document(s) carefully and answer the question "
+                "based on the information provided."
+            )
+        elif category == "summarization":
+            task_instruction = (
+                "Read the following document(s) carefully and provide a concise summary "
+                "capturing the key points."
+            )
+        elif category == "few_shot":
+            task_instruction = (
+                "Study the provided examples carefully and answer the question "
+                "following the same pattern."
+            )
+        elif category == "synthetic":
+            task_instruction = (
+                "Analyze the provided passages carefully and answer the question precisely."
+            )
+        elif category == "code":
+            task_instruction = "Read the provided code context and complete the code as requested."
+        else:
+            task_instruction = "Read the following context carefully and respond to the request."
+
+        parts = [task_instruction, ""]
+
+        if context:
+            parts.append("--- CONTEXT ---")
+            parts.append(context)
+            parts.append("--- END CONTEXT ---")
+            parts.append("")
+
+        parts.append(f"Question/Task: {input_text}")
+
+        return "\n".join(parts)
+
+    async def create_environment(
+        self,
+        task: dict[str, Any],
+        docker_manager: DockerEnvironmentManager,
+    ) -> TaskEnvironment:
+        """Create environment for LongBench task.
+
+        LongBench doesn't require repository setup - creates minimal environment.
+
+        Args:
+            task: LongBench task dictionary.
+            docker_manager: Docker environment manager.
+
+        Returns:
+            TaskEnvironment for the task.
+        """
+        instance_id = task.get("instance_id", "longbench_unknown")
+        temp_task = {
+            "instance_id": instance_id,
+            "repo": "longbench/context",
+            "base_commit": "HEAD",
+        }
+        return await docker_manager.create_environment(temp_task)
+
+    async def evaluate(
+        self,
+        _env: TaskEnvironment,
+        task: dict[str, Any],
+        solution: str,
+    ) -> dict[str, Any]:
+        """Evaluate a solution for LongBench task.
+
+        Uses the appropriate metric based on the task subset:
+        - F1 score for QA tasks
+        - ROUGE-L for summarization tasks
+        - Accuracy for classification tasks
+        - Edit similarity for code tasks
+
+        Args:
+            _env: Task environment (unused; evaluation is offline).
+            task: LongBench task dictionary.
+            solution: Solution to evaluate (agent's response).
+
+        Returns:
+            Dictionary with evaluation results including 'resolved' boolean.
+        """
+        answers = task.get("ground_truth_answers", task.get("answers", []))
+        if not answers:
+            return {"resolved": False, "error": "No ground truth answers available"}
+
+        subset_name = task.get("_subset", task.get("dataset", self.subset))
+        all_classes = task.get("all_classes")
+
+        if subset_name in F1_SUBSETS:
+            score = self._compute_f1_max(solution, answers)
+            metric = "f1"
+        elif subset_name in ROUGE_SUBSETS:
+            score = self._compute_rouge_l_max(solution, answers)
+            metric = "rouge_l"
+        elif subset_name in ACCURACY_SUBSETS:
+            score = self._compute_classification_accuracy(solution, answers, all_classes)
+            metric = "accuracy"
+        elif subset_name in EDIT_SIM_SUBSETS:
+            score = self._compute_edit_similarity_max(solution, answers)
+            metric = "edit_similarity"
+        else:
+            # Default to F1 for unknown subsets
+            score = self._compute_f1_max(solution, answers)
+            metric = "f1"
+
+        # Threshold: score >= 0.5 is considered resolved for F1/ROUGE
+        # For accuracy, exact match (1.0) is required
+        if metric == "accuracy":
+            resolved = score >= 1.0
+        else:
+            resolved = score >= 0.5
+
+        return {
+            "resolved": resolved,
+            "score": score,
+            "metric": metric,
+            "subset": subset_name,
+            "category": SUBSET_TO_CATEGORY.get(subset_name, "unknown"),
+        }
+
+    def _normalize_text(self, text: str) -> str:
+        """Normalize text for scoring by lowercasing, removing punctuation and articles.
+
+        Args:
+            text: Raw text.
+
+        Returns:
+            Normalized text.
+        """
+        text = text.lower().strip()
+        # Remove articles
+        text = re.sub(r"\b(a|an|the)\b", " ", text)
+        # Remove punctuation
+        text = text.translate(str.maketrans("", "", string.punctuation))
+        # Collapse whitespace
+        text = re.sub(r"\s+", " ", text).strip()
+        return text
+
+    def _compute_f1(self, prediction: str, reference: str) -> float:
+        """Compute token-level F1 score between prediction and reference.
+
+        Args:
+            prediction: Predicted answer text.
+            reference: Reference answer text.
+
+        Returns:
+            F1 score between 0.0 and 1.0.
+        """
+        pred_tokens = self._normalize_text(prediction).split()
+        ref_tokens = self._normalize_text(reference).split()
+
+        if not pred_tokens or not ref_tokens:
+            return 1.0 if pred_tokens == ref_tokens else 0.0
+
+        common = Counter(pred_tokens) & Counter(ref_tokens)
+        num_common = sum(common.values())
+
+        if num_common == 0:
+            return 0.0
+
+        precision = num_common / len(pred_tokens)
+        recall = num_common / len(ref_tokens)
+
+        return 2 * precision * recall / (precision + recall)
+
+    def _compute_f1_max(self, prediction: str, references: list[str]) -> float:
+        """Compute maximum F1 score across all reference answers.
+
+        Args:
+            prediction: Predicted answer text.
+            references: List of reference answer texts.
+
+        Returns:
+            Maximum F1 score.
+        """
+        if not references:
+            return 0.0
+        return max(self._compute_f1(prediction, ref) for ref in references)
+
+    def _compute_rouge_l(self, prediction: str, reference: str) -> float:
+        """Compute ROUGE-L (Longest Common Subsequence) F-measure.
+
+        Args:
+            prediction: Predicted text.
+            reference: Reference text.
+
+        Returns:
+            ROUGE-L F-measure between 0.0 and 1.0.
+        """
+        pred_tokens = self._normalize_text(prediction).split()
+        ref_tokens = self._normalize_text(reference).split()
+
+        if not pred_tokens or not ref_tokens:
+            return 1.0 if pred_tokens == ref_tokens else 0.0
+
+        lcs_length = self._lcs_length(pred_tokens, ref_tokens)
+
+        if lcs_length == 0:
+            return 0.0
+
+        precision = lcs_length / len(pred_tokens)
+        recall = lcs_length / len(ref_tokens)
+
+        return 2 * precision * recall / (precision + recall)
+
+    def _lcs_length(self, seq1: list[str], seq2: list[str]) -> int:
+        """Compute length of longest common subsequence.
+
+        Args:
+            seq1: First sequence.
+            seq2: Second sequence.
+
+        Returns:
+            Length of the LCS.
+        """
+        m, n = len(seq1), len(seq2)
+        # Use space-optimized DP with two rows
+        prev = [0] * (n + 1)
+        curr = [0] * (n + 1)
+
+        for i in range(1, m + 1):
+            for j in range(1, n + 1):
+                if seq1[i - 1] == seq2[j - 1]:
+                    curr[j] = prev[j - 1] + 1
+                else:
+                    curr[j] = max(prev[j], curr[j - 1])
+            prev, curr = curr, [0] * (n + 1)
+
+        return prev[n]
+
+    def _compute_rouge_l_max(self, prediction: str, references: list[str]) -> float:
+        """Compute maximum ROUGE-L score across all reference answers.
+
+        Args:
+            prediction: Predicted text.
+            references: List of reference texts.
+
+        Returns:
+            Maximum ROUGE-L score.
+        """
+        if not references:
+            return 0.0
+        return max(self._compute_rouge_l(prediction, ref) for ref in references)
+
+    def _compute_classification_accuracy(
+        self,
+        prediction: str,
+        references: list[str],
+        all_classes: list[str] | None = None,
+    ) -> float:
+        """Compute classification accuracy.
+
+        For classification tasks, checks if the predicted class matches
+        any of the reference answers.
+
+        Args:
+            prediction: Predicted class/answer.
+            references: List of correct answers.
+            all_classes: All possible class labels (unused but available).
+
+        Returns:
+            1.0 if correct, 0.0 otherwise.
+        """
+        _ = all_classes
+        pred_normalized = self._normalize_text(prediction)
+        for ref in references:
+            ref_normalized = self._normalize_text(ref)
+            if ref_normalized == pred_normalized or ref_normalized in pred_normalized:
+                return 1.0
+        return 0.0
+
+    def _compute_edit_similarity(self, prediction: str, reference: str) -> float:
+        """Compute edit similarity between prediction and reference.
+
+        Edit similarity = 1 - (edit_distance / max_length).
+
+        Args:
+            prediction: Predicted code.
+            reference: Reference code.
+
+        Returns:
+            Edit similarity between 0.0 and 1.0.
+        """
+        if not prediction and not reference:
+            return 1.0
+        if not prediction or not reference:
+            return 0.0
+
+        m, n = len(prediction), len(reference)
+
+        # Levenshtein distance using space-optimized DP
+        prev = list(range(n + 1))
+        curr = [0] * (n + 1)
+
+        for i in range(1, m + 1):
+            curr[0] = i
+            for j in range(1, n + 1):
+                if prediction[i - 1] == reference[j - 1]:
+                    curr[j] = prev[j - 1]
+                else:
+                    curr[j] = 1 + min(prev[j], curr[j - 1], prev[j - 1])
+            prev, curr = curr, [0] * (n + 1)
+
+        edit_distance = prev[n]
+        max_length = max(m, n)
+        return 1.0 - (edit_distance / max_length)
+
+    def _compute_edit_similarity_max(self, prediction: str, references: list[str]) -> float:
+        """Compute maximum edit similarity across all reference answers.
+
+        Args:
+            prediction: Predicted code.
+            references: List of reference code snippets.
+
+        Returns:
+            Maximum edit similarity score.
+        """
+        if not references:
+            return 0.0
+        return max(self._compute_edit_similarity(prediction, ref) for ref in references)
+
+    def get_prebuilt_image(self, _task: dict[str, Any]) -> str | None:
+        """Get pre-built Docker image name for LongBench task.
+
+        LongBench doesn't use pre-built images.
+
+        Args:
+            _task: LongBench task dictionary (unused).
+
+        Returns:
+            None (no pre-built images available).
+        """
+        return None
+
+    def get_prompt_template(self) -> str:
+        """Get LongBench prompt template.
+
+        Returns:
+            Prompt template for long-context understanding tasks.
+        """
+        return (
+            "You are given a long-context understanding task.\n\n"
+            "{problem_statement}\n\n"
+            "IMPORTANT INSTRUCTIONS:\n"
+            "- Read the entire context carefully before answering\n"
+            "- Base your answer solely on the information provided in the context\n"
+            "- Be concise and precise in your answer\n"
+            "- For QA tasks, provide a direct answer to the question\n"
+            "- For summarization tasks, capture the key points concisely\n"
+            "- For code tasks, complete the code following the existing patterns"
+        )