bench_tasks.py

import asyncio
from dataclasses import dataclass
import json
import os
import random
from typing import Any, Coroutine, Optional, TypeVar
import re
import traceback

from openai import AsyncClient, InternalServerError, LengthFinishReasonError
from pydantic import BaseModel

from bench_utils import (
    extract_block_spans,
    calculate_latex_command_text_ratio,
    check_span_overlapping,
    get_questions,
)

from arxiv import get_arxiv_title


SYSTEM_PROMPT = "Please read the following paper and answer the question below."


T = TypeVar('T')
BaseModelT = TypeVar('BaseModelT', bound=BaseModel)



class GlobalConfig:
    cot: bool = False
    one_shot_question: Optional[dict] = None
    add_context: bool = True
    invoke_extra_body_kwargs: dict = {}
    temperature: float = 0.0
    is_thinking_model: bool = False
    num_filters = None
    
    stat_counter = 0
    
    dry_run = False
    
    @staticmethod
    def kwargs():
        return dict(
            add_cot=GlobalConfig.cot,
            add_one_shot_question=GlobalConfig.one_shot_question,
            add_context=GlobalConfig.add_context,
        )


def get_invoke_model_name() -> str:
    if (model_name := os.getenv('EXT_MODEL')):
        return model_name
    else:
        return 'custom-model'
    

async def invoke_model(
    async_client: AsyncClient, messages: list[dict], 
    max_tokens: int, temperature = None, response_format = None,
    strip_thinking = True,
) -> Optional[str]:
    
    if GlobalConfig.dry_run:
        return ''
    
    for retry in range(10):
        try:
            response = await async_client.chat.completions.create(
                model=get_invoke_model_name(),
                messages=messages,
                temperature=temperature if temperature is not None else GlobalConfig.temperature,
                max_tokens=max_tokens + (2048 if GlobalConfig.is_thinking_model else 0),
                response_format=response_format,
                extra_body=GlobalConfig.invoke_extra_body_kwargs,
            )
        except InternalServerError as e:
            traceback.print_exception(e)
            return ''
        
        if response.choices is None or response.choices[0].message.content is None:
            print('No choices in response, maybe rate limit')
            await asyncio.sleep(5)
            continue
        
        content = response.choices[0].message.content
        
        if strip_thinking:
            if '</think>' in content:
                content = content[content.find('</think>') + len('</think>'):]
        
            assert '</think>' not in content
        return content
    
    return ''
    
    
async def invole_model_parsed(
    async_client: AsyncClient, messages: list[dict], 
    max_tokens: int, response_format: BaseModelT, temperature = None,
) -> Optional[BaseModelT]:
    if GlobalConfig.dry_run:
        return None
    
    try:
        response = await async_client.beta.chat.completions.parse(
            model=get_invoke_model_name(),
            messages=messages,
            temperature=temperature if temperature is not None else GlobalConfig.temperature,
            max_tokens=max_tokens + (8192 if GlobalConfig.is_thinking_model else 0),
            response_format=response_format,
            extra_body={
                "provider": {"require_parameters": True},
                **GlobalConfig.invoke_extra_body_kwargs,
            },
        )
    except InternalServerError as e:
        traceback.print_exception(e)
        return None
    except LengthFinishReasonError as e:
        traceback.print_exception(e)
        return None
    
    return response.choices[0].message.parsed
    

def check_answer_in_response(outputs, questions) -> int:
    correct = 0
    for output, question in zip(outputs, questions):
        answer = question[-1]
        if answer in output:
            correct += 1
    return correct
    

def extract_figure_labels(latex_content: str) -> list[str]:
    figure_pattern = r'\\begin{figure\*?}(.*?)\\end{figure\*?}'
    label_pattern = r'\\label{([^}]*)}'
    
    labels = []
    
    figure_matches = re.finditer(figure_pattern, latex_content, re.DOTALL)
    
    for match in figure_matches:
        figure_content = match.group(1)
        label_matches = re.finditer(label_pattern, figure_content)
        for label_match in label_matches:
            labels.append(label_match.group(1))
    
    return labels


