### Python: Environment Setup and Model Initialization for AIMo2

Source: https://github.com/imagination-research/aimo2/blob/master/imagination_aimo2/backup_submission.ipynb

This script sets up the Python environment by modifying the PYTHONPATH to include a specific Kaggle input directory. It then initializes a language model actor (`LLMActor`) and a math solver (`MathSolver`) with specified generation and model configurations. Finally, it sets up an inference server for the AIMo2 competition.

```python
import sys
import os

original_pythonpath = os.environ.get("PYTHONPATH", "")
new_path = "/kaggle/input/lmdeploy-package"
merged_pythonpath = f"{new_path}:{original_pythonpath}" if original_pythonpath else new_path
os.environ["PYTHONPATH"] = merged_pythonpath

if __name__ == '__main__':
    
    gen_config = GenerationConfig(
        temperature=0.9,
        min_p=0.1,
        skip_special_tokens=True,
        max_new_tokens=16384,
        top_p=0.95,
        do_sample=True,
        repetition_penalty=1.05,
    )

    model_config = ModelConfig(
        model_path=llm_model_pth_14_3_16, 
        gpu_indices=[0,1,2,3],
        gpu_memory_utilization=0.97, 
        max_model_len=20000, 
        quant_policy=8,
        max_batch_size=max_batch_size,
        num_samples=num_samples,
    )
    
    print("loading model 1...",flush=True)
    actor1 = LLMActor(model_config)
    actor1.is_ready()

    math_solver = MathSolver(actor1, gen_config)
    
    inference_server = kaggle_evaluation.aimo_2_inference_server.AIMO2InferenceServer(predict)

    if os.getenv('KAGGLE_IS_COMPETITION_RERUN'):
        inference_server.serve()
    else:
        inference_server.run_local_gateway(
            (
                '/kaggle/input/ai-mathematical-olympiad-progress-prize-2/reference.csv',
            )
        )
```

--------------------------------

### Kaggle Inference Server Setup and Prediction (Python)

Source: https://context7.com/imagination-research/aimo2/llms.txt

Sets up a Kaggle-compatible inference server using Python. It defines a prediction function that utilizes an EarlyStopActor for generating answers based on questions and IDs. The server can be run for online evaluation or local testing.

```python
import polars as pl
from imagination_aimo2.local_eval_kaggle import EarlyStopActor, predict
import kaggle_evaluation.aimo_2_inference_server

# Global configuration (embedded in submission notebook)
config = {
    "main_model": {...},
    "actor": {...},
    "early_stop_strategy": {...},
}

# Create actor with early stopping
actor = EarlyStopActor(config, **config["actor"], early_stop_strategy=config["early_stop_strategy"])

# Prediction function for Kaggle API
def predict(id_: pl.DataFrame, question: pl.DataFrame, answer: pl.DataFrame = None) -> pl.DataFrame:
    id_val = id_.item(0)
    question_val = question.item(0)

    prediction, *_ = actor.predict_for_question(question_val, id_val)
    return pl.DataFrame({"id": id_val, "answer": prediction})

# Start inference server
inference_server = kaggle_evaluation.aimo_2_inference_server.AIMO2InferenceServer(predict)

import os
if os.getenv("KAGGLE_IS_COMPETITION_RERUN"):
    inference_server.serve()  # Online evaluation
else:
    inference_server.run_local_gateway(("data/reference.csv",))  # Local testing

```

--------------------------------

### SFT and DPO Training (Bash)

Source: https://context7.com/imagination-research/aimo2/llms.txt

Details the two-stage training process using the 360-LLaMA-Factory framework for long-context mathematical reasoning. Stage 1 involves Supervised Fine-Tuning (SFT) on a combined dataset, and Stage 2 uses Direct Preference Optimization (DPO) on filtered pairs, with specific parameters for each stage including model, dataset, epochs, learning rate, and sequence parallelism.

```bash
# Stage 1: Supervised Fine-Tuning (SFT)
# Initialize and update the 360-LLaMA-Factory submodule first
git submodule update --init --recursive
cd 360-LLaMA-Factory

bash ../scripts/run_sft.sh
# Key parameters:
# - Base model: deepseek-r1-distill-qwen-14b
# - Dataset: sft-4k (Light-R1 + LIMO data)
# - Epochs: 8, Learning rate: 1e-5
# - Context length: 18000 tokens

# Stage 2: Direct Preference Optimization (DPO)
bash ../scripts/run_dpo.sh
# Key parameters:
# - Base model: dpsk-14b-sft (from Stage 1)
# - Dataset: dpo-1 (filtered pairs)
# - Epochs: 4, Learning rate: 5e-7
# - Sequence parallelism: 8 for long context
# - Beta: 0.3, Loss: nca_pair
```

--------------------------------

### Initialize BasicActor for Question Solving

Source: https://context7.com/imagination-research/aimo2/llms.txt

Demonstrates initializing the BasicActor with a configuration dictionary that includes model and inference settings, prompt lists, and generation configurations. It then uses the actor to predict an answer for a given question.

