diff --git a/examples/Guru/README.md b/examples/Guru/README.md new file mode 100644 index 00000000..ebf9445c --- /dev/null +++ b/examples/Guru/README.md @@ -0,0 +1,170 @@ +# Instructions on Using Guru Reward Functions in VeRL + +This directory provides an example of launching [DAPO](https://arxiv.org/abs/2503.14476) training with Guru reward functions, using the `reward_fn_import.sh` script. The script demonstrates how to integrate a modular reward library (located in `llm-reasoners/reasoners/reward/`) with the VeRL framework for large language models (LLMs). + +--- + +## 1. Overview + +- **Goal:** + Run DAPO training for LLMs with a custom reward function. +- **Approach:** + - Import a reward function from the local reward library at runtime. + - Pass its path and function name to VeRL via environment variables. + - Launch DAPO training using the `reward_fn_import.sh` script. + +--- + +## 2. Environment Setup + +### 2.1 Create a Conda Environment + +We recommend using Conda to manage dependencies: + +```bash +conda create -n verl python=3.10 -y +conda activate verl +```` + +### 2.2 Install VeRL Dependencies + +After activating the `verl` environment, install VeRL and its core dependencies: + +1. **Run the install script** (choose either Megatron or FSDP backend): + + * **With Megatron (default):** + + ```bash + bash scripts/install_vllm_sglang_mcore.sh + ``` + + * **With FSDP (no Megatron):** + + ```bash + USE_MEGATRON=0 bash scripts/install_vllm_sglang_mcore.sh + ``` + + > If you encounter errors, inspect the script and manually install any missing packages. + +2. **Clone and install VeRL from source:** + + ```bash + git clone https://github.com/volcengine/verl.git + cd verl + pip install --no-deps -e . + ``` + +3. **Install additional Python packages:** + + ```bash + pip install vllm==0.8.5 + pip install tensordict==0.6.2 + pip install datasets transformers numpy polars pandas rich tqdm matplotlib sympy pylatexenc requests + ``` + +--- + +## 3. How the Reward Library Is Used + +VeRL accepts a custom reward function by specifying two arguments: + +* `custom_reward_function.path` + Path to the Python file containing the reward logic. + + ```bash + custom_reward_function.path=llm-reasoners/reasoners/reward/__init__.py + ``` + +* `custom_reward_function.name` + Name of the function to invoke as the reward entry point. + + ```bash + custom_reward_function.name=_default_compute_score + ``` + +You can modify these values in `reward_fn_import.sh` to point to any reward function you prefer. For example, to use a function defined in `naive_dapo.py`: + +```bash +custom_reward_function.path=llm-reasoners/reasoners/reward/naive_dapo.py +custom_reward_function.name=compute_score +``` + +> **Tip:** Ensure your custom function’s signature matches VeRL’s expected interface (see `_default_compute_score` for reference). + +--- + +## 4. How to Run + +1. **Prepare Your Data** + + * Download the training dataset (Parquet format) from [Huggingface](https://huggingface.co/datasets/LLM360/guru_RL_verl). + * Place all `.parquet` files under `data/parquet/` (or update `DATA_DIR` in the script). + +2. **Set the Base Model** + + * By default, `BASE_MODEL` is set to `Qwen/Qwen2.5-7B-Instruct`. + * Open `reward_fn_import.sh` and change `BASE_MODEL` if you want to use a different model. + +3. **Configure Additional Parameters** + + * Inside `reward_fn_import.sh`, adjust any training hyperparameters (batch size, maximum prompt/response length, etc.) as needed. + * Specify your custom reward function via `custom_reward_function.path` and `custom_reward_function.name`. + +4. **Launch Training** + Run the script: + + ```bash + bash reward_fn_import.sh + ``` + + The script will: + + * Set up environment variables + * Collect all Parquet files in `DATA_DIR` + * Import the specified reward function + * Launch DAPO training with VeRL + +--- + +## 5. Example Command + +Simply run: + +```bash +bash reward_fn_import.sh +``` + +This will start DAPO training using: + +* `Qwen/Qwen2.5-7B-Instruct` as the base model +* The reward function defined in `__init__.py` (default entry point `_default_compute_score`) +* All training data in `data/parquet/` + +--- + +## 6. Customization Tips + +* **Switch Reward Functions** + Edit these lines in `reward_fn_import.sh`: + + ```bash + custom_reward_function.path=llm-reasoners/reasoners/reward/your_module.py + custom_reward_function.name=your_compute_function + ``` + +* **Change Base Model** + + ```bash + BASE_MODEL="your_org/your_model-name" + ``` + +* **Adjust Training Parameters** + Modify variables like `BATCH_SIZE`, `MAX_PROMPT_LENGTH`, or `LEARNING_RATE` in the script. + +* **Use a Different Data Directory** + Set: + + ```bash + DATA_DIR="path/to/your/parquets" + ``` + diff --git a/examples/Guru/reward_fn_import.sh b/examples/Guru/reward_fn_import.sh new file mode 100644 index 00000000..90f95240 --- /dev/null +++ b/examples/Guru/reward_fn_import.sh @@ -0,0 +1,114 @@ +#!/bin/bash +export DATA_DIR="data/parquet/" + +export train_files="[$(\ + ls "${DATA_DIR}"/*.parquet \ + | xargs -n1 basename \ + | sed "s|^|'${DATA_DIR}|;s|$|'|" \ + | paste -sd, - +)]" +echo "train_files = $train_files" +export test_dir="data/test/" +export test_files=$train_files +export BASE_MODEL="Qwen/Qwen2.5-7B-Instruct" + +export WANDB_PROJECT="Reward-fn-import" +export WANDB_EXPERIMENT_NAME="local-${BASE_MODEL##*/}-$(date +%s)" + +export VLLM_ATTENTION_BACKEND="XFORMERS" +export HYDRA_FULL_ERROR=1 + +# 4. Clean up existing Ray state +echo "Stopping any existing Ray cluster…" +ray stop || true +rm -rf /tmp/ray/ray_current_cluster + +project_name='DAPO' +exp_name='DAPO-Qwen2.5-7B' + +adv_estimator=grpo + +use_kl_in_reward=False +kl_coef=0.0 +use_kl_loss=False +kl_loss_coef=0.0 + +clip_ratio_low=0.2 +clip_ratio_high=0.28 + +max_prompt_length=1100 +max_response_length=6000 +enable_overlong_buffer=True +overlong_buffer_len=512 +overlong_penalty_factor=1.0 + +loss_agg_mode="token-mean" + +enable_filter_groups=False +max_num_gen_batches=10 +train_prompt_bsz=16 +gen_prompt_bsz=$((train_prompt_bsz * 3)) +n_resp_per_prompt=2 +train_prompt_mini_bsz=32 + +# Paths +MODEL_PATH=${BASE_MODEL} + +# Algorithm +temperature=1.0 +top_p=1.0 +top_k=-1 # 0 for HF rollout, -1 for vLLM rollout + +# Performance Related Parameter +sp_size=8 +use_dynamic_bsz=True +actor_ppo_max_token_len=$((max_prompt_length + max_response_length)) +infer_ppo_max_token_len=$((max_prompt_length + max_response_length)) +offload=True +gen_tp=4 + +custom_reward_function="reward_score/__init__.py" +name="_default_compute_score" + +# 7. Launch training script +echo "Launching training…" +python3 -m verl.trainer.main_ppo \ + algorithm.adv_estimator=grpo \ + data.train_files=$train_files \ + data.val_files=$train_files \ + data.train_batch_size=128 \ + data.max_prompt_length=4096 \ + data.max_response_length=4096 \ + data.filter_overlong_prompts=True \ + data.truncation='error' \ + actor_rollout_ref.model.path=Qwen/Qwen2-7B-Instruct \ + actor_rollout_ref.actor.optim.lr=1e-6 \ + actor_rollout_ref.model.use_remove_padding=True \ + actor_rollout_ref.actor.ppo_mini_batch_size=128 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=40 \ + actor_rollout_ref.actor.use_kl_loss=True \ + actor_rollout_ref.actor.kl_loss_coef=0.001 \ + actor_rollout_ref.actor.kl_loss_type=low_var_kl \ + actor_rollout_ref.actor.entropy_coeff=0 \ + actor_rollout_ref.model.enable_gradient_checkpointing=True \ + actor_rollout_ref.actor.fsdp_config.param_offload=False \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=40 \ + actor_rollout_ref.rollout.tensor_model_parallel_size=1 \ + actor_rollout_ref.rollout.name=vllm \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \ + actor_rollout_ref.rollout.n=5 \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=40 \ + actor_rollout_ref.ref.fsdp_config.param_offload=True \ + algorithm.use_kl_in_reward=False \ + trainer.critic_warmup=0 \ + trainer.logger=['console','wandb'] \ + trainer.project_name='verl_grpo_example_gsm8k' \ + trainer.experiment_name='qwen2_7b_function_rm' \ + trainer.n_gpus_per_node=8 \ + trainer.nnodes=1 \ + trainer.save_freq=20 \ + trainer.test_freq=5 \ + custom_reward_function.path=llm-reasoners/reasoners/reward/__init__.py \ + custom_reward_function.name=_default_compute_score \ + trainer.total_epochs=15 $@ \ No newline at end of file diff --git a/examples/model_filtering/README.md b/examples/model_filtering/README.md new file mode 100644 index 00000000..bfe82c87 --- /dev/null +++ b/examples/model_filtering/README.md @@ -0,0 +1,433 @@ +# **Difficulty Model Filtering Pipeline** + +## Table of Contents + +1. [Overview](#overview) +2. [Resource Requirements](#resource-requirements) +3. [Installation and Setup](#installation-and-setup) +4. [Pipeline Workflow & Scripts](#pipeline-workflow--scripts) + + 1. [`run_inference.py`](#1-run_inferencepy) + 2. [`run_reward.py`](#2-run_rewardpy) + 3. [`run_add_pr.py`](#3-run_add_prpy) + 4. [`run_filter.py`](#4-run_filterpy) + 5. [`run_sample_and_postprocess.py`](#5-run_sample_and_postprocesspy) +5. [Checkpoint & Resumption](#checkpoint--resumption) + + * [Advanced Checkpoint Options](#advanced-checkpoint-options) +6. [Utility Functions](#utility-functions) + + * [Generating Pass Rate Mappings](#generating-pass-rate-mappings) + * [Analyzing Results](#analyzing-results) +7. [Notes](#notes) + +--- + +## Overview + +This pipeline performs difficulty filtering by combining two stages: + +1. **Inference (GPU‐intensive):** Run a large language model (e.g., Qwen2.5-7B-Instruct, Qwen3-30B-A3B) on a Parquet dataset in data-parallel and tensor-parallel modes to generate responses. +2. **Reward Calculation (CPU‐intensive):** Evaluate those responses with reward functions (e.g., pass/fail metrics) to compute pass rates and categorize problem difficulty. + +--- + +## Resource Requirements + +* **GPUs:** 8× H200 (140 GB each) +* **CPUs:** ≥ 64 cores recommended for reward evaluation +* **RAM & Disk:** Sufficient for loading large Parquet files and storing intermediate JSON/Parquet outputs + +--- + +## Installation and Setup + +1. **Create and activate a conda environment** (recommended): + + ```bash + conda create -n model_filtering python=3.10 -y + conda activate model_filtering + ``` + +2. **Install required dependencies:** + + The following Python packages are required for the pipeline scripts: + + - `torch` (PyTorch, version >= 1.13 recommended) + - `vllm==0.8.5` (required version) + - `verl` (internal/private package, see below) + - `datasets` (HuggingFace Datasets) + - `transformers` (HuggingFace Transformers) + - `numpy` + - `polars` + - `pandas` + - `rich` + - `tqdm` + - `matplotlib` (only needed for plotting in some scripts, e.g., test_parquet.py) + + **Install public dependencies with pip:** + + ```bash + pip install torch + pip install vllm==0.8.5 + pip install datasets transformers numpy polars pandas rich tqdm matplotlib + ``` + + > **Important:** The pipeline specifically requires `vllm` version 0.8.5. If you have another version installed, uninstall it first: + > + > ```bash + > pip uninstall vllm + > pip install vllm==0.8.5 + > ``` + + **Install `verl`:** + - `verl` is an internal/private package. Please refer to your organization or project maintainer for installation instructions. Typically, this may involve cloning a private repository and running `pip install -e .` in the `verl` directory. + +3. **Verify GPU drivers and distributed setup** (e.g., NCCL, CUDA, networking) for multi-GPU or multi-node execution. + +--- + +## Pipeline Workflow & Scripts + + +### 1. `run_inference.py` + +**Purpose:** + +* Load a Parquet dataset of prompts and run model inference in distributed (data-parallel + tensor-parallel) fashion. +* Generate *n* completions per prompt and save outputs as batch-indexed JSON files under the specified output directory. + +**Key Arguments:** + +* `--model_path`: Model identifier or local path (e.g., `"Qwen/Qwen2.5-7B-Instruct"`). + +* `--dataset_parquet_path`: Path to input Parquet (e.g., `"data/train/codegen__leetcode2k_2.4k.parquet"`). + +* `--output_dir`: Directory to save inference JSON outputs. + +* `--max_prompt_length`: Max tokens for each prompt (e.g., 4096). + +* `--truncation`: Truncation strategy if prompt exceeds length (`left`, `right`, `error`). + +* `--batch_size`: Batch size per GPU (e.g., 128). + +* `--n`: Number of generated completions per prompt (e.g., 16). + +* `--max_new_tokens`: Maximum tokens to generate per completion (e.g., 4096). + +* Parallel/distributed flags: + + * `--dp_size` (data-parallel world size) + * `--tp_size` (tensor-parallel world size) + * `--node_size`, `--node_rank`, `--master_addr`, `--master_port` + +* `--force_regenerate` (optional): Ignore all existing batch JSONs in `output_dir` and regenerate from scratch. + +**Example A: Qwen2.5-7B-Instruct (\~ 0.09 s per data point on leetcode2k)** + +```bash +python model_filtering/run_inference.py \ + --model_path "Qwen/Qwen2.5-7B-Instruct" \ + --dataset_parquet_path "data/train/codegen__leetcode2k_2.4k.parquet" \ + --output_dir "./diff_filter_output" \ + --max_prompt_length 4096 \ + --truncation "left" \ + --dp_size 8 \ + --tp_size 1 \ + --node_size 1 \ + --node_rank 0 \ + --master_addr "127.0.0.1" \ + --master_port 0 \ + --batch_size 128 \ + --n 16 \ + --max_new_tokens 4096 +``` + +**Example B: Qwen3-30B-A3B (\~ 12 s per data point on leetcode2k)** + +```bash +python model_filtering/run_inference.py \ + --model_path "Qwen/Qwen3-30B-A3B" \ + --dataset_parquet_path "data/train/codegen__leetcode2k_2.4k.parquet" \ + --output_dir "./diff_filter_output" \ + --max_prompt_length 4096 \ + --truncation "left" \ + --dp_size 4 \ + --tp_size 2 \ + --node_size 1 \ + --node_rank 0 \ + --master_addr "127.0.0.1" \ + --master_port 0 \ + --batch_size 128 \ + --n 16 \ + --max_new_tokens 32768 \ + --enable_expert_parallel +``` + +--- + +### 2. `run_reward.py` + +**Purpose:** + +* Read the JSON outputs from `run_inference.py` (one JSON per batch). +* Evaluate each generated response using reward functions (e.g., pass/fail, custom scoring) in parallel on CPU. +* Save reward results (e.g., pass rates and any detailed metrics) as batch-indexed JSON files in the same output directory. + +**Key Arguments:** + +* `--model_path`: Same model identifier used during inference (e.g., `"Qwen/Qwen2.5-7B-Instruct"`). +* `--dataset_parquet_path`: Path to the original input Parquet (for consistency and directory structure). +* `--output_dir`: Directory containing inference JSONs. +* `--reward_workers`: Number of CPU processes for reward computation (e.g., 64). +* `--correct_reward_threshold` (optional): A numeric threshold above which a response counts as correct. +* `--recalculate_rewards` (optional): Force recomputation of all rewards, ignoring existing reward JSONs. +* `--task_timeout` (optional): Timeout (in seconds) for each response's reward computation. + +**Example: Qwen2.5-7B-Instruct Reward Calculation** + +```bash +python model_filtering/run_reward.py \ + --model_path "Qwen/Qwen2.5-7B-Instruct" \ + --dataset_parquet_path "data/train/codegen__leetcode2k_2.4k.parquet" \ + --output_dir "./diff_filter_output" \ + --reward_workers 64 +``` + +--- + +### 3. `run_add_pr.py` + +**Purpose:** + +* Annotate an existing Parquet dataset by adding two new columns—`low_pass_rate` and `high_pass_rate`—based on provided JSON mappings. +* Output a new Parquet file that merges original data with pass-rate metadata. + +**Key Arguments:** + +* `--parquet_in`: Input Parquet file to annotate (e.g., `dataset.parquet`). +* `--low_json`: JSON file mapping `idx → pass_rate` for the low threshold (e.g., `low_pass_rates.json`). +* `--high_json`: JSON file mapping `idx → pass_rate` for the high threshold (e.g., `high_pass_rates.json`). +* `--parquet_out_dir`: Directory where the annotated Parquet file will be saved. + +**Example:** + +```bash +python model_filtering/run_add_pr.py \ + --parquet_in data/train/codegen__leetcode2k_2.4k.parquet \ + --low_json diff_filter_output/Qwen2.5-7B-Instruct/idx_to_passrate_low.json \ + --high_json diff_filter_output/Qwen2.5-7B-Instruct/idx_to_passrate_high.json \ + --parquet_out_dir data/annotated_output +``` + +--- + +### 4. `run_filter.py` + +**Purpose:** + +* Filter a Parquet dataset (e.g., the one produced by `run_add_pr.py`) using two JSON mappings (low-threshold and high-threshold pass rates). +* Drop rows whose pass rates fall outside specified bounds. +* Save the filtered subset(s) as Parquet—and optionally output a summary JSON of filtered indices. + +**Key Arguments:** + +* `--parquet_in`: Input Parquet file to filter (e.g., `annotated_dataset.parquet`). +* `--parquet_out_dir`: Directory to save filtered Parquet files. +* `--low_json`: JSON mapping index→pass\_rate for low threshold. +* `--high_json`: JSON mapping index→pass\_rate for high threshold. +* `--low_thresh`: Drop rows whose `low_pass_rate` ≤ this value (e.g., 0.1). +* `--high_thresh`: Drop rows whose `high_pass_rate` ≥ this value (e.g., 0.9). +* `--json_out` (optional): Path to save a JSON summary of which indices were filtered. + +**Example:** + +```bash +python model_filtering/run_filter.py \ + --parquet_in data/annotated_output/codegen__leetcode2k_2.4k.parquet \ + --parquet_out_dir data/filtered_output \ + --low_json diff_filter_output/Qwen2.5-7B-Instruct/idx_to_passrate_low.json \ + --high_json diff_filter_output/Qwen2.5-7B-Instruct/idx_to_passrate_high.json \ + --low_thresh 0.1 \ + --high_thresh 0.9 \ + --json_out data/filtered_output/filtered_summary.json +``` + +--- + +### 5. `run_sample_and_postprocess.py` + +**Purpose:** + +* Take raw Parquet files (possibly across multiple domains) and apply domain-specific filtering, patching, and sampling to reach a target sample size. +* Optionally enforce a maximum prompt length (in tokens) to filter out overly long examples. + +**Key Arguments:** + +* `--input_data_dir`: Directory containing one or more raw Parquet files (e.g., `data/raw_parquets/`). +* `--input_data_names`: Space-separated list of dataset names (without the `.parquet` extension). +* `--output_data_dir`: Directory where the sampled/filtered Parquet files will be saved. +* `--target_sample_size`: Target number of total examples to select (e.g., 15000). +* `--domain`: Domain type (`math`, `codegen`, `simulation`, `logic`, `table`, or `stem`). +* `--max_prompt_tokens` (optional): Remove any example whose prompt length (token count) exceeds this threshold (e.g., 4096). + +**Example:** + +```bash +python model_filtering/run_sample_and_postprocess.py \ + --input_data_dir data/raw_parquets \ + --input_data_names dataset1 dataset2 \ + --output_data_dir data/processed_parquets \ + --target_sample_size 15000 \ + --domain codegen \ + --max_prompt_tokens 4096 +``` + +--- + +## Checkpoint & Resumption + +All scripts that process data in batches (namely `run_inference.py` and `run_reward.py`) automatically checkpoint by writing batch-level JSON files to the specified output directory. If a run is interrupted, the script: + +1. **Scans** the `output_dir` for existing batch files. +2. **Loads** results from those files. +3. **Continues** from the next unprocessed batch index onward. + +This avoids reprocessing already completed work. + +### Advanced Checkpoint Options + +* **Inference Stage (`run_inference.py`):** + + * `--force_regenerate`: Ignore all existing batch JSON files and start inference from batch 0, overwriting previous outputs. Use this if you change model settings (e.g., temperature) and want fresh generations. + +* **Reward Stage (`run_reward.py`):** + + * `--recalculate_rewards`: Force recomputation of all reward scores even if reward JSON files already exist. Use this if you add or modify reward functions and wish to re-evaluate without re-running inference. + +--- + +## Utility Functions + +After both stages are complete, use these utilities (in `model_filtering/utils.py`) to generate mappings and analyze difficulty distributions. + +--- + +### Generating Pass Rate Mappings + +Creates an `idx_to_passrate.json` for a specific model and dataset—mapping dataset row index → pass rate (0.0–1.0). This is stored under: + +``` +//idx_to_passrate.json +``` + +**Usage:** + +```bash +python -m model_filtering.utils map \ + --output_dir "./diff_filter_output" \ + --dataset "codegen__leetcode2k_2.4k" \ + --model "Qwen2.5-7B-Instruct" +``` + +* **Flags:** + + * `--output_dir`: Root output directory containing model subdirectories. + * `--dataset`: Dataset basename (without `.parquet`). + * `--model`: Model identifier (subdirectory name). + +* **Output:** + A JSON of the form: + + ```json + { + "0": 0.75, + "1": 0.20, + // … + } + ``` + + where each key is the zero-based index of the example in the original Parquet and each value is its pass rate. + +--- + +### Analyzing Results + +Categorizes problems into one of seven difficulty buckets based on pass rates. Optionally aggregates across multiple datasets. + +**Difficulty Buckets:** + +1. **Impossible:** pass rate == 0.0 +2. **Very Hard:** (0.0, 0.2) +3. **Hard:** \[0.2, 0.4) +4. **Medium:** \[0.4, 0.6) +5. **Easy:** \[0.6, 0.8) +6. **Very Easy:** (0.8, 1.0) +7. **Perfect:** pass rate == 1.0 + +**Single‐Dataset Example:** + +```bash +python -m model_filtering.utils analyze \ + --output_dir "./diff_filter_output" \ + --datasets "codegen__leetcode2k_2.4k" \ + --model "Qwen2.5-7B-Instruct" +``` + +**Multi‐Dataset Example (same model):** + +```bash +python -m model_filtering.utils analyze \ + --output_dir "./diff_filter_output" \ + --datasets barc_1.1k_chunk_00 barc_1.1k_chunk_01 barc_1.1k_chunk_02 \ + --model "Qwen3-30B-A3B" +``` + +* **Optional Flags:** + + * `--regenerate`: Force regeneration of `idx_to_passrate.json` mappings before analysis. + * `--save_combined`: Save a combined mapping and aggregated statistics when analyzing multiple datasets. + +* **Outputs:** + + * A summary (to stdout or saved file) listing counts and percentages for each difficulty bucket. + * If `--save_combined` is used, a combined JSON mapping of all indices across specified datasets. + +--- + +## Notes + +* **Python Version:** ≥ 3.8 (recommend ≥ 3.10). +* **Dependencies:** See main `README.md` for full list (e.g., `verl`, `vllm==0.8.5`, `pandas`, `pyarrow`). +* **Distributed Setup:** Ensure NCCL/CUDA networking is correctly configured when using multiple GPUs or nodes. +* **Help Flags:** For detailed options, run: + + ```bash + python model_filtering/run_inference.py --help + python model_filtering/run_reward.py --help + # …etc. + ``` +* **Output Directory Structure Example:** + + ``` + diff_filter_output/ + ├─ Qwen2.5-7B-Instruct/ + │ ├─ inference_batch_000.json + │ ├─ inference_batch_001.json + │ ├─ ⋯ + │ ├─ reward_batch_000.json + │ ├─ reward_batch_001.json + │ ├─ ⋯ + │ └─ idx_to_passrate.json + └─ Qwen3-30B-A3B/ + ├─ inference_batch_000.json + ├─ ⋯ + └─ reward_batch_000.json + ``` +* **Error Recovery:** + + * If a batch JSON is corrupted or missing, manually remove it and rerun the corresponding script with `--force_regenerate` (for inference) or `--recalculate_rewards` (for reward). + +--- + +*End of README* diff --git a/examples/model_filtering/run_add_pr.py b/examples/model_filtering/run_add_pr.py new file mode 100644 index 00000000..cbcdfa3b --- /dev/null +++ b/examples/model_filtering/run_add_pr.py @@ -0,0 +1,96 @@ +#!/usr/bin/env python3 +""" +Annotate a Parquet dataset by adding two columns per row: +- `low_pass_rate`: pass rate from the low_json mapping (default 0) +- `high_pass_rate`: pass rate from the high_json mapping (default 0) + +Example +------- +python annotate_pass_rates.py \ + --parquet_in webinstruct_le30.parquet \ + --low_json low_pass_rates.json \ + --high_json high_pass_rates.json \ + --parquet_out webinstruct_with_rates.parquet +""" +import argparse +import json +from pathlib import Path +import os +from datasets import Dataset, load_dataset +from rich.console import Console +from rich.table import Table + +console = Console() + +def load_pass_rates(json_path: Path) -> dict[str, float]: + """ + Load a JSON mapping of idx→pass_rate. + Returns an empty dict on missing file. + """ + if not json_path.exists(): + console.print(f"[yellow]Warning:[/] Pass-rate file not found: {json_path}") + return {} + try: + data = json.loads(json_path.read_text(encoding="utf-8")) + # ensure keys are strings + return {str(k): float(v) for k, v in data.items()} + except Exception as e: + console.print(f"[red]Error loading {json_path}:[/] {e}") + return {} + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Annotate Parquet with low/high pass rates" + ) + parser.add_argument( + "--parquet_in", required=True, help="Input Parquet file" + ) + parser.add_argument("--parquet_out_dir", required=True, help="Directory to save filtered Parquets") + parser.add_argument( + "--low_json", required=True, help="JSON mapping idx→pass_rate for lows" + ) + parser.add_argument( + "--high_json", required=True, help="JSON mapping idx→pass_rate for highs" + ) + args = parser.parse_args() + + console.rule("[bold blue] Loading dataset") + ds = load_dataset("parquet", data_files=args.parquet_in, split="train") + console.print(f"Loaded [bold]{len(ds):,}[/bold] rows from [cyan]{args.parquet_in}[/cyan]") + + console.rule("[bold blue] Loading pass-rate JSONs") + low_rates = load_pass_rates(Path(args.low_json)) + high_rates = load_pass_rates(Path(args.high_json)) + console.print(f"Low rates: {len(low_rates):,} entries; High rates: {len(high_rates):,} entries") + + console.rule("[bold blue] Annotating rows") + def annotate(example: dict) -> dict: + extra = example.get("extra_info", {}) + + for field in ("idx", "index", "id"): + if field in extra: + idx = field + break + + idx = str(extra.get(idx, "")) + # default to 0.0 if not found + example["low_pass_rate"] = low_rates.get(idx, 0.0) + example["high_pass_rate"] = high_rates.get(idx, 0.0) + return example + + annotated = ds.map(annotate) + console.print("Annotation complete.") + + console.rule("[bold blue] Saving annotated dataset") + base = Path(args.parquet_in).stem + out_name = f"{base}_pr.parquet" + out_path = Path(args.parquet_out_dir) / out_name + os.makedirs(args.parquet_out_dir, exist_ok=True) + + annotated.to_parquet(str(out_path)) + print(annotate) + console.print(f"[bold green]Saved[/bold green] annotated dataset to [cyan]{out_path}[/cyan]") + +if __name__ == "__main__": + main() diff --git a/examples/model_filtering/run_filter.py b/examples/model_filtering/run_filter.py new file mode 100644 index 00000000..214c0694 --- /dev/null +++ b/examples/model_filtering/run_filter.py @@ -0,0 +1,107 @@ +#!/usr/bin/env python3 +""" +Filter a Parquet dataset by pass-rate thresholds with rich output. + +Example +------- +python filter_by_pass_rate.py \ + --parquet_in webinstruct_le30.parquet \ + --parquet_out_dir ./filtered_parquets \ + --json_out webinstruct_le30_by_idx.json \ + --low_json low_pass_rates.json \ + --high_json high_pass_rates.json \ + --low_thresh 0.1 \ + --high_thresh 0.9 +""" +import argparse +import json +import os +from pathlib import Path + +from datasets import Dataset, load_dataset +from rich.console import Console +from rich.table import Table + +console = Console() + +def build_idx_map(ds: Dataset) -> dict[str, dict]: + """Return a mapping {idx: row_dict} where idx comes from row['extra_info'] keys.""" + idx_map: dict[str, dict] = {} + for row in ds: + extra = row.get("extra_info", {}) + for field in ("idx", "index", "id"): + if field in extra: + idx_map[str(extra[field])] = row + break + return idx_map + + +def load_idx_set(json_path: Path, low_thresh: float = None, high_thresh: float = None) -> set[str]: + """Return the set of idx keys to drop based on thresholds in the JSON mapping.""" + data = json.loads(json_path.read_text(encoding="utf-8")) + drop_set: set[str] = set() + for key, rate in data.items(): + if low_thresh is not None and rate <= low_thresh: + drop_set.add(str(key)) + if high_thresh is not None and rate >= high_thresh: + drop_set.add(str(key)) + return drop_set + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Filter Parquet by pass-rate thresholds with rich output" + ) + parser.add_argument("--parquet_in", required=True, help="Input Parquet file") + parser.add_argument("--parquet_out_dir", required=True, help="Directory to save filtered Parquets") + parser.add_argument("--json_out", help="Optional: dump idx→row JSON for inspection") + parser.add_argument("--low_json", required=True, help="JSON idx→pass_rate for low-pass filters") + parser.add_argument("--high_json", required=True, help="JSON idx→pass_rate for high-pass filters") + parser.add_argument("--low_thresh", type=float, required=True, help="Drop if pass_rate <= this") + parser.add_argument("--high_thresh", type=float, required=True, help="Drop if pass_rate >= this") + args = parser.parse_args() + + console.rule("[bold blue] Loading Parquet Dataset") + ds = load_dataset("parquet", data_files=args.parquet_in, split="train") + total = len(ds) + console.print(f"Loaded [bold]{total:,}[/bold] rows from [cyan]{args.parquet_in}[/cyan]") + + console.rule("[bold blue] Building Index Map") + idx2row = build_idx_map(ds) + console.print(f"Built map with [yellow]{len(idx2row):,}[/yellow] entries") + + if args.json_out: + console.print(f"Writing JSON output to [cyan]{args.json_out}[/cyan]") + with open(args.json_out, "w", encoding="utf-8") as f: + json.dump(idx2row, f, ensure_ascii=False, indent=2) + + console.rule("[bold blue] Loading Pass-Rate Filters") + low_drop = load_idx_set(Path(args.low_json), low_thresh=args.low_thresh) + high_drop = load_idx_set(Path(args.high_json), high_thresh=args.high_thresh) + drop_set = low_drop | high_drop + + table = Table(show_header=True, header_style="bold magenta") + table.add_column("Filter") + table.add_column("Threshold", justify="right") + table.add_column("# IDs to Drop", justify="right") + table.add_row("Low Pass", f"<= {args.low_thresh}", f"{len(low_drop):,}") + table.add_row("High Pass", f">= {args.high_thresh}", f"{len(high_drop):,}") + table.add_row("Total Drop", "-", f"{len(drop_set):,}") + console.print(table) + + console.rule("[bold blue] Filtering Rows") + kept_rows = [row for idx, row in idx2row.items() if idx not in drop_set] + console.print(f"Kept [bold green]{len(kept_rows):,}[/bold green] rows out of [bold]{len(idx2row):,}[/bold]") + + console.rule("[bold blue] Saving Filtered Dataset") + # construct output file name based on input name and thresholds + base = Path(args.parquet_in).stem + out_name = f"{base}_l{args.low_thresh}_h{args.high_thresh}.parquet" + out_path = Path(args.parquet_out_dir) / out_name + os.makedirs(args.parquet_out_dir, exist_ok=True) + + Dataset.from_list(kept_rows).to_parquet(str(out_path)) + console.print(f"[bold green]Saved[/bold green] filtered dataset to [cyan]{out_path}[/cyan]") + +if __name__ == "__main__": + main() diff --git a/examples/model_filtering/run_inference.py b/examples/model_filtering/run_inference.py new file mode 100644 index 00000000..04a8b027 --- /dev/null +++ b/examples/model_filtering/run_inference.py @@ -0,0 +1,179 @@ +#!/usr/bin/env python3 +import argparse +import os +import sys +import signal +import multiprocessing +from multiprocessing import Process +import torch +import vllm +from vllm.utils import get_open_port +from datetime import datetime +from rich.panel import Panel + +from reasoners.tools.model_filtering.utils import console +from reasoners.tools.model_filtering.pipeline import DifficultyFilterPipeline +from datasets import Dataset + +# --------------------------------------------------------------------------- # +# Data-parallel worker # +# --------------------------------------------------------------------------- # +def run_dp_worker(args, dp_rank, dp_size, dataset): + + os.environ.update( + { + "VLLM_DP_RANK": str(dp_rank), + "VLLM_DP_RANK_LOCAL": str(dp_rank), + "VLLM_DP_SIZE": str(dp_size), + "VLLM_DP_MASTER_IP": args.master_addr, + "VLLM_DP_MASTER_PORT": str(args.master_port), + "VLLM_WORKER_MULTIPROC_METHOD": "spawn", + "VLLM_USE_TORCH_DIST": "1", + } + ) + console.print(f"[DP rank {dp_rank}] sees {torch.cuda.device_count()} visible GPU(s).") + + + args.dp_rank = dp_rank + args.dp_size = dp_size + + console.rule(f"[bold]Worker Configuration DP {dp_rank}/{dp_size-1}", style="cyan") + console.print( + Panel( + f"[bold]Model:[/bold] {args.model_path}\n" + f"[bold]Dataset:[/bold] {args.dataset_parquet_path}\n" + f"[bold]Batch size:[/bold] {args.batch_size}\n" + f"[bold]Generations:[/bold] {args.n}\n" + f"[bold]Max tokens:[/bold] {args.max_new_tokens}\n" + f"[bold]Tensor parallel:[/bold] {args.tp_size}\n" + f"[bold]GPU Device:[/bold] {torch.cuda.current_device()}", + title="📋 Configuration", + border_style="cyan", + ) + ) + + # ---------- Inference only (no reward) --------------------------------- # + pipeline = DifficultyFilterPipeline(args, dataset) + pipeline.run_inference() + +# --------------------------------------------------------------------------- # +# CLI / launcher # +# --------------------------------------------------------------------------- # +def main(): + def handle_signal(signum, _): + console.print(f"\n⚠️ [warning]Received signal {signum}, shutting down…[/warning]") + sys.exit(1) + + signal.signal(signal.SIGINT, handle_signal) + signal.signal(signal.SIGTERM, handle_signal) + + parser = argparse.ArgumentParser(description="Inference-only pipeline (rewards run later)") + parser.add_argument("--model_path", type=str, required=True) + parser.add_argument("--dataset_parquet_path", type=str, required=True) + parser.add_argument("--output_dir", type=str, default="./diff_filter_output") + + parser.add_argument("--max_prompt_length", type=int, default=4096) + parser.add_argument( + "--truncation", type=str, default="error", choices=["left", "right", "error"] + ) + parser.add_argument("--default_data_source", type=str, default="None") + + parser.add_argument("--batch_size", type=int, default=4) + parser.add_argument("--n", type=int, default=16) + parser.add_argument("--max_new_tokens", type=int, default=2048) + parser.add_argument("--top_k", type=int, default=100) + parser.add_argument("--top_p", type=float, default=0.9) + parser.add_argument("--temperature", type=float, default=1.0) + parser.add_argument("--repetition_penalty", type=float, default=1.0) + + parser.add_argument("--enable_expert_parallel", action="store_true", + help="Enable expert parallel mode for model initialization (default: True)") + + parser.add_argument("--debug", action="store_true") + + parser.add_argument("--force_regenerate", action="store_true", + help="Force regeneration of outputs for all data, ignoring previously saved results") + + parser.add_argument("--skip_percentage", type=float, default=0.0, + help="Skip to this percentage (0.0-1.0) of the dataset (e.g., 0.6 for 60%)") + + parser.add_argument("--dp_size", type=int, default=1) + parser.add_argument("--tp_size", type=int, default=1) + parser.add_argument("--node_size", type=int, default=1) + parser.add_argument("--node_rank", type=int, default=0) + parser.add_argument("--master_addr", type=str, default="127.0.0.1") + parser.add_argument("--master_port", type=int, default=0) + + args = parser.parse_args() + + console.rule("[bold]Difficulty Filter — Inference-only", style="cyan") + console.print(f"⏰ Start time: {datetime.now():%Y-%m-%d %H:%M:%S}") + console.rule(style="cyan") + + + os.makedirs(args.output_dir, exist_ok=True) + multiprocessing.set_start_method("spawn", force=True) + os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + + if args.node_size == 1: + args.master_port = get_open_port() + else: + assert args.master_addr != "127.0.0.1" + assert args.master_port > 0 + + assert args.dp_size % args.node_size == 0 + dp_per_node = args.dp_size // args.node_size + + #---Load dataset at main process + if "livecodebench" in args.dataset_parquet_path: + import polars as pl + console.print(f"🐞 [DEBUG] Loading livecodebench dataset using polars") + pd_data = pl.read_parquet(args.dataset_parquet_path).to_pandas() + dataset = Dataset.from_pandas(pd_data) + else: + dataset = Dataset.from_parquet(args.dataset_parquet_path) + + #---slice the data and dispatch to each dp worker + if args.dp_size == 1: + run_dp_worker(args, dp_rank=0, dp_size=1, dataset=dataset) + else: + procs = [] + console.print(f"🔄 Starting {dp_per_node} worker(s) on node {args.node_rank}/{args.node_size-1}") + for local_rank in range(dp_per_node): + global_rank = args.node_rank * dp_per_node + local_rank + + # ── DP split + if args.dp_size > 1: + total = len(dataset) + per_rank = total // args.dp_size + start = global_rank * per_rank + end = start + per_rank if global_rank != args.dp_size - 1 else total + console.print("start:", start) + console.print("end:",end) + dp_dataset = dataset.select(range(start, end)) + console.print( + f"🔢 DP rank [highlight]{global_rank}[/highlight] " + f"processing [highlight]{len(dp_dataset)}[/highlight] / {total}" + ) + else: + console.print(f"📊 Dataset loaded with [highlight]{len(dataset)}[/highlight] samples") + + p = Process(target=run_dp_worker, args=(args, global_rank, args.dp_size, dp_dataset)) + p.start() + procs.append(p) + + exit_code = 0 + for proc in procs: + proc.join() + if proc.exitcode is None: + print(f"Killing process {proc.pid}") + proc.kill() + exit_code = 1 + elif proc.exitcode: + exit_code = proc.exitcode + + exit(exit_code) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/examples/model_filtering/run_reward.py b/examples/model_filtering/run_reward.py new file mode 100644 index 00000000..7710186a --- /dev/null +++ b/examples/model_filtering/run_reward.py @@ -0,0 +1,300 @@ +#!/usr/bin/env python3 +""" +Reward-only pass for batched inference with preserved response ordering +""" + +# Standard library imports +import argparse +import glob +import json +import os +import signal +import time +from datetime import datetime, timedelta +from multiprocessing import cpu_count +from multiprocessing.pool import Pool +from typing import Optional, Dict, List, Tuple, Any, Union + +# Third party imports +from tqdm import tqdm + +# Local imports +from reasoners.reward import _default_compute_score +from reasoners.tools.model_filtering.utils import console, json_default + +# --------------------------------------------------------------------------- # +# Globals set once per worker via Pool.initializer # +# --------------------------------------------------------------------------- # +_TASK_TIMEOUT = None # seconds + + +def _init_pool_processes(task_timeout: int): + """Initializer so every worker knows the hard timeout value.""" + global _TASK_TIMEOUT + _TASK_TIMEOUT = task_timeout + + +# --------------------------------------------------------------------------- # +# Utility helpers # +# --------------------------------------------------------------------------- # +def format_time(seconds: float) -> str: + return str(timedelta(seconds=int(seconds))) + + +def _timeout_handler(signum, frame): # noqa: D401, N802 + """SIGALRM handler that simply raises a TimeoutError.""" + raise TimeoutError("Hard task timeout hit") + + +def compute_single_reward(arg_tuple): + """ + Compute reward for one (response, ground-truth) pair. + + * arg_tuple = (gid, response, data_source, ground_truth, extra_info, resp_idx) + * Returns (gid, detailed_dict, resp_idx) + """ + gid, response, data_source, ground_truth, extra_info, resp_idx = arg_tuple + + if _TASK_TIMEOUT is not None: + signal.signal(signal.SIGALRM, _timeout_handler) + signal.alarm(_TASK_TIMEOUT) + + try: + result = _default_compute_score( + data_source=data_source, + solution_str=response, + ground_truth=ground_truth, + extra_info=extra_info, + ) + if isinstance(result, dict): + detailed = result + score = detailed.get("score", 0.0) + else: # float + detailed = {"score": float(result)} + score = float(result) + detailed["score"] = score + return gid, detailed, resp_idx + + except TimeoutError: + return gid, {"score": 0.0, "error": f"task_timeout>{_TASK_TIMEOUT}s"}, resp_idx + except Exception as e: + return gid, {"score": 0.0, "error": str(e)}, resp_idx + finally: + if _TASK_TIMEOUT is not None: + signal.alarm(0) + + +# --------------------------------------------------------------------------- # +# Per-rank scoring # +# --------------------------------------------------------------------------- # +def score_rank_dir(rank_dir: str, args, reward_pool: Optional[Pool]): + batch_files = ( + [os.path.join(rank_dir, "batch_00000.json")] + if args.debug + else sorted(glob.glob(os.path.join(rank_dir, "batch_*.json"))) + ) + if not batch_files: + console.print(f"⚠️ [warning]No batch files found in {rank_dir} — skipping[/warning]") + return 0.0 + + rank_start = time.time() + rank_results = {} + total_responses = 0 + + for batch_file in tqdm(batch_files, desc=f"💯 Scoring {os.path.basename(rank_dir)}", position=0): + with open(batch_file, "r") as f: + batch = json.load(f) + + tasks, lookup, gid = [], {}, 0 + batch_name = os.path.basename(batch_file) + + # Track the number of responses per sample to pre-allocate results arrays + response_counts = {} + + for s_idx, sample in batch.items(): + if not args.recalculate_rewards and sample.get("scores"): + rank_results[f"{batch_name}_{s_idx}"] = sample + continue + + # Store the count of responses for this sample + response_counts[s_idx] = len(sample["responses"]) + + for resp_idx, raw_resp in enumerate(sample["responses"]): + stripped = raw_resp.split("", 1)[1] if "" in raw_resp else raw_resp + # Include the resp_idx in the task tuple + tasks.append((gid, stripped, sample["source"], sample["ground_truth"], sample["extra_info"], resp_idx)) + lookup[gid] = s_idx # Maps back to sample index + gid += 1 + + total_responses += len(tasks) + + if tasks: + results_iter = ( + reward_pool.imap_unordered(compute_single_reward, tasks, chunksize=1) + if reward_pool + else map(compute_single_reward, tasks) + ) + + inner_pbar = tqdm( + total=len(tasks), + desc=f"🧮 Responses in {batch_name}", + position=1, + leave=False, + ncols=80, + bar_format='{l_bar}{bar}| {n_fmt}/{total_fmt}', + ) + + # Initialize results containers with the right size for each sample + detailed_by_sample: Dict[str, List[Optional[Dict]]] = { + s_idx: [None] * count for s_idx, count in response_counts.items() + } + + # Process results as they come in, maintaining order + for gidx, detailed, resp_idx in results_iter: + s_idx = lookup[gidx] + # Place the result at the correct position based on resp_idx + detailed_by_sample[s_idx][resp_idx] = detailed + inner_pbar.update(1) + + inner_pbar.close() + + for s_idx, sample in batch.items(): + if s_idx not in detailed_by_sample: + continue + + detailed_list = detailed_by_sample[s_idx] + + # Verify all positions were filled + if None in detailed_list: + missing_indices = [i for i, d in enumerate(detailed_list) if d is None] + console.print(f"⚠️ Missing results for sample {s_idx} at positions {missing_indices}") + # Fill any missing results with error placeholders + for idx in missing_indices: + detailed_list[idx] = {"score": 0.0, "error": "result_missing"} + + scores = [d["score"] for d in detailed_list] + pass_cnt = sum(s >= args.correct_reward_threshold for s in scores) + sample.update( + { + "detailed_scores": detailed_list, + "scores": scores, + "pass_rate": pass_cnt / len(scores) if scores else 0.0, + } + ) + rank_results[f"{batch_name}_{s_idx}"] = sample + + # Calculate and print overall pass rate for this batch + batch_pass_rates = [sample["pass_rate"] for sample in batch.values() if "pass_rate" in sample] + if batch_pass_rates: + avg_batch_pass_rate = sum(batch_pass_rates) / len(batch_pass_rates) + console.print(f"📊 {batch_name} pass rate: {avg_batch_pass_rate:.2%}") + else: + console.print(f"ℹ️ All samples in {batch_name} already scored; skipping computation") + + with open(batch_file, "w") as f: + json.dump(batch, f, indent=2, default=json_default) + + elapsed = time.time() - rank_start + final_path = os.path.join(rank_dir, "final_results.json") + with open(final_path, "w") as f: + json.dump( + { + "results": rank_results, + "errors": {}, + "metrics": { + "total_time": elapsed, + "num_responses": total_responses, + "avg_reward_time": elapsed / max(1, total_responses), + }, + }, + f, + indent=2, + default=json_default, + ) + console.print(f"✅ Saved summary to [highlight]{final_path}[/highlight]") + return elapsed + + +# --------------------------------------------------------------------------- # +# Main # +# --------------------------------------------------------------------------- # +def main(): + parser = argparse.ArgumentParser(description="Reward-only pass (efficient, balanced) with preserved ordering") + parser.add_argument("--model_path", type=str, required=True) + parser.add_argument("--dataset_parquet_path", type=str, required=True, + help="Only used to locate the dataset-named output directory") + parser.add_argument("--output_dir", type=str, default="./diff_filter_output") + + parser.add_argument("--reward_workers", type=int, default=16, + help="Upper bound on CPU processes; auto-downscales.") + parser.add_argument("--correct_reward_threshold", type=float, default=1.0) + parser.add_argument("--recalculate_rewards", action="store_true", + help="Recompute even if sample already has scores") + parser.add_argument("--debug", action="store_true", + help="Process only batch_00000.json per rank") + + parser.add_argument("--task_timeout", type=int, default=35, # a little longer than 30s timeout of code_exec_local + help="Hard wall-clock timeout per response (s)") + parser.add_argument("--maxtasks_per_child", type=int, default=100, + help="Recycle worker after N tasks") + parser.add_argument("--join_timeout", type=int, default=30, + help="Max seconds to wait for pool shutdown before force-kill") + + args = parser.parse_args() + + avail = cpu_count() or 1 + reserved = max(1, avail // 8) + workers = ( + 1 if args.reward_workers <= 1 else min(args.reward_workers, max(1, avail - reserved)) + ) + + console.rule("[bold]Difficulty Filter — Ordered Reward pass", style="cyan") + console.print(f"⏰ Start : {datetime.now():%Y-%m-%d %H:%M:%S}") + console.print(f"🖥️ CPUs : available={avail}, using={workers}") + console.print(f"⏱️ Hard per-response timeout : {args.task_timeout}s") + console.rule(style="cyan") + + dataset_name = os.path.basename(args.dataset_parquet_path).rsplit(".parquet", 1)[0] + model_name = args.model_path.split("/")[-1] + root_dir = os.path.join(args.output_dir, dataset_name, model_name) + + rank_dirs = sorted(glob.glob(os.path.join(root_dir, "dp*"))) + if not rank_dirs: + console.print(f"❌ [error]No dp* directories under {root_dir}") + return + + reward_pool = ( + Pool( + processes=workers, + initializer=_init_pool_processes, + initargs=(args.task_timeout,), + maxtasksperchild=args.maxtasks_per_child, + ) + if workers > 1 + else None + ) + + total_elapsed = 0.0 + try: + for rd in rank_dirs: + total_elapsed += score_rank_dir(rd, args, reward_pool) + finally: + if reward_pool: + reward_pool.close() + + deadline = time.time() + args.join_timeout + while any(p.is_alive() for p in reward_pool._pool) and time.time() < deadline: + time.sleep(0.25) + + if any(p.is_alive() for p in reward_pool._pool): + console.print("[yellow]Pool shutdown took too long — forcing terminate()[/yellow]") + reward_pool.terminate() + + reward_pool.join() + + console.rule(style="cyan") + console.print(f"🏁 Finished scoring in {format_time(total_elapsed)}") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/examples/model_filtering/run_sample_and_postprocess.py b/examples/model_filtering/run_sample_and_postprocess.py new file mode 100644 index 00000000..8c834959 --- /dev/null +++ b/examples/model_filtering/run_sample_and_postprocess.py @@ -0,0 +1,185 @@ +import argparse +import os +import datasets +import json +import transformers +import numpy as np +from reasoners.tools.model_filtering.filter import LengthFilter +from reasoners.tools.model_filtering.utils import save_dataset + +GLOBAL_SEED = 42 +np.random.seed(GLOBAL_SEED) + +def filter_dataset_by_length(dataset, max_length) -> datasets.Dataset: + """ + Filter the dataset by the prompt length. + """ + tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B") + prompt_length_filter = LengthFilter(tokenizer=tokenizer, max_length=args.max_prompt_tokens) + print(f"Before filtering by prompt length {max_length}: num_items = {len(dataset)}") + dataset = dataset.filter(lambda x: prompt_length_filter.check(x), num_proc=64) + print(f"After filtering by prompt length {max_length}: num_items = {len(dataset)}") + return dataset + + +def postprocess_dataset(dataset, dataset_name) -> datasets.Dataset: + """ + Postprocess (patch) the dataset for the given domain. + """ + # Normalize the data source for 'logic' domain + if "zebra_puzzle" in dataset_name or "ordering_puzzle" in dataset_name or "graph" in dataset_name: + dataset = dataset.map(lambda item: { + "data_source": "logic__" + item["data_source"] + }) + + # Refine the prompt for 'table' domain + if "table" in dataset_name: + original_instruction = "Please output the final answer within \\boxed{}. If there are multiple answers, please output them separated by |." + refined_instruction = "Please output the final answer in \\boxed{}. If there are multiple answers, include them in a single box, separated by |." + dataset = dataset.map(lambda item: { + "prompt": [ + {"role": "user", "content": item["prompt"][0]["content"].replace(original_instruction, refined_instruction)}, + ] + }) + + # Add 'question' key for 'stem' domain as it's required in the reward + if "stem" in dataset_name: + # add 'question' to 'extra_info' + dataset = dataset.map(lambda item: { + "extra_info": { + **item["extra_info"], + "question": item['raw_prompt'] + } + }) + + # Keep 8 unit tests, otherwise the reward model will be too slow + # Truncate (right) too long input prompts, usually with very long unit test io pairs. + if "codegen" in dataset_name: + def _map_func(item): + + MAX_PROMPT_CHARS = 4096 * 5 # Use char count instead of token count as it's faster; 4096 is max input tokens and 5 is avg. chars per token + if len(item['prompt'][0]['content']) > MAX_PROMPT_CHARS: + print(f"Truncating {item['id']} from {dataset_name} to {MAX_PROMPT_CHARS} chars") + item['prompt'][0]['content'] = item['prompt'][0]['content'][:MAX_PROMPT_CHARS] + + MAX_UNIT_TEST_CHARS = 1024 * 128 # max unit test input/output length is 128KB + ground_truth = json.loads(item['reward_model']['ground_truth']) + if 'inputs' in ground_truth and len(ground_truth['inputs']) > 0: + # Store complete inputs/outputs in ground truth + complete_inputs = ground_truth['inputs'] + complete_outputs = ground_truth['outputs'] + + complete_inputs_after_filtered_too_long = [] + complete_outputs_after_filtered_too_long = [] + for ut_in, ut_out in zip(complete_inputs, complete_outputs): + if len(ut_in) > MAX_UNIT_TEST_CHARS: + pass + print(f"Unit test input is too long, {len(ut_in)} chars") + elif len(ut_out) > MAX_UNIT_TEST_CHARS: + pass + print(f"Unit test output is too long, {len(ut_out)} chars") + else: + complete_inputs_after_filtered_too_long.append(ut_in) + complete_outputs_after_filtered_too_long.append(ut_out) + + # Get 8 random indices from this test case's inputs/outputs + n_tests = len(complete_inputs_after_filtered_too_long) + sample_indices = np.random.choice(n_tests, min(8, n_tests), replace=False) + + # Update ground truth with sampled data while preserving complete data + ground_truth['complete_inputs'] = complete_inputs_after_filtered_too_long + ground_truth['complete_outputs'] = complete_outputs_after_filtered_too_long + ground_truth['inputs'] = [complete_inputs_after_filtered_too_long[i] for i in sample_indices] + ground_truth['outputs'] = [complete_outputs_after_filtered_too_long[i] for i in sample_indices] + + item['reward_model']['ground_truth'] = json.dumps(ground_truth) + + if n_tests > 8: + print(f"Filtering {n_tests} unit tests to 8 for {item['id']} from {dataset_name}") + print(type(item['reward_model']['ground_truth'])) + print(item['reward_model']['ground_truth']) + + return item + + dataset = dataset.map(_map_func, num_proc=16) + return dataset + + +def sample_dataset(dataset, target_sample_size) -> datasets.Dataset: + """ + Sample the dataset to the target sample size. + """ + if target_sample_size != -1: + subset_sample_size = int(target_sample_size * subset_sample_sizes[input_data_name] / total_sample_size) + + sampled_dataset = dataset.shuffle(seed=GLOBAL_SEED).select(range(subset_sample_size)) + print(f"{len(sampled_dataset)} items sampled from {input_data_name}") + else: + sampled_dataset = dataset + subset_sample_size = len(dataset) + return sampled_dataset, subset_sample_size + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--input_data_dir", type=str, required=True) + parser.add_argument("--input_data_names", type=str, nargs="+") + parser.add_argument("--output_data_dir", type=str, required=True) + parser.add_argument("--target_sample_size", type=int, required=True) + parser.add_argument("--domain", type=str, default="math", choices=["math", "codegen", "simulation", "logic", "table", "stem"]) + parser.add_argument("--max_prompt_tokens", type=int, default=None) + args = parser.parse_args() + + # Load data + print(args.input_data_names) + total_sample_size = 0 + subset_sample_sizes = {} + dataset_list = [] + for input_data_name in args.input_data_names: + data_path = os.path.join(args.input_data_dir, f"{input_data_name}.parquet") + print(f"Loading {input_data_name} from {data_path}") + try: + dataset = datasets.Dataset.from_parquet(data_path) + except Exception as e: + # for nested columns, e.g., livecodebench + import polars as pl + dataframe = pl.read_parquet(data_path).to_pandas() + dataset = datasets.Dataset.from_pandas(dataframe) + print(f"Loaded {input_data_name} from {data_path}, with {len(dataset)} items") + subset_sample_sizes[input_data_name] = len(dataset) + total_sample_size += len(dataset) + dataset_list.append(dataset) + + # Postprocessing and sample + os.makedirs(args.output_data_dir, exist_ok=True) + for input_data_name, dataset in zip(args.input_data_names, dataset_list): + + # Filter by length + if args.max_prompt_tokens is not None: + dataset = filter_dataset_by_length(dataset, args.max_prompt_tokens) + # Misc postprocess (patch some issues not covered in previous data processing) + dataset = postprocess_dataset(dataset, input_data_name) + # Sample to `target_sample_size` + sampled_dataset, subset_sample_size = sample_dataset(dataset, args.target_sample_size) + + # Print the first 5 items + for idx, item in enumerate(sampled_dataset): + if idx < 5: + print(f"========== {input_data_name} item#{idx} ==========") + try: + print(json.dumps(item, indent=4)) + except: + print(item) + else: + break + + # Save the sampled dataset + output_path = save_dataset( + sampled_dataset, + args.output_data_dir, + f"{args.domain}__{input_data_name.split('__')[1]}_sampled", + subset_sample_size + ) + print(f"Saving {output_path} with {subset_sample_size} items") + + print(f"Saved {args.output_data_dir} with {args.target_sample_size} items") diff --git a/reasoners/reward/__init__.py b/reasoners/reward/__init__.py new file mode 100644 index 00000000..fa356829 --- /dev/null +++ b/reasoners/reward/__init__.py @@ -0,0 +1,91 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import json +from typing import Dict, Any +from reward_score.utils import _deserialise_extra, _decompress_str + +def _default_compute_score(data_source, solution_str, ground_truth, extra_info=None): + compressed = True + reward_metric = None + if type(ground_truth) is dict: + ground_truth = _deserialise_extra(_decompress_str(ground_truth["compressed"]))["ground_truth"] + + if type(extra_info) is dict: + extra_info = _deserialise_extra(_decompress_str(extra_info)) + reward_metric = extra_info.get("reward_metric", None) + compressed = False + + # math + if data_source.startswith("math"): + if reward_metric == "prime_math": + from . import prime_math + res = prime_math.compute_score(data_source=data_source, model_output=solution_str, ground_truth=ground_truth, extra_info=extra_info, compressed=compressed) + elif reward_metric == "math_llm_judge": + from . import math_llm_judge + res = math_llm_judge.compute_score( + data_source=data_source, model_output=solution_str, ground_truth=ground_truth, extra_info=extra_info, compressed=compressed + ) + else: + # Default + from . import naive_dapo + res = naive_dapo.compute_score(data_source=data_source, solution_str=solution_str, ground_truth=ground_truth, extra_info=extra_info, compressed=compressed) + # code generation + elif data_source.startswith('codegen'): + from . import coder1 + res = coder1.compute_score(data_source=data_source, solution_str=solution_str, ground_truth=ground_truth, extra_info=extra_info, compressed=compressed) + # simulation (code) + elif data_source.startswith("simulation__codeio"): + from . import codeio + res = codeio.compute_score(data_source=data_source, model_output=solution_str, ground_truth=ground_truth, extra_info=extra_info, compressed=compressed) + elif data_source.startswith("simulation__cruxeval"): + from . import cruxeval + res = cruxeval.compute_score(data_source=data_source, model_output=solution_str, ground_truth=ground_truth, extra_info=extra_info, compressed=compressed) + # logic + elif data_source.startswith("simulation__arcagi") or data_source.startswith("simulation__barc"): + from . import arcagi + res = arcagi.compute_score(data_source=data_source, model_output=solution_str, ground_truth=ground_truth, extra_info=extra_info, compressed=compressed) + elif data_source.startswith("logic__zebra_puzzle"): + from . import zebra_puzzle + res = zebra_puzzle.compute_score(data_source=data_source, solution_str=solution_str, ground_truth=ground_truth, extra_info=extra_info, compressed=compressed) + elif data_source.startswith("logic__ordering_puzzle"): + from . import puzzles_dataset + res = puzzles_dataset.compute_score(data_source=data_source, solution_str=solution_str, ground_truth=ground_truth, extra_info=extra_info, compressed=compressed) + elif data_source.startswith("logic__graph"): + from . import graph_dataset + res = graph_dataset.compute_score(data_source=data_source, solution_str=solution_str, ground_truth=ground_truth, extra_info=extra_info, compressed=compressed) + # table + elif data_source.startswith("table"): + # TODO: tmp placeholder using math_verify + from . import tablereason + res = tablereason.compute_score(data_source=data_source, model_output=solution_str, ground_truth=ground_truth, extra_info=extra_info, compressed=compressed) + elif data_source.startswith('stem__gpqa'): + from . import gpqa + res = gpqa.compute_score(data_source=data_source, solution_str=solution_str, ground_truth=ground_truth, extra_info=extra_info, compressed=compressed) + elif data_source.startswith('stem__supergpqa'): + from . import supergpqa + res = supergpqa.compute_score(data_source=data_source, solution_str=solution_str, ground_truth=ground_truth, extra_info=extra_info, compressed=compressed) + elif data_source.startswith('stem_web'): + from . import stem_llm_judge + res = stem_llm_judge.compute_score(data_source=data_source, model_output=solution_str, ground_truth=ground_truth, extra_info=extra_info, compressed=compressed) + else: + raise NotImplementedError(f"Reward function is not implemented for {data_source=}") + + if isinstance(res, dict): + if "score" in res: + res["score"] = float(res["score"]) + return res + elif isinstance(res, (int, float, bool)): + return float(res) + else: + return float(res[0]) diff --git a/reasoners/reward/arcagi.py b/reasoners/reward/arcagi.py new file mode 100644 index 00000000..915c96a7 --- /dev/null +++ b/reasoners/reward/arcagi.py @@ -0,0 +1,101 @@ +import re +import ast +import numpy as np + +def extract_solution(solution_str): + answer_pattern = r'(.*?)' + matches = list(re.finditer(answer_pattern, solution_str, flags=re.DOTALL)) + if matches: + final_answer = matches[-1].group(1).strip() + final_answer = final_answer.replace("\n", "") + final_answer = final_answer.replace("...", "-1") + try: + # Find the part of the text that looks like a nested list + start = final_answer.index('[[') + end = final_answer.index(']]', start) + 2 + array_str = final_answer[start:end] + # Use ast.literal_eval to safely evaluate the string as a Python expression + array = ast.literal_eval(array_str) + if all(isinstance(i, list) for i in array): + return array + else: + return [[0]] + except Exception as e: + return [[0]] + else: + return [[0]] + +def pad_array_with_value(array, target_shape, pad_value): + """ + Pad the given array to the target shape with the specified pad value. + + This function pads the original array to fit the target shape by adding additional + pixels at the ends. This method ensures that the smaller array is placed at the + top-left corner of the target shape, making sense of the number of correct pixels + during comparison. + + Note: + Depending on how you pad the arrays, the number of correct pixels might vary. + For example, placing the smaller array in the center versus adding pixels at the ends + can yield different results. Here, we pad by adding pixels at the ends. + + Parameters: + array (list): The original array to be padded. + target_shape (tuple): The desired shape of the padded array (rows, columns). + pad_value (int): The value to use for padding the array. + + Returns: + np.ndarray: A padded array with the specified target shape and pad value. + """ + padded_array = np.full(target_shape, pad_value, dtype=int) + try: + array = np.stack(array).astype(int) + except Exception as e: + array = np.array([[0]]) + original_shape = array.shape + padded_array[:original_shape[0], :original_shape[1]] = array + return padded_array + + +def compare_solutions_with_padding(generated_output, correct_output, pad_value=-1): + """ + Compare the generated output with the correct output, using padding to align their shapes. + + Parameters: + generated_output (list): The generated solution array. + correct_output (list): The correct solution array. + pad_value (int, optional): The value to use for padding. Default is -1. The colour value -1 should not be present in the solutions. + + Returns: + tuple: A tuple containing: + - is_correct (bool): True if the solutions match exactly, False otherwise. + - correct_percentage (float): The percentage of correctly matched pixels. + """ + max_rows = max(len(generated_output), len(correct_output)) + max_cols = max(len(generated_output[0]), len(correct_output[0])) + target_shape = (max_rows, max_cols) + + padded_generated = pad_array_with_value(generated_output, target_shape, pad_value) + padded_correct = pad_array_with_value(correct_output, target_shape, pad_value) + total_pixels = max_rows * max_cols + correct_pixels = np.sum((padded_generated == padded_correct) & (padded_generated != pad_value) & (padded_correct != pad_value)) + correct_percentage = (correct_pixels / total_pixels) + is_correct = float(correct_pixels == total_pixels) + return is_correct, correct_percentage + + + +def compute_score(data_source: str, + model_output: str, + ground_truth: np.ndarray, + extra_info: any = None, + compressed: bool = False) -> float: + if compressed: + from reward_score.utils import _deserialise_extra, _decompress_str + extra_info = _deserialise_extra(_decompress_str(extra_info)) + ground_truth = _deserialise_extra(_decompress_str(ground_truth["compressed"]))["ground_truth"] + print(f"ground_truth: {ground_truth}") + model_output = str(model_output) + final_answer = extract_solution(model_output) + is_correct, correct_percentage = compare_solutions_with_padding(final_answer, ground_truth) + return float(is_correct) \ No newline at end of file diff --git a/reasoners/reward/codeio.py b/reasoners/reward/codeio.py new file mode 100644 index 00000000..a9f32c17 --- /dev/null +++ b/reasoners/reward/codeio.py @@ -0,0 +1,190 @@ +import json +import ast +import re +from typing import Dict, Any, Tuple, List + +def normalize(obj: Any) -> Any: + """ + Recursively normalize objects so that semantically‑equivalent + values compare equal. + + Handles: + • "true"/"false" (any case) → bool + • strings that are JSON objects/arrays → parsed & normalized + • lists / dicts → element‑wise normalize + """ + # ---------- primitives ---------- # + if isinstance(obj, str): + s = obj.strip() + + # try JSON string first + if (s.startswith("{") and s.endswith("}")) or (s.startswith("[") and s.endswith("]")): + try: + parsed = json.loads(s) + return normalize(parsed) # recurse on parsed + except json.JSONDecodeError: + pass # fall through if not valid JSON + + # bool strings + low = s.lower() + if low in {"true", "false"}: + return low == "true" + return obj + + # ---------- containers ---------- # + if isinstance(obj, list): + return [normalize(v) for v in obj] + if isinstance(obj, dict): + return {k: normalize(v) for k, v in obj.items()} + + # everything else untouched + return obj +# ------------------------------------------------------------------ # + +def extract_last_complete_json(s: str): + """ + Extract the last complete JSON object from a string. + """ + stack, last_json_start, last_json_str = [], None, None + for i, ch in enumerate(s): + if ch == "{": + stack.append(i) + if last_json_start is None: + last_json_start = i + elif ch == "}": + if stack: + stack.pop() + if not stack: + last_json_str = s[last_json_start:i+1] + last_json_start = None + if last_json_str: + try: + return json.loads(last_json_str.replace("\n", "")) + except json.JSONDecodeError: + pass + return None + +def extract_json_from_code_block(text: str): + """ + Extract JSON content from text that might contain code blocks. + """ + text = str(text) + m = re.findall(r"```(?:json)?\s*([\s\S]*?)```", text) + return (m[0] if m else text).strip() + +def extract_answer(output_str: str): + """ + Extract the predicted answer from the output string. + """ + if output_str is None: + return None + clean = extract_json_from_code_block(output_str) + try: + obj = json.loads(clean) + return obj.get("output", obj.get("input", obj)) + except json.JSONDecodeError: + last = extract_last_complete_json(clean) + if last is None: + return None + return last.get("output", last.get("input", last)) + except Exception: + return None + +def convert_gt_to_object(gt: str): + """ + Convert ground truth string to Python objects. + """ + if isinstance(gt, str) and ('"input":' in gt or '"output":' in gt): + if not (gt.strip().startswith("{") and gt.strip().endswith("}")): + gt = "{" + gt + "}" + try: + obj = json.loads(gt) + return obj.get("input", obj.get("output", obj)) + except json.JSONDecodeError: + pass + try: + obj = json.loads(gt) + return obj.get("output", obj.get("input", obj)) + except Exception: + try: + return ast.literal_eval(gt) + except Exception: + return gt + +def check_accuracy(raw_pred: str, gt: str, any_order: bool = False) -> Tuple[bool, bool]: + """ + Check if the prediction matches the ground truth. + """ + pred = normalize(extract_answer(raw_pred)) + truth = normalize(convert_gt_to_object(gt)) + + # top‑level bool convenience shortcut + tf_groups = [["True", "true", True], ["False", "false", False]] + for g in tf_groups: + if pred in g and gt in g: + return True, False + + no_answer = pred is None + if not any_order: + return pred == truth, no_answer + + # order‑agnostic list comparison + if not isinstance(pred, list) or not isinstance(truth, list): + return False, no_answer + if len(pred) != len(truth): + return False, no_answer + return all(item in truth for item in pred), no_answer + +def compute_score(data_source: str, + model_output: str, + ground_truth: str, + extra_info: any = None, + compressed: bool = False) -> float: + """ + Compute score dict for evaluation harness. + """ + if compressed: + from reward_score.utils import _deserialise_extra, _decompress_str + extra_info = _deserialise_extra(_decompress_str(extra_info)) + ground_truth = _deserialise_extra(_decompress_str(ground_truth["compressed"]))["ground_truth"] + print(f"ground_truth: {ground_truth}") + correct, _ = check_accuracy(str(model_output), str(ground_truth), any_order=False) + return float(correct) + + +# --------------------------- test --------------------------- # +if __name__ == "__main__": + # Example 1 + model_out1 = '''```json +{"input": {"upper_limit": 2924}} +```''' + ground_truth1 = '"input": {"upper_limit": 2719}' + result1 = compute_score(model_out1, ground_truth1) + print(f"Example 1 result: {result1}") + + # Example 2 + model_out2 = '''```json +{ + "preorder": "[97]", + "inorder": "[97]" +} +```''' + ground_truth2 = '"input": {"preorder": "[97]", "inorder": "[97]"}' + result2 = compute_score(model_out2, ground_truth2) + print(f"Example 2 result: {result2}") + + # Example 3 - testing "output" wrapper + model_out3 = '''```json +{"output": {"result": 42}} +```''' + ground_truth3 = '"output": {"result": 42}' + result3 = compute_score(model_out3, ground_truth3) + print(f"Example 3 result: {result3}") + + # Example 4 - previously failing case, should now pass + model_out4 = '''```json +{"answer": true} +```''' + ground_truth4 = '"output": {"answer": "true"}' + result4 = compute_score(model_out4, ground_truth4) + print(f"Example 4 result: {result4}") diff --git a/reasoners/reward/coder1/__init__.py b/reasoners/reward/coder1/__init__.py new file mode 100644 index 00000000..347963ca --- /dev/null +++ b/reasoners/reward/coder1/__init__.py @@ -0,0 +1,301 @@ +import re +from typing import Tuple, Optional +import time +import json +import time +import os +from concurrent.futures import ThreadPoolExecutor, as_completed + +import numpy as np + +from reasoners.reward.coder1.utils import _ERROR_MSG_PREFIX + +_MAX_CHAR_DISPLAY = 2048 + +# CODER1_EXEC = os.environ.get("CODER1_EXEC", "bwrap") +CODER1_EXEC = os.environ.get("CODER1_EXEC", "unsafe_local") + +if CODER1_EXEC == "docker": + from reasoners.reward.coder1.docker_exec import code_exec_docker + + code_exec = code_exec_docker +elif CODER1_EXEC == "firejail": + from reasoners.reward.coder1.firejail_exec import code_exec_firejail + + code_exec = code_exec_firejail +elif CODER1_EXEC == "ces": + from reasoners.reward.coder1.ces_exec import remote_code_exec_ces + + code_exec = remote_code_exec_ces +elif CODER1_EXEC == "kira": + from reasoners.reward.coder1.kira_exec import remote_code_exec_kira + + code_exec = remote_code_exec_kira +elif CODER1_EXEC == "bwrap": + from reasoners.reward.coder1.bwrap_exec import code_exec_bwrap + + code_exec = code_exec_bwrap +elif CODER1_EXEC == "unsafe_local": + from reasoners.reward.coder1.unsafe_local_exec import code_exec_local + + code_exec = code_exec_local +else: + raise ValueError(f"Unknown CODER1_EXEC: {CODER1_EXEC}") + + +CODE_PATTERN = re.compile(r"```(?:\w+)?\n(.*?)\n```", re.DOTALL) +ANSWER_PATTERN = re.compile(r"(.*)", re.DOTALL) + +def remote_check_stdio(code, stdin, stdout): + succ, output = code_exec(code=code, stdin=stdin) + return succ, output, stdin, stdout + + +# def validate_response_structure(processed_str: str) -> bool: +# pattern = re.compile(r".*.*", re.DOTALL) +# return bool(pattern.search(processed_str.strip())) + + +def try_extract_solution(solution_str: str) -> str: + match = re.search(ANSWER_PATTERN, solution_str) + + if match: + return match.group(1).strip() + + return solution_str + + +def extract_code_from_string(solution_str): + solution_str = try_extract_solution(solution_str) + code_blocks = CODE_PATTERN.findall(solution_str) + return "\n".join(code_blocks).strip() + + +def fuzzy_equal(actual: str, expected: str, tolerance: float = 1e-6, verbose=True) -> bool: + """ + Compare two outputs line by line and element by element for approximate equality. + Handles: + 1. Integer and floating-point number comparison with tolerance + 2. Case-insensitive comparison for yes/no + + Args: + actual: The actual output from code execution + expected: The expected output + tolerance: Tolerance for floating point number comparison + + Returns: + bool: True if outputs are approximately equal + """ + # Save original values for debugging + original_actual = actual + original_expected = expected + + # Normalize line endings + actual = actual.strip().replace('\r\n', '\n') + expected = expected.strip().replace('\r\n', '\n') + + # If exact match after normalization, return early + if actual == expected: + return True + + # Split into lines + actual_lines = actual.split('\n') + expected_lines = expected.split('\n') + + # If different number of lines, they're definitely not equal + if len(actual_lines) != len(expected_lines): + return False + + # Track fuzzy matches for debugging + fuzzy_match_reasons = [] + + # Compare each line + for i, (actual_line, expected_line) in enumerate(zip(actual_lines, expected_lines)): + # If lines match exactly, continue + if actual_line == expected_line: + continue + + # Split into tokens by whitespace + actual_tokens = actual_line.split() + expected_tokens = expected_line.split() + + # If different number of tokens, they're not equal + if len(actual_tokens) != len(expected_tokens): + return False + + # Compare each token + for j, (actual_token, expected_token) in enumerate(zip(actual_tokens, expected_tokens)): + # If tokens match exactly, continue + if actual_token == expected_token: + continue + + # For yes/no, use case-insensitive comparison + if actual_token.lower() in ["yes", "no"] and expected_token.lower() in ["yes", "no"]: + if actual_token.lower() == expected_token.lower(): + fuzzy_match_reasons.append(f"Line {i+1}, Token {j+1}: Case-insensitive yes/no match '{actual_token}' ≈ '{expected_token}'") + continue + else: + return False + + # Try numeric comparison + try: + actual_num = float(actual_token) + expected_num = float(expected_token) + diff = abs(actual_num - expected_num) + + if diff <= tolerance: + fuzzy_match_reasons.append(f"Line {i+1}, Token {j+1}: Numeric match '{actual_token}' ≈ '{expected_token}' (diff: {diff})") + continue + else: + return False + except ValueError: + # Not numeric values + return False + + # Output fuzzy match information if any occurred + if fuzzy_match_reasons and verbose: + print(f"🔍 FUZZY MATCH - Outputs approximately equal:") + print(f" Expected: {repr(original_expected)}") + print(f" Actual: {repr(original_actual)}") + print(f" Reasons for fuzzy matching:") + for reason in fuzzy_match_reasons: + print(f" • {reason}") + + # If we made it here, all lines are approximately equal + return True + +def _compute_score( + solution_str, ground_truth, extra_info, format_reward=0.0, answer_reward=1.0 +): + reward_log = [] + + # ground_truth is not code, but tests + # pass_fmt = validate_response_structure(solution_str) + solution_code = extract_code_from_string(solution_str) + + # NameError: name 'List' is not defined. Did you mean: 'list'? + # NameError: name 'defaultdict' is not defined + # reference solutions fail due to imports not being present + + if ( + # not pass_fmt or len(solution_code) == 0 + len(solution_code) == 0 + ): # only print full output when there is an error + reward_log.append("-" * 16 + "Bad format detected!" + "-" * 16) + reward_log.append("-" * 16 + "Original Model Output" + "-" * 16) + reward_log.append(solution_str) + # return -answer_reward - format_reward, "\n".join(reward_log) + return 0.0, "\n".join(reward_log) + + reward_log.append("-" * 16 + "Extracted Code to Execute" + "-" * 16) + ground_truth = json.loads(ground_truth) + + t_start = time.time() + + # log code + if "functional" in ground_truth: + if "prefix" in extra_info and extra_info["prefix"] != None: + solution_code = extra_info["prefix"] + "\n" + solution_code + reward_log.append(solution_code + "\n" + ground_truth["functional"]) + else: + reward_log.append(solution_code) + + if ( + "pytest" in ground_truth + or "functional" in ground_truth + or "solution_file" in ground_truth + ): + if "functional" in ground_truth: + if "prefix" in extra_info and extra_info["prefix"] != None: + solution_code = extra_info["prefix"] + "\n" + solution_code + succ, output = code_exec(solution_code + "\n" + ground_truth["functional"]) + elif "solution_file" in ground_truth: + succ, output = code_exec( + solution_code, solution=ground_truth["solution_file"] + ) + else: # pytest + succ, output = code_exec(solution_code, pytest=ground_truth["pytest"]) + if not succ: + # reward_log.append( + # "!" * 16 + # + f"⚠️ Test Execution Failed in {time.time() - t_start:.1f}s" + # + "!" * 16 + # ) + # reward_log.append(output[:_MAX_CHAR_DISPLAY]) + # reward_log.append("-" * 16 + "Failed Prompt" + "-" * 16) + # reward_log.append(extra_info["original_prompt"].replace("\n\n", "\n")) + return format_reward, "\n".join(reward_log) + elif "inputs" in ground_truth and "outputs" in ground_truth: + stdin_list: str = ground_truth["inputs"] + stdout_list: str = ground_truth["outputs"] + + # Add parallelism + # ok i see why they had concurrency issues. + # they call this one by one by one for each of the inputs + # spawns a separate process for each input as opposed to just being a single file + with ThreadPoolExecutor(max_workers=min(8, len(stdin_list))) as executor: + futures = [ + executor.submit(remote_check_stdio, solution_code, stdin, stdout) + for stdin, stdout in zip(stdin_list, stdout_list) + ] + for future in as_completed(futures): + succ, output, stdin, stdout = future.result() + if not succ or not fuzzy_equal(output.strip(), stdout.strip()): + output = output[:_MAX_CHAR_DISPLAY] # truncate output to print + # reward_log.append( + # "!" * 16 + # + f"⚠️ Test Execution Failed in {time.time() - t_start:.1f}s" + # + "!" * 16 + # ) + # reward_log.append(f"🔎Input: {repr(stdin)}") + # reward_log.append(f"✅Expected: {repr(stdout.strip())}") + # reward_log.append( + # f"❌Actual: {output if output.startswith(_ERROR_MSG_PREFIX) else repr(output.strip())}" + # ) + # reward_log.append("-" * 16 + "Failed Prompt" + "-" * 16) + # reward_log.append(extra_info["original_prompt"].replace("\n\n", "\n")) + return format_reward, "\n".join(reward_log) + else: + raise ValueError( + "Current supports for ground-truth are ['pytest', 'functional', 'solution_file', 'inputs/outputs'] -- No idea what's: {ground_truth = }" + ) + + reward_log.append("+" * 16 + "Test Execution Passed! (Output)" + "+" * 16) + reward_log.append(output) + return format_reward + answer_reward, "\n".join(reward_log) + + +def compute_score(data_source: str, + solution_str: str, + ground_truth: str, + extra_info: dict, + compressed: bool = False, + format_reward=0.0, + answer_reward=1.0): + if compressed: + from reward_score.utils import _deserialise_extra, _decompress_str + extra_info = _deserialise_extra(_decompress_str(extra_info)) + ground_truth = _deserialise_extra(_decompress_str(ground_truth["compressed"]))["ground_truth"] + + if isinstance(extra_info, np.ndarray): + extra_info = extra_info.item() + score, reward_log = _compute_score( + solution_str, + ground_truth, + extra_info=extra_info, + format_reward=format_reward, + answer_reward=answer_reward, + ) + marker = "✅" if score == (format_reward + answer_reward) else "❌" + reward_log = ( + marker * 16 + + "Reward Calculation" + + marker * 16 + + "\n" + + reward_log + + "\n" + + marker * 16 + + f"Final Rward = {score}" + + marker * 16 + ) + return float(score) diff --git a/reasoners/reward/coder1/bwrap_exec.py b/reasoners/reward/coder1/bwrap_exec.py new file mode 100644 index 00000000..b1b2f87c --- /dev/null +++ b/reasoners/reward/coder1/bwrap_exec.py @@ -0,0 +1,125 @@ +import os +import subprocess + +from tempfile import NamedTemporaryFile, TemporaryDirectory + +from .utils import _ERROR_MSG_PREFIX, _DEFAULT_TIMEOUT_SECONDS + +CLI_ARG_SIZE_LIMIT = 1024 * 3 + + +def code_exec_bwrap( + code, + stdin: str = None, + timeout=5, + pytest: str = None, + solution: str = None, + python_env: str = os.environ.get("CONDA_BIN_PATH", None), +): + env = os.environ.copy() + env["OPENBLAS_NUM_THREADS"] = "1" + if "PYTHONPATH" in env: + del env["PYTHONPATH"] # avoid importing wrong stuff + + if python_env is None: + python_executable = "/usr/bin/python3" + else: + python_executable = os.path.join(python_env, "python3") + + command = [ + "timeout", + str(timeout), + "bwrap", + "--unshare-all", + # "--ro-bind", "/usr/bin/python3", "/usr/bin/python3", + "--ro-bind", + python_executable, + "/sandbox/bin/python3", + "--ro-bind", + "/usr/lib", + "/usr/lib", + "--ro-bind", + "/usr/lib64", + "/usr/lib64", + "--ro-bind", + "/lib", + "/lib", + "--ro-bind", + "/lib64", + "/lib64", + "--dev", + "/dev", + "--proc", + "/proc", + # "--new-session", # causes issues w. timeout + ] + if ( + solution + ): # not necessarily a pytest file. test solution is just in a separate file that imports solution + with TemporaryDirectory() as tmpdir: + assert stdin is None, "STDIN is not supported with solution_file" + with open(os.path.join(tmpdir, "solution.py"), "w") as f: + f.write(code) + with open(os.path.join(tmpdir, "test_solution.py"), "w") as f: + f.write(solution) + command.extend(["--ro-bind", tmpdir, tmpdir]) + command.extend( + ["/sandbox/bin/python3", os.path.join(tmpdir, "test_solution.py")] + ) + result = subprocess.run( + command, + cwd=tmpdir, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + env=env, + check=False, + ) + elif pytest: # none of the reference solutions in code-r1-12k are pytest based + with TemporaryDirectory() as tmpdir: + assert stdin is None, "STDIN is not supported with pytest" + with open(os.path.join(tmpdir, "solution.py"), "w") as f: + f.write(code) + with open(os.path.join(tmpdir, "test_solution.py"), "w") as f: + f.write(pytest) + command.extend(["--ro-bind", tmpdir, tmpdir]) + command.extend(["/sandbox/bin/python3", "-m", "pytest", tmpdir]) + result = subprocess.run( + command, + cwd=tmpdir, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + env=env, + check=False, + ) + else: + if len(code) < CLI_ARG_SIZE_LIMIT: + command.extend(["/sandbox/bin/python3", "-c", code]) + result = subprocess.run( + command, + input=stdin.encode() if stdin else None, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + env=env, + check=False, + ) + else: + with NamedTemporaryFile() as tmp: + tmp.write(code.encode()) + tmp.flush() + command.extend(["--ro-bind", tmp.name, tmp.name]) + command.extend(["/sandbox/bin/python3", tmp.name]) + result = subprocess.run( + command, + input=stdin.encode() if stdin else None, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + env=env, + check=False, + ) + + stderr = result.stderr.decode().strip() + stdout = result.stdout.decode() + + if result.returncode == 0: + return True, stdout + return False, _ERROR_MSG_PREFIX + f"STDOUT:\n{stdout}\n\nSTDERR:\n{stderr}" diff --git a/reasoners/reward/coder1/ces_exec.py b/reasoners/reward/coder1/ces_exec.py new file mode 100644 index 00000000..34f694e0 --- /dev/null +++ b/reasoners/reward/coder1/ces_exec.py @@ -0,0 +1,52 @@ +from traceback import format_exc +import requests +import time +import json +import random +from typing import Tuple + +from .utils import check_executor_alive, _ERROR_MSG_PREFIX, _DEFAULT_TIMEOUT_SECONDS + + +# https://github.com/cassanof/code_exec_server +def remote_code_exec_ces(code: str, stdin: str = None) -> Tuple[bool, str]: + _DEFAULT_CES_URL = "http://localhost:8000" + timeout = _DEFAULT_TIMEOUT_SECONDS + max_retry_on_timeout: int = 2 + cur_retry_on_timeout = 0 + while True: # loop for server downtime + try: + t_start = time.time() + headers = {"Content-Type": "application/json"} + r = requests.post( + _DEFAULT_CES_URL + "/py_exec", + data=json.dumps({ + "code": code, + "timeout": timeout, + "stdin": stdin + }), + headers=headers, + ) + resp, outs = r.text.split("\n", 1) + succ_exit = resp == "0" + + # sometimes the code fail without output despite it should + # it could because the server is busy and it hits the timeout + # for such case we should retry. + # FIXME: also mark this as a potential bug in code exec server or prime + if not succ_exit and outs == "" and time.time() - t_start > timeout: # likely server timeout + note_message = f"🚨 Failed with timeout for {timeout}s X {max_retry_on_timeout} times" + if cur_retry_on_timeout >= max_retry_on_timeout: + return False, _ERROR_MSG_PREFIX + note_message + cur_retry_on_timeout += 1 + print(note_message) + time.sleep(random.randint(5, 20)) # sleep randomly to avoid all clients retry at the same time + continue + return succ_exit, outs + except Exception: + if not check_executor_alive(_DEFAULT_CES_URL): # check if the server is alive + print("Request rejected, waiting 3 seconds and then retrying...") + time.sleep(3) + continue + + return False, _ERROR_MSG_PREFIX + format_exc() diff --git a/reasoners/reward/coder1/docker_exec.py b/reasoners/reward/coder1/docker_exec.py new file mode 100644 index 00000000..80ed07b9 --- /dev/null +++ b/reasoners/reward/coder1/docker_exec.py @@ -0,0 +1,78 @@ +from traceback import format_exc + +import docker +from docker.types import Ulimit +import tempfile +from pathlib import Path + +from .utils import _ERROR_MSG_PREFIX, _DEFAULT_TIMEOUT_SECONDS + +CLI_ARG_SIZE_LIMIT = 1024 * 3 + + +# trying the fifth option as firejail can be tricky to configure in some other test beds +def code_exec_docker(code, stdin: str = None, timeout=_DEFAULT_TIMEOUT_SECONDS): + image = "jupyter/scipy-notebook" # python 3.11 + ulimits = [ + Ulimit(name="nofile", soft=16, hard=16), + Ulimit(name="fsize", soft=524288, hard=524288), # 512 KB + Ulimit(name="as", soft=4096 * 1024 * 1024, hard=4096 * 1024 * 1024) # 4096 MB in bytes RAM + ] + + client = docker.from_env() + container_kwargs = { + "image": image, + "environment": { + "OPENBLAS_NUM_THREADS": "1" + }, + "ulimits": ulimits, + "stdin_open": stdin is not None, + "pids_limit": 16, + } + + tmpdir = None + if len(code) < CLI_ARG_SIZE_LIMIT: + pycmd = ["python", "-c", code] + else: + tmpdir = tempfile.TemporaryDirectory() + container_kwargs["volumes"] = {tmpdir.name: {"bind": "/tmp"}} + script_path = Path(tmpdir.name) / "script.py" + script_path.write_text(code) + pycmd = ["python", "/tmp/script.py"] + + if stdin: + bash_cmd = ["sh", "-c"] + if "-c" in pycmd: + pycmd[-1] = f"'{pycmd[-1]}'" + if len(stdin) < CLI_ARG_SIZE_LIMIT: + bash_cmd.append(f"echo {repr(stdin)} | {' '.join(pycmd)}") + else: + if tmpdir is None: + tmpdir = tempfile.TemporaryDirectory() + container_kwargs["volumes"][stdin_path.parent] = {"bind": "/tmp"} + + stdin_path = Path(tmpdir.name) / "stdin.txt" + stdin_path.write_text(stdin) + bash_cmd.append(f"cat /tmp/stdin.txt | {' '.join(pycmd)}") + container_kwargs["command"] = bash_cmd + else: + container_kwargs["command"] = pycmd + + container = None + try: + container = client.containers.run(**container_kwargs, detach=True) + result = container.wait(timeout=timeout) + exit_code = result.get("StatusCode", 1) + logs = container.logs(stdout=True, stderr=True).decode() + except Exception: + return False, _ERROR_MSG_PREFIX + f"{format_exc()}" + finally: + if container: + container.remove(force=True) + if tmpdir: + tmpdir.cleanup() + + if exit_code == 0: + return True, logs + else: + return False, _ERROR_MSG_PREFIX + f"Exit code: {exit_code}\nLogs:\n{logs}" diff --git a/reasoners/reward/coder1/firejail_exec.py b/reasoners/reward/coder1/firejail_exec.py new file mode 100644 index 00000000..872eca10 --- /dev/null +++ b/reasoners/reward/coder1/firejail_exec.py @@ -0,0 +1,97 @@ +import os +import subprocess + +from tempfile import NamedTemporaryFile, TemporaryDirectory + +from .utils import _ERROR_MSG_PREFIX, _DEFAULT_TIMEOUT_SECONDS + +# So I tried 4 approaches for code execution (after a few all-nighters...): +# 1. _remote_code_exec_ces -- Directly using https://github.com/cassanof/code_exec_server +# - Is fast but leads to unreasonable false positives of timeouts +# - I tried to alleviate this by (i) restarting the server frequently; (ii) bigger timeout; (iii) lower concurrency +# - Still feels 10% false positives and bad concurrency +# 2. _remote_code_exec_kira -- Extending https://github.com/cassanof/code_exec_server to support my format and use some OS features for isolation +# - Less unreasonable timeouts but the concurrency is very bad, stucking at create temp dirs all the time +# 3. https://e2b.dev/ +# - Concurrency is fine +# - Does not support STDIN by default - needs some hack to support it +# - I don't want to pay other servers when I have 192 physical CPUs... +# 4. _code_exec_firejail -- Using firejail (https://github.com/netblue30/firejail) +# - User space isolation (some ulimit/rlimit features) +# - Drop-in OS isolation via seccomp (blocking socket, etc.) +# - Concurrency is the best so far +# - This is not the safest - but docker is not safe either :L. Looks good enough for my dataset anyways. +# sudo add-apt-repository ppa:deki/firejail +# sudo apt-get update +# sudo apt-get install firejail firejail-profiles + +CLI_ARG_SIZE_LIMIT = 1024 * 3 + + +def code_exec_firejail(code, stdin: str = None, timeout=_DEFAULT_TIMEOUT_SECONDS, pytest: str = None): + env = os.environ.copy() + env["OPENBLAS_NUM_THREADS"] = "1" + if "PYTHONPATH" in env: + del env["PYTHONPATH"] # avoid importing wrong stuff + + # Build the firejail command with resource limits and cleanup options + command = [ + "firejail", + "--private", + "--quiet", + "--seccomp=socket", + "--profile=pip", + "--rlimit-nproc=32", + "--rlimit-nofile=32", + "--rlimit-fsize=2m", # Limit file size + "--rlimit-as=4096m", + f"--timeout=00:00:{timeout}", + ] + + if pytest: + # solution is in {tmpdir}/solution.py + with TemporaryDirectory() as tmpdir: + assert stdin is None, "STDIN is not supported with pytest" + # Write the solution to a file + with open(os.path.join(tmpdir, "solution.py"), "w") as f: + f.write(code) + with open(os.path.join(tmpdir, "test_solution.py"), "w") as f: + f.write(pytest) + command.insert(4, f"--whitelist={tmpdir}") + command.extend(["python3", "-m", "pytest", tmpdir]) + result = subprocess.run( + command, + cwd=tmpdir, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + env=env, + check=False, + ) + else: + if len(code) < CLI_ARG_SIZE_LIMIT: + command.extend(["python3", "-c", code]) + result = subprocess.run(command, + input=stdin.encode() if stdin else None, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + env=env, + check=False) + else: + with NamedTemporaryFile() as tmp: + tmp.write(code.encode()) + tmp.flush() + command.insert(4, f"--whitelist={tmp.name}") + command.extend(["python3", tmp.name]) + result = subprocess.run(command, + input=stdin.encode() if stdin else None, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + env=env, + check=False) + + stderr = result.stderr.decode().strip() + stdout = result.stdout.decode() + + if result.returncode == 0: + return True, stdout + return False, _ERROR_MSG_PREFIX + f"STDOUT:\n{stdout}\n\nSTDERR:\n{stderr}" diff --git a/reasoners/reward/coder1/kira_exec.py b/reasoners/reward/coder1/kira_exec.py new file mode 100644 index 00000000..0d4288a9 --- /dev/null +++ b/reasoners/reward/coder1/kira_exec.py @@ -0,0 +1,44 @@ +from traceback import format_exc +import requests +import time +import json +from typing import Tuple + +from .utils import check_executor_alive, _ERROR_MSG_PREFIX + + +# Modified version of https://github.com/FlorianWoelki/kira +def remote_code_exec_kira(code: str, stdin: str = None) -> Tuple[bool, str]: + """Returns either (True, output) or (False, error)""" + _DEFAULT_KIRA_URL = "http://localhost:9090" + while True: # loop for server downtime + try: + headers = {"Content-Type": "application/json"} + r = requests.post( + _DEFAULT_KIRA_URL + "/execute", + data=json.dumps({ + "language": "python", + "content": code, + "tests": [{ + "stdin": stdin + }] + }), + headers=headers, + ) + response = r.json() + result = response["testOutput"]["results"] + if not result: + return False, _ERROR_MSG_PREFIX + f"{response}" + result = result[0] + stdout = result["received"] + stderr = result["runError"] + if len(stderr) == 0: + return True, stdout + return False, stderr + except Exception: + if not check_executor_alive(_DEFAULT_KIRA_URL): # check if the server is alive + print("Request rejected, waiting 3 seconds and then retrying...") + time.sleep(3) + continue + + return False, _ERROR_MSG_PREFIX + format_exc() diff --git a/reasoners/reward/coder1/unsafe_local_exec.py b/reasoners/reward/coder1/unsafe_local_exec.py new file mode 100644 index 00000000..8c01d262 --- /dev/null +++ b/reasoners/reward/coder1/unsafe_local_exec.py @@ -0,0 +1,113 @@ +import os +import subprocess +import shlex +from tempfile import NamedTemporaryFile, TemporaryDirectory + +from .utils import _ERROR_MSG_PREFIX, _DEFAULT_TIMEOUT_SECONDS + +CLI_ARG_SIZE_LIMIT = 1024 * 3 +MEMORY_LIMIT_KB = 10 * 1024 * 1024 # 10GB in KB + +def wrap_command_with_ulimit(command): + cmd_str = ' '.join(shlex.quote(c) for c in command) + return ["bash", "-c", f"ulimit -v {MEMORY_LIMIT_KB}; exec {cmd_str}"] + +def code_exec_local( + code, + stdin: str = None, + timeout=30, + pytest: str = None, + solution: str = None, + python_env: str = os.environ.get("CONDA_BIN_PATH", None), +): + env = os.environ.copy() + env["OPENBLAS_NUM_THREADS"] = "1" + if "PYTHONPATH" in env: + del env["PYTHONPATH"] + + if python_env is None: + python_executable = "/usr/bin/python3" + else: + python_executable = os.path.join(python_env, "python3") + + if solution: + with TemporaryDirectory() as tmpdir: + assert stdin is None, "STDIN is not supported with solution_file" + with open(os.path.join(tmpdir, "solution.py"), "w") as f: + f.write(code) + with open(os.path.join(tmpdir, "test_solution.py"), "w") as f: + f.write(solution) + + command = [ + "timeout", + str(timeout), + python_executable, + os.path.join(tmpdir, "test_solution.py"), + ] + command = wrap_command_with_ulimit(command) + result = subprocess.run( + command, + cwd=tmpdir, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + env=env, + check=False, + ) + elif pytest: + with TemporaryDirectory() as tmpdir: + assert stdin is None, "STDIN is not supported with pytest" + with open(os.path.join(tmpdir, "solution.py"), "w") as f: + f.write(code) + with open(os.path.join(tmpdir, "test_solution.py"), "w") as f: + f.write(pytest) + + command = [ + "timeout", + str(timeout), + python_executable, + "-m", + "pytest", + tmpdir, + ] + command = wrap_command_with_ulimit(command) + result = subprocess.run( + command, + cwd=tmpdir, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + env=env, + check=False, + ) + else: + if len(code) < CLI_ARG_SIZE_LIMIT: + command = ["timeout", str(timeout), python_executable, "-c", code] + command = wrap_command_with_ulimit(command) + result = subprocess.run( + command, + input=(stdin.encode() if stdin else None), + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + env=env, + check=False, + ) + else: + with NamedTemporaryFile() as tmp: + tmp.write(code.encode()) + tmp.flush() + command = ["timeout", str(timeout), python_executable, tmp.name] + command = wrap_command_with_ulimit(command) + result = subprocess.run( + command, + input=(stdin.encode() if stdin else None), + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + env=env, + check=False, + ) + + stderr = result.stderr.decode().strip() + stdout = result.stdout.decode() + + if result.returncode == 0: + return True, stdout + return False, _ERROR_MSG_PREFIX + f"STDOUT:\n{stdout}\n\nSTDERR:\n{stderr}" \ No newline at end of file diff --git a/reasoners/reward/coder1/utils.py b/reasoners/reward/coder1/utils.py new file mode 100644 index 00000000..39f5cba1 --- /dev/null +++ b/reasoners/reward/coder1/utils.py @@ -0,0 +1,11 @@ +import requests + +_ERROR_MSG_PREFIX = "Failed to execute program: " +_DEFAULT_TIMEOUT_SECONDS = 30 + + +def check_executor_alive(executor): + try: + return requests.get(executor + "/").status_code in [200, 404] + except Exception: + return False diff --git a/reasoners/reward/cruxeval/__init__.py b/reasoners/reward/cruxeval/__init__.py new file mode 100644 index 00000000..9455fc00 --- /dev/null +++ b/reasoners/reward/cruxeval/__init__.py @@ -0,0 +1,33 @@ +import re +from reasoners.reward.cruxeval.utils import check_correctness + + +def compute_score(data_source: str, + model_output: str, + ground_truth: str, + extra_info: any = None, + compressed: bool = False) -> float: + if compressed: + from reward_score.utils import _deserialise_extra, _decompress_str + extra_info = _deserialise_extra(_decompress_str(extra_info)) + ground_truth = _deserialise_extra(_decompress_str(ground_truth["compressed"]))["ground_truth"] + + model_output = str(model_output) + # print(f">>> {model_output}") + try: + if "" in model_output: + # remove content until + model_output = re.split(r"", model_output)[1] + else: + model_output = model_output + # remove content between ```python and ``` + model_output = re.split(r"```python", model_output)[1] + model_output = re.split(r"```", model_output)[0] + except: + model_output = model_output + + full_code = eval(ground_truth)["functional"] + "\n" + model_output + # print(f">>> {full_code}") + is_correct = 1 if check_correctness(full_code) else 0 + # print(f">>> {is_correct}") + return float(is_correct) diff --git a/reasoners/reward/cruxeval/utils.py b/reasoners/reward/cruxeval/utils.py new file mode 100644 index 00000000..9eec77cd --- /dev/null +++ b/reasoners/reward/cruxeval/utils.py @@ -0,0 +1,219 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. + +# This code is adapted from OpenAI's release +# https://github.com/openai/human-eval/blob/master/human_eval/execution.py + +import contextlib +import faulthandler +import io +import multiprocessing +import os +import platform +import signal +import tempfile + + +def check_correctness(check_program, timeout=3): + """ + Evaluates the functional correctness of a completion by running the test + suite provided in the problem. + + :param completion_id: an optional completion ID so we can match + the results later even if execution finishes asynchronously. + """ + manager = multiprocessing.Manager() + result = manager.list() + + p = multiprocessing.Process(target=unsafe_execute, args=(check_program, result, timeout)) + p.start() + p.join(timeout=timeout + 1) + if p.is_alive(): + p.kill() + + if not result: + result.append("timed out") + + return result[0] == "passed" + + +def unsafe_execute(check_program, result, timeout): + + with create_tempdir(): + + # These system calls are needed when cleaning up tempdir. + import os + import shutil + + rmtree = shutil.rmtree + rmdir = os.rmdir + chdir = os.chdir + + # Disable functionalities that can make destructive changes to the test. + reliability_guard() + + # Run program. + try: + exec_globals = {} + with swallow_io(): + with time_limit(timeout): + exec(check_program, exec_globals) + result.append("passed") + except TimeoutException: + result.append("timed out") + except BaseException as e: + result.append(f"failed: {e}") + + # Needed for cleaning up. + shutil.rmtree = rmtree + os.rmdir = rmdir + os.chdir = chdir + + +@contextlib.contextmanager +def time_limit(seconds): + def signal_handler(signum, frame): + raise TimeoutException("Timed out!") + + signal.setitimer(signal.ITIMER_REAL, seconds) + signal.signal(signal.SIGALRM, signal_handler) + try: + yield + finally: + signal.setitimer(signal.ITIMER_REAL, 0) + + +@contextlib.contextmanager +def swallow_io(): + stream = WriteOnlyStringIO() + with contextlib.redirect_stdout(stream): + with contextlib.redirect_stderr(stream): + with redirect_stdin(stream): + yield + + +@contextlib.contextmanager +def create_tempdir(): + with tempfile.TemporaryDirectory() as dirname: + with chdir(dirname): + yield dirname + + +class TimeoutException(Exception): + pass + + +class WriteOnlyStringIO(io.StringIO): + """StringIO that throws an exception when it's read from""" + + def read(self, *args, **kwargs): + raise OSError + + def readline(self, *args, **kwargs): + raise OSError + + def readlines(self, *args, **kwargs): + raise OSError + + def readable(self, *args, **kwargs): + """Returns True if the IO object can be read.""" + return False + + +class redirect_stdin(contextlib._RedirectStream): # type: ignore + _stream = "stdin" + + +@contextlib.contextmanager +def chdir(root): + if root == ".": + yield + return + cwd = os.getcwd() + os.chdir(root) + try: + yield + except BaseException as exc: + raise exc + finally: + os.chdir(cwd) + + +def reliability_guard(maximum_memory_bytes=None): + """ + This disables various destructive functions and prevents the generated code + from interfering with the test (e.g. fork bomb, killing other processes, + removing filesystem files, etc.) + + WARNING + This function is NOT a security sandbox. Untrusted code, including, model- + generated code, should not be blindly executed outside of one. See the + Codex paper for more information about OpenAI's code sandbox, and proceed + with caution. + """ + + if maximum_memory_bytes is not None: + import resource + + resource.setrlimit(resource.RLIMIT_AS, (maximum_memory_bytes, maximum_memory_bytes)) + resource.setrlimit(resource.RLIMIT_DATA, (maximum_memory_bytes, maximum_memory_bytes)) + if not platform.uname().system == "Darwin": + resource.setrlimit(resource.RLIMIT_STACK, (maximum_memory_bytes, maximum_memory_bytes)) + + faulthandler.disable() + + import builtins + + builtins.exit = None + builtins.quit = None + + import os + + os.environ["OMP_NUM_THREADS"] = "1" + + os.kill = None + os.system = None + os.putenv = None + os.remove = None + os.removedirs = None + os.rmdir = None + os.fchdir = None + os.setuid = None + os.fork = None + os.forkpty = None + os.killpg = None + os.rename = None + os.renames = None + os.truncate = None + os.replace = None + os.unlink = None + os.fchmod = None + os.fchown = None + os.chmod = None + os.chown = None + os.chroot = None + os.fchdir = None + os.lchflags = None + os.lchmod = None + os.lchown = None + os.getcwd = None + os.chdir = None + + import shutil + + shutil.rmtree = None + shutil.move = None + shutil.chown = None + + import subprocess + + subprocess.Popen = None # type: ignore + + __builtins__["help"] = None + + import sys + + sys.modules["ipdb"] = None + sys.modules["joblib"] = None + sys.modules["resource"] = None + sys.modules["psutil"] = None + sys.modules["tkinter"] = None \ No newline at end of file diff --git a/reasoners/reward/gpqa.py b/reasoners/reward/gpqa.py new file mode 100644 index 00000000..916e4890 --- /dev/null +++ b/reasoners/reward/gpqa.py @@ -0,0 +1,78 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import re + +def extract_solution(solution_str, method='strict'): + """ + Extract the final answer choice from an LLM's response to a multiple-choice GPQA question. + + Args: + solution_str (str): The full text response from the LLM + method (str): 'strict' for exact format matching, 'flexible' for more lenient matching + + Returns: + str: The extracted answer choice (A, B, C, or D) or None if not found + """ + assert method in ['strict', 'flexible'] + + if method == 'strict': + # Use regex from OpenAI simple-eval https://github.com/openai/simple-evals/blob/main/gpqa_eval.py + solution = re.search(r"(?i)Answer[ \t]*:[ \t]*\$?([A-D])\$?", solution_str) + if solution is None: + final_answer = None + else: + final_answer = solution.group(1) + elif method == 'flexible': + answer = re.findall(r"\(([A-D])\)", solution_str) + final_answer = None + if len(answer) == 0: + # No answer choices found in parentheses + pass + else: + invalid_str = [''] + # Find the last letter that is a valid answer choice + for final_answer in reversed(answer): + if final_answer not in invalid_str: + break + + return final_answer + + +def compute_score(solution_str, ground_truth, method='strict', format_score=0., score=1., extra_info=None, compressed: bool = False): + """The scoring function for GPQA. + + Reference: Trung, Luong, et al. "Reft: Reasoning with reinforced fine-tuning." Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). 2024. + + Args: + solution_str: the solution text + ground_truth: the ground truth + method: the method to extract the solution, choices are 'strict' and 'flexible' + format_score: the score for the format + score: the score for the correct answer + """ + if compressed: + from reward_score.utils import _deserialise_extra, _decompress_str + extra_info = _deserialise_extra(_decompress_str(extra_info)) + ground_truth = _deserialise_extra(_decompress_str(ground_truth["compressed"]))["ground_truth"] + print(f"ground_truth: {ground_truth}") + answer = extract_solution(solution_str=solution_str, method=method) + if answer is None: + return 0 + else: + if answer == ground_truth: + return score + else: + return format_score \ No newline at end of file diff --git a/reasoners/reward/graph_dataset.py b/reasoners/reward/graph_dataset.py new file mode 100644 index 00000000..4d0ec08c --- /dev/null +++ b/reasoners/reward/graph_dataset.py @@ -0,0 +1,83 @@ +import re +import random +import ast +import operator +import signal +import contextlib + +class TimeoutException(Exception): + pass + +@contextlib.contextmanager +def time_limit(seconds: float): + def signal_handler(signum, frame): + raise TimeoutException("Timed out!") + + signal.setitimer(signal.ITIMER_REAL, seconds) + signal.signal(signal.SIGALRM, signal_handler) + try: + yield + finally: + signal.setitimer(signal.ITIMER_REAL, 0) + +def extract_solution(solution_str): + + answer_pattern = r'(.*?)' + match = re.finditer(answer_pattern, solution_str) + matches = list(match) + + if matches: + final_answer = matches[-1].group(1).strip() + if re.search(r'^[A-Za-z]+$', final_answer): + return final_answer + else: + return None + +def compute_score(data_source: str, + solution_str: str, + ground_truth: str, + extra_info: any = None, + compressed: bool = False, + timeout: float = 10.0): + """The scoring function for graph dataset task. + + Args: + solution_str: the solution text + ground_truth: the correct answer + timeout: maximum time in seconds to allow for computation + """ + if compressed: + from reward_score.utils import _deserialise_extra, _decompress_str + extra_info = _deserialise_extra(_decompress_str(extra_info)) + ground_truth = _deserialise_extra(_decompress_str(ground_truth["compressed"]))["ground_truth"] + print(f"ground_truth: {ground_truth}") + try: + with time_limit(timeout): + if not isinstance(ground_truth, str): + ground_truth = str(ground_truth) + + target = ground_truth.lower() + solution = extract_solution(solution_str) + + if solution: + solution = solution.lower() + else: + score = 0.0 + + try: + if target == solution: + score = 1.0 + else: + score = 0.0 + + except Exception as e: + score = 0.0 + + except TimeoutException: + print("Computation timed out in graph_dataset") + score = 0.0 + except Exception as e: + print(f"Error in compute_score in graph_dataset: {e}") + score = 0.0 + + return float(score) diff --git a/reasoners/reward/math.py b/reasoners/reward/math.py new file mode 100644 index 00000000..50792aa6 --- /dev/null +++ b/reasoners/reward/math.py @@ -0,0 +1,227 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# Adapted from https://github.com/EleutherAI/lm-evaluation-harness/blob/main/lm_eval/tasks/hendrycks_math/utils.py + + +def compute_score(solution_str, ground_truth) -> float: + retval = 0. + try: + string_in_last_boxed = last_boxed_only_string(solution_str) + if string_in_last_boxed is not None: + answer = remove_boxed(string_in_last_boxed) + if is_equiv(answer, ground_truth): + retval = 1. + except Exception as e: + print(e) + + return retval + + +# string normalization from https://github.com/EleutherAI/lm-evaluation-harness/blob/master/lm_eval/tasks/hendrycks_math.py +def is_equiv(str1, str2, verbose=False): + if str1 is None and str2 is None: + print("WARNING: Both None") + return True + if str1 is None or str2 is None: + return False + + try: + ss1 = strip_string(str1) + ss2 = strip_string(str2) + if verbose: + print(ss1, ss2) + return ss1 == ss2 + except Exception: + return str1 == str2 + + +def remove_boxed(s): + if "\\boxed " in s: + left = "\\boxed " + assert s[:len(left)] == left + return s[len(left):] + + left = "\\boxed{" + + assert s[:len(left)] == left + assert s[-1] == "}" + + return s[len(left):-1] + + +def last_boxed_only_string(string): + idx = string.rfind("\\boxed") + if "\\boxed " in string: + return "\\boxed " + string.split("\\boxed ")[-1].split("$")[0] + if idx < 0: + idx = string.rfind("\\fbox") + if idx < 0: + return None + + i = idx + right_brace_idx = None + num_left_braces_open = 0 + while i < len(string): + if string[i] == "{": + num_left_braces_open += 1 + if string[i] == "}": + num_left_braces_open -= 1 + if num_left_braces_open == 0: + right_brace_idx = i + break + i += 1 + + if right_brace_idx is None: + retval = None + else: + retval = string[idx:right_brace_idx + 1] + + return retval + + +def fix_fracs(string): + substrs = string.split("\\frac") + new_str = substrs[0] + if len(substrs) > 1: + substrs = substrs[1:] + for substr in substrs: + new_str += "\\frac" + if substr[0] == "{": + new_str += substr + else: + try: + assert len(substr) >= 2 + except AssertionError: + return string + a = substr[0] + b = substr[1] + if b != "{": + if len(substr) > 2: + post_substr = substr[2:] + new_str += "{" + a + "}{" + b + "}" + post_substr + else: + new_str += "{" + a + "}{" + b + "}" + else: + if len(substr) > 2: + post_substr = substr[2:] + new_str += "{" + a + "}" + b + post_substr + else: + new_str += "{" + a + "}" + b + string = new_str + return string + + +def fix_a_slash_b(string): + if len(string.split("/")) != 2: + return string + a = string.split("/")[0] + b = string.split("/")[1] + try: + a = int(a) + b = int(b) + assert string == "{}/{}".format(a, b) + new_string = "\\frac{" + str(a) + "}{" + str(b) + "}" + return new_string + except AssertionError: + return string + + +def remove_right_units(string): + # "\\text{ " only ever occurs (at least in the val set) when describing units + if "\\text{ " in string: + splits = string.split("\\text{ ") + assert len(splits) == 2 + return splits[0] + else: + return string + + +def fix_sqrt(string): + if "\\sqrt" not in string: + return string + splits = string.split("\\sqrt") + new_string = splits[0] + for split in splits[1:]: + if split[0] != "{": + a = split[0] + new_substr = "\\sqrt{" + a + "}" + split[1:] + else: + new_substr = "\\sqrt" + split + new_string += new_substr + return new_string + + +def strip_string(string): + # linebreaks + string = string.replace("\n", "") + + # remove inverse spaces + string = string.replace("\\!", "") + + # replace \\ with \ + string = string.replace("\\\\", "\\") + + # replace tfrac and dfrac with frac + string = string.replace("tfrac", "frac") + string = string.replace("dfrac", "frac") + + # remove \left and \right + string = string.replace("\\left", "") + string = string.replace("\\right", "") + + # Remove circ (degrees) + string = string.replace("^{\\circ}", "") + string = string.replace("^\\circ", "") + + # remove dollar signs + string = string.replace("\\$", "") + + # remove units (on the right) + string = remove_right_units(string) + + # remove percentage + string = string.replace("\\%", "") + string = string.replace("\%", "") # noqa: W605 + + # " 0." equivalent to " ." and "{0." equivalent to "{." Alternatively, add "0" if "." is the start of the string + string = string.replace(" .", " 0.") + string = string.replace("{.", "{0.") + # if empty, return empty string + if len(string) == 0: + return string + if string[0] == ".": + string = "0" + string + + # to consider: get rid of e.g. "k = " or "q = " at beginning + if len(string.split("=")) == 2: + if len(string.split("=")[0]) <= 2: + string = string.split("=")[1] + + # fix sqrt3 --> sqrt{3} + string = fix_sqrt(string) + + # remove spaces + string = string.replace(" ", "") + + # \frac1b or \frac12 --> \frac{1}{b} and \frac{1}{2}, etc. Even works with \frac1{72} (but not \frac{72}1). Also does a/b --> \\frac{a}{b} + string = fix_fracs(string) + + # manually change 0.5 --> \frac{1}{2} + if string == "0.5": + string = "\\frac{1}{2}" + + # NOTE: X/Y changed to \frac{X}{Y} in dataset, but in simple cases fix in case the model output is X/Y + string = fix_a_slash_b(string) + + return string diff --git a/reasoners/reward/math_llm_judge/__init__.py b/reasoners/reward/math_llm_judge/__init__.py new file mode 100644 index 00000000..8ac4ec51 --- /dev/null +++ b/reasoners/reward/math_llm_judge/__init__.py @@ -0,0 +1,458 @@ +# Copyright 2024 PRIME team and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +input_template = """You are a teacher and your task is to grade the student's answer with the reference answer. + +Question: {QUESTION} + +Student's Answer: {STUDENT_ANSWER} + +Reference Answer: {REFERENCE_ANSWER} + +You only need to refer to the reference answer to grade the student's answer. Sometimes the student's answer is expressed in a different way from the reference answer, but the meaning is the same, and you should still consider it correct. If they are not equivalent in mathematical sense, you should consider it incorrect. + +Note 1: You don't need to solve the problem yourself. Just grade the student's answer based on the reference answer. + +Note 2: If the reference answer is a range, please make sure the student's answer is strictly identical, including the open or closed interval. + +Note 3: If the reference answer is an expression and it looks like the student's answer is equivalent to the reference answer, you should present the derivation process to check if they are equivalent. + +Note 4: If the reference answer includes multiple solutions, please make sure the student's answer covers all of them. + +Please provide a brief explanation (a few sentences) of your grading process and put your final grade in the following format: + +Final Grade: CORRECT or INCORRECT +""" + +import re +import sympy +from pylatexenc import latex2text +from sympy.parsing import sympy_parser +import os +from . import math_normalize +from .grader import math_equal + +import requests + +# import math_normalize +# from grader import math_equal + +# sympy might hang -- we don't care about trying to be lenient in these cases +BAD_SUBSTRINGS = ["^{", "^("] +BAD_REGEXES = ["\^[0-9]+\^", "\^[0-9][0-9]+"] +TUPLE_CHARS = "()[]" + + +def timeout(timeout_seconds: int = 8): + if os.name == "posix": + import signal + + def decorator(func): + + def handler(signum, frame): + raise TimeoutError("Operation timed out!") + + def wrapper(*args, **kwargs): + old_handler = signal.getsignal(signal.SIGALRM) + signal.signal(signal.SIGALRM, handler) + signal.alarm(timeout_seconds) + + try: + return func(*args, **kwargs) + finally: + signal.alarm(0) + signal.signal(signal.SIGALRM, old_handler) + + return wrapper + + return decorator + else: + raise NotImplementedError(f"Unsupported OS: {os.name}") + + +def _sympy_parse(expr: str): + """Parses an expression with sympy.""" + py_expr = expr.replace("^", "**") + return sympy_parser.parse_expr( + py_expr, + transformations=(sympy_parser.standard_transformations + (sympy_parser.implicit_multiplication_application,)), + ) + + +def _parse_latex(expr: str) -> str: + """Attempts to parse latex to an expression sympy can read.""" + expr = expr.replace("\\tfrac", "\\frac") + expr = expr.replace("\\dfrac", "\\frac") + expr = expr.replace("\\frac", " \\frac") # Play nice with mixed numbers. + expr = latex2text.LatexNodes2Text().latex_to_text(expr) + + # Replace the specific characters that this parser uses. + expr = expr.replace("√", "sqrt") + expr = expr.replace("π", "pi") + expr = expr.replace("∞", "inf") + expr = expr.replace("∪", "U") + expr = expr.replace("·", "*") + expr = expr.replace("×", "*") + + return expr.strip() + + +def _is_float(num: str) -> bool: + try: + float(num) + return True + except ValueError: + return False + + +def _is_int(x: float) -> bool: + try: + return abs(x - int(round(x))) <= 1e-7 + except: + return False + + +def _is_frac(expr: str) -> bool: + return bool(re.search(r"^-?[0-9]+.?/0*[1-9][0-9]*.?$", expr)) + + +def _str_is_int(x: str) -> bool: + try: + x = _strip_properly_formatted_commas(x) + x = float(x) + return abs(x - int(round(x))) <= 1e-7 + except: + return False + + +def _str_to_int(x: str) -> bool: + x = x.replace(",", "") + x = float(x) + return int(x) + + +def _inject_implicit_mixed_number(step: str): + """ + Automatically make a mixed number evalable + e.g. 7 3/4 => 7+3/4 + """ + p1 = re.compile("([0-9]) +([0-9])") + step = p1.sub("\\1+\\2", step) ## implicit mults + return step + + +def _strip_properly_formatted_commas(expr: str): + # We want to be careful because we don't want to strip tuple commas + p1 = re.compile("(\d)(,)(\d\d\d)($|\D)") + while True: + next_expr = p1.sub("\\1\\3\\4", expr) + if next_expr == expr: + break + expr = next_expr + return next_expr + + +def _normalize(expr: str) -> str: + """Normalize answer expressions.""" + if expr is None: + return None + + # Remove enclosing `\text{}`. + m = re.search("^\\\\text\{(?P.+?)\}$", expr) + if m is not None: + expr = m.group("text") + + expr = expr.replace("\\%", "%") + expr = expr.replace("\\$", "$") + expr = expr.replace("$", "") + expr = expr.replace("%", "") + expr = expr.replace(" or ", " , ") + expr = expr.replace(" and ", " , ") + + expr = expr.replace("million", "*10^6") + expr = expr.replace("billion", "*10^9") + expr = expr.replace("trillion", "*10^12") + + for unit in [ + "degree", + "cm", + "centimeter", + "meter", + "mile", + "second", + "minute", + "hour", + "day", + "week", + "month", + "year", + "foot", + "feet", + "inch", + "yard", + "liter", + ]: + expr = re.sub(f"{unit}(es)?(s)? *(\^[0-9]+)?", "", expr) + expr = re.sub(f"\^ *\\\\circ", "", expr) + + if len(expr) > 0 and expr[0] == "{" and expr[-1] == "}": + expr = expr[1:-1] + + expr = re.sub(",\\\\! *", "", expr) + if _is_float(expr) and _is_int(float(expr)): + expr = str(int(round(float(expr)))) + if "\\" in expr: + try: + expr = _parse_latex(expr) + except: + pass + + # edge case with mixed numbers and negative signs + expr = re.sub("- *", "-", expr) + + expr = _inject_implicit_mixed_number(expr) + + # don't be case sensitive for text answers + expr = expr.lower() + + if _str_is_int(expr): + expr = str(_str_to_int(expr)) + + return expr + + +def count_unknown_letters_in_expr(expr: str): + expr = expr.replace("sqrt", "") + expr = expr.replace("frac", "") + letters_in_expr = set([x for x in expr if x.isalpha()]) + return len(letters_in_expr) + + +def should_allow_eval(expr: str): + # we don't want to try parsing unknown text or functions of more than two variables + if count_unknown_letters_in_expr(expr) > 2: + return False + + for bad_string in BAD_SUBSTRINGS: + if bad_string in expr: + return False + + for bad_regex in BAD_REGEXES: + if re.search(bad_regex, expr) is not None: + return False + + return True + + +@timeout(timeout_seconds=10) +def are_equal_under_sympy(ground_truth_normalized: str, given_normalized: str): + are_equal = False + try: + expr = f"({ground_truth_normalized})-({given_normalized})" + if should_allow_eval(expr): + sympy_diff = _sympy_parse(expr) + simplified = sympy.simplify(sympy_diff) + if simplified == 0: + are_equal = True + except: + pass + return are_equal + + +def split_tuple(expr: str): + """ + Split the elements in a tuple/interval, while handling well-formatted commas in large numbers + """ + expr = _strip_properly_formatted_commas(expr) + if len(expr) == 0: + return [] + if (len(expr) > 2 and expr[0] in TUPLE_CHARS and expr[-1] in TUPLE_CHARS and + all([ch not in expr[1:-1] for ch in TUPLE_CHARS])): + elems = [elem.strip() for elem in expr[1:-1].split(",")] + else: + elems = [expr] + return elems + + +def grade_answer(given_answer: str, ground_truth: str) -> bool: + """ + The answer will be considered correct if: + (a) it normalizes to the same string as the ground truth answer + OR + (b) sympy can simplify the difference between the expressions to 0 + """ + if given_answer is None: + return False + + ground_truth_normalized_mathd = math_normalize.normalize_answer(ground_truth) + given_answer_normalized_mathd = math_normalize.normalize_answer(given_answer) + + # be at least as lenient as mathd + if ground_truth_normalized_mathd == given_answer_normalized_mathd: + return True + + ground_truth_normalized = _normalize(ground_truth) + given_normalized = _normalize(given_answer) + + if ground_truth_normalized is None: + return False + + if ground_truth_normalized == given_normalized: + return True + + if len(given_normalized) == 0: + return False + + ground_truth_elems = split_tuple(ground_truth_normalized) + given_elems = split_tuple(given_normalized) + + if len(ground_truth_elems) > 1 and (ground_truth_normalized[0] != given_normalized[0] or + ground_truth_normalized[-1] != given_normalized[-1]): + is_correct = False + elif len(ground_truth_elems) != len(given_elems): + is_correct = False + else: + for ground_truth_elem, given_elem in zip(ground_truth_elems, given_elems): + if _is_frac(ground_truth_elem) and _is_frac(given_elem): + # if fractions aren't reduced, then shouldn't be marked as correct + # so, we don't want to allow sympy.simplify in this case + is_correct = ground_truth_elem == given_elem + elif _str_is_int(ground_truth_elem) != _str_is_int(given_elem): + # if the ground truth answer is an integer, we require the given answer to be a strict match (no sympy.simplify) + is_correct = False + else: + is_correct = are_equal_under_sympy(ground_truth_elem, given_elem) + if not is_correct: + break + + return is_correct + + +def remove_boxed(s): + left = "\\boxed{" + try: + assert s[:len(left)] == left + assert s[-1] == "}" + return s[len(left):-1] + except: + return None + + +def _last_boxed_only_string(string): + idx = string.rfind("\\boxed") + if idx < 0: + idx = string.rfind("\\fbox") + if idx < 0: + return None + + i = idx + left_brace_idx = None + right_brace_idx = None + num_left_braces_open = 0 + while i < len(string): + if string[i] == "{": + num_left_braces_open += 1 + if left_brace_idx is None: + left_brace_idx = i + elif string[i] == "}": + num_left_braces_open -= 1 + if num_left_braces_open == 0: + right_brace_idx = i + break + + i += 1 + + if left_brace_idx is None or right_brace_idx is None: + return None + + return string[left_brace_idx + 1:right_brace_idx].strip() + + +def match_answer(response): + is_matched = False + response = response.split("")[-1] + + # Find boxed + ans_boxed = _last_boxed_only_string(response) + if ans_boxed: + is_matched = True + response = ans_boxed + + return is_matched, response + + +import math + +def llm_check_answer(model_output: str, ground_truth: str, question: str) -> bool: + # use llm to check if the answer is correct + + # url = "http://176.56.200.81:30000/v1/chat/completions" + url = os.getenv("MATH_LLM_JUDGE_URL") + if not url: + raise ValueError("MATH_LLM_JUDGE_URL is not set") + + prompt = input_template.format(QUESTION=question, STUDENT_ANSWER=model_output, REFERENCE_ANSWER=ground_truth) + + data = { + "model": "Qwen/Qwen2.5-32B-Instruct", + "messages": [{"role": "user", "content": prompt}], + } + response = requests.post(url, json=data) + eval_result = not "INCORRECT" in response.json()['choices'][0]['message']['content'] \ + and "CORRECT" in response.json()['choices'][0]['message']['content'] + # print({ + # "model_output": model_output, + # "ground_truth": ground_truth, + # "question": question, + # "response": response.json()['choices'][0]['message']['content'], + # "eval_result": eval_result, + # }) + return eval_result + +def compute_score(data_source: str, + model_output: str, + ground_truth: str, + extra_info: dict, + compressed: bool = False) -> float: + if compressed: + from reward_score.utils import _deserialise_extra, _decompress_str + extra_info = _deserialise_extra(_decompress_str(extra_info["compressed"])) + ground_truth = _deserialise_extra(_decompress_str(ground_truth["compressed"]))["ground_truth"] + + question = extra_info["question"] + model_output = str(model_output) + ground_truth = str(ground_truth) + + is_matched, extracted_model_output = match_answer(model_output) + + # grade simple algebra questions. if succeeded, return; otherwise, proceed to more complex grading + if grade_answer(extracted_model_output, ground_truth): + return 1.0 + + try: + if "\\pi" in extracted_model_output or "\\pi" in ground_truth: + equivs = [] + for pi in [math.pi, 3.14]: + equivs.append(math_equal(extracted_model_output, ground_truth, timeout=True, pi=pi)) + is_correct = any(equivs) + else: + is_correct = math_equal(extracted_model_output, ground_truth, timeout=True) + except: + is_correct = False + + if is_matched and not is_correct: + # use llm to check if the answer is correct + is_correct = llm_check_answer(extracted_model_output, ground_truth, question) + + return float(is_correct) diff --git a/reasoners/reward/math_llm_judge/grader.py b/reasoners/reward/math_llm_judge/grader.py new file mode 100644 index 00000000..34f0c7f5 --- /dev/null +++ b/reasoners/reward/math_llm_judge/grader.py @@ -0,0 +1,380 @@ +# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Copyright (c) Microsoft Corporation. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE + +# Copyright (c) 2023 OpenAI +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: + +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +# Copyright (c) 2021 Dan Hendrycks +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +# Copyright 2024 PRIME team and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +This logic is largely copied from the Hendrycks' MATH release (math_equivalence), and borrowed from: +- https://github.com/microsoft/ToRA/blob/main/src/eval/grader.py +- https://github.com/microsoft/ProphetNet/tree/master/CRITIC +- https://github.com/openai/prm800k +""" + +import contextlib +import re +import signal +import math +from math import isclose +from typing import Union + +from sympy import N, simplify +from sympy.parsing.latex import parse_latex +from sympy.parsing.sympy_parser import parse_expr + + +def is_digit(s): + try: + if "{,}" in str(s): + num = float(str(s).replace("{,}", "")) + return True, num + + num = float(str(s).replace(",", "")) + return True, num + except ValueError: + return False, None + + +def normalize(answer, pi) -> str: + # checking if answer is $ and removing $ in that case to compare + if isinstance(answer, str) and bool(re.match(r'\$\d+(\.\d+)?', answer)): + return answer[1:] + + # checking if answer is % or \\% and removing % + if isinstance(answer, str) and (bool(re.match(r'^\d+(\.\d+)?%$', answer)) or + bool(re.match(r'^\d+(\.\d+)?\\%$', answer))): + return answer.replace("\\%", "").replace("%", "") + + # handle base + answer = handle_base(answer) + + # handle pi + answer = handle_pi(answer, pi) + + return answer + + +def handle_base(x) -> str: + if isinstance(x, str) and "_" in x: + # Due to base + x = x.split("_")[0] + x = float(x) + return int(x) + return x + + +def handle_pi(string, pi): + if isinstance(string, str) and "\\pi" in string: + # Find the first occurrence of "\pi" + idx = string.find("\\pi") + + # Iterate over the string and find all occurrences of "\pi" with a valid previous character + while idx != -1: + + if idx > 0 and string[idx - 1].isdigit(): + # Replace "\pi" with "*math.pi" if the previous character is a digit + string = string[:idx] + f"*{pi}" + string[idx + 3:] + else: + # Replace "\pi" with "1*math.pi" if the previous character is not a digit + string = string[:idx] + f"1*{pi}" + string[idx + 3:] + + # Find the next occurrence of "\pi" + idx = string.find("\\pi", idx + 1) + + # Evaluate the expression using eval() function + try: + string = eval(string) + except: + pass + + return string + + +def math_equal(prediction: Union[bool, float, str], + reference: Union[float, str], + include_percentage: bool = True, + tolerance: float = 1e-4, + timeout: float = 10.0, + pi: float = math.pi) -> bool: + """ + Exact match of math if and only if: + 1. numerical equal: both can convert to float and are equal + 2. symbolic equal: both can convert to sympy expression and are equal + """ + + prediction = normalize(prediction, pi) + reference = normalize(reference, pi) + + if isinstance(prediction, str) and len(prediction) > 1000: # handling weird corner-cases + prediction = prediction[:1000] + + # 0. string comparison + if isinstance(prediction, str) and isinstance(reference, str): + if prediction.strip().lower() == reference.strip().lower(): + return True + if prediction.replace(" ", "") == reference.replace(" ", ""): + return True + + try: # 1. numerical equal + if is_digit(prediction)[0] and is_digit(reference)[0]: + prediction = is_digit(prediction)[1] + reference = is_digit(reference)[1] + # number questions + if include_percentage: + gt_result = [reference / 100, reference, reference * 100] + else: + gt_result = [reference] + for item in gt_result: + try: + if isclose(item, prediction, rel_tol=tolerance): + return True + except Exception: + continue + return False + except Exception: + pass + + if not prediction and prediction not in [0, False]: + return False + + # 2. symbolic equal + reference = str(reference).strip() + prediction = str(prediction).strip() + + ## deal with [], (), {} + prediction = format_intervals(prediction) + + pred_str, ref_str = prediction, reference + if (prediction.startswith("[") and prediction.endswith("]") and + not reference.startswith("(")) or (prediction.startswith("(") and prediction.endswith(")") and + not reference.startswith("[")): + pred_str = pred_str.strip("[]()") + ref_str = ref_str.strip("[]()") + for s in ["{", "}", "(", ")"]: + ref_str = ref_str.replace(s, "") + pred_str = pred_str.replace(s, "") + if pred_str == ref_str: + return True + + ## [a, b] vs. [c, d], return a==c and b==d + if (prediction and reference and prediction[0] in "([" and prediction[-1] in ")]" and + prediction[0] == reference[0] and prediction[-1] == reference[-1]): + pred_parts = prediction[1:-1].split(",") + ref_parts = reference[1:-1].split(",") + if len(pred_parts) == len(ref_parts): + if all([ + math_equal(pred_pt, ref_pt, include_percentage, tolerance) + for pred_pt, ref_pt in zip(pred_parts, ref_parts) + ]): + return True + + if "," in prediction and "," in reference: + pred_parts = [item.strip() for item in prediction.split(",")] + ref_parts = [item.strip() for item in reference.split(",")] + + if len(pred_parts) == len(ref_parts): + if all([ + math_equal(pred_parts[i], ref_parts[i], include_percentage, tolerance) + for i in range(len(pred_parts)) + ]): + return True + else: + return False + + # if we have point == tuple of values + if prediction.startswith("Point") and reference[0] == "(" and reference[-1] == ")": + pred_parts = prediction[prediction.find("(") + 1:-1].split(",") + ref_parts = reference[1:-1].split(",") + if len(pred_parts) == len(ref_parts): + if all([ + math_equal(pred_pt, ref_pt, include_percentage, tolerance) + for pred_pt, ref_pt in zip(pred_parts, ref_parts) + ]): + return True + + # if reference is a matrix + if "\\\begin{pmatrix}" in reference and prediction.startswith("Matrix"): + try: + pred_matrix = parse_expr(prediction) + ref_matrix_items = reference.split()[1:-1:2] + if len(pred_matrix) == len(ref_matrix_items): + if all([ + math_equal(pred, ref, include_percentage, tolerance) + for ref, pred in zip(ref_matrix_items, pred_matrix) + ]): + return True + except Exception: + pass + elif "\\begin{pmatrix}" in reference and prediction.startswith("[") and prediction.endswith("]"): + if isinstance(eval(prediction), list): + try: + pred_matrix = eval(prediction) + # ref_matrix_items = reference.split()[1:-1:2] + ref_matrix_items = reference.lstrip("\\begin{pmatrix}").lstrip("\begin{pmatrix}").rstrip( + "\\end{pmatrix}").rstrip("\\end{pmatrix}") + ref_matrix_items = ref_matrix_items.split("\\") + ref_matrix_items = [row.split("&") if "&" in row else row for row in ref_matrix_items] + if len(pred_matrix) == len(ref_matrix_items): + if all([ + math_equal(pred, ref, include_percentage, tolerance) + for ref, pred in zip(ref_matrix_items, pred_matrix) + ]): + return True + except Exception: + pass + + return symbolic_equal(prediction, reference, tolerance, timeout) + + +def symbolic_equal(a, b, tolerance, timeout=10.0): + + def _parse(s): + for f in [parse_expr, parse_latex]: + try: + with time_limit(timeout): + return f(s) + except Exception: + pass + return s + + a = _parse(a) + b = _parse(b) + + try: + with time_limit(timeout): + if simplify(a - b) == 0: + return True + except Exception: + pass + + try: + with time_limit(timeout): + if isclose(N(a), N(b), rel_tol=tolerance): + return True + except Exception: + pass + return False + + +class TimeoutException(Exception): + pass + + +@contextlib.contextmanager +def time_limit(seconds: float): + + def signal_handler(signum, frame): + raise TimeoutException("Timed out!") + + signal.setitimer(signal.ITIMER_REAL, seconds) + signal.signal(signal.SIGALRM, signal_handler) + try: + yield + finally: + signal.setitimer(signal.ITIMER_REAL, 0) + + +def format_intervals(prediction): + patterns = { + "Interval(": r"^Interval\((.*)\)$", + "Interval.Ropen(": r"^Interval\.Ropen\((.*)\)$", + "Interval.Lopen(": r"^Interval\.Lopen\((.*)\)$", + "Interval.open(": r"^Interval\.open\((.*)\)$", + } + + for key, pattern in patterns.items(): + match = re.match(pattern, prediction) + if match: + inner_content = match.group(1) + + if key == "Interval(": # Intarval(a, b) == [a, b] + return f"[{inner_content}]" + elif key == "Interval.Ropen(": # Intarval.Ropen(a, b) == [a, b) + return f"[{inner_content})" + elif key == "Interval.Lopen(": # Intarval.Lopen(a, b) == (a, b] + return f"({inner_content}]" + elif key == "Interval.open(": # Intarval.open(a, b) == (a, b) + return f"({inner_content})" + + return prediction diff --git a/reasoners/reward/math_llm_judge/math_normalize.py b/reasoners/reward/math_llm_judge/math_normalize.py new file mode 100644 index 00000000..e9921a5a --- /dev/null +++ b/reasoners/reward/math_llm_judge/math_normalize.py @@ -0,0 +1,191 @@ +# Copyright 2024 PRIME team and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Copyright (c) 2021 Dan Hendrycks +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +""" +This logic is largely copied from the Hendrycks' MATH release (math_equivalence). + +From: https://github.com/openai/prm800k/blob/main/prm800k/grading/math_normalize.py +""" +import re +from typing import Optional + + +def normalize_answer(answer: Optional[str]) -> Optional[str]: + if answer is None: + return None + answer = answer.strip() + try: + # Remove enclosing `\text{}`. + m = re.search("^\\\\text\{(?P.+?)\}$", answer) + if m is not None: + answer = m.group("text").strip() + return _strip_string(answer) + except: + return answer + + +def _fix_fracs(string): + substrs = string.split("\\frac") + new_str = substrs[0] + if len(substrs) > 1: + substrs = substrs[1:] + for substr in substrs: + new_str += "\\frac" + if substr[0] == "{": + new_str += substr + else: + try: + assert len(substr) >= 2 + except: + return string + a = substr[0] + b = substr[1] + if b != "{": + if len(substr) > 2: + post_substr = substr[2:] + new_str += "{" + a + "}{" + b + "}" + post_substr + else: + new_str += "{" + a + "}{" + b + "}" + else: + if len(substr) > 2: + post_substr = substr[2:] + new_str += "{" + a + "}" + b + post_substr + else: + new_str += "{" + a + "}" + b + string = new_str + return string + + +def _fix_a_slash_b(string): + if len(string.split("/")) != 2: + return string + a = string.split("/")[0] + b = string.split("/")[1] + try: + a = int(a) + b = int(b) + assert string == "{}/{}".format(a, b) + new_string = "\\frac{" + str(a) + "}{" + str(b) + "}" + return new_string + except: + return string + + +def _remove_right_units(string): + # "\\text{ " only ever occurs (at least in the val set) when describing units + if "\\text{ " in string: + splits = string.split("\\text{ ") + assert len(splits) == 2 + return splits[0] + else: + return string + + +def _fix_sqrt(string): + if "\\sqrt" not in string: + return string + splits = string.split("\\sqrt") + new_string = splits[0] + for split in splits[1:]: + if split[0] != "{": + a = split[0] + new_substr = "\\sqrt{" + a + "}" + split[1:] + else: + new_substr = "\\sqrt" + split + new_string += new_substr + return new_string + + +def _strip_string(string): + # linebreaks + string = string.replace("\n", "") + + # remove inverse spaces + string = string.replace("\\!", "") + + # replace \\ with \ + string = string.replace("\\\\", "\\") + + # replace tfrac and dfrac with frac + string = string.replace("tfrac", "frac") + string = string.replace("dfrac", "frac") + + # remove \left and \right + string = string.replace("\\left", "") + string = string.replace("\\right", "") + + # Remove circ (degrees) + string = string.replace("^{\\circ}", "") + string = string.replace("^\\circ", "") + + # remove dollar signs + string = string.replace("\\$", "") + + # remove units (on the right) + string = _remove_right_units(string) + + # remove percentage + string = string.replace("\\%", "") + string = string.replace("\%", "") + + # " 0." equivalent to " ." and "{0." equivalent to "{." Alternatively, add "0" if "." is the start of the string + string = string.replace(" .", " 0.") + string = string.replace("{.", "{0.") + # if empty, return empty string + if len(string) == 0: + return string + if string[0] == ".": + string = "0" + string + + # to consider: get rid of e.g. "k = " or "q = " at beginning + if len(string.split("=")) == 2: + if len(string.split("=")[0]) <= 2: + string = string.split("=")[1] + + # fix sqrt3 --> sqrt{3} + string = _fix_sqrt(string) + + # remove spaces + string = string.replace(" ", "") + + # \frac1b or \frac12 --> \frac{1}{b} and \frac{1}{2}, etc. Even works with \frac1{72} (but not \frac{72}1). Also does a/b --> \\frac{a}{b} + string = _fix_fracs(string) + + # manually change 0.5 --> \frac{1}{2} + if string == "0.5": + string = "\\frac{1}{2}" + + # NOTE: X/Y changed to \frac{X}{Y} in dataset, but in simple cases fix in case the model output is X/Y + string = _fix_a_slash_b(string) + + return string \ No newline at end of file diff --git a/reasoners/reward/naive_dapo.py b/reasoners/reward/naive_dapo.py new file mode 100644 index 00000000..a4f5c30a --- /dev/null +++ b/reasoners/reward/naive_dapo.py @@ -0,0 +1,526 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# Adapted from https://github.com/EleutherAI/lm-evaluation-harness/blob/main/lm_eval/tasks/hendrycks_math/utils.py + +import re +import signal +from typing import Optional + +import sympy +from pylatexenc import latex2text +from sympy.parsing import sympy_parser +import os +from reward_score.prime_math import math_normalize +from reward_score.prime_math.grader import math_equal + + +class timeout: + + def __init__(self, seconds=1, error_message="Timeout"): + self.seconds = seconds + self.error_message = error_message + + def handle_timeout(self, signum, frame): + raise TimeoutError(self.error_message) + + def __enter__(self): + signal.signal(signal.SIGALRM, self.handle_timeout) + signal.alarm(self.seconds) + + def __exit__(self, type, value, traceback): + signal.alarm(0) + + +# Constants for normalization +SUBSTITUTIONS = [ + ("an ", ""), + ("a ", ""), + (".$", "$"), + ("\\$", ""), + (r"\ ", ""), + (" ", ""), + ("mbox", "text"), + (",\\text{and}", ","), + ("\\text{and}", ","), + ("\\text{m}", "\\text{}"), +] + +REMOVED_EXPRESSIONS = [ + "square", + "ways", + "integers", + "dollars", + "mph", + "inches", + "hours", + "km", + "units", + "\\ldots", + "sue", + "points", + "feet", + "minutes", + "digits", + "cents", + "degrees", + "cm", + "gm", + "pounds", + "meters", + "meals", + "edges", + "students", + "childrentickets", + "multiples", + "\\text{s}", + "\\text{.}", + "\\text{\ns}", + "\\text{}^2", + "\\text{}^3", + "\\text{\n}", + "\\text{}", + r"\mathrm{th}", + r"^\circ", + r"^{\circ}", + r"\;", + r",\!", + "{,}", + '"', + "\\dots", +] + + +def normalize_final_answer(final_answer: str) -> str: + """Normalize a final answer to a quantitative reasoning question. + + Args: + final_answer: The answer string to normalize + + Returns: + Normalized answer string + """ + final_answer = final_answer.split("=")[-1] + + # Apply substitutions and removals + for before, after in SUBSTITUTIONS: + final_answer = final_answer.replace(before, after) + for expr in REMOVED_EXPRESSIONS: + final_answer = final_answer.replace(expr, "") + + # Extract and normalize LaTeX math + final_answer = re.sub(r"(.*?)(\$)(.*?)(\$)(.*)", "$\\3$", final_answer) + final_answer = re.sub(r"(\\text\{)(.*?)(\})", "\\2", final_answer) + final_answer = re.sub(r"(\\textbf\{)(.*?)(\})", "\\2", final_answer) + final_answer = re.sub(r"(\\overline\{)(.*?)(\})", "\\2", final_answer) + final_answer = re.sub(r"(\\boxed\{)(.*)(\})", "\\2", final_answer) + + # Normalize shorthand TeX: + # \fracab -> \frac{a}{b} + # \frac{abc}{bef} -> \frac{abc}{bef} + # \fracabc -> \frac{a}{b}c + # \sqrta -> \sqrt{a} + # \sqrtab -> sqrt{a}b + final_answer = re.sub(r"(frac)([^{])(.)", "frac{\\2}{\\3}", final_answer) + final_answer = re.sub(r"(sqrt)([^{])", "sqrt{\\2}", final_answer) + final_answer = final_answer.replace("$", "") + + # Normalize numbers + if final_answer.replace(",", "").isdigit(): + final_answer = final_answer.replace(",", "") + + return final_answer.strip() + + +# sympy might hang -- we don't care about trying to be lenient in these cases +BAD_SUBSTRINGS = ["^{", "^("] +BAD_REGEXES = ["\^[0-9]+\^", "\^[0-9][0-9]+"] +TUPLE_CHARS = "()[]" + + +def timeout(timeout_seconds: int = 8): + if os.name == "posix": + import signal + + def decorator(func): + + def handler(signum, frame): + raise TimeoutError("Operation timed out!") + + def wrapper(*args, **kwargs): + old_handler = signal.getsignal(signal.SIGALRM) + signal.signal(signal.SIGALRM, handler) + signal.alarm(timeout_seconds) + + try: + return func(*args, **kwargs) + finally: + signal.alarm(0) + signal.signal(signal.SIGALRM, old_handler) + + return wrapper + + return decorator + else: + raise NotImplementedError(f"Unsupported OS: {os.name}") + + +def _sympy_parse(expr: str): + """Parses an expression with sympy.""" + py_expr = expr.replace("^", "**") + return sympy_parser.parse_expr( + py_expr, + transformations=(sympy_parser.standard_transformations + (sympy_parser.implicit_multiplication_application,)), + ) + + +def _parse_latex(expr: str) -> str: + """Attempts to parse latex to an expression sympy can read.""" + expr = expr.replace("\\tfrac", "\\frac") + expr = expr.replace("\\dfrac", "\\frac") + expr = expr.replace("\\frac", " \\frac") # Play nice with mixed numbers. + expr = latex2text.LatexNodes2Text().latex_to_text(expr) + + # Replace the specific characters that this parser uses. + expr = expr.replace("√", "sqrt") + expr = expr.replace("π", "pi") + expr = expr.replace("∞", "inf") + expr = expr.replace("∪", "U") + expr = expr.replace("·", "*") + expr = expr.replace("×", "*") + + return expr.strip() + + +def _is_float(num: str) -> bool: + try: + float(num) + return True + except ValueError: + return False + + +def _is_int(x: float) -> bool: + try: + return abs(x - int(round(x))) <= 1e-7 + except: + return False + + +def _is_frac(expr: str) -> bool: + return bool(re.search(r"^-?[0-9]+.?/0*[1-9][0-9]*.?$", expr)) + + +def _str_is_int(x: str) -> bool: + try: + x = _strip_properly_formatted_commas(x) + x = float(x) + return abs(x - int(round(x))) <= 1e-7 + except: + return False + + +def _str_to_int(x: str) -> bool: + x = x.replace(",", "") + x = float(x) + return int(x) + + +def _inject_implicit_mixed_number(step: str): + """ + Automatically make a mixed number evalable + e.g. 7 3/4 => 7+3/4 + """ + p1 = re.compile("([0-9]) +([0-9])") + step = p1.sub("\\1+\\2", step) ## implicit mults + return step + + +def _strip_properly_formatted_commas(expr: str): + # We want to be careful because we don't want to strip tuple commas + p1 = re.compile("(\d)(,)(\d\d\d)($|\D)") + while True: + next_expr = p1.sub("\\1\\3\\4", expr) + if next_expr == expr: + break + expr = next_expr + return next_expr + + +def _normalize(expr: str) -> str: + """Normalize answer expressions.""" + if expr is None: + return None + + # Remove enclosing `\text{}`. + m = re.search("^\\\\text\{(?P.+?)\}$", expr) + if m is not None: + expr = m.group("text") + + expr = expr.replace("\\%", "%") + expr = expr.replace("\\$", "$") + expr = expr.replace("$", "") + expr = expr.replace("%", "") + expr = expr.replace(" or ", " , ") + expr = expr.replace(" and ", " , ") + + expr = expr.replace("million", "*10^6") + expr = expr.replace("billion", "*10^9") + expr = expr.replace("trillion", "*10^12") + + for unit in [ + "degree", + "cm", + "centimeter", + "meter", + "mile", + "second", + "minute", + "hour", + "day", + "week", + "month", + "year", + "foot", + "feet", + "inch", + "yard", + "liter", + ]: + expr = re.sub(f"{unit}(es)?(s)? *(\^[0-9]+)?", "", expr) + expr = re.sub(f"\^ *\\\\circ", "", expr) + + if len(expr) > 0 and expr[0] == "{" and expr[-1] == "}": + expr = expr[1:-1] + + expr = re.sub(",\\\\! *", "", expr) + if _is_float(expr) and _is_int(float(expr)): + expr = str(int(round(float(expr)))) + if "\\" in expr: + try: + expr = _parse_latex(expr) + except: + pass + + # edge case with mixed numbers and negative signs + expr = re.sub("- *", "-", expr) + + expr = _inject_implicit_mixed_number(expr) + + # don't be case sensitive for text answers + expr = expr.lower() + + if _str_is_int(expr): + expr = str(_str_to_int(expr)) + + return expr + + +def count_unknown_letters_in_expr(expr: str): + expr = expr.replace("sqrt", "") + expr = expr.replace("frac", "") + letters_in_expr = set([x for x in expr if x.isalpha()]) + return len(letters_in_expr) + + +def should_allow_eval(expr: str): + # we don't want to try parsing unknown text or functions of more than two variables + if count_unknown_letters_in_expr(expr) > 2: + return False + + for bad_string in BAD_SUBSTRINGS: + if bad_string in expr: + return False + + for bad_regex in BAD_REGEXES: + if re.search(bad_regex, expr) is not None: + return False + + return True + + +@timeout(timeout_seconds=10) +def are_equal_under_sympy(ground_truth_normalized: str, given_normalized: str): + are_equal = False + try: + expr = f"({ground_truth_normalized})-({given_normalized})" + if should_allow_eval(expr): + sympy_diff = _sympy_parse(expr) + simplified = sympy.simplify(sympy_diff) + if simplified == 0: + are_equal = True + except: + pass + return are_equal + + +def split_tuple(expr: str): + """ + Split the elements in a tuple/interval, while handling well-formatted commas in large numbers + """ + expr = _strip_properly_formatted_commas(expr) + if len(expr) == 0: + return [] + if (len(expr) > 2 and expr[0] in TUPLE_CHARS and expr[-1] in TUPLE_CHARS and + all([ch not in expr[1:-1] for ch in TUPLE_CHARS])): + elems = [elem.strip() for elem in expr[1:-1].split(",")] + else: + elems = [expr] + return elems + + +def grade_answer(given_answer: str, ground_truth: str) -> tuple[bool, str]: + """ + The answer will be considered correct if: + (a) it normalizes to the same string as the ground truth answer + OR + (b) sympy can simplify the difference between the expressions to 0 + """ + if given_answer is None: + return False + + ground_truth_normalized_mathd = math_normalize.normalize_answer(ground_truth) + given_answer_normalized_mathd = math_normalize.normalize_answer(given_answer) + + # be at least as lenient as mathd + if ground_truth_normalized_mathd == given_answer_normalized_mathd: + return True, given_answer_normalized_mathd + + ground_truth_normalized = _normalize(ground_truth) + given_normalized = _normalize(given_answer) + + if ground_truth_normalized is None: + return False, given_normalized + + if ground_truth_normalized == given_normalized: + return True, given_normalized + + if len(given_normalized) == 0: + return False, given_normalized + + ground_truth_elems = split_tuple(ground_truth_normalized) + given_elems = split_tuple(given_normalized) + + if len(ground_truth_elems) > 1 and (ground_truth_normalized[0] != given_normalized[0] or + ground_truth_normalized[-1] != given_normalized[-1]): + is_correct = False + elif len(ground_truth_elems) != len(given_elems): + is_correct = False + else: + for ground_truth_elem, given_elem in zip(ground_truth_elems, given_elems): + if _is_frac(ground_truth_elem) and _is_frac(given_elem): + # if fractions aren't reduced, then shouldn't be marked as correct + # so, we don't want to allow sympy.simplify in this case + is_correct = ground_truth_elem == given_elem + elif _str_is_int(ground_truth_elem) != _str_is_int(given_elem): + # if the ground truth answer is an integer, we require the given answer to be a strict match (no sympy.simplify) + is_correct = False + else: + is_correct = are_equal_under_sympy(ground_truth_elem, given_elem) + if not is_correct: + break + + return is_correct, given_normalized + +def _last_boxed_only_string(string): + idx = string.rfind("\\boxed") + if idx < 0: + idx = string.rfind("\\fbox") + if idx < 0: + return None + + i = idx + left_brace_idx = None + right_brace_idx = None + num_left_braces_open = 0 + while i < len(string): + if string[i] == "{": + num_left_braces_open += 1 + if left_brace_idx is None: + left_brace_idx = i + elif string[i] == "}": + num_left_braces_open -= 1 + if num_left_braces_open == 0: + right_brace_idx = i + break + + i += 1 + + if left_brace_idx is None or right_brace_idx is None: + return None + + return string[left_brace_idx + 1:right_brace_idx].strip() + + +def match_answer(response): + is_matched = False + response = response.split("")[-1] + + # Find boxed + ans_boxed = _last_boxed_only_string(response) + if ans_boxed: + is_matched = True + response = ans_boxed + + return is_matched, response + +import math + +def compute_score(data_source: str, + solution_str: str, + ground_truth: str, + extra_info: dict, + compressed: bool = False) -> float: + """Compute the reward score for a solution. This draws heavily from the LLM-as-judge and PRIME reward functions + + Args: + solution_str: The solution string + ground_truth: The ground truth answer + extra_info: dict with additional info for the score computation + + Returns: + Reward score (1.0 for correct, -1.0 for incorrect) + """ + if compressed: + from .utils import _deserialise_extra, _decompress_str + extra_info = _deserialise_extra(_decompress_str(extra_info)) + ground_truth = _deserialise_extra(_decompress_str(ground_truth["compressed"]))["ground_truth"] + + model_output = str(solution_str) + ground_truth = str(ground_truth) + + # Extract answer from generated output + is_matched, extracted_model_output = match_answer(model_output) + + # TWK NOTE: WE REMOVED THE RESPONSE TRUNCATION FROM math_dapo.compute_score + + # Verify the solution, first check simple comparisons. + correct, pred = grade_answer(extracted_model_output, ground_truth) + + if not correct: + try: + if "\\pi" in extracted_model_output or "\\pi" in ground_truth: + equivs = [] + for pi in [math.pi, 3.14]: + equivs.append(math_equal(extracted_model_output, ground_truth, tiemout=True, pi=pi)) + correct = any(equivs) + else: + correct = math_equal(extracted_model_output, ground_truth, timeout=True) + except: + correct = False + + + # reward = 1.0 if correct else -1.0 + reward = 1.0 if correct else 0. + acc = correct + return reward \ No newline at end of file diff --git a/reasoners/reward/prime_math/__init__.py b/reasoners/reward/prime_math/__init__.py new file mode 100644 index 00000000..01390e5b --- /dev/null +++ b/reasoners/reward/prime_math/__init__.py @@ -0,0 +1,437 @@ +# Copyright 2024 PRIME team and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Answer checker API that uses sympy to simplify expressions and check for equality. + +Call grade_answer(given_answer: str, ground_truth: str). + +FROM: https://github.com/openai/prm800k/blob/main/prm800k/grading/grader.py +""" +import re +import sympy +from pylatexenc import latex2text +from sympy.parsing import sympy_parser +import os +from . import math_normalize +from .grader import math_equal +from reward_score.utils import _decompress_str, _deserialise_extra +# import math_normalize +# from grader import math_equal + +# sympy might hang -- we don't care about trying to be lenient in these cases +BAD_SUBSTRINGS = ["^{", "^("] +BAD_REGEXES = ["\^[0-9]+\^", "\^[0-9][0-9]+"] +TUPLE_CHARS = "()[]" + + +def timeout(timeout_seconds: int = 8): + if os.name == "posix": + import signal + + def decorator(func): + + def handler(signum, frame): + raise TimeoutError("Operation timed out!") + + def wrapper(*args, **kwargs): + old_handler = signal.getsignal(signal.SIGALRM) + signal.signal(signal.SIGALRM, handler) + signal.alarm(timeout_seconds) + + try: + return func(*args, **kwargs) + finally: + signal.alarm(0) + signal.signal(signal.SIGALRM, old_handler) + + return wrapper + + return decorator + else: + raise NotImplementedError(f"Unsupported OS: {os.name}") + + +def _sympy_parse(expr: str): + """Parses an expression with sympy.""" + py_expr = expr.replace("^", "**") + return sympy_parser.parse_expr( + py_expr, + transformations=(sympy_parser.standard_transformations + (sympy_parser.implicit_multiplication_application,)), + ) + + +def _parse_latex(expr: str) -> str: + """Attempts to parse latex to an expression sympy can read.""" + expr = expr.replace("\\tfrac", "\\frac") + expr = expr.replace("\\dfrac", "\\frac") + expr = expr.replace("\\frac", " \\frac") # Play nice with mixed numbers. + expr = latex2text.LatexNodes2Text().latex_to_text(expr) + + # Replace the specific characters that this parser uses. + expr = expr.replace("√", "sqrt") + expr = expr.replace("π", "pi") + expr = expr.replace("∞", "inf") + expr = expr.replace("∪", "U") + expr = expr.replace("·", "*") + expr = expr.replace("×", "*") + + return expr.strip() + + +def _is_float(num: str) -> bool: + try: + float(num) + return True + except ValueError: + return False + + +def _is_int(x: float) -> bool: + try: + return abs(x - int(round(x))) <= 1e-7 + except: + return False + + +def _is_frac(expr: str) -> bool: + return bool(re.search(r"^-?[0-9]+.?/0*[1-9][0-9]*.?$", expr)) + + +def _str_is_int(x: str) -> bool: + try: + x = _strip_properly_formatted_commas(x) + x = float(x) + return abs(x - int(round(x))) <= 1e-7 + except: + return False + + +def _str_to_int(x: str) -> bool: + x = x.replace(",", "") + x = float(x) + return int(x) + + +def _inject_implicit_mixed_number(step: str): + """ + Automatically make a mixed number evalable + e.g. 7 3/4 => 7+3/4 + """ + p1 = re.compile("([0-9]) +([0-9])") + step = p1.sub("\\1+\\2", step) ## implicit mults + return step + + +def _strip_properly_formatted_commas(expr: str): + # We want to be careful because we don't want to strip tuple commas + p1 = re.compile("(\d)(,)(\d\d\d)($|\D)") + while True: + next_expr = p1.sub("\\1\\3\\4", expr) + if next_expr == expr: + break + expr = next_expr + return next_expr + + +def _normalize(expr: str) -> str: + """Normalize answer expressions.""" + if expr is None: + return None + + # Remove enclosing `\text{}`. + m = re.search("^\\\\text\{(?P.+?)\}$", expr) + if m is not None: + expr = m.group("text") + + expr = expr.replace("\\%", "%") + expr = expr.replace("\\$", "$") + expr = expr.replace("$", "") + expr = expr.replace("%", "") + expr = expr.replace(" or ", " , ") + expr = expr.replace(" and ", " , ") + + expr = expr.replace("million", "*10^6") + expr = expr.replace("billion", "*10^9") + expr = expr.replace("trillion", "*10^12") + + for unit in [ + "degree", + "cm", + "centimeter", + "meter", + "mile", + "second", + "minute", + "hour", + "day", + "week", + "month", + "year", + "foot", + "feet", + "inch", + "yard", + "liter", + ]: + expr = re.sub(f"{unit}(es)?(s)? *(\^[0-9]+)?", "", expr) + expr = re.sub(f"\^ *\\\\circ", "", expr) + + if len(expr) > 0 and expr[0] == "{" and expr[-1] == "}": + expr = expr[1:-1] + + expr = re.sub(",\\\\! *", "", expr) + if _is_float(expr) and _is_int(float(expr)): + expr = str(int(round(float(expr)))) + if "\\" in expr: + try: + expr = _parse_latex(expr) + except: + pass + + # edge case with mixed numbers and negative signs + expr = re.sub("- *", "-", expr) + + expr = _inject_implicit_mixed_number(expr) + + # don't be case sensitive for text answers + expr = expr.lower() + + if _str_is_int(expr): + expr = str(_str_to_int(expr)) + + return expr + + +def count_unknown_letters_in_expr(expr: str): + expr = expr.replace("sqrt", "") + expr = expr.replace("frac", "") + letters_in_expr = set([x for x in expr if x.isalpha()]) + return len(letters_in_expr) + + +def should_allow_eval(expr: str): + # we don't want to try parsing unknown text or functions of more than two variables + if count_unknown_letters_in_expr(expr) > 2: + return False + + for bad_string in BAD_SUBSTRINGS: + if bad_string in expr: + return False + + for bad_regex in BAD_REGEXES: + if re.search(bad_regex, expr) is not None: + return False + + return True + + +@timeout(timeout_seconds=10) +def are_equal_under_sympy(ground_truth_normalized: str, given_normalized: str): + are_equal = False + try: + expr = f"({ground_truth_normalized})-({given_normalized})" + if should_allow_eval(expr): + sympy_diff = _sympy_parse(expr) + simplified = sympy.simplify(sympy_diff) + if simplified == 0: + are_equal = True + except: + pass + return are_equal + + +def split_tuple(expr: str): + """ + Split the elements in a tuple/interval, while handling well-formatted commas in large numbers + """ + expr = _strip_properly_formatted_commas(expr) + if len(expr) == 0: + return [] + if (len(expr) > 2 and expr[0] in TUPLE_CHARS and expr[-1] in TUPLE_CHARS and + all([ch not in expr[1:-1] for ch in TUPLE_CHARS])): + elems = [elem.strip() for elem in expr[1:-1].split(",")] + else: + elems = [expr] + return elems + + +def grade_answer(given_answer: str, ground_truth: str) -> bool: + """ + The answer will be considered correct if: + (a) it normalizes to the same string as the ground truth answer + OR + (b) sympy can simplify the difference between the expressions to 0 + """ + if given_answer is None: + return False + + ground_truth_normalized_mathd = math_normalize.normalize_answer(ground_truth) + given_answer_normalized_mathd = math_normalize.normalize_answer(given_answer) + + # be at least as lenient as mathd + if ground_truth_normalized_mathd == given_answer_normalized_mathd: + return True + + ground_truth_normalized = _normalize(ground_truth) + given_normalized = _normalize(given_answer) + + if ground_truth_normalized is None: + return False + + if ground_truth_normalized == given_normalized: + return True + + if len(given_normalized) == 0: + return False + + ground_truth_elems = split_tuple(ground_truth_normalized) + given_elems = split_tuple(given_normalized) + + if len(ground_truth_elems) > 1 and (ground_truth_normalized[0] != given_normalized[0] or + ground_truth_normalized[-1] != given_normalized[-1]): + is_correct = False + elif len(ground_truth_elems) != len(given_elems): + is_correct = False + else: + for ground_truth_elem, given_elem in zip(ground_truth_elems, given_elems): + if _is_frac(ground_truth_elem) and _is_frac(given_elem): + # if fractions aren't reduced, then shouldn't be marked as correct + # so, we don't want to allow sympy.simplify in this case + is_correct = ground_truth_elem == given_elem + elif _str_is_int(ground_truth_elem) != _str_is_int(given_elem): + # if the ground truth answer is an integer, we require the given answer to be a strict match (no sympy.simplify) + is_correct = False + else: + is_correct = are_equal_under_sympy(ground_truth_elem, given_elem) + if not is_correct: + break + + return is_correct + + +def remove_boxed(s): + left = "\\boxed{" + try: + assert s[:len(left)] == left + assert s[-1] == "}" + return s[len(left):-1] + except: + return None + + +def _last_boxed_only_string(string): + idx = string.rfind("\\boxed") + if idx < 0: + idx = string.rfind("\\fbox") + if idx < 0: + return None + + i = idx + left_brace_idx = None + right_brace_idx = None + num_left_braces_open = 0 + while i < len(string): + if string[i] == "{": + num_left_braces_open += 1 + if left_brace_idx is None: + left_brace_idx = i + elif string[i] == "}": + num_left_braces_open -= 1 + if num_left_braces_open == 0: + right_brace_idx = i + break + + i += 1 + + if left_brace_idx is None or right_brace_idx is None: + return None + + return string[left_brace_idx + 1:right_brace_idx].strip() + + +def match_answer(response): + is_matched = False + response = response.split("")[-1] + # for ans_marker in ['answer:', "answer is", "answers are"]: + # ans_idx = response.lower().rfind(ans_marker) + # if ans_idx != -1: + # is_matched = True + # response = response[ans_idx + len(ans_marker):].strip() + # if response.endswith("\n"): + # response = response[:-2] + + # for ans_marker in ["is answer", "is the answer", "are answers", "are the answers"]: + # ans_idx = response.lower().rfind(ans_marker) + # if ans_idx != -1: + # is_matched = True + # response = response[:ans_idx].strip() + # if response.endswith("\n"): + # response = response[:-2] + + # Find boxed + ans_boxed = _last_boxed_only_string(response) + if ans_boxed: + is_matched = True + response = ans_boxed + + # if ". " in response: + # dot_idx = response.lower().rfind(". ") + # if dot_idx != -1: + # response = response[:dot_idx].strip() + + # for ans_marker in ['be ', "is ", "are ", "=", ": ", "get ", 'be\n', "is\n", "are\n", ":\n", "get\n"]: + # ans_idx = response.lower().rfind(ans_marker) + # if ans_idx != -1: + # is_matched = True + # response = response[ans_idx + len(ans_marker):].strip() + # if response.endswith("\n"): + # response = response[:-2] + + # is_matched = is_matched if any([c.isdigit() for c in response]) else False # answer must have a digit + # # Grade + return is_matched, response + + +import math + + +def compute_score(data_source: str, model_output: str, ground_truth: str, extra_info: dict, compressed: bool = False) -> float: + if compressed: + extra_info = _deserialise_extra(_decompress_str(extra_info["compressed"])) + ground_truth = _deserialise_extra(_decompress_str(ground_truth["compressed"]))["ground_truth"] + + model_output = str(model_output) + ground_truth = str(ground_truth) + + is_matched, extracted_model_output = match_answer(model_output) + format_correctness = "Step 2:" in model_output and "\\box" in model_output + + # grade simple algebra questions. if succeeded, return; otherwise, proceed to more complex grading + if grade_answer(extracted_model_output, ground_truth): + # return True, True, extracted_model_output + return 1.0 + + try: + if "\\pi" in extracted_model_output or "\\pi" in ground_truth: + equivs = [] + for pi in [math.pi, 3.14]: + equivs.append(math_equal(extracted_model_output, ground_truth, timeout=True, pi=pi)) + is_correct = any(equivs) + else: + is_correct = math_equal(extracted_model_output, ground_truth, timeout=True) + except: + is_correct = False + + # return is_correct, format_correctness, extracted_model_output + return float(is_correct) diff --git a/reasoners/reward/prime_math/grader.py b/reasoners/reward/prime_math/grader.py new file mode 100644 index 00000000..e3e1686f --- /dev/null +++ b/reasoners/reward/prime_math/grader.py @@ -0,0 +1,380 @@ +# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Copyright (c) Microsoft Corporation. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE + +# Copyright (c) 2023 OpenAI +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: + +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +# Copyright (c) 2021 Dan Hendrycks +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +# Copyright 2024 PRIME team and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +This logic is largely copied from the Hendrycks' MATH release (math_equivalence), and borrowed from: +- https://github.com/microsoft/ToRA/blob/main/src/eval/grader.py +- https://github.com/microsoft/ProphetNet/tree/master/CRITIC +- https://github.com/openai/prm800k +""" + +import contextlib +import re +import signal +import math +from math import isclose +from typing import Union + +from sympy import N, simplify +from sympy.parsing.latex import parse_latex +from sympy.parsing.sympy_parser import parse_expr + + +def is_digit(s): + try: + if "{,}" in str(s): + num = float(str(s).replace("{,}", "")) + return True, num + + num = float(str(s).replace(",", "")) + return True, num + except ValueError: + return False, None + + +def normalize(answer, pi) -> str: + # checking if answer is $ and removing $ in that case to compare + if isinstance(answer, str) and bool(re.match(r'\$\d+(\.\d+)?', answer)): + return answer[1:] + + # checking if answer is % or \\% and removing % + if isinstance(answer, str) and (bool(re.match(r'^\d+(\.\d+)?%$', answer)) or + bool(re.match(r'^\d+(\.\d+)?\\%$', answer))): + return answer.replace("\\%", "").replace("%", "") + + # handle base + answer = handle_base(answer) + + # handle pi + answer = handle_pi(answer, pi) + + return answer + + +def handle_base(x) -> str: + if isinstance(x, str) and "_" in x: + # Due to base + x = x.split("_")[0] + x = float(x) + return int(x) + return x + + +def handle_pi(string, pi): + if isinstance(string, str) and "\pi" in string: + # Find the first occurrence of "\pi" + idx = string.find("\pi") + + # Iterate over the string and find all occurrences of "\pi" with a valid previous character + while idx != -1: + + if idx > 0 and string[idx - 1].isdigit(): + # Replace "\pi" with "*math.pi" if the previous character is a digit + string = string[:idx] + f"*{pi}" + string[idx + 3:] + else: + # Replace "\pi" with "1*math.pi" if the previous character is not a digit + string = string[:idx] + f"1*{pi}" + string[idx + 3:] + + # Find the next occurrence of "\pi" + idx = string.find("\pi", idx + 1) + + # Evaluate the expression using eval() function + try: + string = eval(string) + except: + pass + + return string + + +def math_equal(prediction: Union[bool, float, str], + reference: Union[float, str], + include_percentage: bool = True, + tolerance: float = 1e-4, + timeout: float = 10.0, + pi: float = math.pi) -> bool: + """ + Exact match of math if and only if: + 1. numerical equal: both can convert to float and are equal + 2. symbolic equal: both can convert to sympy expression and are equal + """ + + prediction = normalize(prediction, pi) + reference = normalize(reference, pi) + + if isinstance(prediction, str) and len(prediction) > 1000: # handling weird corner-cases + prediction = prediction[:1000] + + # 0. string comparison + if isinstance(prediction, str) and isinstance(reference, str): + if prediction.strip().lower() == reference.strip().lower(): + return True + if prediction.replace(" ", "") == reference.replace(" ", ""): + return True + + try: # 1. numerical equal + if is_digit(prediction)[0] and is_digit(reference)[0]: + prediction = is_digit(prediction)[1] + reference = is_digit(reference)[1] + # number questions + if include_percentage: + gt_result = [reference / 100, reference, reference * 100] + else: + gt_result = [reference] + for item in gt_result: + try: + if isclose(item, prediction, rel_tol=tolerance): + return True + except Exception: + continue + return False + except Exception: + pass + + if not prediction and prediction not in [0, False]: + return False + + # 2. symbolic equal + reference = str(reference).strip() + prediction = str(prediction).strip() + + ## deal with [], (), {} + prediction = format_intervals(prediction) + + pred_str, ref_str = prediction, reference + if (prediction.startswith("[") and prediction.endswith("]") and + not reference.startswith("(")) or (prediction.startswith("(") and prediction.endswith(")") and + not reference.startswith("[")): + pred_str = pred_str.strip("[]()") + ref_str = ref_str.strip("[]()") + for s in ["{", "}", "(", ")"]: + ref_str = ref_str.replace(s, "") + pred_str = pred_str.replace(s, "") + if pred_str == ref_str: + return True + + ## [a, b] vs. [c, d], return a==c and b==d + if (prediction and reference and prediction[0] in "([" and prediction[-1] in ")]" and + prediction[0] == reference[0] and prediction[-1] == reference[-1]): + pred_parts = prediction[1:-1].split(",") + ref_parts = reference[1:-1].split(",") + if len(pred_parts) == len(ref_parts): + if all([ + math_equal(pred_pt, ref_pt, include_percentage, tolerance) + for pred_pt, ref_pt in zip(pred_parts, ref_parts) + ]): + return True + + if "," in prediction and "," in reference: + pred_parts = [item.strip() for item in prediction.split(",")] + ref_parts = [item.strip() for item in reference.split(",")] + + if len(pred_parts) == len(ref_parts): + if all([ + math_equal(pred_parts[i], ref_parts[i], include_percentage, tolerance) + for i in range(len(pred_parts)) + ]): + return True + else: + return False + + # if we have point == tuple of values + if prediction.startswith("Point") and reference[0] == "(" and reference[-1] == ")": + pred_parts = prediction[prediction.find("(") + 1:-1].split(",") + ref_parts = reference[1:-1].split(",") + if len(pred_parts) == len(ref_parts): + if all([ + math_equal(pred_pt, ref_pt, include_percentage, tolerance) + for pred_pt, ref_pt in zip(pred_parts, ref_parts) + ]): + return True + + # if reference is a matrix + if "\begin{pmatrix}" in reference and prediction.startswith("Matrix"): + try: + pred_matrix = parse_expr(prediction) + ref_matrix_items = reference.split()[1:-1:2] + if len(pred_matrix) == len(ref_matrix_items): + if all([ + math_equal(pred, ref, include_percentage, tolerance) + for ref, pred in zip(ref_matrix_items, pred_matrix) + ]): + return True + except Exception: + pass + elif "\begin{pmatrix}" in reference and prediction.startswith("[") and prediction.endswith("]"): + if isinstance(eval(prediction), list): + try: + pred_matrix = eval(prediction) + # ref_matrix_items = reference.split()[1:-1:2] + ref_matrix_items = reference.lstrip("\\begin{pmatrix}").lstrip("\begin{pmatrix}").rstrip( + "\\end{pmatrix}").rstrip("\end{pmatrix}") + ref_matrix_items = ref_matrix_items.split("\\") + ref_matrix_items = [row.split("&") if "&" in row else row for row in ref_matrix_items] + if len(pred_matrix) == len(ref_matrix_items): + if all([ + math_equal(pred, ref, include_percentage, tolerance) + for ref, pred in zip(ref_matrix_items, pred_matrix) + ]): + return True + except Exception: + pass + + return symbolic_equal(prediction, reference, tolerance, timeout) + + +def symbolic_equal(a, b, tolerance, timeout=10.0): + + def _parse(s): + for f in [parse_expr, parse_latex]: + try: + with time_limit(timeout): + return f(s) + except Exception: + pass + return s + + a = _parse(a) + b = _parse(b) + + try: + with time_limit(timeout): + if simplify(a - b) == 0: + return True + except Exception: + pass + + try: + with time_limit(timeout): + if isclose(N(a), N(b), rel_tol=tolerance): + return True + except Exception: + pass + return False + + +class TimeoutException(Exception): + pass + + +@contextlib.contextmanager +def time_limit(seconds: float): + + def signal_handler(signum, frame): + raise TimeoutException("Timed out!") + + signal.setitimer(signal.ITIMER_REAL, seconds) + signal.signal(signal.SIGALRM, signal_handler) + try: + yield + finally: + signal.setitimer(signal.ITIMER_REAL, 0) + + +def format_intervals(prediction): + patterns = { + "Interval(": r"^Interval\((.*)\)$", + "Interval.Ropen(": r"^Interval\.Ropen\((.*)\)$", + "Interval.Lopen(": r"^Interval\.Lopen\((.*)\)$", + "Interval.open(": r"^Interval\.open\((.*)\)$", + } + + for key, pattern in patterns.items(): + match = re.match(pattern, prediction) + if match: + inner_content = match.group(1) + + if key == "Interval(": # Intarval(a, b) == [a, b] + return f"[{inner_content}]" + elif key == "Interval.Ropen(": # Intarval.Ropen(a, b) == [a, b) + return f"[{inner_content})" + elif key == "Interval.Lopen(": # Intarval.Lopen(a, b) == (a, b] + return f"({inner_content}]" + elif key == "Interval.open(": # Intarval.open(a, b) == (a, b) + return f"({inner_content})" + + return prediction diff --git a/reasoners/reward/prime_math/math_normalize.py b/reasoners/reward/prime_math/math_normalize.py new file mode 100644 index 00000000..e9921a5a --- /dev/null +++ b/reasoners/reward/prime_math/math_normalize.py @@ -0,0 +1,191 @@ +# Copyright 2024 PRIME team and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Copyright (c) 2021 Dan Hendrycks +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +""" +This logic is largely copied from the Hendrycks' MATH release (math_equivalence). + +From: https://github.com/openai/prm800k/blob/main/prm800k/grading/math_normalize.py +""" +import re +from typing import Optional + + +def normalize_answer(answer: Optional[str]) -> Optional[str]: + if answer is None: + return None + answer = answer.strip() + try: + # Remove enclosing `\text{}`. + m = re.search("^\\\\text\{(?P.+?)\}$", answer) + if m is not None: + answer = m.group("text").strip() + return _strip_string(answer) + except: + return answer + + +def _fix_fracs(string): + substrs = string.split("\\frac") + new_str = substrs[0] + if len(substrs) > 1: + substrs = substrs[1:] + for substr in substrs: + new_str += "\\frac" + if substr[0] == "{": + new_str += substr + else: + try: + assert len(substr) >= 2 + except: + return string + a = substr[0] + b = substr[1] + if b != "{": + if len(substr) > 2: + post_substr = substr[2:] + new_str += "{" + a + "}{" + b + "}" + post_substr + else: + new_str += "{" + a + "}{" + b + "}" + else: + if len(substr) > 2: + post_substr = substr[2:] + new_str += "{" + a + "}" + b + post_substr + else: + new_str += "{" + a + "}" + b + string = new_str + return string + + +def _fix_a_slash_b(string): + if len(string.split("/")) != 2: + return string + a = string.split("/")[0] + b = string.split("/")[1] + try: + a = int(a) + b = int(b) + assert string == "{}/{}".format(a, b) + new_string = "\\frac{" + str(a) + "}{" + str(b) + "}" + return new_string + except: + return string + + +def _remove_right_units(string): + # "\\text{ " only ever occurs (at least in the val set) when describing units + if "\\text{ " in string: + splits = string.split("\\text{ ") + assert len(splits) == 2 + return splits[0] + else: + return string + + +def _fix_sqrt(string): + if "\\sqrt" not in string: + return string + splits = string.split("\\sqrt") + new_string = splits[0] + for split in splits[1:]: + if split[0] != "{": + a = split[0] + new_substr = "\\sqrt{" + a + "}" + split[1:] + else: + new_substr = "\\sqrt" + split + new_string += new_substr + return new_string + + +def _strip_string(string): + # linebreaks + string = string.replace("\n", "") + + # remove inverse spaces + string = string.replace("\\!", "") + + # replace \\ with \ + string = string.replace("\\\\", "\\") + + # replace tfrac and dfrac with frac + string = string.replace("tfrac", "frac") + string = string.replace("dfrac", "frac") + + # remove \left and \right + string = string.replace("\\left", "") + string = string.replace("\\right", "") + + # Remove circ (degrees) + string = string.replace("^{\\circ}", "") + string = string.replace("^\\circ", "") + + # remove dollar signs + string = string.replace("\\$", "") + + # remove units (on the right) + string = _remove_right_units(string) + + # remove percentage + string = string.replace("\\%", "") + string = string.replace("\%", "") + + # " 0." equivalent to " ." and "{0." equivalent to "{." Alternatively, add "0" if "." is the start of the string + string = string.replace(" .", " 0.") + string = string.replace("{.", "{0.") + # if empty, return empty string + if len(string) == 0: + return string + if string[0] == ".": + string = "0" + string + + # to consider: get rid of e.g. "k = " or "q = " at beginning + if len(string.split("=")) == 2: + if len(string.split("=")[0]) <= 2: + string = string.split("=")[1] + + # fix sqrt3 --> sqrt{3} + string = _fix_sqrt(string) + + # remove spaces + string = string.replace(" ", "") + + # \frac1b or \frac12 --> \frac{1}{b} and \frac{1}{2}, etc. Even works with \frac1{72} (but not \frac{72}1). Also does a/b --> \\frac{a}{b} + string = _fix_fracs(string) + + # manually change 0.5 --> \frac{1}{2} + if string == "0.5": + string = "\\frac{1}{2}" + + # NOTE: X/Y changed to \frac{X}{Y} in dataset, but in simple cases fix in case the model output is X/Y + string = _fix_a_slash_b(string) + + return string \ No newline at end of file diff --git a/reasoners/reward/puzzles_dataset.py b/reasoners/reward/puzzles_dataset.py new file mode 100644 index 00000000..135013ea --- /dev/null +++ b/reasoners/reward/puzzles_dataset.py @@ -0,0 +1,137 @@ +import re +import random +import ast +import operator +import signal +import contextlib + +class TimeoutException(Exception): + pass + +@contextlib.contextmanager +def time_limit(seconds: float): + def signal_handler(signum, frame): + raise TimeoutException("Timed out!") + + signal.setitimer(signal.ITIMER_REAL, seconds) + signal.signal(signal.SIGALRM, signal_handler) + try: + yield + finally: + signal.setitimer(signal.ITIMER_REAL, 0) + +def extract_solution(solution_str): + + # Find the answer tag content + answer_pattern = r'(.*?)' + match = re.finditer(answer_pattern, solution_str) + matches = list(match) + + if matches: + final_answer = matches[-1].group(1).strip() + + # Use regex to safely extract the bird names without eval + bird_pattern = r'\[[\s\'\"]*([^\],]+)[\s\'\"]*(?:,[\s\'\"]*([^\],]+)[\s\'\"]*)*\]' + bird_match = re.search(bird_pattern, final_answer) + + if bird_match: + # Extract all bird names from the list + bird_list_str = bird_match.group(0) + # Extract individual birds, handling different formats + birds = re.findall(r'[\'"]?([\w\s]+)[\'"]?', bird_list_str) + # Clean up the extracted birds (removing list brackets, quotes, etc.) + birds = [bird.lower().strip() for bird in birds if bird.strip() and bird.strip() not in ['[', ']']] + return birds + else: + # If we can't extract using regex, try a safer approach + try: + # Add quotes around unquoted words to make it valid Python + fixed_str = re.sub(r'(\[|\s*,\s*)(\w+)(\s*,|\])', r'\1"\2"\3', final_answer) + return eval(fixed_str) + except: + # Last resort: just return as text + return None + else: + return None + + +def compute_edit_distance(list1, list2): + """ + Calculate edit distance between two lists. + Returns the minimum number of operations (insertions, deletions, substitutions) + required to transform list1 into list2. + """ + # Create a matrix of size (len(list1)+1) x (len(list2)+1) + dp = [[0 for _ in range(len(list2) + 1)] for _ in range(len(list1) + 1)] + + # Initialize the first row and column + for i in range(len(list1) + 1): + dp[i][0] = i + for j in range(len(list2) + 1): + dp[0][j] = j + + # Fill the matrix + for i in range(1, len(list1) + 1): + for j in range(1, len(list2) + 1): + if list1[i-1] == list2[j-1]: + dp[i][j] = dp[i-1][j-1] + else: + dp[i][j] = 1 + min(dp[i-1][j], # deletion + dp[i][j-1], # insertion + dp[i-1][j-1]) # substitution + + return dp[len(list1)][len(list2)] + +# granular reward function +def compute_score(data_source: str, + solution_str: str, + ground_truth: str, + extra_info: any = None, + compressed: bool = False, + method: str = 'strict', + timeout: float = 10.0): + """The scoring function for bird puzzles task. + + Args: + solution_str: the solution text + ground_truth: dictionary containing target number and available numbers + method: the method to extract the solution + timeout: maximum time in seconds to allow for computation + """ + if compressed: + from reward_score.utils import _deserialise_extra, _decompress_str + extra_info = _deserialise_extra(_decompress_str(extra_info)) + ground_truth = _deserialise_extra(_decompress_str(ground_truth["compressed"]))["ground_truth"] + print(f"ground_truth: {ground_truth}") + + try: + with time_limit(timeout): + target = ground_truth.tolist() if not isinstance(ground_truth,list) else ground_truth + predicted_arrangement = extract_solution(solution_str=solution_str) + + if predicted_arrangement is None: + score = 0.0 + + # Evaluate equation + try: + if isinstance(predicted_arrangement, list) and isinstance(target, list): + edit_distance = compute_edit_distance(predicted_arrangement, target) + max_possible_dist = max(len(predicted_arrangement), len(target)) + result = predicted_arrangement == target + if result: + score = 1.0 + elif method != 'strict': + score = max(1.0 - (edit_distance / max_possible_dist)) + else: + score = 0.0 + except Exception as e: + score = 0.0 + + except TimeoutException: + print("Computation timed out in puzzles_dataset") + score = 0.0 + except Exception as e: + print(f"Error in compute_score in puzzles_dataset: {e}") + score = 0.0 + + return float(score) diff --git a/reasoners/reward/stem_llm_judge/__init__.py b/reasoners/reward/stem_llm_judge/__init__.py new file mode 100644 index 00000000..dcc4e79b --- /dev/null +++ b/reasoners/reward/stem_llm_judge/__init__.py @@ -0,0 +1,182 @@ +""" +Minimal STEM judge: every answer is graded by an external LLM. + +Prerequisite: + - Set env var STEM_LLM_JUDGE_URL to an OpenAI-compatible /v1/chat/completions. + - Launch the service with your preferred model beforehand, e.g. + + vllm serve TIGER-Lab/general-verifier + export STEM_LLM_JUDGE_URL=http://127.0.0.1:8000/v1/chat/completions +""" + +import os, re, requests +from typing import Tuple + +# # ------------ Prompt template ------------------------------------------------ +# JUDGE_TEMPLATE = """\ +# You are a strict grader for university-level STEM problems. + +# Question: +# {QUESTION} + +# Reference Answer: +# {REFERENCE_ANSWER} + +# Student Answer: +# {STUDENT_ANSWER} + +# Carefully think and check whether the student answer is equivalent to the reference answer. +# You only need to refer to the reference answer to grade the student's answer. Sometimes the student's answer is expressed in a different way from the reference answer, but the meaning is the same, and you should still consider it correct. If they are not equivalent in mathematical sense, you should consider it incorrect. + +# : CORRECT or INCORRECT + +# """ + + +# ------------ Core LLM call -------------------------------------------------- +def _llm_judge(question: str, student: str, reference: str, verbose: bool = False) -> bool: + url_base = os.getenv("STEM_LLM_JUDGE_URL") + if not url_base: + raise EnvironmentError("STEM_LLM_JUDGE_URL not set") + url = url_base.rstrip("/") + "/v1/chat/completions" + + # prompt = JUDGE_TEMPLATE.format( + # QUESTION=question, + # STUDENT_ANSWER=student, + # REFERENCE_ANSWER=reference, + # ) + + prompt = ( + f"User: ### Question: {question}\n\n" + f"### Ground Truth Answer: {reference}\n\n" + f"### Student Answer: {student}\n\n" + "For the above question, please verify if the student's answer is equivalent to the ground truth answer.\n" + "Do not solve the question by yourself; just check if the student's answer is equivalent to the ground truth answer.\n" + "If the student's answer is correct, output \"Final Decision: Yes\". If the student's answer is incorrect, output \"Final Decision: No\". Assistant:" + ) + + payload = { + "model": "TIGER-Lab/general-verifier", + "messages": [{"role": "user", "content": prompt}], + "chat_template_kwargs": {"enable_thinking": False}, + "temperature": 0.0 + } + + resp = requests.post(url, json=payload, timeout=60) + resp.raise_for_status() + data = resp.json() + + text = data["choices"][0]["message"]["content"] + score = float("Final Decision: Yes" in text) + + marker = "✅" if score == 1. else "❌" + + if verbose: + print(marker*50 ) + print("student answer: ", student) + print("gt: ", reference) + import json as _json + print(_json.dumps(data, indent=2, ensure_ascii=False)) + print(marker*16 + " LLM Judge CONTENT " + marker*16 ) + print(text) + print(marker*16 + "End of LLM Judge Reply \n"+ marker*16 ) + + return score + +def _last_boxed_only_string(string): + idx = string.rfind("\\boxed") + if idx < 0: + idx = string.rfind("\\fbox") + if idx < 0: + return None + + i = idx + left_brace_idx = None + right_brace_idx = None + num_left_braces_open = 0 + while i < len(string): + if string[i] == "{": + num_left_braces_open += 1 + if left_brace_idx is None: + left_brace_idx = i + elif string[i] == "}": + num_left_braces_open -= 1 + if num_left_braces_open == 0: + right_brace_idx = i + break + + i += 1 + + if left_brace_idx is None or right_brace_idx is None: + return None + + return string[left_brace_idx + 1:right_brace_idx].strip() + + +def match_answer(response): + is_matched = False + + # Find boxed + ans_boxed = _last_boxed_only_string(response) + if ans_boxed: + is_matched = True + response = ans_boxed + + return is_matched, response + + +# ------------ Public API ----------------------------------------------------- +def compute_score(data_source: str, + model_output: str, + ground_truth: str, + extra_info: dict, + compressed: bool = False) -> Tuple[bool, float, str]: + """ + Arguments + --------- + model_output : str – agent's raw answer + ground_truth : str – reference answer + extra_info : dict – MUST contain key "question" + + Returns + ------- + (is_correct, score, normalized_student_answer) + score is 1.0 if correct, else 0.0 + """ + if compressed: + from reward_score.utils import _deserialise_extra, _decompress_str + extra_info = _deserialise_extra(_decompress_str(extra_info)) + ground_truth = _deserialise_extra(_decompress_str(ground_truth["compressed"]))["ground_truth"] + print(f" stem_llm_judge ground_truth: {ground_truth}") + + model_output = str(model_output) + ground_truth = str(ground_truth) + is_matched, extracted_model_output = match_answer(model_output) + question = extra_info["question"] + if is_matched==False: + return 0. + else: + try: + is_correct = _llm_judge(question, extracted_model_output, ground_truth, verbose=False) + except Exception as e: + print(f"[judge-error] {e}") + return 0. + return float(is_correct) + +# ---------------- Demo ------------------------------------------------------- +if __name__ == "__main__": + demo_item = { + "question": "A cash-strapped young professional offers to buy your car " + "with four equal annual payments of $3,000, beginning 2 " + "years from today. Assuming you can invest at 10% and want " + "to receive $9,000 for the car, should you accept?", + "answer": "$8,645.09" + } + agent_reply = "The answer is\\boxed{$8,645.09}" # pretend this comes from your agent + score = compute_score( + data_source="", + model_output=agent_reply, + ground_truth=demo_item["answer"], + extra_info={"question": demo_item["question"]} + ) + print("score =", score) diff --git a/reasoners/reward/supergpqa.py b/reasoners/reward/supergpqa.py new file mode 100644 index 00000000..6c3ef69c --- /dev/null +++ b/reasoners/reward/supergpqa.py @@ -0,0 +1,71 @@ +import re +import random +def extract_last_boxed(text): + pattern = r'\\boxed\{((?:[^{}]|\{(?:[^{}]|\{[^{}]*\})*\})*)\}' + matches = list(re.finditer(pattern, text)) + if matches: + return matches[-1].group(1) + return None + +def extract_last_final_answer(text): + pattern1 = r'Final Answer:((?:[^<]|<[^<])*?)\n' + pattern2 = r'The answer is:((?:[^<]|<[^<])*?)\n' + matches1 = list(re.finditer(pattern1, text)) + matches2 = list(re.finditer(pattern2, text)) + if matches1: + return matches1[-1].group(1) + elif matches2: + return matches2[-1].group(1) + return None + +def extract_solution(solution_str): + if '<|im_start|>user' in solution_str: + model_output = re.sub(r'^.*?<\|im_start\|>assistant', '<|im_start|>assistant', solution_str, flags=re.DOTALL, count=1) + elif 'Assistant:' in solution_str: + model_output = solution_str.split('Assistant:')[-1].strip() + else: + model_output = solution_str + + stop_words = ["", "<|im_end|>", "<|endoftext|>"] + for stop_word in stop_words: + if stop_word in model_output: + model_output = model_output.split(stop_word)[0].strip() + + extract_boxed_answer = extract_last_boxed(model_output) + if extract_boxed_answer: + return extract_boxed_answer + else: + return extract_last_final_answer(model_output) + + +def get_prediction(output): + solution = extract_solution(output) + if solution is None: + return random.choice(['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J']) + for option in ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J']: + if option in solution: + return option + return random.choice(['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J']) + +def compute_score(data_source: str, + solution_str: str, + ground_truth: str, + extra_info: any = None, + compressed: bool = False, + method: str = 'strict', + format_score: float = 0., + score: float = 1.): + if compressed: + from reward_score.utils import _deserialise_extra, _decompress_str + extra_info = _deserialise_extra(_decompress_str(extra_info)) + ground_truth = _deserialise_extra(_decompress_str(ground_truth["compressed"]))["ground_truth"] + print(f"ground_truth: {ground_truth}") + answer = get_prediction(solution_str) + if answer is None: + return 0 + else: + if answer == ground_truth: + return score + else: + return format_score + \ No newline at end of file diff --git a/reasoners/reward/tablereason.py b/reasoners/reward/tablereason.py new file mode 100644 index 00000000..55fcd74d --- /dev/null +++ b/reasoners/reward/tablereason.py @@ -0,0 +1,62 @@ +from reasoners.reward.prime_math.grader import math_equal +from reasoners.reward import math +import re + + +def _check_single_answer(answer: str, ground_truth: str) -> bool: + try: + nanswer = answer.replace(",", "").replace("%", " / 100").replace("$", "").replace(":", "/").replace("\\", "") + nanswer = float(eval(nanswer)) + return math_equal(nanswer, ground_truth, tolerance=1e-3) + except: + # If the answer is not a number, use the original answer for full string match + return math.is_equiv(answer, ground_truth) + +def drop_latex_text(answer: str) -> str: + # Remove \\text{} from "20 \\text{to} 39". There could be multiple \\text{} in the answer. + # Replace \text{something} with something + answer = re.sub(r'\\\\text\{([^}]*)\}', r'\1', answer) + answer = re.sub(r'\\\\', r'', answer) + return answer + + +def compute_score(data_source: str, model_output: str, ground_truth: str, extra_info: any = None, compressed: bool = False) -> float: + if compressed: + from reward_score.utils import _deserialise_extra, _decompress_str + extra_info = _deserialise_extra(_decompress_str(extra_info)) + ground_truth = _deserialise_extra(_decompress_str(ground_truth["compressed"]))["ground_truth"] + print(f"ground_truth: {ground_truth}") + + model_output = str(model_output).lower() + ground_truth = str(ground_truth).lower() + + solution_str = model_output.split("")[-1] + answer_str = math.last_boxed_only_string(solution_str) + if answer_str is not None: + answer = math.remove_boxed(answer_str) + answer = drop_latex_text(answer) + else: + answer = solution_str + + # print(f">>> {answer}, {ground_truth}") + if "|" not in ground_truth: + # Single numeric answer + score = _check_single_answer(answer, ground_truth) + else: + # Multiple answers, in format "ans1|ans2|ans3" + try: + ground_truth = sorted([ans.strip() for ans in ground_truth.split("|")]) + answer = sorted([ans.strip() for ans in answer.split("|")]) + if len(ground_truth) != len(answer): + score = 0 + else: + score = 1 + for gt, res in zip(ground_truth, answer): + score = _check_single_answer(res, gt) + if not score: + break + except Exception as e: + print(f"Error postprocessing ground truth or response in tablereason.py: {e}") + return 0.0 + + return float(score) diff --git a/reasoners/reward/utils.py b/reasoners/reward/utils.py new file mode 100644 index 00000000..aa86dd2a --- /dev/null +++ b/reasoners/reward/utils.py @@ -0,0 +1,37 @@ +import json +import base64 +import gzip +from typing import Dict, Any + +def _deserialise_extra(extra_str: str) -> Dict[str, Any]: + """Opposite of :func:`_serialise_extra`. Returns an empty dict if input is + not valid JSON (so we fail gracefully).""" + try: + val = json.loads(extra_str) + return val if isinstance(val, dict) else {} + except Exception: + return {} + +def _decompress_str(b64) -> str: + """Base64 解码 + gzip 解压，支持 dict 十片合并后的处理""" + if isinstance(b64, str): + # 单块字符串，直接解码 + gz = base64.b64decode(b64.encode('ascii')) + return gzip.decompress(gz).decode('utf-8') + + elif isinstance(b64, dict): + # 1) 按 part_1, part_2, … part_N 顺序取片段 + total = [] + for i in range(1, len(b64) + 1): + key = f'part_{i}' + total.append(b64[key]) + full_b64 = ''.join(total) + + # 2) 一次性 Base64 解码 + gz_bytes = base64.b64decode(full_b64.encode('ascii')) + # 3) gzip 解压 + return gzip.decompress(gz_bytes).decode('utf-8') + + else: + raise ValueError(f"Unsupported type for decompress: {type(b64)}") + diff --git a/reasoners/reward/zebra_puzzle.py b/reasoners/reward/zebra_puzzle.py new file mode 100644 index 00000000..ed42cb60 --- /dev/null +++ b/reasoners/reward/zebra_puzzle.py @@ -0,0 +1,102 @@ +import re +import random +import ast +import operator +import json +import signal +import contextlib + +class TimeoutException(Exception): + pass + +@contextlib.contextmanager +def time_limit(seconds: float): + def signal_handler(signum, frame): + raise TimeoutException("Timed out!") + + signal.setitimer(signal.ITIMER_REAL, seconds) + signal.signal(signal.SIGALRM, signal_handler) + try: + yield + finally: + signal.setitimer(signal.ITIMER_REAL, 0) + +def extract_solution(solution_str): + + answer_pattern = r'(.*?)' + match = re.finditer(answer_pattern, solution_str) + matches = list(match) + + if matches: + final_answer = matches[-1].group(1).strip() + try: + solution = ast.literal_eval(final_answer) + return solution + except (SyntaxError, ValueError): + try: + solution = json.loads(final_answer) + return solution + except json.JSONDecodeError: + return None + except Exception as e: + print(f"Error extracting solution: {e}") + return None + else: + return None + + +def compute_accuracy(answer, ground_truth): + """ + compare grid level accuracy of the final answer w the ground truth + """ + if not isinstance(answer, dict): + return 0 + + # num_objects + num_rows = len(ground_truth["rows"]) + #num_attributes + num_cols = len(ground_truth["header"]) + + #total_correct_cells + correct_cells = 0 + for i in range(num_rows): + for j in range(num_cols): + if answer["rows"][i][j] == ground_truth["rows"][i][j]: + correct_cells += 1 + #accuracy + accuracy = correct_cells / (num_rows * num_cols) + return accuracy + +def compute_score(data_source: str, + solution_str: str, + ground_truth: str, + extra_info: any = None, + compressed: bool = False, + method: str = 'strict', + timeout: float = 10.0): + if compressed: + from reward_score.utils import _deserialise_extra, _decompress_str + extra_info = _deserialise_extra(_decompress_str(extra_info)) + ground_truth = _deserialise_extra(_decompress_str(ground_truth["compressed"]))["ground_truth"] + print(f"ground_truth: {ground_truth}") + try: + with time_limit(timeout): + predicted_arrangement = extract_solution(solution_str) + + if predicted_arrangement is None: + score = 0.0 + else: + try: + accuracy = compute_accuracy(predicted_arrangement, ground_truth) + score = accuracy + except Exception as e: + score = 0.0 + + except TimeoutException: + print("Computation timed out in zebra_puzzle") + score = 0.0 + except Exception as e: + print(f"Error in compute_score in zebra_puzzle: {e}") + score = 0.0 + + return float(score) diff --git a/reasoners/tools/model_filtering/__init__.py b/reasoners/tools/model_filtering/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/reasoners/tools/model_filtering/filter.py b/reasoners/tools/model_filtering/filter.py new file mode 100644 index 00000000..a8b40647 --- /dev/null +++ b/reasoners/tools/model_filtering/filter.py @@ -0,0 +1,40 @@ +import transformers +import abc +from typing import Dict + + +class Filter(abc.ABC): + """ + Filter class for filtering data. + """ + def __init__(self): + pass + + @abc.abstractmethod + def check(self, data_entry: Dict) -> bool: + pass + + +class LengthFilter(Filter): + """ + Filter class for filtering data by length. + """ + def __init__(self, tokenizer: transformers.PreTrainedTokenizer = None, min_length: int = 0, max_length: int = 2048, length_tolerance: int = 100): + if tokenizer is None: + self.tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B") + else: + self.tokenizer = tokenizer + self.min_length = min_length + self.max_length = max_length + self.length_tolerance = length_tolerance + + def check(self, data_entry: Dict) -> bool: + if data_entry["prompt"]: + + prompt_tokens = self.tokenizer.tokenize(self.tokenizer.apply_chat_template(data_entry["prompt"], tokenize=False)) + elif data_entry["raw_prompt"]: + prompt_tokens = self.tokenizer.tokenize(data_entry["raw_prompt"]) + else: + raise ValueError("No prompt found in data") + # print(f"Prompt length: {len(prompt_tokens)}") + return self.min_length <= len(prompt_tokens) <= self.max_length - self.length_tolerance diff --git a/reasoners/tools/model_filtering/pipeline.py b/reasoners/tools/model_filtering/pipeline.py new file mode 100644 index 00000000..299d969f --- /dev/null +++ b/reasoners/tools/model_filtering/pipeline.py @@ -0,0 +1,433 @@ +#!/usr/bin/env python3 + +# Standard library imports +import datetime +import json +import os +import time +import traceback +from datetime import timedelta + +# Third party imports +import numpy as np +from rich.panel import Panel +from rich.table import Table +from torch.utils.data import DataLoader +from tqdm import tqdm +from transformers import AutoTokenizer +from vllm import LLM, SamplingParams +import traceback +# Local imports +from model_filtering.utils import console, json_default, custom_collate_fn + +class DifficultyFilterPipeline: + def __init__(self, args, dataset): + self.args = args + self.dataset = dataset + self.tokenizer = None + self.model = None + self.sampling_params = None + + self.start_time = time.time() + self.gen_times = [] + self.args.skip_percentage = getattr(self.args, 'skip_percentage', 0.0) + + @staticmethod + def format_time(seconds): + return str(timedelta(seconds=int(seconds))) + + # ------------- component init ------------------------------------------ # + def initialize_components(self): + console.print(f"🔄 Loading tokenizer from [highlight]{self.args.model_path}[/highlight]...") + self.tokenizer = AutoTokenizer.from_pretrained(self.args.model_path) + + console.print( + f"🚀 Initializing model from [highlight]{self.args.model_path}[/highlight]" + f"with TP={self.args.tp_size}..." + ) + console.print(f"Warning: enable_expert_parallel is set to {self.args.enable_expert_parallel}") + console.print("If you are using MOE models, set enable_expert_parallel to True") + + self.model = LLM( + self.args.model_path, + tensor_parallel_size=self.args.tp_size, + enable_expert_parallel=self.args.enable_expert_parallel, + enforce_eager=True, + ) + self.sampling_params = SamplingParams( + n=self.args.n, + max_tokens=self.args.max_new_tokens, + top_k=self.args.top_k, + top_p=self.args.top_p, + temperature=self.args.temperature, + repetition_penalty=self.args.repetition_penalty, + skip_special_tokens=True, + truncate_prompt_tokens=self.args.max_prompt_length + ) + console.print("✅ Model initialization [success]complete[/success].") + + # ------------- dataset -------------------------------------------------- # + def prepare_dataset(self): + console.print(f"📂 Loading dataset from [highlight]{self.args.dataset_parquet_path}[/highlight]...") + console.print(f"🐞 [DEBUG] Dataset loaded with [highlight]{len(self.dataset)}[/highlight] samples") + + # TODO: replace None values with empty strings in dataset columns and nested dictionaries + # TODO: Below sometimes causes stucks in the process + # # Replace None values with empty strings in dataset columns and nested dictionaries + # def replace_none_with_empty(obj, max_depth=3): + # if obj is None: + # return "" + # elif isinstance(obj, dict) and max_depth > 0: + # return {k: replace_none_with_empty(v, max_depth-1) for k, v in obj.items()} + # elif isinstance(obj, list) and max_depth > 0: + # return [replace_none_with_empty(item, max_depth-1) for item in obj] + # return obj + + # def process_example(example): + # return {k: replace_none_with_empty(v) for k, v in example.items()} + + # # Use batched processing with multiple workers + # dataset = dataset.map( + # process_example, + # num_proc=min(8, os.cpu_count()), + # desc="Processing None values" + # ) + # console.print("🧹 Replaced None values with empty strings using parallel processing") + + # ── debug slice + if self.args.debug: + self.dataset = self.dataset.select(range(min(48, len(self.dataset)))) + console.print(f"🐞 [DEBUG] Using first [highlight]{len(self.dataset)}[/highlight] samples") + + # ── KEEP-ONLY columns actually referenced downstream ---------------- # + required_cols = { + "prompt", + "reward_model", + "apply_chat_template", + "data_source", + "extra_info", + } + cols_to_drop = [c for c in self.dataset.column_names if c not in required_cols] + if cols_to_drop: + self.dataset = self.dataset.remove_columns(cols_to_drop) + console.print(f"🧹 Dropped {len(cols_to_drop)} column(s) for easier processing: {', '.join(cols_to_drop)}") + + # ------------- output directory / checkpoint handling ------------------ # + def get_rank_output_dir(self): + """Get the output directory for the current rank""" + model_name = self.args.model_path.split("/")[-1] + dataset_name = os.path.basename(self.args.dataset_parquet_path).rsplit(".parquet", 1)[0] + rank_output_dir = os.path.join(self.args.output_dir, dataset_name, model_name, f"dp{self.args.dp_rank}") + os.makedirs(rank_output_dir, exist_ok=True) + return rank_output_dir + + def batch_exists(self, batch_idx): + """Check if a specific batch has already been processed""" + rank_output_dir = self.get_rank_output_dir() + batch_file = os.path.join(rank_output_dir, f"batch_{batch_idx:05d}.json") + return os.path.exists(batch_file) and not self.args.force_regenerate + + def get_processed_batches_count(self): + """Count the number of already processed batches""" + rank_output_dir = self.get_rank_output_dir() + batch_files = [f for f in os.listdir(rank_output_dir) + if f.startswith("batch_") and f.endswith(".json")] + return len(batch_files) + + # ------------- batch serialization ------------------------------------ # + def process_batch_outputs(self, outputs, batch_dict): + batch_results = {} + console.print(f"DP{self.args.dp_rank} Begin to process {len(outputs)} samples") + + tokenizer = self.model.get_tokenizer() + MAX_LEN = self.args.max_prompt_length + + for i in range(len(outputs)): + messages = batch_dict["prompt"][i] + text = "".join(m["content"] for m in messages) + token_len = len(tokenizer(text).input_ids) + + # build extra_info, adding flags and lengths + extra_info = {k: batch_dict["extra_info"][k][i] + for k in batch_dict["extra_info"]} + extra_info["too_long"] = (token_len > MAX_LEN) + extra_info["token_length"] = token_len + try: + # collect model outputs + full_responses = [r.text for r in outputs[i].outputs] + + # compute mean response length in tokens + if full_responses: + resp_token_counts = [ + len(tokenizer(resp).input_ids) for resp in full_responses + ] + mean_resp_len = sum(resp_token_counts) / len(resp_token_counts) + else: + mean_resp_len = 0.0 + extra_info["response_mean_length"] = mean_resp_len + + data_source = batch_dict["data_source"][i] + if "zebra" in self.args.dataset_parquet_path: + ground_truth = "" + else: + ground_truth = batch_dict["reward_model"]["ground_truth"][i] + question = next((m["content"] for m in messages if m["role"] == "user"), + "No question found") + + batch_results[i] = { + "messages": messages, + "question": question, + "ground_truth": ground_truth, + "source": data_source, + "responses": full_responses, + "extra_info": extra_info, + } + + if token_len > MAX_LEN: + console.print(f"⚠️ Sample {i} is over-length ({token_len} tokens); marked in extra_info.") + + except Exception as e: + console.print(f"❌ Error processing sample {i}: {e}") + console.print(traceback.format_exc()) + + console.print(f"✅ Finished processing. Kept {len(batch_results)} / {len(outputs)} samples.") + return batch_results + + # ------------- progress ------------------------------------------------- # + def print_progress_stats(self, idx, total_batches): + elapsed = time.time() - self.start_time + eta = "calculating..." + if idx > 0 and self.gen_times: + remain = total_batches - idx - 1 + eta_batch = np.mean(self.gen_times[-10:]) + eta = self.format_time(remain * eta_batch) + + console.print() + console.print( + Panel( + f"[bold]Progress: [metric]{idx+1}/{total_batches}[/metric] " + f"({(idx+1)/total_batches*100:.1f}%)[/bold]", + title="📊 Summary", + border_style="cyan", + ) + ) + console.print(f"⏱️ Elapsed: [time]{self.format_time(elapsed)}[/time] | ETA: [time]{eta}[/time]") + if self.gen_times: + console.print(f"⚡ Generation avg (last 10): [metric]{np.mean(self.gen_times[-10:]):.2f}s[/metric]") + + def check_all_ranks_completed(self, base_output_dir): + """ + Check if all DP ranks have completed their processing. + + Args: + base_output_dir: Base directory for all DP ranks output + + Returns: + tuple: (is_complete, completed_ranks, missing_ranks) + """ + if self.args.dp_size <= 1: + return True, [0], [] + + completed_ranks = [] + missing_ranks = [] + + for rank in range(self.args.dp_size): + rank_dir = os.path.join(base_output_dir, f"dp{rank}") + rank_completion_file = os.path.join(rank_dir, "RANK_COMPLETE") + if os.path.exists(rank_completion_file): + completed_ranks.append(rank) + else: + missing_ranks.append(rank) + + return len(completed_ranks) == self.args.dp_size, completed_ranks, missing_ranks + + def busy_wait_for_all_ranks(self, base_output_dir, dataloader): + """ + Wait for all DP ranks to complete their processing while keeping the process active. + Instead of sleeping, this will continue running inference without saving results. + + Args: + base_output_dir: Base directory for all DP ranks output + dataloader: The dataloader to use for busy waiting + """ + if self.args.dp_size <= 1: + return + + console.print(f"\n⏳ [DP-{self.args.dp_rank}] Waiting for all DP ranks to complete with busy waiting...\n") + + check_interval = 60 # Check status every 60 seconds + last_check_time = time.time() + + # Create an iterator over the dataloader that loops continuously + busy_iter = iter(dataloader) + + while True: + # Check if all ranks have completed (but not too frequently) + current_time = time.time() + if current_time - last_check_time >= check_interval: + all_ready, completed_ranks, missing_ranks = self.check_all_ranks_completed(base_output_dir) + last_check_time = current_time + + if all_ready: + # Create a nice table summarizing all ranks + table = Table(title=f"[DP-{self.args.dp_rank}] DP Ranks Completion Status") + table.add_column("Rank", style="cyan") + table.add_column("Status", style="green") + + for rank in range(self.args.dp_size): + table.add_row(f"DP-{rank}", "✅ Complete") + + console.print(table) + console.print(f"\n✅ [DP-{self.args.dp_rank}] All {self.args.dp_size} DP ranks have completed their work") + break + else: + # Create a table showing status of all ranks + table = Table(title=f"[DP-{self.args.dp_rank}] DP Ranks Completion Status") + table.add_column("Rank", style="cyan") + table.add_column("Status", style="green") + + for rank in range(self.args.dp_size): + if rank in completed_ranks: + table.add_row(f"DP-{rank}", "✅ Complete") + else: + table.add_row(f"DP-{rank}", "⏳ Waiting") + + console.print(table) + console.print(f"\n⏳ [DP-{self.args.dp_rank}] Waiting for {len(missing_ranks)} rank(s) to complete... ({len(completed_ranks)}/{self.args.dp_size} done)") + console.print(f" [DP-{self.args.dp_rank}] Missing: {', '.join(f'DP-{r}' for r in missing_ranks)}") + console.print(f" [DP-{self.args.dp_rank}] Continuing to run inference without saving to stay active") + + # Do some work to keep the process active + try: + batch_dict = next(busy_iter) + except StopIteration: + # Reset the iterator if we reach the end + busy_iter = iter(dataloader) + batch_dict = next(busy_iter) + + # Run inference on a batch but don't save the results + console.print(f"\n🔄 [DP-{self.args.dp_rank}] Running busy-wait inference to stay active (results won't be saved)...") + outputs = self.model.chat( + batch_dict["prompt"], + sampling_params=self.sampling_params, + use_tqdm=False + ) + console.print(f"✅ [DP-{self.args.dp_rank}] Completed busy-wait inference") + + # ------------- main loop ------------------------------------------------ # + def run_inference(self): + self.initialize_components() + self.prepare_dataset() + + dataloader = DataLoader( + self.dataset, + batch_size=self.args.batch_size, + pin_memory=True, + num_workers=min(4, os.cpu_count() // max(1, self.args.dp_size)), + prefetch_factor=8, + collate_fn=custom_collate_fn + ) + + total_batches = len(dataloader) + + # Get the rank output directory + model_name = self.args.model_path.split("/")[-1] + dataset_name = os.path.basename(self.args.dataset_parquet_path).rsplit(".parquet", 1)[0] + rank_output_dir = self.get_rank_output_dir() + + # Determine starting batch index + if self.args.skip_percentage > 0: + start_batch_idx = int(total_batches * self.args.skip_percentage) + console.print(f"⏩ [DP-{self.args.dp_rank}] Skipping to {self.args.skip_percentage*100:.0f}% of batches (batch {start_batch_idx}/{total_batches})") + else: + start_batch_idx = 0 + + # Count already processed batches for logging + processed_count = self.get_processed_batches_count() + if processed_count > 0 and not self.args.force_regenerate: + console.print(f"📋 [DP-{self.args.dp_rank}] Found {processed_count} already processed batches") + elif self.args.force_regenerate: + console.print(f"🔄 [DP-{self.args.dp_rank}] [warning]Force regeneration requested - ignoring existing batch files[/warning]") + + progress_bar = tqdm( + total=len(dataloader), + initial=start_batch_idx, + desc=f"🔄 DP{self.args.dp_rank} Processing", + position=0, + ) + + # Skip to the starting batch + batch_iter = iter(dataloader) + for _ in range(start_batch_idx): + try: + next(batch_iter) + except StopIteration: + console.print(f"⚠️ [DP-{self.args.dp_rank}] [warning]Skip index {start_batch_idx} exceeds dataset size[/warning]") + break + + # Process each batch with smart checkpoint checking + for idx, batch_dict in enumerate(batch_iter, start=start_batch_idx): + # Check if this specific batch has already been processed + if self.batch_exists(idx): + console.print(f"\n⏩ [DP-{self.args.dp_rank}] Batch {idx} already exists, skipping...") + progress_bar.update(1) + continue + + # Batch doesn't exist, process it + batch_size = len(batch_dict["reward_model"]["ground_truth"]) + console.print(f"\n🔄 [DP-{self.args.dp_rank}] Generating for batch {idx}/{len(dataloader)-1} ({batch_size} samples)…") + + # enforce apply_chat_template==True + if not all(batch_dict.get("apply_chat_template", [True] * batch_size)): + raise RuntimeError( + "Encountered apply_chat_template=False but raw_prompt column is removed. " + "Please ensure all samples set apply_chat_template=True." + ) + + gen_start = time.time() + outputs = self.model.chat( + batch_dict["prompt"], + sampling_params=self.sampling_params + ) + + self.gen_times.append(time.time() - gen_start) + console.print(f"⏱️ [DP-{self.args.dp_rank}] Generation took [time]{self.gen_times[-1]:.2f}s[/time]") + + # ----------- store outputs ------------------------- # + batch_results = self.process_batch_outputs(outputs, batch_dict) + + batch_out = os.path.join(rank_output_dir, f"batch_{idx:05d}.json") + with open(batch_out, "w") as f: + json.dump(batch_results, f, indent=2, default=json_default) + console.print(f"💾 [DP-{self.args.dp_rank}] Saved batch results to [highlight]{batch_out}[/highlight]") + + self.print_progress_stats(idx, len(dataloader)) + progress_bar.update(1) + console.rule(style="cyan") + + progress_bar.close() + elapsed_total = time.time() - self.start_time + + console.print() + console.print(Panel(f"[bold][DP-{self.args.dp_rank}] Inference completed!", title="🏁 Finished", border_style="green")) + console.print(f"⏱️ [DP-{self.args.dp_rank}] Total time: [time]{self.format_time(elapsed_total)}[/time]") + console.print( + f"ℹ️ [DP-{self.args.dp_rank}] All per-batch JSON files are ready. " + "Run the separate reward script to score and assemble final results." + ) + + # Create a completion marker for this rank + model_name = self.args.model_path.split("/")[-1] + dataset_name = os.path.basename(self.args.dataset_parquet_path).rsplit(".parquet", 1)[0] + base_output_dir = os.path.join(self.args.output_dir, dataset_name, model_name) + completion_file = os.path.join(rank_output_dir, "RANK_COMPLETE") + with open(completion_file, 'w') as f: + f.write(f"Completed at {datetime.datetime.now().isoformat()}") + + console.print(f"✅ [DP-{self.args.dp_rank}] Created completion marker for DP rank {self.args.dp_rank}") + + # Use busy waiting instead of idle sleep for waiting for all ranks + if self.args.dp_size > 1: + self.busy_wait_for_all_ranks(base_output_dir, dataloader) + + console.print(f"🏁 [DP-{self.args.dp_rank}] Process complete!") \ No newline at end of file diff --git a/reasoners/tools/model_filtering/utils.py b/reasoners/tools/model_filtering/utils.py new file mode 100644 index 00000000..b91f0f97 --- /dev/null +++ b/reasoners/tools/model_filtering/utils.py @@ -0,0 +1,519 @@ +#!/usr/bin/env python3 +import os +import json +import glob +from typing import Dict, List, Optional, Set, Tuple +import numpy as np +from collections import defaultdict +import torch +from datetime import datetime +import argparse +from rich.console import Console +from rich.theme import Theme +from rich.table import Table +from rich.progress import Progress, SpinnerColumn, TimeElapsedColumn +from torch.utils.data.dataloader import default_collate +import random + +def add_suffix(filename, sample_size): + if sample_size < 1000: + size_str = f"{sample_size}" + elif (sample_size / 1000) % 1 != 0: + size_str = f"{sample_size / 1000:.1f}k" + else: + size_str = f"{sample_size // 1000}k" + return f"{filename}_{size_str}" + + +def set_seed(seed): + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + + +def sample_dataset(dataset, sample_size): + """ + Sample a dataset to a given size. + """ + if sample_size is not None: + indices = list(range(len(dataset))) + random.shuffle(indices) + indices = indices[:min(sample_size, len(dataset))] + dataset = dataset.select(indices) + return dataset + + +def save_dataset(dataset, output_dir, filename_prefix, sample_size=None): + """ + Save a dataset to a parquet file with appropriate naming. + + Args: + dataset: The dataset to save + output_dir: Directory to save the dataset + filename_prefix: Base filename to use + sample_size: Sample size to add as suffix to filename + + Returns: + str: Path to the saved file + """ + # Add suffix based on actual dataset size if sample_size is None + if sample_size is None: + sample_size = len(dataset) + + # Create filename with appropriate suffix + filename = add_suffix(filename_prefix, sample_size) + output_path = os.path.join(output_dir, f"{filename}.parquet") + + # Save dataset + dataset.to_parquet(output_path) + + return output_path + +# --------------------------------------------------------------------------- # +# Rich console setup # +# --------------------------------------------------------------------------- # +custom_theme = Theme({ + "info": "cyan", + "success": "green", + "warning": "yellow", + "error": "bold red", + "highlight": "bold magenta", + "metric": "bold cyan", + "time": "bold blue", +}) +console = Console(theme=custom_theme) + +# --------------------------------------------------------------------------- # +# Helper: make anything JSON serialisable # +# --------------------------------------------------------------------------- # +def json_default(obj): + """Fallback encoder for json.dump.""" + if isinstance(obj, torch.Tensor): + return obj.item() if obj.dim() == 0 else obj.tolist() + if isinstance(obj, (np.ndarray, np.generic)): + return obj.item() if np.ndim(obj) == 0 else obj.tolist() + if isinstance(obj, datetime): + return obj.isoformat() + raise TypeError(f"{type(obj).__name__} is not JSON serialisable") + +# --------------------------------------------------------------------------- # +# Preserves dictionary and list formats without converting to tensors # +# --------------------------------------------------------------------------- # +def custom_collate_fn(batch): + """ + Custom collate function that preserves dictionary and list formats without converting to tensors + """ + elem = batch[0] + if isinstance(elem, dict): + # For dictionaries, process each key separately + result = {} + for key in elem: + values = [d[key] for d in batch] + # Recursively process values for each key + result[key] = custom_collate_fn(values) + return result + elif isinstance(elem, list): + # For lists, return original list directly + return batch + elif isinstance(elem, tuple): + # For tuples, process each element separately + transposed = zip(*batch) + result = [] + for samples in transposed: + result.append(custom_collate_fn(samples)) + return tuple(result) + else: + # For other types, use default collate + try: + return default_collate(batch) + except: + # If default collate fails, return original data + return batch + +# --------------------------------------------------------------------------- # +# Data loading, concatenation, and analysis # +# --------------------------------------------------------------------------- # +def read_json(file_path: str) -> Dict: + """Read a single JSON file.""" + try: + with open(file_path, 'r') as f: + return json.load(f) + except Exception as e: + console.print(f"[error]Error reading {file_path}: {e}[/error]") + return {} + +# --------------------------------------------------------------------------- # +# Extract idx to pass_rate mapping # +# --------------------------------------------------------------------------- # +def extract_idx_to_passrate(output_dir: str, dataset_name: str, model_name: str) -> Dict[str, float]: + """ + Extract mapping from idx to pass_rate for a specific model and dataset. + + Args: + output_dir: Base output directory + dataset_name: Name of the dataset + model_name: Name of the model + + Returns: + Dictionary mapping unique idx to pass_rate + """ + model_dir = os.path.join(output_dir, dataset_name, model_name) + if not os.path.exists(model_dir): + console.print(f"[error]Directory not found: {model_dir}[/error]") + return {} + + idx_to_pass_rate = {} + seen_idx_set = set() + dp_dirs = glob.glob(os.path.join(model_dir, "dp*")) + + if not dp_dirs: + console.print(f"[warning]No DP directories found in {model_dir}[/warning]") + return {} + + console.print(f"Found {len(dp_dirs)} DP directories in {model_dir}") + + for dp_dir in dp_dirs: + # Extract dp_rank from directory name + try: + dp_rank = int(os.path.basename(dp_dir).replace('dp', '')) + except ValueError: + console.print(f"[warning]Could not extract dp_rank from {dp_dir}[/warning]") + continue + + final_results_path = os.path.join(dp_dir, "final_results.json") + if not os.path.exists(final_results_path): + console.print(f"[warning]No final_results.json found in {dp_dir}[/warning]") + continue + + data = read_json(final_results_path) + if not data or "results" not in data: + console.print(f"[warning]Invalid results format in {final_results_path}[/warning]") + continue + + # Extract idx to pass_rate mapping + for key, value in data["results"].items(): + idx = None + if "extra_info" in value: + # Check if extra_info contains an index field + for index_field in ["index", "idx", "id"]: + if index_field in value["extra_info"]: + idx = value["extra_info"][index_field] + break + + # Assert that we found an index + assert idx is not None, f"No index field found in extra_info for sample {key}" + + # Assert that source exists in the value + assert "source" in value, f"No 'source' field found in value for sample {key}" + + # Create combined id with source and index + idx = f"{value['source']}_{idx}" + + if idx is not None: + # Check for duplicate idx + if idx in seen_idx_set: + raise ValueError(f"Duplicate idx '{idx}' found in dataset {dataset_name}, model {model_name}") + + seen_idx_set.add(idx) + pass_rate = value["pass_rate"] + idx_to_pass_rate[idx] = pass_rate + else: + console.print(f"[warning]Missing index in extra_info for sample {key} in {final_results_path}[/warning]") + + return idx_to_pass_rate + +# --------------------------------------------------------------------------- # +# Generate idx to pass_rate mapping for a dataset # +# --------------------------------------------------------------------------- # +def generate_idx_mapping(output_dir: str, dataset_name: str, model_name: str) -> Dict[str, float]: + """ + Generate idx to pass_rate mapping for a single model in a dataset. + + Args: + output_dir: Base output directory + dataset_name: Name of the dataset + model_name: Name of the model to process + + Returns: + Dictionary mapping unique idx to pass_rate + """ + dataset_dir = os.path.join(output_dir, dataset_name) + if not os.path.exists(dataset_dir): + console.print(f"[error]Dataset directory not found: {dataset_dir}[/error]") + return {} + + console.print(f"[bold]Processing Model: {model_name}[/bold]") + + # Extract idx to pass_rate mapping for this model + idx_to_pass_rate = extract_idx_to_passrate(output_dir, dataset_name, model_name) + + if not idx_to_pass_rate: + console.print(f"[warning]No valid idx to pass_rate mapping found for {model_name}[/warning]") + return {} + + # Save individual model mapping + model_dir = os.path.join(dataset_dir, model_name) + model_mapping_path = os.path.join(model_dir, "idx_to_passrate.json") + try: + with open(model_mapping_path, 'w') as f: + json.dump(idx_to_pass_rate, f, indent=2, default=json_default) + console.print(f"[success]Saved idx to pass_rate mapping for {model_name} ({len(idx_to_pass_rate)} samples)[/success]") + except Exception as e: + console.print(f"[error]Failed to save idx to pass_rate mapping for {model_name}: {e}[/error]") + + return idx_to_pass_rate + +# --------------------------------------------------------------------------- # +# Combine mappings from multiple datasets # +# --------------------------------------------------------------------------- # +def combine_mappings(output_dir: str, dataset_names: List[str], model_name: str, regenerate: bool = False) -> Dict[str, float]: + """ + Combine idx to pass_rate mappings from multiple datasets for a single model. + + Args: + output_dir: Base output directory + dataset_names: List of dataset names to combine + model_name: Model name to use for all datasets + regenerate: Whether to regenerate mappings + + Returns: + Combined mapping of unique idx to pass_rate + """ + combined_mapping = {} + all_seen_idx = set() + + for dataset_name in dataset_names: + console.print(f"\n[bold]Processing dataset: {dataset_name}[/bold]") + model_dir = os.path.join(output_dir, dataset_name, model_name) + + if not os.path.exists(model_dir): + console.print(f"[warning]Directory {model_dir} does not exist, skipping[/warning]") + continue + + mapping_path = os.path.join(model_dir, "idx_to_passrate.json") + idx_to_pass_rate = {} + + # Generate or load mapping + if regenerate: + console.print(f"[bold]Regenerating mapping for {dataset_name} with {model_name}[/bold]") + idx_to_pass_rate = generate_idx_mapping(output_dir, dataset_name, model_name) + elif not os.path.exists(mapping_path): + console.print(f"[error]Mapping file not found: {mapping_path}[/error]") + console.print(f"[info]Please run the 'map' command first[/info]") + return + else: + idx_to_pass_rate = read_json(mapping_path) + console.print(f"[info]Loaded existing mapping for {dataset_name} with {len(idx_to_pass_rate)} samples[/info]") + + # Check for duplicate idx across datasets + for idx in idx_to_pass_rate: + if idx in all_seen_idx: + raise ValueError(f"Duplicate idx '{idx}' found across multiple datasets") + all_seen_idx.add(idx) + + # Add to combined mapping + combined_mapping.update(idx_to_pass_rate) + + # Save combined mapping + combined_dir = os.path.join(output_dir, "combined") + os.makedirs(combined_dir, exist_ok=True) + + combined_path = os.path.join(combined_dir, f"{model_name}_combined.json") + try: + with open(combined_path, 'w') as f: + json.dump(combined_mapping, f, indent=2, default=json_default) + console.print(f"[success]Saved combined mapping for {model_name} with {len(combined_mapping)} samples from {len(dataset_names)} datasets[/success]") + except Exception as e: + console.print(f"[error]Failed to save combined mapping: {e}[/error]") + + return combined_mapping + +# --------------------------------------------------------------------------- # +# Analyze dataset difficulty based on idx_to_passrate mapping # +# --------------------------------------------------------------------------- # +def analyze_dataset_difficulty(idx_to_pass_rate: Dict[str, float]) -> None: + """ + Analyze difficulty distribution based on idx_to_passrate mapping. + + Args: + idx_to_pass_rate: Dictionary mapping unique idx to pass_rate + """ + if not idx_to_pass_rate: + console.print(f"[error]Empty mapping provided for analysis[/error]") + return + + # Extract pass rates + pass_rates = list(idx_to_pass_rate.values()) + + # Calculate statistics + mean_pass_rate = float(np.mean(pass_rates)) + median_pass_rate = float(np.median(pass_rates)) + std_pass_rate = float(np.std(pass_rates)) + + # Print basic statistics + stats_table = Table(show_header=True, header_style="bold magenta") + stats_table.add_column("Metric") + stats_table.add_column("Value") + + stats = { + "Total samples": len(pass_rates), + "Mean pass rate": f"{mean_pass_rate:.3f}", + "Median pass rate": f"{median_pass_rate:.3f}", + "Std pass rate": f"{std_pass_rate:.3f}", + } + + for metric, value in stats.items(): + stats_table.add_row(metric, str(value)) + + console.print(stats_table) + + # Print difficulty distribution + dist_table = Table(show_header=True, header_style="bold magenta") + dist_table.add_column("Difficulty") + dist_table.add_column("Pass Rate Range") + dist_table.add_column("Count") + dist_table.add_column("Percentage") + + # Count exact 0s and 1s + exact_zeros = sum(1 for r in pass_rates if r == 0.0) + exact_ones = sum(1 for r in pass_rates if r == 1.0) + + # Add special categories for exactly 0 and exactly 1 + dist_table.add_row( + "Impossible", + "Exactly 0.0", + str(exact_zeros), + f"{(exact_zeros / len(pass_rates)) * 100:.1f}%" + ) + + # Update bins to exclude exact 0s and 1s + bins = [(0, 0.2, "Very Hard", True, False), + (0.2, 0.4, "Hard", False, False), + (0.4, 0.6, "Medium", False, False), + (0.6, 0.8, "Easy", False, False), + (0.8, 1.0, "Very Easy", False, True)] + + for bin_start, bin_end, difficulty, exclude_start, exclude_end in bins: + # Count items in this bin, respecting exclusions + count = sum(1 for r in pass_rates if + (bin_start < r if exclude_start else bin_start <= r) and + (r < bin_end if exclude_end else r <= bin_end)) + + percentage = (count / len(pass_rates)) * 100 + + # Show range notation with appropriate brackets + range_str = f"{'(' if exclude_start else '['}{bin_start:.1f}-{bin_end:.1f}{')' if exclude_end else ']'}" + + dist_table.add_row( + difficulty, + range_str, + str(count), + f"{percentage:.1f}%" + ) + + # Add special category for exactly 1 + dist_table.add_row( + "Perfect", + "Exactly 1.0", + str(exact_ones), + f"{(exact_ones / len(pass_rates)) * 100:.1f}%" + ) + + console.print(dist_table) + +def main(): + parser = argparse.ArgumentParser(description="Dataset difficulty analysis tools") + subparsers = parser.add_subparsers(dest='command', help='Command to execute') + + # Map command to generate idx to pass_rate mapping + map_parser = subparsers.add_parser('map', help='Generate idx to pass_rate mapping') + map_parser.add_argument('--output_dir', type=str, required=True, help='Base output directory') + map_parser.add_argument('--dataset', type=str, required=True, help='Dataset name') + map_parser.add_argument('--model', type=str, required=True, help='Model name to process') + + # Analyze command with support for multiple datasets + analyze_parser = subparsers.add_parser('analyze', help='Analyze dataset difficulty') + analyze_parser.add_argument('--output_dir', type=str, required=True, help='Base output directory') + analyze_parser.add_argument('--datasets', type=str, nargs='+', required=True, help='Dataset name(s) to analyze') + analyze_parser.add_argument('--model', type=str, required=True, help='Model name to analyze') + analyze_parser.add_argument('--regenerate', action='store_true', help='Regenerate idx to pass_rate mapping') + analyze_parser.add_argument('--save_combined', action='store_true', help='Save combined mapping (default: False)') + + args = parser.parse_args() + + if args.command == 'map': + generate_idx_mapping(args.output_dir, args.dataset, args.model) + elif args.command == 'analyze': + # Handle multiple datasets + combined_mapping = {} + all_seen_idx = set() + + for dataset_name in args.datasets: + console.print(f"\n[bold]Processing dataset: {dataset_name}[/bold]") + model_dir = os.path.join(args.output_dir, dataset_name, args.model) + + if not os.path.exists(model_dir): + console.print(f"[warning]Directory {model_dir} does not exist, skipping[/warning]") + continue + + mapping_path = os.path.join(model_dir, "idx_to_passrate.json") + idx_to_pass_rate = {} + + # Generate or load mapping + if args.regenerate: + console.print(f"[bold]Regenerating mapping for {dataset_name} with {args.model}[/bold]") + idx_to_pass_rate = generate_idx_mapping(args.output_dir, dataset_name, args.model) + elif not os.path.exists(mapping_path): + console.print(f"[error]Mapping file not found: {mapping_path}[/error]") + console.print(f"[info]Generating mapping for {dataset_name} with {args.model}[/info]") + idx_to_pass_rate = generate_idx_mapping(args.output_dir, dataset_name, args.model) + else: + idx_to_pass_rate = read_json(mapping_path) + console.print(f"[info]Loaded existing mapping for {dataset_name} with {len(idx_to_pass_rate)} samples[/info]") + + # Check for duplicate idx across datasets + for idx in idx_to_pass_rate: + if idx in all_seen_idx: + console.print(f"[warning]Duplicate idx '{idx}' found across multiple datasets, last occurrence will be used[/warning]") + all_seen_idx.add(idx) + + # Add to combined mapping + combined_mapping.update(idx_to_pass_rate) + + # Save combined mapping if requested + if args.save_combined and len(args.datasets) > 1: + # Determine base dataset name by removing "_chunk_XX" from dataset names + dataset_bases = set() + for dataset_name in args.datasets: + if "_chunk" in dataset_name: + base_name = dataset_name.split("_chunk")[0] + dataset_bases.add(base_name) + else: + dataset_bases.add(dataset_name) + + # Use the base name if all datasets have the same base, otherwise use "combined" + if len(dataset_bases) == 1: + combined_name = list(dataset_bases)[0] + else: + combined_name = "combined" + + # Create directory if it doesn't exist + combined_dir = os.path.join(args.output_dir, combined_name) + os.makedirs(combined_dir, exist_ok=True) + + combined_path = os.path.join(combined_dir, f"{args.model}_combined.json") + try: + with open(combined_path, 'w') as f: + json.dump(combined_mapping, f, indent=2, default=json_default) + console.print(f"[success]Saved combined mapping for {args.model} with {len(combined_mapping)} samples to {combined_path}[/success]") + except Exception as e: + console.print(f"[error]Failed to save combined mapping: {e}[/error]") + # Always analyze the mapping (whether single dataset or combined) + dataset_label = args.datasets[0] if len(args.datasets) == 1 else f"{len(args.datasets)} combined datasets" + console.print(f"\n[bold]Analyzing {dataset_label} for {args.model}[/bold]") + analyze_dataset_difficulty(combined_mapping) + else: + parser.print_help() + +if __name__ == "__main__": + main() \ No newline at end of file