@dataclass
class BenchTask:
    arxiv_id: str
    json_name: str
    add_one_shot_question: Optional[dict] = None
    add_context: bool = True
    
    def __post_init__(self):
        self.json_output_name = self.json_name
        
        if self.add_one_shot_question:
            if self.add_one_shot_question['type'] == 'bench':
                full_text = self.add_one_shot_question['bench']['text']
                question = get_questions(self.add_one_shot_question['bench'][self.json_name])[0]
                self.json_output_name += '_one_shot'
                self.one_shot_messages = [
                    {"role": "user", "content": self.build_query(full_text, question)},
                    {"role": "assistant", "content": self.build_answer(full_text, question)},
                ]
            else:
                assert 0
        else:
            self.one_shot_messages = []
        
        if not self.add_context:
            self.json_output_name += '_no_context'
    
    def get_paper_title_string(self) -> str:
        return f'[{self.arxiv_id}] {get_arxiv_title("https://arxiv.org/abs/" + self.arxiv_id)}\n'

    def generate_bench(self, **kwargs) -> list[tuple]:
        raise NotImplementedError()
    
    def run_task_cached(self, cache: dict, client: AsyncClient, full_text: str, questions: list[tuple]) -> list[Coroutine[Any, Any, T]]:
        cache_data = cache.get(self.json_output_name)
        if cache_data is None or GlobalConfig.dry_run:
            return self.run_task(client, full_text, questions)
        else:
            assert full_text == cache.get('text')
            assert questions == cache_data.get('questions')
            async def warp(x):
                return x
            return [warp(x) for x in cache_data['outputs']]
    
    def run_task(self, client: AsyncClient, full_text: str, questions: list[tuple]) -> list[Coroutine[Any, Any, T]]:
        raise NotImplementedError()
    
    def dump_task(self, text: str, questions: list[tuple], outputs: list[T]) -> dict:
        raise NotImplementedError()
    
    def build_query(self, full_text: str, question: tuple) -> str:
        raise NotImplementedError()
    
    def build_answer(self, full_text: str, question: tuple) -> str:
        raise NotImplementedError()


@dataclass
class SectionCountTask(BenchTask):
    json_name: str = 'section_count'
    
    def __post_init__(self):
        super().__post_init__()
        self.json_output_name += '_v2'
    
    def generate_bench(self, full_text: str, **kwargs) -> list[tuple]:
        pattern = r'\\section\b'
        matches = list(re.finditer(pattern, full_text, re.DOTALL))
        count = len(matches)
        return [(count,)]
    
    def run_task(self, client: AsyncClient, full_text: str, questions: list[tuple]) -> list[Coroutine[Any, Any, int]]:
        return [
            self.eval(client, full_text, questions[0])
        ]
    
    def build_query(self, full_text: str, question: tuple) -> str:
        if self.add_context:
            query = f'''\
```latex
{full_text}
```
'''
        else:
            query = self.get_paper_title_string()
            query += '''
Question: In the paper above, how many sections are there (with the "\section" command, including "\section*")? Please think step by step and give the answer at last. Put the answer number inside the <answer> </answer> tag.
'''
            
        return query
    
    def build_answer(self, full_text: str, question: tuple):
        pattern = r'\\section(?:\*|\[(?:[^]]*)\])?\{([^}]*)\}'
        matches = list(re.finditer(pattern, full_text, re.DOTALL))
        count = len(matches)
        
        res = "Let's count how many sections are there in the paper:\n"
        for idx, match in enumerate(matches, start=1):
            res += f'{idx}. {match.group(0)}\n'
        res += '\n'
        res += f'So, there are <answer>{count}</answer> sections are there.'
        return res
    
    async def eval(self, client: AsyncClient, full_text: str, question: tuple) -> int:
        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            *self.one_shot_messages,
            {"role": "user", "content": self.build_query(full_text, question)}
        ]
        
        response = await invoke_model(
            client,
            messages,
            max_tokens=1024,
            strip_thinking=False,
        )
        # print(messages)
        # print(response)
        match = re.search(r'<answer>\D*(\d+).*</answer>', response)
        if match is not None:
            section_count = int(match.group(1))
            return section_count, response
        else:
            return -1, response

    def dump_task(self, text, questions, outputs):
        pattern = r'\\section(?:\*|\[(?:[^]]*)\])?\{([^}]*)\}'
        matches = list(re.finditer(pattern, text, re.DOTALL))
        count = len(matches)
        output = outputs[0]
        if isinstance(output, tuple):
            answer, response = output
            kwargs = {'_response': response}
        else:
            answer = output
            kwargs = {}
        correct = answer == count
        return {
            'correct': int(correct),
            'accuracy': int(correct),
            'total': 1,
            'questions': questions,
            'outputs': [answer],
            **kwargs,
        }