```python
import yaml
from imagination_aimo2.local_eval import BasicActor

# Load configuration from YAML
config = {
    "main_model": {
        "model_cfg": {"model_path": "models/dpsk-qwen-14b-awq"},
        "inference_cfg": {
            "max_batch_size": 32,
            "quant_policy": 8,
            "max_prefill_token_num": 8192,
            "gen_max_new_tokens": 32768,
        }
    },
    "actor": {
        "model_dict": {"main_model": "main_model"},
        "prompt_list_combine": "interleave",
        "prompt_list": [
            {
                "system": "You are a helpful math assistant. Please reason step by step to put the answer in \boxed{}.",
                "user_suffix": "\nYou excel at reasoning.\nYou must put the final answer in \boxed{}.",
                "number": 8,
            },
            {
                "system": "You are a helpful math assistant. Please provide the python code to solve the math problem.",
                "user_suffix": "\nYou excel at coding.\nProvide python code.\nImport necessary libraries.",
                "number": 8,
            }
        ],
        "gen_cfg": {
            "temperature": 0.9,
            "min_p": 0.1,
            "max_new_tokens": 16384,
        }
    }
}

# Initialize actor and solve a question
actor = BasicActor(config, **config["actor"])

question = "Find the sum of all positive integers n such that n^2 - 19n + 99 is a perfect square."
result = actor.predict_for_question(question, id_="q1")

aggregated_answer, cot_answers, code_answers, out_lens, python_codes, code_errors, outputs, duration = result
print(f"Answer: {aggregated_answer}, Time: {duration:.2f}s")
```

--------------------------------

### Execute SFT and DPO Training Pipelines

Source: https://github.com/imagination-research/aimo2/blob/master/README.md

Shell scripts to manage the training lifecycle of the model. These scripts handle Supervised Fine-Tuning (SFT) and Direct Preference Optimization (DPO) data filtering and training.

```bash
# Stage 1 SFT training
bash scripts/run_sft.sh

# Stage 2 DPO data filtering
bash scripts/run_filter.sh

# Stage 2 DPO training
bash scripts/run_dpo.sh
```

--------------------------------

### Initialize EarlyStopActor with Dynamic Strategies

Source: https://context7.com/imagination-research/aimo2/llms.txt

Initializes the EarlyStopActor, inheriting from BasicActor, with additional configurations for early stopping strategies. This includes sample-level and question-level stopping rules, consistency checks, and dynamic speed adjustments based on remaining time.

```python
from imagination_aimo2.local_eval_kaggle import EarlyStopActor

config = {
    "main_model": {
        "model_cfg": {"model_path": "models/dpsk-qwen-14b-awq"},
        "inference_cfg": {"max_batch_size": 32, "quant_policy": 8, "gen_max_new_tokens": 32768}
    },
    "actor": {
        "model_dict": {"main_model": "main_model"},
        "prompt_list": [
            {"system": "Reason step by step...", "number": 8},
            {"system": "Provide python code...", "number": 8}
        ],
        "gen_cfg": {"temperature": 0.9, "max_new_tokens": 16384}
    },
    "early_stop_strategy": {
        "every_n": 20,  # Check every 20 tokens
        "consistency_rules": [
            {"max_answers": 4, "min_repeats": 4, "message": "4/4 answers agree"},
            {"recent": 5, "min_repeats": 4, "message": "4/5 recent answers agree"},
            {"max_answers": 7, "min_repeats": 5, "message": "5/7 answers agree"},
        ],
        "speed_settings": {
            "very_fast": {"min_answers": 6, "timeout_minutes": 9.5},
            "fast": {"min_answers": 7, "timeout_minutes": 9.5},
            "normal": {"min_answers": 8, "timeout_minutes": 11},
            "slow": {"min_answers": 9, "timeout_minutes": 12},
        }
    }
}

actor = EarlyStopActor(
    config,
    **config["actor"],
    early_stop_strategy=config["early_stop_strategy"]
)

# Actor automatically adjusts speed based on remaining time
# Speed 1 (very_fast): 10 samples, Speed 3 (normal): 15 samples, Speed 5 (slow): 17 samples

```

--------------------------------

### Initialize Inference Environment and Model Paths

Source: https://github.com/imagination-research/aimo2/blob/master/imagination_aimo2/backup_submission.ipynb

Sets up the Python environment by configuring environment variables, importing necessary libraries like lmdeploy and vllm, and defining paths for various DeepSeek-R1 model variants. It also initializes a random seed and sets a global cutoff time for inference tasks.

```python
import os
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
import sys
sys.path.append('/kaggle/input/lmdeploy-package')
from lmdeploy import pipeline, TurbomindEngineConfig, GenerationConfig
from transformers import set_seed
import time
import warnings
import pandas as pd

warnings.simplefilter('ignore')
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ['PYTHONUNBUFFERED'] = '1'

random_seed = int(time.time()) % 10000
set_seed(random_seed)

llm_model_pth_14 = "/kaggle/input/deepseek-r1-14b-awq-turobmind/other/default/1/deepseek-14b"
llm_model_pth_14_3_13 = "/kaggle/input/dpsk-14b-sft-dpo-3-13-awq-tb/keras/default/1/dpsk-14b-sft-dpo4-3-13-awq-tb"
llm_model_pth_14_3_16 = "/kaggle/input/dpsk-14b-sft-dpo-3-16-awq-tb/keras/default/1/dpsk-14b-sft-dpo2-3-16-awq-tb"
```

--------------------------------

### Analysis and Visualization Tools (Bash)

Source: https://context7.com/imagination-research/aimo2/llms.txt

Provides command-line tools for analyzing experiment results, including visualizing early-stopping effectiveness and generating comparison tables. These scripts process directories containing experiment results to produce plots and markdown tables summarizing performance metrics.

