evalstone/evalscope/evalscope/benchmarks/aa_lcr/aa_lcr_adapter.py

# Copyright (c) Alibaba, Inc. and its affiliates.
# flake8: noqa: E501
import zipfile
from pathlib import Path
from typing import Any, Dict, List, Literal

from pydantic import BaseModel

from evalscope.api.benchmark import BenchmarkMeta, DefaultDataAdapter
from evalscope.api.dataset import Sample
from evalscope.api.judge import (
    CaseVerdict,
    JudgeCase,
    JudgeContext,
    JudgeDefinition,
    JudgeRequest,
    OutputContract,
    ReducedVerdict,
)
from evalscope.api.messages import ChatMessageUser
from evalscope.api.metric import Score
from evalscope.api.registry import register_benchmark
from evalscope.constants import DEFAULT_EVALSCOPE_CACHE_DIR, ScoringPolicy, Tags
from evalscope.utils.download_utils import download_url
from evalscope.utils.logger import get_logger

logger = get_logger()


# The judge prompt ends with "Reply only with CORRECT or INCORRECT."
class Grade(BaseModel):
    reasoning: str = ''
    verdict: Literal['CORRECT', 'INCORRECT']


GRADE_CONTRACT = OutputContract(schema_model=Grade)

# Default judge prompt template
JUDGE_PROMPT = """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: {correct_answer}
CANDIDATE ANSWER TO ASSESS: {response}

Reply only with CORRECT or INCORRECT."""

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
"""

# New constants for auto-download
DOWNLOAD_URL: str = (
    'https://modelscope.cn/datasets/evalscope/AA-LCR/resolve/master/extracted_text/AA-LCR_extracted-text.zip'
)
DEFAULT_CACHE_SUBDIR: str = 'aa_lcr'
DEFAULT_ZIP_NAME: str = 'AA-LCR_extracted-text.zip'
DEFAULT_EXTRACTED_DIR_NAME: str = 'lcr'


@register_benchmark(
    BenchmarkMeta(
        name='aa_lcr',
        pretty_name='AA-LCR',
        tags=[Tags.KNOWLEDGE, Tags.REASONING, Tags.LONG_CONTEXT],
        description="""
## Overview

AA-LCR (Artificial Analysis Long Context Retrieval) is a benchmark for evaluating long-context retrieval and reasoning capabilities of language models. It requires models to find and synthesize information across multiple documents.

## Task Description

- **Task Type**: Long-Context Question Answering
- **Input**: Multiple documents + question requiring cross-document reasoning
- **Output**: Answer synthesized from document information
- **Context**: Very long context (multiple documents concatenated)

## Key Features

- Tests long-context retrieval abilities
- Multiple document understanding
- Cross-document reasoning required
- LLM-based judging for answer correctness
- Auto-download of document corpus

## Evaluation Notes