@dataclass
class ParagraphOrderingTask(BenchTask):
    json_name: str = 'paragraph_ordering'
    
    def generate_bench(self, full_text: str, **kwargs) -> list[tuple]:
        block_spans = extract_block_spans(full_text)
        
        questions = []
        pattern = r'(.{20,}?)\n\s*\n+(.{20,}?)\n\s*\n+(.{20,}?)\n\s*\n+(.{20,}?)\n\s*\n+'
        
        for match in re.finditer(pattern, full_text, re.DOTALL):
            span = match.span(0)
            if not check_span_overlapping([span], block_spans) and \
                calculate_latex_command_text_ratio(match.group(1)) < 0.5 and \
                calculate_latex_command_text_ratio(match.group(2)) < 0.5 and \
                calculate_latex_command_text_ratio(match.group(3)) < 0.5 and \
                calculate_latex_command_text_ratio(match.group(4)) < 0.5 and \
                True:
                paragraphs = [
                    match.group(1),
                    match.group(2),
                    match.group(3),
                    match.group(4),
                ]
                indices = list(range(len(paragraphs)))
                random.shuffle(indices)
                questions.append((span, paragraphs, indices))
        
        random.shuffle(questions)

        return questions[:50]
    
    def run_task(self, client: AsyncClient, full_text: str, questions: list[tuple]) -> list[Coroutine[Any, Any, str]]:
        return [
            self.eval(client, full_text, *t)
            for t in questions
        ]
    
    def build_query(self, full_text: str, question: tuple) -> str:
        span, paragraphs, indices = question
        
        if self.add_context:
            modified_text = full_text[:span[0]] + '__MISSING_PARAGRAPHS__' + full_text[span[1]:]
            query = f'''\
```latex
{modified_text}
```
'''
        else:
            query = self.get_paper_title_string()

        query += f'''
Question: In the paper above, there is a __MISSING_PARAGRAPHS__ blank, where 4 paragraphs are taken out and shuffled. Please put them back in the correct order and answer in the exact form of 'X X X X', where each 'X' is the letter representing a paragraph.

The paragraphs are as follows:
A.
{paragraphs[indices[0]]}
 
B.
{paragraphs[indices[1]]}

C.
{paragraphs[indices[2]]}

D.
{paragraphs[indices[3]]}

Please think step by step and give the answer at last.
'''
        return query
    
    def build_answer(self, full_text: str, question: tuple) -> str:
        span, paragraphs, indices = question
        inv_indices = {v: k for k, v in enumerate(indices)}
        return ' '.join(chr(ord('A') + inv_indices[i]) for i in range(4))
    
    async def eval(
        self, client: AsyncClient, full_text: str,
        span, paragraphs, indices
    ) -> int:
        query = self.build_query(full_text, (span, paragraphs, indices))

        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            *self.one_shot_messages,
            {"role": "user", "content": query}
        ]
        
        response_content = await invoke_model(
            client,
            messages=messages,
            max_tokens=512,
        )
        return response_content

    def dump_task(self, text, questions, outputs):
        correct = 0
        for question, output in zip(questions, outputs):
            inv_indices = {v: k for k, v in enumerate(question[2])}
            seq = ' '.join(chr(ord('A') + inv_indices[i]) for i in range(4))
            correct += int(seq in output)
        return {
            'correct': correct,
            'accuracy': correct / len(questions) if len(questions) > 0 else float('nan'),
            'total': len(questions),
            'questions': questions,
            'outputs': outputs,
        }