```bash
# Analyze early stopping effectiveness with visualizations
python scripts/analyze_early_stop.py results/exp1 results/exp2 results/exp3

# Generates for each result directory:
# - answer_refine_vis.pdf: Stack bar plot showing answer refinement patterns
# - answer_token_length.pdf: Token length comparison for early stopping
# - outputs_per_question/<idx>_answer_status.pdf: Per-question heatmaps

# Generate markdown comparison table
python scripts/plot_table_local.py results/exp1 results/exp2 results/exp3

# Output includes:
# | Model | Quantization | Total solving time | Aggregated correct | Code error breakdown |
```

--------------------------------

### Configure LMDeploy TurbomindEngine with AWQ Quantization

Source: https://context7.com/imagination-research/aimo2/llms.txt

Sets up configuration for LMDeploy's TurbomindEngine, specifying model path, GPU indices, inference parameters like batch size and quantization policy (AWQ), and generation configuration for streaming inference.

```python
from lmdeploy import TurbomindEngineConfig, GenerationConfig

# Model configuration for AWQ-quantized model
model_cfg = {
    "model_path": "models/dpsk-qwen-14b-finetune-v1-epoch4-awq",
    "gpu_indices": [0, 1],  # Use specific GPUs
}

inference_cfg = {
    "max_batch_size": 32,
    "quant_policy": 8,  # 8-bit KV cache quantization (0 for fp16)
    "enable_prefix_caching": True,
    "cache_max_entry_count": 0.95,
    "max_prefill_token_num": 8192,
    "gen_max_new_tokens": 32768,
}

# Generation configuration
gen_cfg = GenerationConfig(
    temperature=0.9,
    min_p=0.1,
    top_p=0.95,
    do_sample=True,
    skip_special_tokens=True,
    max_new_tokens=16384,
    repetition_penalty=1.05,
)
```

--------------------------------

### Analyze Model Performance and Early Stopping

Source: https://github.com/imagination-research/aimo2/blob/master/README.md

Python scripts to analyze evaluation results, generate performance tables, and visualize token lengths or correctness to inform early-stopping strategies.

```python
# Generate Markdown summary table
python scripts/plot_table_local.py results/experiment_dir

# Analyze early stopping metrics and generate PDFs
python scripts/analyze_early_stop.py results/experiment_dir
```

--------------------------------

### LLM Inference Workflow with lmdeploy

Source: https://github.com/imagination-research/aimo2/blob/master/kaggle_forum.md

This snippet outlines the overall inference workflow using the lmdeploy framework. It involves prompt preparation (CoT and Code prompts), batch generation of multiple samples using `stream_infer`, and continuous extraction and aggregation of answers with early stopping mechanisms based on sample and question-level checks.

```python
# Assume 'lmdeploy' is imported and initialized
# Assume 'question' is the input query
# Assume 'cot_prompt_template' and 'code_prompt_template' are defined

# Prompt Preparation Task
num_samples = 15

cot_samples = []
code_samples = []
for _ in range(num_samples):
    # Prepare CoT prompt
    cot_prompt = cot_prompt_template.format(user_input=question)
    cot_samples.append(cot_prompt)
    # Prepare Code prompt
    code_prompt = code_prompt_template.format(user_input=question)
    code_samples.append(code_prompt)

all_prompts = cot_samples + code_samples

# LLM Generation Task
time_limit = get_remaining_time()
adjusted_num_samples = adjust_samples_based_on_time(num_samples, time_limit)

# Use lmdeploy's stream_infer for batch generation
# The iterator yields results for each sample
response_iterator = lmdeploy.create_engine().stream_infer(
    prompts=all_prompts[:adjusted_num_samples],
    # other parameters like temperature, top_p, etc.
)

# Answer Extraction and Aggregation (within the loop or after)
aggregated_answer = None
for i, output in enumerate(response_iterator):
    # Sample-level checking and early stopping
    if should_early_stop_sample(output, N_yields):
        # Handle early stopped sample
        continue

    # Extract answer from the current sample's output
    extracted_answer = extract_answer(output)

    # Question-level checking and aggregation
    if is_sample_complete(output):
        if aggregated_answer is None:
            aggregated_answer = aggregate_answers([extracted_answer]) # Initial aggregation
        else:
            aggregated_answer = aggregate_answers([extracted_answer], aggregated_answer)

        # Early stop all remaining samples if question is answered
        if should_early_stop_question(aggregated_answer):
            break

# Return the final aggregated answer
return aggregated_answer

```

--------------------------------

### Code Prompt Configuration for Math Assistant

Source: https://github.com/imagination-research/aimo2/blob/master/README.md

Defines the system and user suffix for a Code prompt, instructing the model to act as a math assistant, provide Python code to solve the problem, and format the final answer within a \boxed{} tag. It specifies that the answer must be an integer, less than or equal to 1000 (modulo 1000), and requires importing necessary libraries.

```yaml
# Code prompt
- system: "You are a helpful math assistant. Please provide the python code to solve the math problem and also put the final answer in \boxed{}."
  user_suffix: "\nYou excel at coding
You must provide the python code, avoid redundant analysis.
If the final answer is greater than 1000, then take the modulo of 1000.
The answer must be integer.
There is only one answer for each question.
Import necessary libraries."
```

--------------------------------

### Perform Local Evaluation with Configuration

Source: https://github.com/imagination-research/aimo2/blob/master/README.md

Helper script to execute local evaluation using a specific configuration file and dataset. Results are automatically saved to a structured directory based on the configuration and seed.

```bash
SEED=42 bash scripts/run_cfg.sh cfgs/example.yaml data/aime_2025_30.csv
```

--------------------------------

### CoT Prompt Configuration for Math Assistant

Source: https://github.com/imagination-research/aimo2/blob/master/README.md

