From 934f93c2fbf9682b70318b301fe04134ad043a32 Mon Sep 17 00:00:00 2001 From: ssiq Date: Fri, 17 Jul 2026 17:52:26 +0800 Subject: [PATCH] Add AA-LCR dataset integration --- dataset-index.yml | 6 + examples/eval_aa_lcr_gpt55_xhigh.py | 84 ++++++ .../configs/datasets/aa_lcr/__init__.py | 1 + .../configs/datasets/aa_lcr/aa_lcr_gen.py | 58 +++++ opencompass/datasets/__init__.py | 1 + opencompass/datasets/aa_lcr.py | 239 ++++++++++++++++++ tests/datasets/test_aa_lcr.py | 167 ++++++++++++ 7 files changed, 556 insertions(+) create mode 100644 examples/eval_aa_lcr_gpt55_xhigh.py create mode 100644 opencompass/configs/datasets/aa_lcr/__init__.py create mode 100644 opencompass/configs/datasets/aa_lcr/aa_lcr_gen.py create mode 100644 opencompass/datasets/aa_lcr.py create mode 100644 tests/datasets/test_aa_lcr.py diff --git a/dataset-index.yml b/dataset-index.yml index 04b03026d..0dc551a16 100644 --- a/dataset-index.yml +++ b/dataset-index.yml @@ -40,6 +40,12 @@ paper: https://arxiv.org/pdf/2406.10149 configpath: opencompass/configs/datasets/babilong configpath_llmjudge: '' +- aa_lcr: + name: AA-LCR + category: Long Context + paper: https://huggingface.co/datasets/ArtificialAnalysis/AA-LCR + configpath: opencompass/configs/datasets/aa_lcr/aa_lcr_gen.py + configpath_llmjudge: opencompass/configs/datasets/aa_lcr/aa_lcr_gen.py - bigcodebench: name: BigCodeBench category: Code diff --git a/examples/eval_aa_lcr_gpt55_xhigh.py b/examples/eval_aa_lcr_gpt55_xhigh.py new file mode 100644 index 000000000..3e4c1348d --- /dev/null +++ b/examples/eval_aa_lcr_gpt55_xhigh.py @@ -0,0 +1,84 @@ +from copy import deepcopy + +from mmengine.config import read_base + +from opencompass.models import OpenAISDK, OpenAISDKResponse +from opencompass.partitioners import NaivePartitioner +from opencompass.runners import LocalRunner +from opencompass.tasks import OpenICLEvalTask, OpenICLInferTask +from opencompass.utils.text_postprocessors import extract_non_reasoning_content + +with read_base(): + from opencompass.configs.datasets.aa_lcr.aa_lcr_gen import \ + aa_lcr_datasets + + +AA_LCR_NUM_RUNS = 3 +AA_LCR_MAX_OUT_LEN = 128000 +AA_LCR_MAX_SEQ_LEN = 1050000 +AA_LCR_TEST_RANGE = None + +work_dir = './output/aa_lcr_gpt55_xhigh' + +datasets = deepcopy(aa_lcr_datasets) +for dataset in datasets: + dataset['n'] = AA_LCR_NUM_RUNS + dataset['eval_cfg']['evaluator']['dataset_cfg']['n'] = 1 + + if AA_LCR_TEST_RANGE: + dataset['reader_cfg']['test_range'] = AA_LCR_TEST_RANGE + dataset['eval_cfg']['evaluator']['dataset_cfg']['reader_cfg'][ + 'test_range'] = AA_LCR_TEST_RANGE + + judge_cfg = dict( + abbr='qwen3-235b-a22b-instruct-2507', + type=OpenAISDK, + path='qwen3-235b-a22b-instruct-2507', + key='ENV', + query_per_second=1, + batch_size=10, + temperature=0, + tokenizer_path='gpt-4o-2024-05-13', + max_out_len=1024, + max_seq_len=131072, + retry=5, + timeout=3600, + ) + dataset['eval_cfg']['evaluator']['judge_cfg'] = judge_cfg + +models = [ + dict( + abbr='gpt-5.5-xhigh', + type=OpenAISDKResponse, + path='gpt-5.5', + key='ENV', + query_per_second=1, + batch_size=100, + tokenizer_path='gpt-4o-2024-05-13', + max_out_len=AA_LCR_MAX_OUT_LEN, + max_seq_len=AA_LCR_MAX_SEQ_LEN, + max_workers=20, + retry=5, + timeout=3600, + openai_extra_kwargs=dict(reasoning=dict(effort='xhigh')), + pred_postprocessor=dict(type=extract_non_reasoning_content), + ) +] + +infer = dict( + partitioner=dict(type=NaivePartitioner), + runner=dict( + type=LocalRunner, + max_num_workers=10, + task=dict(type=OpenICLInferTask), + ), +) + +eval = dict( + partitioner=dict(type=NaivePartitioner), + runner=dict( + type=LocalRunner, + max_num_workers=10, + task=dict(type=OpenICLEvalTask), + ), +) diff --git a/opencompass/configs/datasets/aa_lcr/__init__.py b/opencompass/configs/datasets/aa_lcr/__init__.py new file mode 100644 index 000000000..7db43f116 --- /dev/null +++ b/opencompass/configs/datasets/aa_lcr/__init__.py @@ -0,0 +1 @@ +from .aa_lcr_gen import aa_lcr_datasets # noqa: F401, F403 diff --git a/opencompass/configs/datasets/aa_lcr/aa_lcr_gen.py b/opencompass/configs/datasets/aa_lcr/aa_lcr_gen.py new file mode 100644 index 000000000..1829b4053 --- /dev/null +++ b/opencompass/configs/datasets/aa_lcr/aa_lcr_gen.py @@ -0,0 +1,58 @@ +from opencompass.datasets import AALCRDataset, aa_lcr_llmjudge_postprocess +from opencompass.evaluator import GenericLLMEvaluator +from opencompass.openicl.icl_inferencer import GenInferencer +from opencompass.openicl.icl_raw_prompt_template import RawPromptTemplate +from opencompass.openicl.icl_retriever import ZeroRetriever + + +AA_LCR_JUDGE_TEMPLATE = """Assess whether the following CANDIDATE ANSWER is \ +CORRECT or INCORRECT. +For the CANDIDATE ANSWER to be correct, it must be consistent with the \ +OFFICIAL ANSWER. + +The question, for reference only: {question} +The OFFICIAL ANSWER: {answer} +CANDIDATE ANSWER TO ASSESS: {prediction} + +Reply only with CORRECT or INCORRECT.""" + + +aa_lcr_reader_cfg = dict(input_columns=['prompt'], output_column='answer') + +aa_lcr_infer_cfg = dict( + prompt_template=dict( + type=RawPromptTemplate, + messages=[dict(role='user', content='{prompt}')], + ), + retriever=dict(type=ZeroRetriever), + inferencer=dict(type=GenInferencer), +) + +aa_lcr_eval_cfg = dict( + evaluator=dict( + type=GenericLLMEvaluator, + prompt_template=dict( + type=RawPromptTemplate, + messages=[dict(role='user', content=AA_LCR_JUDGE_TEMPLATE)], + ), + dataset_cfg=dict( + type=AALCRDataset, + path='ArtificialAnalysis/AA-LCR', + reader_cfg=aa_lcr_reader_cfg, + ), + judge_cfg=dict(), + dict_postprocessor=dict(type=aa_lcr_llmjudge_postprocess), + ), + pred_role='BOT', +) + +aa_lcr_datasets = [ + dict( + abbr='AA-LCR', + type=AALCRDataset, + path='ArtificialAnalysis/AA-LCR', + reader_cfg=aa_lcr_reader_cfg, + infer_cfg=aa_lcr_infer_cfg, + eval_cfg=aa_lcr_eval_cfg, + ) +] diff --git a/opencompass/datasets/__init__.py b/opencompass/datasets/__init__.py index a93690e23..a3596b664 100644 --- a/opencompass/datasets/__init__.py +++ b/opencompass/datasets/__init__.py @@ -1,3 +1,4 @@ +from .aa_lcr import AALCRDataset, aa_lcr_llmjudge_postprocess # noqa: F401 from .advancedIF import AdvancedIFDataset # noqa: F401 from .advancedIF import advancedif_rubric_postprocess # noqa: F401 from .advglue import * # noqa: F401, F403 diff --git a/opencompass/datasets/aa_lcr.py b/opencompass/datasets/aa_lcr.py new file mode 100644 index 000000000..642e2590d --- /dev/null +++ b/opencompass/datasets/aa_lcr.py @@ -0,0 +1,239 @@ +import csv +import os +import posixpath +import re +import unicodedata +import zipfile +from collections import defaultdict +from typing import Dict, List, Optional, Tuple + +from datasets import Dataset, DatasetDict + +from opencompass.registry import DICT_POSTPROCESSORS, LOAD_DATASET + +from .base import BaseDataset + +DEFAULT_REPO_ID = 'ArtificialAnalysis/AA-LCR' +DEFAULT_CSV_FILENAME = 'AA-LCR_Dataset.csv' +DEFAULT_ZIP_FILENAME = 'extracted_text/AA-LCR_extracted-text.zip' + +PROMPT_TEMPLATE = """BEGIN INPUT DOCUMENTS + +{documents_text} + +END INPUT DOCUMENTS + +Answer the following question using the input documents provided above. + +START QUESTION + +{question} + +END QUESTION +""" + + +def _split_semicolon_list(value: str) -> List[str]: + return [item.strip() for item in value.split(';') if item.strip()] + + +def _resolve_local_file(root: str, filename: str) -> Optional[str]: + candidates = [ + os.path.join(root, filename), + os.path.join(root, os.path.basename(filename)), + ] + for candidate in candidates: + if os.path.isfile(candidate): + return candidate + return None + + +def _resolve_data_files(path: str, csv_filename: str, + zip_filename: str) -> Tuple[str, str]: + if os.path.isdir(path): + csv_path = _resolve_local_file(path, csv_filename) + zip_path = _resolve_local_file(path, zip_filename) + if csv_path and zip_path: + return csv_path, zip_path + + if os.path.isfile(path): + csv_path = path + zip_path = _resolve_local_file(os.path.dirname(path), zip_filename) + if zip_path: + return csv_path, zip_path + + from huggingface_hub import hf_hub_download + + csv_path = hf_hub_download(repo_id=path, + filename=csv_filename, + repo_type='dataset') + zip_path = hf_hub_download(repo_id=path, + filename=zip_filename, + repo_type='dataset') + return csv_path, zip_path + + +def _doc_name(category: str, set_id: str, filename: str) -> str: + return posixpath.join('lcr', category, set_id, filename) + + +def _recover_zip_name(filename: str) -> str: + try: + return filename.encode('cp437').decode('utf-8') + except (UnicodeEncodeError, UnicodeDecodeError): + return filename + + +def _normalize_doc_name(filename: str) -> str: + return unicodedata.normalize('NFC', _recover_zip_name(filename)) + + +def _read_doc_from_zip(zip_file: zipfile.ZipFile, category: str, set_id: str, + filename: str) -> str: + doc_name = _doc_name(category, set_id, filename) + if doc_name in zip_file.namelist(): + with zip_file.open(doc_name) as f: + return f.read().decode('utf-8') + + recovered_names = { + _normalize_doc_name(zip_name): zip_name + for zip_name in zip_file.namelist() + } + zip_name = recovered_names.get(_normalize_doc_name(doc_name)) + if zip_name is not None: + with zip_file.open(zip_name) as f: + return f.read().decode('utf-8') + + raise FileNotFoundError(f'Cannot find AA-LCR document {filename!r} for ' + f'{category}/{set_id} in {zip_file.filename}.') + + +def _build_prompt(documents: List[str], question: str) -> str: + documents_text = '\n\n'.join( + 'BEGIN DOCUMENT {}:\n{}\nEND DOCUMENT {}'.format( + idx + 1, doc, idx + 1) for idx, doc in enumerate(documents)) + return PROMPT_TEMPLATE.format(documents_text=documents_text, + question=question) + + +@LOAD_DATASET.register_module() +class AALCRDataset(BaseDataset): + + @staticmethod + def load(path: str = DEFAULT_REPO_ID, + csv_filename: str = DEFAULT_CSV_FILENAME, + zip_filename: str = DEFAULT_ZIP_FILENAME): + csv_path, zip_path = _resolve_data_files(path, csv_filename, + zip_filename) + + raw_data = [] + with open(csv_path, encoding='utf-8', newline='') as f, \ + zipfile.ZipFile(zip_path) as zip_file: + reader = csv.DictReader(f) + for row_idx, row in enumerate(reader): + data_source_filenames = _split_semicolon_list( + row['data_source_filenames']) + data_source_urls = _split_semicolon_list( + row.get('data_source_urls', '')) + category = row['document_category'] + set_id = row['document_set_id'] + documents = [ + _read_doc_from_zip(zip_file, category, set_id, filename) + for filename in data_source_filenames + ] + question = row['question'] + answer = _split_semicolon_list(row['answer']) + + raw_data.append({ + 'id': int(row.get('', row_idx)), + 'question_id': int(row['question_id']), + 'document_category': category, + 'document_set_id': set_id, + 'question': question, + 'answer': answer, + 'data_source_filenames': data_source_filenames, + 'data_source_urls': data_source_urls, + 'input_tokens': int(row['input_tokens']), + 'prompt': _build_prompt(documents, question), + }) + + dataset = Dataset.from_list(raw_data) + return DatasetDict({'test': dataset, 'train': dataset}) + + +def _parse_aa_lcr_judgement(judgement: str) -> str: + normalized = judgement.strip().upper() + if re.search(r'\bINCORRECT\b', normalized): + return 'INCORRECT' + if re.search(r'\bCORRECT\b', normalized): + return 'CORRECT' + return 'UNKNOWN' + + +def _get_test_samples(dataset) -> List[Dict]: + if dataset is None: + return [] + try: + return list(dataset.reader.dataset['test']) + except AttributeError: + return [] + + +@DICT_POSTPROCESSORS.register_module() +def aa_lcr_llmjudge_postprocess(output: dict, + output_path: str, + dataset=None) -> dict: + samples = _get_test_samples(dataset) + details = [] + category_stats = defaultdict(lambda: {'correct': 0, 'total': 0}) + correct = 0 + incorrect = 0 + unknown = 0 + + for key, value in sorted(output.items(), key=lambda item: int(item[0])): + idx = int(key) + sample = samples[idx] if idx < len(samples) else {} + judge_response = value.get('prediction', '') + judgement = _parse_aa_lcr_judgement(judge_response) + + if judgement == 'CORRECT': + correct += 1 + elif judgement == 'INCORRECT': + incorrect += 1 + else: + unknown += 1 + + category = sample.get('document_category', 'unknown') + category_stats[category]['total'] += 1 + if judgement == 'CORRECT': + category_stats[category]['correct'] += 1 + + details.append({ + 'id': idx, + 'question_id': sample.get('question_id'), + 'document_category': category, + 'document_set_id': sample.get('document_set_id'), + 'question': sample.get('question'), + 'gold': sample.get('answer'), + 'candidate_answer': sample.get('prediction'), + 'judge_response': judge_response, + 'judgement': judgement, + 'correct': judgement == 'CORRECT', + }) + + total = len(details) + results = { + 'accuracy': correct / total * 100 if total else 0, + 'correct_count': correct, + 'incorrect_count': incorrect, + 'unknown_count': unknown, + 'total': total, + 'details': details, + } + for category, stats in category_stats.items(): + safe_category = re.sub(r'\W+', '_', category).strip('_') + results[f'accuracy_{safe_category}'] = (stats['correct'] / + stats['total'] * + 100 if stats['total'] else 0) + + return results diff --git a/tests/datasets/test_aa_lcr.py b/tests/datasets/test_aa_lcr.py new file mode 100644 index 000000000..71ebdbac3 --- /dev/null +++ b/tests/datasets/test_aa_lcr.py @@ -0,0 +1,167 @@ +import csv +import zipfile + +from datasets import Dataset + +from opencompass.configs.datasets.aa_lcr.aa_lcr_gen import \ + AA_LCR_JUDGE_TEMPLATE +from opencompass.datasets.aa_lcr import (AALCRDataset, + aa_lcr_llmjudge_postprocess) +from opencompass.openicl.icl_raw_prompt_template import RawPromptTemplate + + +def test_aa_lcr_local_loader_preserves_document_order(tmp_path): + csv_path = tmp_path / 'AA-LCR_Dataset.csv' + zip_dir = tmp_path / 'extracted_text' + zip_dir.mkdir() + zip_path = zip_dir / 'AA-LCR_extracted-text.zip' + + with csv_path.open('w', encoding='utf-8', newline='') as f: + writer = csv.DictWriter( + f, + fieldnames=[ + '', + 'document_category', + 'document_set_id', + 'question_id', + 'question', + 'answer', + 'data_source_filenames', + 'data_source_urls', + 'input_tokens', + ]) + writer.writeheader() + writer.writerow({ + '': '0', + 'document_category': 'Test_Category', + 'document_set_id': 'test_set', + 'question_id': '1', + 'question': 'What is the answer?', + 'answer': 'Second document wins.', + 'data_source_filenames': 'first.txt;second.txt', + 'data_source_urls': 'https://example.com/1;https://example.com/2', + 'input_tokens': '42', + }) + + with zipfile.ZipFile(zip_path, 'w') as zf: + zf.writestr('lcr/Test_Category/test_set/first.txt', 'first text') + zf.writestr('lcr/Test_Category/test_set/second.txt', 'second text') + + dataset = AALCRDataset.load(path=str(tmp_path))['test'] + + assert len(dataset) == 1 + sample = dataset[0] + assert sample['answer'] == ['Second document wins.'] + assert sample['data_source_filenames'] == ['first.txt', 'second.txt'] + assert sample['prompt'].index('BEGIN DOCUMENT 1:\nfirst text') < ( + sample['prompt'].index('BEGIN DOCUMENT 2:\nsecond text')) + assert sample['prompt'].strip().endswith('END QUESTION') + + +def test_aa_lcr_local_loader_recovers_zip_filename_mojibake(tmp_path): + csv_path = tmp_path / 'AA-LCR_Dataset.csv' + zip_dir = tmp_path / 'extracted_text' + zip_dir.mkdir() + zip_path = zip_dir / 'AA-LCR_extracted-text.zip' + + with csv_path.open('w', encoding='utf-8', newline='') as f: + writer = csv.DictWriter( + f, + fieldnames=[ + '', + 'document_category', + 'document_set_id', + 'question_id', + 'question', + 'answer', + 'data_source_filenames', + 'data_source_urls', + 'input_tokens', + ]) + writer.writeheader() + writer.writerow({ + '': '0', + 'document_category': 'Test_Category', + 'document_set_id': 'test_set', + 'question_id': '1', + 'question': 'What is the answer?', + 'answer': 'Recovered document.', + 'data_source_filenames': 'No cheat-codes needed – guide.txt', + 'data_source_urls': '', + 'input_tokens': '42', + }) + + with zipfile.ZipFile(zip_path, 'w') as zf: + zf.writestr( + 'lcr/Test_Category/test_set/' + 'No cheat-codes needed ΓÇô guide.txt', + 'recovered text', + ) + + sample = AALCRDataset.load(path=str(tmp_path))['test'][0] + + assert 'BEGIN DOCUMENT 1:\nrecovered text' in sample['prompt'] + + +def test_aa_lcr_llmjudge_postprocess(): + test_set = Dataset.from_list([ + { + 'question_id': 1, + 'document_category': 'Category A', + 'document_set_id': 'set_a', + 'question': 'Question A', + 'answer': 'Answer A', + 'prediction': 'Answer A', + }, + { + 'question_id': 2, + 'document_category': 'Category A', + 'document_set_id': 'set_a', + 'question': 'Question B', + 'answer': 'Answer B', + 'prediction': 'Wrong', + }, + ]) + + class _Reader: + pass + + _Reader.dataset = {'test': test_set} + + class _JudgeDataset: + reader = _Reader() + + scores = aa_lcr_llmjudge_postprocess( + { + '0': { + 'prediction': 'CORRECT' + }, + '1': { + 'prediction': 'INCORRECT' + }, + }, + output_path='unused.json', + dataset=_JudgeDataset(), + ) + + assert scores['accuracy'] == 50 + assert scores['correct_count'] == 1 + assert scores['incorrect_count'] == 1 + assert scores['accuracy_Category_A'] == 50 + + +def test_aa_lcr_judge_template_uses_opencompass_fields(): + template = RawPromptTemplate( + messages=[dict(role='user', content=AA_LCR_JUDGE_TEMPLATE)]) + + messages = template.generate_item({ + 'question': 'Which document wins?', + 'answer': ['Second document wins.'], + 'prediction': 'The second document wins.', + }) + content = messages[0]['content'] + + assert '{official_answer}' not in content + assert '{candidate_answer}' not in content + assert 'Second document wins.' in content + assert 'The second document wins.' in content