@dataclass
class XRefTask(BenchTask):
    json_name: str = 'xref'
    
    def __post_init__(self):
        super().__post_init__()
        self.json_output_name += '_fast'
    
    def generate_bench(self, full_text: str, **kwargs) -> list[tuple]:
        figure_labels = extract_figure_labels(full_text)
        
        xref_questions = []
        for line in re.split(r'\n\s*\n+', full_text):
            matches = re.finditer(r'\\(?:auto|c|)ref{([a-zA-Z0-9\-_:.]+?)}', line)
            for match in matches:
                replaced = line[:match.start(1)] + '______' + line[match.end(1):]
                answer = match.group(1)
                
                if 'fig' in answer or answer in figure_labels:
                    continue
                
                if answer in full_text.replace(line, replaced):
                    xref_questions.append((line, replaced, answer))
        return xref_questions
    
    def run_task(self, client: AsyncClient, full_text: str, questions: list[tuple]) -> list[Coroutine[Any, Any, str]]:
        if self.add_one_shot_question and self.add_one_shot_question['type'] == 'same_paper':
            if len(questions) < 1:
                return []
            original, replaced = questions[0][:2]
            modified_text = full_text.replace(original, replaced)
            questions = [question for question in questions if question[0] in modified_text]
            return [
                self.eval_xref(client, full_text, *q[:2], (original, replaced))
                for q in questions
            ]
        else:
            return [
                self.eval_xref(client, full_text, *q[:2])
                for q in questions
            ]
    
    def build_query(self, full_text: str, question: tuple, add_text=True) -> str:
        original, replaced = question[:2]
        
        if self.add_context:
            if add_text:
                modified_text = full_text.replace(original, replaced)
                modified_text = re.sub(r'(\\(?:auto|c|)ref)\{([a-zA-Z0-9\-_:.]+?)\}', r'\1\{xxx\}', modified_text)
                query = f'''\
```latex
{modified_text}
```

'''
            else:
                assert full_text == '' 
                query = ''
        else:
            query = self.get_paper_title_string() + '\n'

        query += f'''\
Question: In the following paragraph, which reference should be filled in the blank (______)?
```latex
{replaced}
```'''
        return query
    
    def build_answer(self, full_text: str, question: tuple) -> str:
        return question[-1]
    
    async def eval_xref(self, async_client: AsyncClient, text: str, original: str, replaced: str, same_paper_one_shot_question = None) -> str:
        if same_paper_one_shot_question is not None:
            original_a, replaced_a = same_paper_one_shot_question[:2]
            modified_text = text.replace(original, replaced)
            query_a = self.build_query(modified_text, (original_a, replaced_a))
            answer_a = self.build_answer(modified_text, same_paper_one_shot_question)
            query_b = self.build_query('', (original, replaced), add_text=False)
            messages = [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": query_a},
                {"role": "assistant", "content": answer_a},
                {"role": "user", "content": query_b}
            ]
        else:
            query = self.build_query(text, (original, replaced))

            messages = [
                {"role": "system", "content": SYSTEM_PROMPT},
                *self.one_shot_messages,
                {"role": "user", "content": query}
            ]
        
        response_content = await invoke_model(
            async_client,
            messages=messages,
            max_tokens=512,
        )
        return response_content
    
    def dump_task(self, text, questions, outputs):
        questions = questions[-len(outputs):]
        
        correct = check_answer_in_response(outputs, questions)
        
        return {
            'correct': correct,
            'total': len(questions),
            'accuracy': correct / len(questions) if len(questions) > 0 else float('nan'),
            'questions': questions,
            'outputs': outputs,
        }


@dataclass
class CitationListTask(BenchTask):
    json_name: str = 'citation_list_v2'
    
    def __post_init__(self):
        super().__post_init__()
        self.json_output_name += '_fast'
    
    def generate_bench(self, full_text: str, ref_map: dict, **kwargs) -> list[tuple]:
        if len(ref_map) < 10:
            return []
        
        cite_pattern = r'\\(?:cite|citep|citet|citealp|citeauthor|citeyear)(?:\[.*?\])?\s*\{([^}]+)\}'
        
        questions = []
        
        sampled_keys = list(ref_map.keys())
        
        cite_info = json.dumps(
            [
                {
                    'ID': ID,
                    'title': ref_map[ID]['title'],
                    'abstract': ref_map[ID]['abstract'],
                }
                for ID in sampled_keys
            ],
            indent=4,
            ensure_ascii=False,
        )

        for line in re.split(r'\n\s*\n+', full_text):
            matches = re.finditer(cite_pattern, line, re.IGNORECASE)
            
            unique_keys = set()
            for match in matches:
                keys = match.group(1).split(',')
                for key in keys:
                    unique_keys.add(key.strip())
            
            for key in unique_keys:
                if key in ref_map:
                    replaced = line.replace(key, '______')
                    answer = key
                    questions.append((line, replaced, cite_info, answer))
        
        return questions
    
    def run_task(self, client: AsyncClient, full_text: str, questions: list[tuple]) -> list[Coroutine[Any, Any, str]]:
        return [
            self.eval_citations(client, full_text, *q[:3])
            for q in questions
        ]
    
    def build_query(self, full_text: str, question: tuple[str, str, str, str]) -> str:
        original, replaced, cite_info, _ = question
        
        if self.add_context:
            modified_text = full_text.replace(original, replaced)
            cite_pattern = r'(\\(?:cite|citep|citet|citealp|citeauthor|citeyear)(?:\[.*?\])?\s*)\{([^}]+)\}'
            modified_text = re.sub(cite_pattern, r'\1\{xxx\}', modified_text)
            query = f'''\
```latex
{modified_text}
```

'''
        else:
            query = self.get_paper_title_string() + '\n'

        query += f'''\
Here is the information of some possible citations:
```json
{cite_info}
```

Question: In the following paragraph, which citation id should be filled in the blank (______)?
```latex
{replaced}
``` 
'''
        
        return query
    
    def build_answer(self, full_text: str, question: tuple[str, str, str, str]) -> str:
        class CitationAnswer(BaseModel):
            citation_id: str | None = None
        return CitationAnswer(citation_id=question[3]).model_dump_json(indent=4)
    
    async def eval_citations(self, async_client: AsyncClient, text: str, original: str, replaced: str, cite_info: str) -> str:
        query = self.build_query(text, (original, replaced, cite_info, None))

        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            *self.one_shot_messages,
            {"role": "user", "content": query}
        ]
        
        response_content = await invoke_model(
            async_client,
            messages=messages,
            max_tokens=512,
        )
        return response_content
    
    def dump_task(self, text, questions, outputs):
        correct = check_answer_in_response(outputs, questions)
        
        return {
            'correct': correct,
            'total': len(questions),
            'accuracy': correct / len(questions) if len(questions) > 0 else float('nan'),
            'questions': questions,
            'outputs': outputs,
        }


