chore: import upstream snapshot with attribution

python/sglang/test/simple_eval_longbench_v2.py

@@ -0,0 +1,344 @@
+# Adapted from https://github.com/openai/simple-evals/
+
+"""
+LongBench v2: Towards Deeper Understanding and Reasoning on Realistic Long-Context Multitasks
+Yushi Bai, Shangqing Tu, Jiajie Zhang, Hao Peng, Xiaozhi Wang, Xin Lv, Shulin Cao, Jiazheng Xu, Lei Hou, Yuxiao Dong, Jie Tang, Juanzi Li
+https://arxiv.org/abs/2412.15204
+"""
+
+import csv
+import json
+import os
+import re
+from typing import Any, Dict, List, Optional
+
+from transformers import AutoTokenizer
+
+from sglang.test import simple_eval_common as common
+from sglang.test.simple_eval_common import (
+    ANSWER_PATTERN_MULTICHOICE,
+    HTML_JINJA,
+    Eval,
+    EvalResult,
+    SamplerBase,
+    SingleEvalResult,
+)
+
+# LongBench-v2 task categories
+TASK_CATEGORIES = {
+    "single_document_qa",
+    "multi_document_qa",
+    "long_in_context_learning",
+    "long_dialogue_history",
+    "code_repo_understanding",
+    "long_structured_data",
+}
+
+DEFAULT_DATASET = "THUDM/LongBench-v2"
+DEFAULT_DATASET_SPLIT = "train"
+
+
+def format_longbench_v2_question(row: dict) -> str:
+    """Format a LongBench-v2 question using the official template."""
+    context = row.get("context", "")
+    question = row.get("question", "")
+
+    # Handle both standard format (A, B, C, D) and alternative format (choices list)
+    if "choices" in row:
+        choices = row["choices"]
+        choice_A = choices[0] if len(choices) > 0 else ""
+        choice_B = choices[1] if len(choices) > 1 else ""
+        choice_C = choices[2] if len(choices) > 2 else ""
+        choice_D = choices[3] if len(choices) > 3 else ""
+    else:
+        choice_A = row.get("A", row.get("choice_A", ""))
+        choice_B = row.get("B", row.get("choice_B", ""))
+        choice_C = row.get("C", row.get("choice_C", ""))
+        choice_D = row.get("D", row.get("choice_D", ""))
+
+    # Official LongBench-v2 template
+    prompt = f"""
+Please read the following text and answer the question below.
+<text>
+{context.strip()}
+</text>
+
+What is the correct answer to this question: {question.strip()}
+Choices:
+(A) {choice_A.strip()}
+(B) {choice_B.strip()}
+(C) {choice_C.strip()}
+(D) {choice_D.strip()}
+
+Format your response as follows: "The correct answer is (insert answer here)"."""
+
+    return prompt
+
+
+def extract_longbench_v2_answer(response: str) -> Optional[str]:
+    """Extract answer from model response using official LongBench-v2 method."""
+    response = response.replace("*", "")
+
+    # First try: "The correct answer is (A)"
+    match = re.search(r"The correct answer is \(([A-D])\)", response, re.IGNORECASE)
+    if match:
+        return match.group(1).upper()
+
+    # Second try: "The correct answer is A"
+    match = re.search(r"The correct answer is ([A-D])", response, re.IGNORECASE)
+    if match:
+        return match.group(1).upper()
+
+    # Fallback: Standard SGLang multichoice pattern
+    match = re.search(ANSWER_PATTERN_MULTICHOICE, response)
+    if match:
+        return match.group(1).upper()
+
+    # Generic fallback when model says "answer is A"
+    match = re.search(r"answer\s+is\s*\(?([A-D])\)?", response, re.IGNORECASE)
+    if match:
+        return match.group(1).upper()
+
+    return None
+
+
+class LongBenchV2Eval(Eval):
+    """
+    Evaluation utility for LongBench-v2 dataset.
+
+    LongBench-v2 is designed to assess the ability of LLMs to handle long-context problems
+    requiring deep understanding and reasoning across real-world multitasks.
+    """
+
+    def __init__(
+        self,
+        model: str = None,
+        data_source: str = DEFAULT_DATASET,
+        num_examples: Optional[int] = None,
+        num_threads: int = 1,
+        n_repeats: int = 1,
+        categories: Optional[List[str]] = None,
+        max_context_length: Optional[int] = None,
+        min_context_length: Optional[int] = None,
+    ):
+        """
+        Initialize LongBench-v2 evaluation.
+
+        Args:
+            data_source: HuggingFace dataset name, local file path (CSV/JSON)
+            num_examples: Number of examples to evaluate (None for all)
+            num_threads: Number of threads for parallel processing
+            n_repeats: Number of times to repeat evaluation for error bars
+            categories: List of task categories to include (None for all)
+            max_context_length: Maximum context length in characters
+            min_context_length: Minimum context length in characters
+        """
+        self.tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=True)
+        self.min_context_length = min_context_length
+        self.max_context_length = max_context_length
+        # Load dataset based on data source type
+        examples = self._load_dataset(data_source)
+
+        # Apply filtering
+        if categories:
+            examples = [ex for ex in examples if ex.get("category") in categories]
+
+        # Sample examples if specified
+        if num_examples:
+            assert n_repeats == 1, "n_repeats only supported when not sampling examples"
+            examples = examples[: min(num_examples, len(examples))]
+
+        # Repeat examples for multiple runs
+        examples = examples * n_repeats
+
+        if not examples:
+            raise ValueError(
+                "No examples available for LongBench-v2 evaluation after filtering"
+            )
+
+        self.examples = examples
+        self.n_repeats = n_repeats
+        self.num_threads = num_threads
+
+        print(f"Loaded {len(self.examples)} examples from LongBench-v2")
+        if categories:
+            print(f"Filtered to categories: {categories}")
+        if min_context_length or max_context_length:
+            print(
+                f"Context length filter: {min_context_length}-{max_context_length} characters"
+            )
+
+    def _load_dataset(self, data_source: str) -> List[Dict[str, Any]]:
+        """Load dataset from HuggingFace hub or local files."""
+
+        if not data_source:
+            data_source = DEFAULT_DATASET
+
+        if os.path.exists(data_source):
+            raw_examples = self._load_local_file(data_source)
+        else:
+            raw_examples = self._load_hf_dataset(data_source)
+
+        return [self._normalize_example(example) for example in raw_examples]
+
+    def _load_local_file(self, path: str) -> List[Dict[str, Any]]:
+        """Load examples from a local CSV/JSON/JSONL file."""
+
+        suffix = os.path.splitext(path)[1].lower()
+        if suffix in {".json", ".jsonl"}:
+            with open(path, "r", encoding="utf-8") as fh:
+                if suffix == ".jsonl":
+                    data = [json.loads(line) for line in fh if line.strip()]
+                else:
+                    data = json.load(fh)
+        elif suffix == ".csv":
+            with open(path, "r", encoding="utf-8") as fh:
+                reader = csv.DictReader(fh)
+                data = list(reader)
+        else:
+            # Try JSON, then CSV as fallback
+            try:
+                with open(path, "r", encoding="utf-8") as fh:
+                    data = json.load(fh)
+            except json.JSONDecodeError:
+                with open(path, "r", encoding="utf-8") as fh:
+                    reader = csv.DictReader(fh)
+                    data = list(reader)
+
+        if isinstance(data, dict):
+            data = data.get("data", [])
+
+        if not isinstance(data, list):
+            raise ValueError("Expected list of examples from local file")
+
+        return data
+
+    def _load_hf_dataset(self, identifier: str) -> List[Dict[str, Any]]:
+        """Load the dataset from HuggingFace Hub."""
+
+        parts = identifier.split(":", maxsplit=1)
+        dataset_name = parts[0]
+        split = parts[1] if len(parts) == 2 else DEFAULT_DATASET_SPLIT
+
+        try:
+            from datasets import load_dataset  # type: ignore
+        except ImportError as exc:
+            raise ImportError(
+                "Please install the 'datasets' package to load LongBench-v2 from HuggingFace: pip install datasets"
+            ) from exc
+
+        dataset = load_dataset(dataset_name, split=split)
+        return [dict(row) for row in dataset]
+
+    def _normalize_example(self, example: Dict[str, Any]) -> Dict[str, Any]:
+        """Ensure each example exposes the expected keys."""
+
+        normalized = dict(example)
+
+        for letter in ["A", "B", "C", "D"]:
+            choice_key = f"choice_{letter}"
+            if letter not in normalized and choice_key in normalized:
+                normalized[letter] = normalized[choice_key]
+
+        if "category" not in normalized and "domain" in normalized:
+            normalized["category"] = normalized["domain"]
+
+        answer = normalized.get("answer")
+        if isinstance(answer, str):
+            normalized["answer"] = answer.strip().upper()
+        elif isinstance(answer, int) and 0 <= answer < 4:
+            normalized["answer"] = ["A", "B", "C", "D"][answer]
+
+        return normalized
+
+    def _check_context_length(
+        self,
+        formatted_question: str,
+        tokenizer: AutoTokenizer,
+        min_length: Optional[int],
+        max_length: Optional[int],
+    ) -> bool:
+        """Filter examples by context length measured in characters."""
+        input_ids = tokenizer.encode(formatted_question)
+        context_length = len(input_ids)
+
+        if min_length is not None and context_length < min_length:
+            return False
+        if max_length is not None and context_length > max_length:
+            return False
+
+        return True
+
+    def __call__(self, sampler: SamplerBase) -> EvalResult:
+        """Run the evaluation."""
+
+        def fn(row: dict):
+            # Format the question using official template
+            formatted_question = format_longbench_v2_question(row)
+
+            if self.min_context_length or self.max_context_length:
+                if not self._check_context_length(
+                    formatted_question,
+                    self.tokenizer,
+                    self.min_context_length,
+                    self.max_context_length,
+                ):
+                    # Skip this example
+                    return None
+
+            prompt_messages = [
+                sampler._pack_message(content=formatted_question, role="user")
+            ]
+
+            # Get model response
+            response_text = sampler(prompt_messages)
+            if response_text is None:
+                response_text = ""
+
+            # Extract answer using official method
+            extracted_answer = extract_longbench_v2_answer(response_text)
+
+            # Get correct answer
+            correct_answer = row.get("answer", "")
+            if isinstance(correct_answer, str):
+                correct_answer = correct_answer.strip().upper()
+            elif isinstance(correct_answer, int) and 0 <= correct_answer < 4:
+                correct_answer = ["A", "B", "C", "D"][correct_answer]
+
+            # Calculate score
+            score = 1.0 if extracted_answer == correct_answer else 0.0
+
+            # Generate HTML report
+            html = common.jinja_env.from_string(HTML_JINJA).render(
+                prompt_messages=prompt_messages,
+                next_message=dict(content=response_text, role="assistant"),
+                score=score,
+                correct_answer=correct_answer,
+                extracted_answer=extracted_answer,
+            )
+
+            # Build conversation
+            convo = prompt_messages + [dict(content=response_text, role="assistant")]
+
+            # Prepare metrics
+            metrics = {"chars": len(response_text)}
+
+            # Add category-specific metrics
+            category = row.get("category", row.get("domain", "unknown"))
+            if category in TASK_CATEGORIES:
+                metrics[category] = score
+
+            difficulty = row.get("difficulty")
+            if isinstance(difficulty, str) and difficulty:
+                metrics[f"difficulty_{difficulty.lower()}"] = score
+
+            return SingleEvalResult(
+                html=html,
+                score=score,
+                convo=convo,
+                metrics=metrics,
+            )
+
+        # Run evaluation with progress tracking
+        results = common.map_with_progress(fn, self.examples, self.num_threads)
+        return common.aggregate_results(results)