Add InfiniteBenchSum scenario and run spec (#3409)

Author: liamjxu

src/helm/benchmark/run_specs/experimental_run_specs.py

@@ -192,3 +192,34 @@ def get_czech_bank_qa_spec(config_name: str = "berka_queries_1024_2024_12_18") -
         annotators=[AnnotatorSpec("helm.benchmark.annotation.czech_bank_qa_annotator.CzechBankQAAnnotator")],
         groups=["czech_bank_qa"],
     )
+
+
+@run_spec_function("infinite_bench_sum")
+def get_infinite_bench_sum_spec(min_num_word: float = 0.0, max_num_word: float = 1e7) -> RunSpec:
+
+    scenario_spec = ScenarioSpec(
+        class_name="helm.benchmark.scenarios.infinite_bench_sum_scenario.InfiniteBenchSumScenario",
+        args={
+            "min_num_word": min_num_word,
+            "max_num_word": max_num_word,
+        },
+    )
+
+    adapter_spec = AdapterSpec(
+        method=ADAPT_GENERATION,
+        input_prefix="",
+        output_prefix="",
+        max_tokens=2000,  # No official number, the average output token is 1.1k according to the paper
+        num_outputs=1,
+        temperature=0.0,
+    )
+
+    metric_specs = get_basic_metric_specs(["rouge_l"])
+
+    return RunSpec(
+        name="infinite_bench_sum",
+        scenario_spec=scenario_spec,
+        adapter_spec=adapter_spec,
+        metric_specs=metric_specs,
+        groups=["infinite_bench_sum"],
+    )

src/helm/benchmark/scenarios/infinite_bench_sum_scenario.py

@@ -0,0 +1,82 @@
+import os
+import re
+from typing import List
+from datasets import load_dataset, Features, Value, Sequence, Dataset
+from helm.benchmark.scenarios.scenario import (
+    Scenario,
+    Instance,
+    Input,
+    Reference,
+    Output,
+    CORRECT_TAG,
+    TEST_SPLIT,
+)
+from helm.common.general import ensure_directory_exists
+
+
+class InfiniteBenchSumScenario(Scenario):
+    """InfiniteBench Sum
+
+    InfiniteBench is a benchmark tailored for evaluating the capabilities of language models to process,
+    understand, and reason over super long contexts (100k+ tokens). InfiniteBench Sum is a subset of
+    InfiniteBench that requires models to generate a concise summary of the novel. The subset is referred
+    to as "En.Sum" in the original paper.
+    """
+
+    name = "infinite_bench_sum"
+    description = "Summarize a novel from InfiniteBench"
+    tags = ["summarization"]
+
+    def __init__(self, min_num_word: float = 0.0, max_num_word: float = 100e6):
+        self.min_num_word = min_num_word
+        self.max_num_word = max_num_word
+        super().__init__()
+
+    def get_instances(self, output_path: str) -> List[Instance]:
+        # Get InfiniteBench from HuggingFace
+        cache_dir = os.path.join(output_path, "data")
+        ensure_directory_exists(cache_dir)
+
+        # Define the features schema
+        ft = Features(
+            {
+                "id": Value("int64"),
+                "context": Value("string"),
+                "input": Value("string"),
+                "answer": Sequence(Value("string")),
+                "options": Sequence(Value("string")),
+            }
+        )
+
+        # Load the dataset with the specified features
+        dataset = load_dataset(
+            "xinrongzhang2022/InfiniteBench",
+            split="longbook_sum_eng",
+            features=ft,
+            revision="90f0394333616266d9fe85824ceaf505093cbaa5",
+        )
+
+        assert isinstance(dataset, Dataset)
+
+        def count_words(text: str) -> int:
+            return len(re.split(r"\s+", text.strip()))
+
+        dataset = dataset.map(
+            lambda example: {"prompt_wc": count_words(example["context"]) + count_words(example["input"])}
+        ).filter(lambda example: self.min_num_word <= example["prompt_wc"] <= self.max_num_word)
+
+        # Read all instances
+        instances: List[Instance] = []
+        for row in dataset:
+            id = row["id"]
+            input = Input(text=row["context"] + "\n\n" + row["input"])
+            instance = Instance(
+                id=id,
+                input=input,
+                references=[Reference(Output(text=row["answer"][0]), tags=[CORRECT_TAG])],
+                split=TEST_SPLIT,
+                extra_data={"word_count": row["prompt_wc"]},
+            )
+            instances.append(instance)
+
+        return instances

src/helm/benchmark/scenarios/test_infinite_bench_sum_scenario.py

@@ -0,0 +1,46 @@
+import pytest
+import re
+from tempfile import TemporaryDirectory
+from helm.benchmark.scenarios.infinite_bench_sum_scenario import InfiniteBenchSumScenario
+from helm.benchmark.scenarios.scenario import CORRECT_TAG
+
+
+def count_words(text: str) -> int:
+    return len(re.split(r"\s+", text.strip()))
+
+
+@pytest.mark.scenarios
+def test_infinite_bench_sum_scenario():
+    with TemporaryDirectory() as tmpdir:
+        scenario = InfiniteBenchSumScenario(min_num_word=0, max_num_word=1e7)
+        instances = scenario.get_instances(tmpdir)
+        assert len(instances) == 103
+        assert instances[0].split == "test"
+        assert len(instances[0].input.text) == 1745528
+        assert instances[0].extra_data
+        assert instances[0].extra_data["word_count"] == 308762
+        references = instances[0].references
+        assert len(references[0].output.text) == 2865
+        assert references[0].tags == [CORRECT_TAG]
+
+        scenario = InfiniteBenchSumScenario(min_num_word=0, max_num_word=100e3)
+        instances = scenario.get_instances(tmpdir)
+        assert len(instances) == 48
+        assert instances[0].split == "test"
+        assert len(instances[0].input.text) == 381778
+        assert instances[0].extra_data
+        assert instances[0].extra_data["word_count"] == 69458
+        references = instances[0].references
+        assert len(references[0].output.text) == 4217
+        assert references[0].tags == [CORRECT_TAG]
+
+        scenario = InfiniteBenchSumScenario(min_num_word=30e3, max_num_word=80e3)
+        instances = scenario.get_instances(tmpdir)
+        assert len(instances) == 32
+        assert instances[0].split == "test"
+        assert len(instances[1].input.text) == 383396
+        assert instances[1].extra_data
+        assert instances[1].extra_data["word_count"] == 68482
+        references = instances[1].references
+        assert len(references[0].output.text) == 5667
+        assert references[0].tags == [CORRECT_TAG]