Defines the system and user suffix for a Chain-of-Thought (CoT) prompt, instructing the model to act as a math assistant, reason step-by-step, and format the final answer within a \boxed{} tag. It includes a condition to take the modulo 1000 if the answer exceeds this value.

```yaml
- system: "You are a helpful math assistant. Please reason step by step to put the answer in \boxed{}."
  user_suffix: "\nYou excel at reasoning.
You must put the final answer in \boxed{}.
If the final answer is greater than 1000, then take the modulo of 1000.
Think carefully and thoroughly, avoid duplication."
```

--------------------------------

### CoT Prompt Template for Math Reasoning

Source: https://github.com/imagination-research/aimo2/blob/master/kaggle_forum.md

This YAML snippet defines a CoT (Chain-of-Thought) prompt template for a math assistant. It includes a system message instructing the model to reason step-by-step and place the final answer in a \boxed{} tag. It also specifies a modulo operation for answers exceeding 1000 and encourages thorough thinking.

```yaml
- system: "You are a helpful math assistant. Please reason step by step to put the answer in \boxed{}."
  user_suffix: "\nYou excel at reasoning.\nYou must put the final answer in \boxed{}.\nIf the final answer is greater than 1000, then take the modulo of 1000.\nThink carefully and thoroughly, avoid duplication."
```

--------------------------------

### Initialize LLMActor with Model Configuration

Source: https://github.com/imagination-research/aimo2/blob/master/imagination_aimo2/backup_submission.ipynb

Initializes the LLMActor with a given ModelConfig. This involves setting up environment variables for GPU visibility, configuring the Turbomind engine, and loading the specified model.

```python
class LLMActor:
    """LLM interaction class"""
    
    def __init__(self, config: ModelConfig):
        os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(map(str, config.gpu_indices))
        
        self._ready = True
        self.model_path = config.model_path
        self.backend_config = TurbomindEngineConfig(
            tp=len(config.gpu_indices),
            download_dir=config.cache_dir,
            max_batch_size=config.max_batch_size,
            enable_prefix_caching=config.enable_prefix_caching,
            cache_max_entry_count=config.gpu_memory_utilization,
            session_len=config.max_model_len,
            max_prefill_token_num=config.max_prefill_token_num,
            quant_policy=8,
        )
        self.pipe = pipeline(
            config.model_path,
            self.backend_config
        )
    
```

--------------------------------

### Configure 8-bit KV Cache Quantization with lmdeploy

Source: https://github.com/imagination-research/aimo2/blob/master/kaggle_forum.md

This configuration snippet demonstrates how to enable 8-bit KV Cache quantization within the lmdeploy framework. This is achieved by setting the 'main_model.inference_cfg.quant_policy' parameter to 8. This quantization method is beneficial for reducing memory usage and potentially increasing inference speed.

```yaml
main_model:
  inference_cfg:
    quant_policy: 8
```

--------------------------------

### Initialize System Path and Global Variables (Python)

Source: https://github.com/imagination-research/aimo2/blob/master/imagination_aimo2/backup_submission.ipynb

Initializes the Python system path to include a specific directory and sets up global variables for tracking rounds, scores, and averages. This is typically used at the beginning of a script or module to configure the environment.

```Python
import sys
sys.path.append('/kaggle/input/lmdeploy-package')
max_round = 1
g_score = 0
g_count = 0
total_avg_score = 0.0
total_avg_length = 0.0
total_solving_time = 0
```

--------------------------------

### AWQ Model Quantization (Bash and Python)

Source: https://context7.com/imagination-research/aimo2/llms.txt

Performs 4-bit AWQ (Activation-aware Weight Quantization) on models for efficient inference. This can be done via a bash script or programmatically using Python with the `awq` library. The process involves specifying input and output model paths, calibration sequence length, and quantization configuration.

```bash
# Quantize a model using AWQ
python scripts/quant_awq.py \
    /path/to/model/dpsk-14b-sft-dpo \
    --out-model-path /path/to/model/dpsk-14b-sft-dpo-awq \
    --calib-seq-len 128
```

```python
# Equivalent Python usage
from transformers import AutoTokenizer
from awq import AutoAWQForCausalLM

quant_config = {
    "zero_point": True,
    "q_group_size": 128,
    "w_bit": 4,
    "version": "GEMM",
}

model = AutoAWQForCausalLM.from_pretrained("models/dpsk-14b-sft-dpo", device_map="auto")
tokenizer = AutoTokenizer.from_pretrained("models/dpsk-14b-sft-dpo", trust_remote_code=True)

model.quantize(tokenizer, quant_config=quant_config, max_calib_seq_len=128, max_calib_samples=1024)
model.save_quantized("models/dpsk-14b-sft-dpo-awq")
tokenizer.save_pretrained("models/dpsk-14b-sft-dpo-awq")
```

--------------------------------

### Generate Text Completions

Source: https://github.com/imagination-research/aimo2/blob/master/imagination_aimo2/backup_submission.ipynb

Generates text completions for a list of input texts using the configured LLM. It takes generation configuration as input and returns a list of generated text strings.

```python
def generate(self, texts, gen_config: GenerationConfig):
    """Generate completions for the given texts"""
    response = self.pipe(
        texts,
        gen_config
    )
    return [r.text for r in response]
```

--------------------------------

### Python: Stream Generation with Early Stopping

Source: https://github.com/imagination-research/aimo2/blob/master/imagination_aimo2/backup_submission.ipynb

Manages the streaming of responses from a language model, processing tokens, and implementing early stopping based on various criteria. It handles multiple prompts concurrently and tracks output completion.