@dataclass
class FillNumberTask(BenchTask):
    json_name: str = 'fill_number'
    
    def __post_init__(self):
        super().__post_init__()
        self.json_output_name += '_v2'
    
    def generate_bench(self, full_text: str, **kwargs) -> list[tuple[int, int]]:
        figure_table_spans = extract_block_spans(full_text)
        result_spans = self.extract_result_spans(full_text)
        
        questions = []
        number_pattern = r'\b(-?\d+(?:\.\d+)?)\b'
        
        for match in re.finditer(number_pattern, full_text):
            span = match.span(1)
            if not check_span_overlapping([span], figure_table_spans) and \
                check_span_overlapping([span], result_spans):
                questions.append(span)
        
        random.shuffle(questions)

        return questions[:50]
    
    def extract_result_spans(self, latex_content: str) -> list[tuple[int, int]]:
        spans = []
        
        pattern = r'(\\section{[^}]*(?:Result)[^}]*}.*?)(?:\\section|$)'
        matches = re.finditer(pattern, latex_content, re.DOTALL)
        for match in matches:
            span = match.span(1)
            spans.append(span)
        
        pattern = r'(\\subsection{[^}]*(?:Result)[^}]*}.*?)(?:\\subsection|$)'
        matches = re.finditer(pattern, latex_content, re.DOTALL)
        for match in matches:
            span = match.span(1)
            spans.append(span)
        
        return spans
    
    def run_task(self, client: AsyncClient, full_text: str, questions: list[tuple[int, int]]) -> list[Coroutine[Any, Any, float]]:
        return [
            self.eval_fill_number(client, full_text, q)
            for q in questions
        ]
    
    def build_query(self, full_text: str, question: tuple[int, int]) -> str:
        if self.add_context:
            modified_text = full_text[:question[0]] + '___FILL_HERE___' + full_text[question[1]:]
            query = f'''\
```latex
{modified_text}
```'''
        else:
            query = self.get_paper_title_string()
        
            query += '''
Question: In the paper above, there is a blank number with placeholder (___FILL_HERE___). Please think step by step and fill in the blank with the most appropriate number. You should put the answer number between the "<answer> </answer>" tag, such as "<answer>2.0</asnwer>".
'''
        return query
    
    def build_answer(self, full_text: str, question: tuple[int, int]) -> str:
        number = float(full_text[question[0]:question[1]])
        return f"It seems that the number should be filled is <answer>{number}</answer>."
    
    async def eval_fill_number(self, async_client: AsyncClient, text: str, span: tuple[int, int]) -> float:
        query = self.build_query(text, span)

        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            *self.one_shot_messages,
            {"role": "user", "content": query}
        ]
        
        response = await invoke_model(
            async_client,
            messages,
            max_tokens=1024,
        )
        # print(messages)
        # print(response)
        match = re.search(r'<answer>\D*\b(-?\d+(?:\.\d+)?)\b.*</answer>', response)
        if match is not None:
            number = float(match.group(1))
            return number
        else:
            return ''
    
    def dump_task(self, text, questions, outputs):
        correct = self.compare_numbers(
            outputs,
            [text[start:end] for start, end in questions]
        )
        
        return {
            'correct': correct,
            'total': len(questions),
            'accuracy': correct / len(questions) if len(questions) > 0 else float('nan'),
            'questions': questions,
            'outputs': outputs,
        }
    
    def compare_numbers(self, outputs, answers) -> int:
        correct = 0
        for output, answer in zip(outputs, answers):
            if output == float(answer):
                correct += 1
        return correct