- Default configuration uses **0-shot** evaluation
- Primary metric: **Accuracy** (via LLM judge)
- Evaluates on **test** split
- Documents auto-downloaded if `text_dir` not specified
- Judge prompt compares candidate answer against reference
""",  # noqa: E501
        dataset_id='evalscope/AA-LCR',
        metric_list=['acc'],
        few_shot_num=0,
        train_split=None,
        eval_split='test',
        prompt_template=PROMPT_TEMPLATE,
        extra_params={
            'text_dir': {
                'type': 'str | null',
                'description': 'Local directory containing extracted AA-LCR text files; if null will auto-download & extract.',
                'value': None,
            }
        },
    )
)
class AALCRAdapter(DefaultDataAdapter):
    scoring_policy = ScoringPolicy.JUDGE_ONLY

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        # Get extra parameters
        self.text_dir = self.extra_params.get('text_dir')

    def load(self):
        # Auto download and extract when text_dir is not provided
        if not self.text_dir:
            self.text_dir = self._ensure_text_dir_downloaded()
        elif not Path(self.text_dir).exists():
            raise ValueError(
                'AA-LCR text_dir does not exist: '
                f'{self.text_dir}. Please provide a valid directory or omit text_dir to auto-download.'
            )

        self.text_dir = Path(self.text_dir)
        return super().load()

    def _ensure_text_dir_downloaded(self) -> Path:
        """Ensure AA-LCR extracted texts are available locally; download and extract if missing."""
        cache_root = Path(DEFAULT_EVALSCOPE_CACHE_DIR) / DEFAULT_CACHE_SUBDIR
        extracted_dir = cache_root / DEFAULT_EXTRACTED_DIR_NAME

        if extracted_dir.exists():
            logger.info(f'AA-LCR documents found: {extracted_dir}')
            return extracted_dir

        cache_root.mkdir(parents=True, exist_ok=True)
        zip_path = cache_root / DEFAULT_ZIP_NAME

        try:
            logger.info(f'Downloading AA-LCR documents from {DOWNLOAD_URL} to {zip_path}...')
            download_url(DOWNLOAD_URL, str(zip_path))

            logger.info(f'Extracting {zip_path} to {cache_root}...')
            with zipfile.ZipFile(zip_path, 'r') as zf:
                zf.extractall(cache_root)

            if not extracted_dir.exists():
                raise ValueError(f'Extraction succeeded but target directory not found: {extracted_dir}')

            logger.info(f'AA-LCR documents ready at {extracted_dir}')
            return extracted_dir
        except Exception as e:
            raise ValueError(
                f'Failed to download or extract AA-LCR documents: {e}. '
                'You can also manually download and set extra_params["text_dir"].'
            ) from e
        finally:
            # Best-effort cleanup of the zip file
            try:
                if zip_path.exists():
                    zip_path.unlink()
            except Exception:
                pass

    def _get_context(self, record: Dict[str, Any]) -> str:
        doc_folder = self.text_dir / record['document_category'] / record['document_set_id']

        # Check if the document folder exists
        if not doc_folder.exists() or not doc_folder.is_dir():
            logger.warning(f'Document folder not found: {doc_folder}. Returning empty context.')
            return ''

        doc_blocks = []
        try:
            for file_path in doc_folder.iterdir():
                if file_path.is_file():
                    try:
                        content = file_path.read_text(encoding='utf-8').strip()
                        if content:
                            doc_blocks.append(content)
                    except (IOError, UnicodeDecodeError) as e:
                        logger.warning(f'Could not read file {file_path}, skipping: {e}')
        except OSError as e:
            logger.warning(f'Could not access document folder {doc_folder}: {e}')
            return f'ERROR: Could not read documents for {record["document_category"]}/{record["document_set_id"]}'

        documents_text = '\n\n'.join(
            f'BEGIN DOCUMENT {i + 1}:\n{doc}\nEND DOCUMENT {i + 1}' for i, doc in enumerate(doc_blocks)
        )
        return documents_text

    def record_to_sample(self, record: Dict[str, Any]) -> Sample:
        """Convert a record to a Sample with long-context prompt."""
        context = self._get_context(record)
        prompt = self.prompt_template.format(documents_text=context, question=record['question'])

        return Sample(
            input=[ChatMessageUser(content=prompt)],
            target=record['answer'],
            metadata={
                'question': record['question'],
                'data_source_urls': record['data_source_urls'],
                'input_tokens': record.get('input_tokens', 0),
            },
        )

    def judge_definition(self, context: JudgeContext) -> JudgeDefinition:

        def request(case, placement, completed_cases, judge_context) -> JudgeRequest:
            prompt = (
                JUDGE_PROMPT.format(
                    question=judge_context.task_state.metadata['question'],
                    correct_answer=judge_context.reference,
                    response=judge_context.filtered_prediction,
                )
                + case.output_contract.instruction()
            )
            return JudgeRequest(messages=[ChatMessageUser(content=prompt)])

        def reduce(case_verdicts, judge_context) -> ReducedVerdict:
            return ReducedVerdict(value={'acc': 1.0 if case_verdicts[0].value.verdict == 'CORRECT' else 0.0})

        return JudgeDefinition.workflow(
            cases=[JudgeCase(case_id='grade', output_contract=GRADE_CONTRACT)],
            request=request,
            reduce=reduce,
            main_score_name='acc',
        )