```python
def stream_generate(self, messages, gen_config, current_speed):
        """Stream generation with early stopping based on specific criteria"""
        text_extractor = TextExtractor()
        python_repl = PythonREPL()
        cot_answers = []
        code_answers = []
        valid_answers = []
        outputs = [""] * len(messages)  # Store complete output for each prompt
        token_counts = [0] * len(messages)  # Store token count for each prompt
        completed_status = [False] * len(messages)  # Flag to mark if each prompt is completed
        check_token_markers = [0] * len(messages) 
        start_time = time.time()
        session_id_start = next(self.pipe._session_id)
        print(f"Starting session ID: {session_id_start}", flush=True)

        try:
            for response in self.pipe.stream_infer(messages, gen_config):  
                if all(completed_status):  # Stop if all prompts are completed
                    print("[End]: All outputs completed.", flush=True)
                    break

                # Safely access index with error handling
                try:
                    index = response.index if response is not None else 0  # Get current output index
                    if index >= len(messages):
                        print(f"[Warning]: Received index {index} outside range of messages ({len(messages)})", flush=True)
                        continue
                        
                    session_id = session_id_start + index

                    if not completed_status[index]:
                        # Safely append text with None check
                        if response.text is not None:
                            outputs[index] += response.text  # Append current token to corresponding prompt output
                            token_counts[index] += 1  # Increment token count
                            
                            if self._should_stop_timeout(valid_answers, start_time, current_speed):
                                break     
                            # Early stopping logic
                            # Check if </think> is found in the current output
                            # Check for stopping criteria every 100 tokens
                            if token_counts[index] - check_token_markers[index] >= 20:
                                check_token_markers[index] = token_counts[index]
                                
                                # For chain-of-thought (index % 2 == 0), check for \boxed{num}
                                if index % 2 == 0 or if_only_cot:
                                    import re
                                    # Look for \boxed{num} where num is an integer
                                    boxed_pattern = re.search(r'oxed{(\d+)}', outputs[index])
                                    if boxed_pattern:
                                        asyncio.run(self._stop_one_session(self.pipe, session_id))
                                        completed_status[index] = True
                                        current_time = time.time()
                                        time_consumed = current_time - start_time
                                        speed = token_counts[index]/time_consumed if time_consumed > 0 else 0
                                        print(f"[Early Output] {index} completed (found boxed). Time:{time_consumed:.2f}, Token:{token_counts[index]}, Speed:{speed:.2f}", flush=True)
                                    
                                        # Process the output

```

--------------------------------

### Local Evaluation Pipeline (Bash and Python)

Source: https://context7.com/imagination-research/aimo2/llms.txt

Executes local evaluations on test datasets, providing comprehensive statistics including accuracy breakdowns, timing analysis, and code execution metrics. Evaluation can be initiated using a bash script with a configuration file or programmatically via Python, processing data from a CSV file and generating detailed reports.

```bash
# Run local evaluation using configuration file
SEED=42 bash scripts/run_cfg.sh cfgs/dpsk-qwen-14b-finetune-v1-epoch4-awq.yaml data/aime_2025_30.csv

# Results saved to: results/<cfg_basename>/seed<seed>/<data_basename>/
```

```python
# Programmatic local evaluation
from imagination_aimo2.local_eval import BasicActor, _report_statistics
import pandas as pd

# Load config and create actor
actor = BasicActor(config, **config["actor"])

# Process exam dataset
df = pd.read_csv("data/aime_2025_30.csv")
results = []

for _, row in df.iterrows():
    answer, cot_ans, code_ans, lens, codes, errors, outputs, duration = \
        actor.predict_for_question(row["question"], row["id"])

    results.append({
        "id": row["id"],
        "question": row["question"],
        "answer": answer,
        "cot_answers": cot_ans,
        "code_answers": code_ans,
        "out_lens": lens,
        "correct_answer": row.get("answer"),
        "python_code_map_list": codes,
        "code_exec_error_map_list": errors,
        "question_duration": duration,
    })

# Generate statistics report
stats = _report_statistics(results, log_each_question=True)
print(f"Aggregated correct: {stats['aggregated_correct']}/{stats['num_result']}")
print(f"Total time: {stats['total_question_duration']:.2f}s")
```

--------------------------------

### Check LLM Readiness

Source: https://github.com/imagination-research/aimo2/blob/master/imagination_aimo2/backup_submission.ipynb

Checks if the LLM model is ready to process inference requests. This is a simple boolean check.

```python
def is_ready(self):
    """Check if the model is ready for inference"""
    return self._ready
```

--------------------------------

### Cite AIMO2 Project BibTeX

Source: https://github.com/imagination-research/aimo2/blob/master/README.md

A BibTeX citation snippet for referencing the AIMO2 project in academic or technical publications.

```bibtex
@misc{imaginationaimo2025,
  author = {Yichen You, Xuefei Ning, Zinan Lin},
  title = {AI Mathematical Olympiad - Progress Prize 2 Solution (2nd place, "imagination-research" team)},
  year = {2025},
  publisher = {GitHub},
  journal = {GitHub repository},
  howpublished = {\url{https://github.com/imagination-research/aimo2/}}
}
```

--------------------------------

### Python: Early Stopping Logic for Stream Generation

Source: https://github.com/imagination-research/aimo2/blob/master/imagination_aimo2/backup_submission.ipynb

Implements early stopping conditions for stream generation based on the number of valid answers, answer repetition, and speed. It checks for specific patterns in recent answers to determine if generation should cease.

```python
def _should_stop_timeout(self, valid_answers, start_time, current_speed):
        # Check for answer repetition
        if len(valid_answers) >= 4 and any(valid_answers.count(x) >= 3 for x in valid_answers):
            print("[End]: An answer repeated 3 times in less than 4 valid answers.", flush=True)
            return True
        if len(valid_answers) >= 4 and any(valid_answers.count(x) >= 4 for x in valid_answers):
            print("[End]: An answer repeated 4 times in less than 4 valid answers.", flush=True)
            return True
        if len(valid_answers) >= 5:
            recent_five = valid_answers[-5:]
            if any(recent_five.count(x) >= 4 for x in recent_five):
                print("[End]: An answer repeated 4 times in recent 5 valid answers.", flush=True)
                return True
        if len(valid_answers) <= 6 and any(valid_answers.count(x) >= 4 for x in valid_answers):
            print("[End]: An answer repeated 4 times in less than 6 valid answers.", flush=True)
            return True
        if len(valid_answers) <= 8 and any(valid_answers.count(x) >= 5 for x in valid_answers):
            print("[End]: An answer repeated 5 times in less than 8 valid answers.", flush=True)
            return True
        if len(valid_answers) <= 9 and any(valid_answers.count(x) >= 6 for x in valid_answers):
            print("[End]: An answer repeated 6 times in less than 9 valid answers.", flush=True)
            return True
            
        # Check for enough answers based on speed
        if len(valid_answers) >= min_answers:
            print(f"[End]: Collected {min_answers} answers (speed={current_speed}).", flush=True)
            return True
            
        return False
```

--------------------------------

### Filter DPO Data Script

Source: https://context7.com/imagination-research/aimo2/llms.txt

Imports the `process_directory` function from the `scripts.filter_dpo_data` module, which is used for constructing DPO training pairs by filtering and selecting responses based on correctness and semantic similarity.

```python
from scripts.filter_dpo_data import process_directory
```

--------------------------------

### Define Chain-of-Thought Prompt Templates

Source: https://github.com/imagination-research/aimo2/blob/master/imagination_aimo2/backup_submission.ipynb

Defines lists of prompt templates used to instruct LLMs to perform step-by-step reasoning and format mathematical answers within LaTeX \boxed{} tags. These templates are categorized into standard thoughts and extended thoughts for improved model performance.

```python
thoughts = [
    'Please use chained reasoning to put the answer in \\boxed{}.',
    'Please reflect and verify while reasoning and put the answer in \\boxed{}.',
    'Solve the following problem using concise and clear reasoning by placing the answer in \\boxed{}.'
]

new_thoughts = [
    'You are the smartest maths expert in the world, please spike this question and put the answer in \\boxed{}.',
    "Scrutinize every logical connection, validate intermediate steps, and conclude with a boxed solution: \\boxed{}."
]

thoughts_cot = ("\n You excel at reasoning." 
                "\n You must put the final answer in \\boxed{} before </think>." 
                "\n The final answer should modulo 1000.")
```

--------------------------------

### Filter Parquet Files for DPO Training Data

Source: https://context7.com/imagination-research/aimo2/llms.txt

Processes parquet files from a specified input directory to generate DPO training pairs, saving them to an output JSON file. It uses similarity and length thresholds to filter pairs, with parameters controlling semantic similarity, length ratio, minimum text length, and the number of files to process.

```python
process_directory(
    input_dir="/path/to/OpenR1-Math-220k/data",
    output_file="/path/to/output/dpo_pairs.json",
    similarity_threshold=0.8,      # Max semantic similarity (lower = less similar pairs)
    length_ratio_threshold=0.7,    # Max length ratio (lower = bigger length difference)
    min_length_threshold=8000,     # Minimum text length for both chosen/rejected
    num_files=3                    # Number of parquet files to process
)
```

--------------------------------

### Process Code Output and Execute Python (Python)

Source: https://github.com/imagination-research/aimo2/blob/master/imagination_aimo2/backup_submission.ipynb

Processes output to extract and execute Python code. It uses a text_extractor to find Python code, executes it via a python_repl, and parses the output for numerical results. It handles potential errors during execution and extraction, collecting valid integer answers. Dependencies include text_extractor, python_repl, re, and math modules.

```Python
import re
import math

def _process_code_output(self, index, output, text_extractor, python_repl, token_count, 
                            code_answers, valid_answers):
        """Process code output to extract and potentially execute Python code"""
        code_answer = 0

        # Try to extract and execute Python code
        python_code = text_extractor.extract_python_code(output)
        if python_code:
            python_code, line_count = text_extractor.process_python_code(python_code[0])
            success, output = python_repl(python_code)
            if success:
                pattern = r'(\d+)(?:\.\d+)?'  # Matches integers or decimals
                matches = re.findall(pattern, output)
                if matches:
                    # Convert the last match to an integer
                    try:
                        last_value = float(matches[-1])
                        if math.isinf(last_value):
                            print(f"[Warning]: Infinite value in Python result for Output {index}: {last_value}", flush=True)
                        else:
                            last_match = int(last_value)
                            print(f"<Python> result for Output {index}: {last_match}", flush=True)
                            code_answer = last_match % 1000
                            if code_answer > 0:
                                code_answers.append(code_answer)
                                valid_answers.append(code_answer)
                    except OverflowError as e:
                        print(f"[Error]: OverflowError while processing Python result for Output {index}: {e}", flush=True)
                    except ValueError as e:
                        print(f"[Error]: ValueError while processing Python result for Output {index}: {e}", flush=True)
            else:
                print(f"[Error] code for Output {index}: {output}", flush=True)
        else:
            print(f"[No] code extracted for Output {index}.", flush=True)
        
        # extract boxed answer if present
        boxed_answer = text_extractor.extract_boxed_text(output)
        print(f"(Answer): extracted for Output {index}: {boxed_answer}", flush=True)
        if 0 < int(boxed_answer) < 1000 and int(boxed_answer) == float(boxed_answer) and code_answer <= 0:
            code_answers.append(boxed_answer)
            valid_answers.append(boxed_answer)
```

--------------------------------

### Execute Python Code Securely with PythonREPL

Source: https://context7.com/imagination-research/aimo2/llms.txt

The PythonREPL class provides a secure Python code execution environment with configurable timeouts. It isolates code execution in temporary directories and captures stdout/stderr. This is crucial for running untrusted code snippets safely.

```python
from imagination_aimo2.local_eval import PythonREPL
executor = PythonREPL(timeout=30)  # 30-second timeout

# Execute Python code and get results
code = """
import sympy as sp
x = sp.Symbol('x')
equation = x**2 - 5*x + 6
solutions = sp.solve(equation, x)
print(sum(solutions))
"""

success, output = executor(code)
if success:
    print(f"Execution successful. Output: {output}")  # Output: 5
else:
    print(f"Execution failed: {output}")

# Handle timeout scenarios
long_running_code = """
import time
time.sleep(60)  # Will timeout
print("done")
"""
success, output = executor(long_running_code)
# success=False, output="Execution timed out after 30 seconds."
```

--------------------------------

### Extract and Process Python Code

Source: https://github.com/imagination-research/aimo2/blob/master/imagination_aimo2/backup_submission.ipynb

Provides static methods to extract Python code blocks from raw text using regex and to prepend necessary mathematical libraries. It also includes utility functions to parse boxed numerical answers from text responses.

```python
class TextExtractor:
    """Class to handle various text extraction operations"""
    
    @staticmethod
    def extract_python_code(text: str) -> List[str]:
        pattern = r'```python\s*(.*?)\s*```'
        matches = re.findall(pattern, text, re.DOTALL)
        return [matches[-1]] if matches else []
    
    @staticmethod
    def process_python_code(query: str) -> Tuple[str, int]:
        query = "import math\nimport numpy as np\nimport sympy as sp\n" + query
        current_rows = query.strip().split("\n")
        new_rows = []
        new_rows_codes = []
        
        for row in current_rows:
            stripped_row = row.strip()
            new_rows.append(row)
            if stripped_row and not stripped_row.startswith("#"):
                new_rows_codes.append(stripped_row)
                
        line_count = len(new_rows_codes)
        ans = "\n".join(new_rows)
        return ans, line_count

    @staticmethod
    def extract_boxed_text(text: str) -> int:
        pattern = r'oxed{(.*?)}'
        matches = re.findall(pattern, text)
        if not matches:
            return -1
            
        content = matches[-1]
        if content.isdigit():
            num = int(content)
        else:
            nums = re.findall(r'the final answer is.*?(\d+)', content)
            if not nums:
                return -1
            num = int(nums[-1])
            
        return num % 1000
```

--------------------------------

### Extract Answers and Code using AnswerExtractor

Source: https://context7.com/imagination-research/aimo2/llms.txt

The AnswerExtractor class extracts mathematical answers from LLM outputs, supporting LaTeX \boxed{} notation and Python code from markdown blocks. It also preprocesses Python code for execution and canonicalizes numbers.

```python
from imagination_aimo2.local_eval import AnswerExtractor

extractor = AnswerExtractor()

# Extract answer from LaTeX boxed notation
text_with_boxed = "The solution is \boxed{42} after careful calculation."
answer = extractor.extract_boxed_text(text_with_boxed)
print(f"Extracted answer: {answer}")  # Output: 42

# Extract Python code from markdown blocks
text_with_code = """
Here's the solution:
```python
import math
result = math.factorial(5)
print(result)
```
The answer is 120.
"""
code_blocks = extractor.extract_python_code(text_with_code)
print(f"Extracted code: {code_blocks[0]}")
# Output: import math\nresult = math.factorial(5)\nprint(result)

# Process and prepare Python code for execution
processed_code = extractor.process_python_code(code_blocks[0])
# Automatically prepends: import math, import numpy as np, import sympy as sp

# Canonicalize numbers (handles modulo 1000 for competition format)
raw_answer = extractor.canonicalize_number("1234")
print(f"Canonicalized: {raw_answer}")  # Output: 234 (1234 % 1000)
```

--------------------------------

### Aggregate Answers with Majority Voting

Source: https://context7.com/imagination-research/aimo2/llms.txt

The framework offers answer aggregation strategies like AllVoteMajorityAggregator and AnswerPriorityVoteMajorityAggregator. These methods combine results from multiple inference runs using weighted voting to improve accuracy, with options to prioritize code or CoT answers.

```python
from imagination_aimo2.aggregators import (
    AllVoteMajorityAggregator,
    AnswerPriorityVoteMajorityAggregator
)
from collections import OrderedDict

# Sample answers from multiple inference runs
# Each dict maps token_count -> answer
cot_answers = [
    OrderedDict({100: 42}),
    OrderedDict({150: 42}),
    OrderedDict({120: 37}),
    OrderedDict(),  # No answer extracted
]
code_answers = [
    OrderedDict({200: 42}),
    OrderedDict(),
    OrderedDict({180: 42}),
    OrderedDict({90: 37}),
]

# All-vote majority aggregation with weighted voting
aggregator = AllVoteMajorityAggregator()
final_answer = aggregator.aggregate_answer(cot_answers, code_answers)
print(f"Aggregated answer: {final_answer}")  # Output: 42

# Priority-based aggregation (prioritize code answers)
code_priority_aggregator = AnswerPriorityVoteMajorityAggregator(
    priority="code",      # Prioritize code over CoT
    highp_weight=1,       # Weight for high-priority answer
    lowp_pos_weight=1,    # Weight for low-priority when high doesn't exist
    lowp_neg_weight=0     # Weight for low-priority when high exists
)
final_answer = code_priority_aggregator.aggregate_answer(cot_answers, code_answers)
print(f"Code-priority answer: {final_answer}")
```

--------------------------------

### MathSolver Class for Question Prediction and Speed Adjustment

Source: https://github.com/imagination-research/aimo2/blob/master/imagination_aimo2/backup_submission.ipynb

Implements the MathSolver class responsible for predicting answers to math questions and dynamically adjusting the solving speed. It utilizes chain-of-thought and code execution strategies, and its speed is determined by the estimated time remaining per question.

```python
import sys
sys.path.append('/kaggle/input/lmdeploy-package')

import time
import os

# Assuming these are defined elsewhere or globally
# TOTAL_QUESTIONS, CHECK_AFTER_QUESTIONS, CHECK_INTERVAL, TIME_THRESHOLDS
# SPEED_TO_SAMPLES, thoughts_cot, thoughts_code, new_thoughts
# g_score, g_count, total_solving_time, total_avg_score, total_avg_length
# cutoff_time, speed

class TextExtractor:
    pass # Placeholder for actual implementation

class AnswerSelector:
    def select_answer(self, answers):
        # Placeholder for actual answer selection logic
        if answers:
            return answers[0]
        return None

class MathSolver:
    """Main class to handle math problem solving"""
    
    def __init__(self, actor, gen_config, speed):
        self.actor = actor
        self.gen_config = gen_config
        self.text_extractor = TextExtractor()
        self.answer_selector = AnswerSelector()
        self.current_speed = speed

    def adjust_speed(self):
        """Adjust speed based on progress through questions"""
        global speed, num_samples, g_count
        
        # Only check at specific question counts
        if (g_count >= CHECK_AFTER_QUESTIONS and g_count % CHECK_INTERVAL == 0):
            
            # Calculate average time per question
            avg_time_remain = (cutoff_time - time.time()) / (TOTAL_QUESTIONS - g_count)
            
            # Determine new speed based on estimated time
            new_speed = 3  # Default
            for time_range, speed_value in TIME_THRESHOLDS.items():
                if time_range[0] <= avg_time_remain < time_range[1]:
                    new_speed = speed_value
                    break
            
            # Update speed if it changed
            if new_speed != self.current_speed:
                old_speed = self.current_speed
                self.current_speed = new_speed
                
                # Update sample count based on new speed
                global num_samples
                num_samples = SPEED_TO_SAMPLES[new_speed]
                
                print(f"[SPEED ADJUSTMENT] After {g_count} questions: remaining avg time: {avg_time_remain:.2f} minutes")
                print(f"[SPEED ADJUSTMENT] Changed speed from {old_speed} to {new_speed}, num_samples={num_samples}")
                
                return True
                
        return False
        
    def predict_for_question(self, question: str, id_=None, correct_answer=None) -> int:
        """Predict answer for a single question"""
        global g_score, g_count, total_solving_time
        global total_avg_score, total_avg_length
        
        if time.time() > cutoff_time:
            return 113
        # if not id_ == "480182" and not os.getenv('KAGGLE_IS_COMPETITION_RERUN'):
        #     return 210

        # Adjust speed based on progress
        self.adjust_speed()
        
        # Start timing this question
        question_start_time = time.time()
        
        # Prepare questions with chain-of-thought and code prompts
        question_cot = question + thoughts_cot
        question_code = question + thoughts_code
        questions = [question_cot, question_code]
        # initial_cot = '<think>\nOkay, so I need to solve this problem step by step and put the final answer in \boxed{},
        # initial_code = '<think>\nOkay, so I need to use python code to solve this problem and put the final answer in \boxed{},
        # initial_prompts = [initial_cot, initial_code]
        
        print("correct answer:", correct_answer, flush=True)
        print(f"Current speed setting: {self.current_speed}, num_samples: {num_samples}", flush=True)
        print(questions[0], flush=True)
            
        # Create messages for the model
        list_of_messages = [
            [{"role": "system", "content": new_thoughts[k%2]},
             {"role": "user", "content": questions[k%2]}]
            for k in range(num_samples)
        ]
    
        # Generate and process model outputs
        cot_answers, code_answers, valid_answers = self.actor.stream_generate(
            list_of_messages, self.gen_config, self.current_speed
        )
        
        # Combine and select final answer
        valid_answers = cot_answers + code_answers
        selected_answer = self.answer_selector.select_answer(valid_answers)
        
        # Print debugging information
        print("cot answers:", cot_answers, flush=True)
        print("code answers:", code_answers, flush=True)
        print("all valid answers:", valid_answers, flush=True)
        print("selected answer:", selected_answer, flush=True)
        
        # Calculate and store timing information
        question_end_time = time.time()
        question_duration = question_end_time - question_start_time
        total_solving_time += question_duration
        
        # Print timing information
        print(f"Question {id_} solving time: {question_duration:.2f} seconds", flush=True)
        
        return selected_answer # Assuming selected_answer is an int or convertible to int

```