From 39760685f6f7518db5b915407e8190e13cd30df7 Mon Sep 17 00:00:00 2001 From: Shibo Hao Date: Thu, 29 May 2025 08:07:01 +0000 Subject: [PATCH 1/5] reward function tested, first commit includes data preprocess, model filtering, reeward funtion --- reasoners/data_preprocess/README.md | 104 + reasoners/data_preprocess/__init__.py | 0 ...iew_filtered_mar21_r1_instruction_style.py | 187 + ...th_preview_filtered_mar21_r1_zero_style.py | 190 + .../archived_for_reference/code_utils.py | 241 + .../archived_for_reference/codegen.py | 977 +++ ...deepscaler_preview_r1_zero_style_filter.py | 192 + ...eepscaler_preview_r1_zero_style_sampled.py | 201 + .../archived_for_reference/geo3k.py | 86 + .../archived_for_reference/orz_data.py | 163 + .../orz_data_r1_zero_style.py | 192 + .../data_preprocess/codegen/humaneval.py | 156 + .../data_preprocess/codegen/humanevalplus.py | 157 + .../data_preprocess/codegen/leetcode2k.py | 174 + .../data_preprocess/codegen/livecodebench.py | 230 + reasoners/data_preprocess/codegen/mbpp.py | 169 + reasoners/data_preprocess/codegen/mbppplus.py | 158 + .../data_preprocess/codegen/primeintellect.py | 236 + reasoners/data_preprocess/codegen/taco.py | 246 + reasoners/data_preprocess/logic/gen_proof.py | 781 +++ .../logic/graph_dataset_gen/core/LICENSE | 201 + .../logic/graph_dataset_gen/core/README.md | 21 + .../logic/graph_dataset_gen/core/array.h | 2429 +++++++ .../logic/graph_dataset_gen/core/core.h | 479 ++ .../logic/graph_dataset_gen/core/io.h | 1496 +++++ .../logic/graph_dataset_gen/core/lex.h | 504 ++ .../logic/graph_dataset_gen/core/map.h | 2508 ++++++++ .../logic/graph_dataset_gen/core/random.h | 321 + .../logic/graph_dataset_gen/core/timer.h | 130 + .../logic/graph_dataset_gen/core/utility.h | 799 +++ .../logic/graph_dataset_gen/core/xxhash.h | 5580 +++++++++++++++++ .../logic/graph_dataset_gen/generator.cpp | 2151 +++++++ .../logic/graph_dataset_gen/jobfile.sub | 11 + .../logic/graph_dataset_gen/logic.py | 359 ++ .../logic/process_graph_dataset.py | 159 + .../logic/process_puzzles_dataset.py | 198 + .../logic/process_zebrapuzzle_dataset.py | 143 + reasoners/data_preprocess/logic/puzzle_gen.py | 318 + .../logic/zebrapuzzle_gen/generator.py | 488 ++ .../logic/zebrapuzzle_gen/puzzle_generator.py | 202 + .../math/bigmath_preview_filtered_mar21.py | 159 + .../math/dapo_or1_merge_dedup_apr30.py | 183 + .../math/deepscaler_preview.py | 143 + reasoners/data_preprocess/ood/ifeval.py | 139 + reasoners/data_preprocess/ood/livebench.py | 160 + .../data_preprocess/sft_rlhf/full_hh_rlhf.py | 146 + reasoners/data_preprocess/sft_rlhf/limo.py | 93 + .../data_preprocess/simulation/arcagi.py | 183 + reasoners/data_preprocess/simulation/barc.py | 166 + .../data_preprocess/simulation/codeio.py | 193 + .../data_preprocess/simulation/cruxeval-i.py | 136 + .../data_preprocess/simulation/cruxeval-o.py | 136 + reasoners/data_preprocess/stem/bbeh.py | 0 reasoners/data_preprocess/stem/gpqa.py | 137 + .../data_preprocess/stem/gpqa_diamond.py | 142 + reasoners/data_preprocess/stem/supergpqa.py | 118 + reasoners/data_preprocess/stem_web/README.md | 20 + .../data_preprocess/stem_web/json2_parquet.py | 109 + .../data_preprocess/stem_web/simple_filter.py | 35 + reasoners/data_preprocess/stem_web/stem.py | 133 + reasoners/data_preprocess/table/finqa.py | 170 + reasoners/data_preprocess/table/hitab.py | 263 + reasoners/data_preprocess/table/multihier.py | 274 + reasoners/reward/__init__.py | 92 + reasoners/reward/arcagi.py | 101 + reasoners/reward/codeio.py | 190 + reasoners/reward/coder1/__init__.py | 301 + reasoners/reward/coder1/bwrap_exec.py | 125 + reasoners/reward/coder1/ces_exec.py | 52 + reasoners/reward/coder1/docker_exec.py | 78 + reasoners/reward/coder1/firejail_exec.py | 97 + reasoners/reward/coder1/kira_exec.py | 44 + reasoners/reward/coder1/unsafe_local_exec.py | 113 + reasoners/reward/coder1/utils.py | 11 + reasoners/reward/cruxeval/__init__.py | 33 + reasoners/reward/cruxeval/cruxeval.py | 0 reasoners/reward/cruxeval/utils.py | 219 + reasoners/reward/geo3k.py | 31 + reasoners/reward/gpqa.py | 78 + reasoners/reward/graph_dataset.py | 83 + reasoners/reward/gsm8k.py | 63 + reasoners/reward/ifeval/__init__.py | 47 + reasoners/reward/ifeval/instructions.py | 1611 +++++ .../reward/ifeval/instructions_registry.py | 168 + reasoners/reward/ifeval/instructions_util.py | 1701 +++++ reasoners/reward/livebench/__init__.py | 68 + .../livebench/data_analysis/cta/utils.py | 32 + .../data_analysis/tablejoin/utils.py | 82 + .../data_analysis/tablereformat/utils.py | 240 + .../reasoning/house_traversal/utils.py | 41 + .../livebench/reasoning/spatial/utils.py | 61 + .../reasoning/web_of_lies_v2/utils.py | 79 + .../reasoning/web_of_lies_v3/utils.py | 87 + .../livebench/reasoning/zebra_puzzle/utils.py | 128 + reasoners/reward/livebench/util.py | 46 + .../livebench/writing/connections/utils.py | 108 + .../writing/plot_unscrambling/utils.py | 66 + .../reward/livebench/writing/typos/utils.py | 59 + reasoners/reward/math.py | 227 + reasoners/reward/math_batch.py | 25 + reasoners/reward/math_dapo.py | 294 + reasoners/reward/math_llm_judge/__init__.py | 458 ++ reasoners/reward/math_llm_judge/grader.py | 380 ++ .../reward/math_llm_judge/math_normalize.py | 191 + reasoners/reward/math_verify.py | 33 + reasoners/reward/naive_dapo.py | 526 ++ reasoners/reward/orz/__init__.py | 3 + reasoners/reward/orz/math_utils.py | 448 ++ reasoners/reward/orz/math_utils_sync.py | 442 ++ reasoners/reward/prime_code/__init__.py | 73 + reasoners/reward/prime_code/testing_util.py | 721 +++ reasoners/reward/prime_code/utils.py | 59 + reasoners/reward/prime_math/__init__.py | 437 ++ reasoners/reward/prime_math/grader.py | 380 ++ reasoners/reward/prime_math/math_normalize.py | 191 + reasoners/reward/puzzles_dataset.py | 137 + reasoners/reward/stem_llm_judge/__init__.py | 182 + reasoners/reward/supergpqa.py | 71 + reasoners/reward/tablereason.py | 62 + reasoners/reward/utils.py | 37 + reasoners/reward/zebra_puzzle.py | 102 + reasoners/tools/model_filtering/README.md | 194 + reasoners/tools/model_filtering/__init__.py | 0 reasoners/tools/model_filtering/pipeline.py | 433 ++ reasoners/tools/model_filtering/run_add_pr.py | 96 + reasoners/tools/model_filtering/run_filter.py | 107 + .../tools/model_filtering/run_inference.py | 179 + reasoners/tools/model_filtering/run_reward.py | 300 + .../run_sample_and_postprocess.py | 186 + .../tools/model_filtering/test_parquet.py | 113 + reasoners/tools/model_filtering/utils.py | 464 ++ 131 files changed, 40461 insertions(+) create mode 100644 reasoners/data_preprocess/README.md create mode 100644 reasoners/data_preprocess/__init__.py create mode 100644 reasoners/data_preprocess/archived_for_reference/bigmath_preview_filtered_mar21_r1_instruction_style.py create mode 100644 reasoners/data_preprocess/archived_for_reference/bigmath_preview_filtered_mar21_r1_zero_style.py create mode 100644 reasoners/data_preprocess/archived_for_reference/code_utils.py create mode 100644 reasoners/data_preprocess/archived_for_reference/codegen.py create mode 100644 reasoners/data_preprocess/archived_for_reference/deepscaler_preview_r1_zero_style_filter.py create mode 100644 reasoners/data_preprocess/archived_for_reference/deepscaler_preview_r1_zero_style_sampled.py create mode 100644 reasoners/data_preprocess/archived_for_reference/geo3k.py create mode 100644 reasoners/data_preprocess/archived_for_reference/orz_data.py create mode 100644 reasoners/data_preprocess/archived_for_reference/orz_data_r1_zero_style.py create mode 100644 reasoners/data_preprocess/codegen/humaneval.py create mode 100644 reasoners/data_preprocess/codegen/humanevalplus.py create mode 100644 reasoners/data_preprocess/codegen/leetcode2k.py create mode 100644 reasoners/data_preprocess/codegen/livecodebench.py create mode 100644 reasoners/data_preprocess/codegen/mbpp.py create mode 100644 reasoners/data_preprocess/codegen/mbppplus.py create mode 100644 reasoners/data_preprocess/codegen/primeintellect.py create mode 100644 reasoners/data_preprocess/codegen/taco.py create mode 100644 reasoners/data_preprocess/logic/gen_proof.py create mode 100644 reasoners/data_preprocess/logic/graph_dataset_gen/core/LICENSE create mode 100644 reasoners/data_preprocess/logic/graph_dataset_gen/core/README.md create mode 100644 reasoners/data_preprocess/logic/graph_dataset_gen/core/array.h create mode 100644 reasoners/data_preprocess/logic/graph_dataset_gen/core/core.h create mode 100644 reasoners/data_preprocess/logic/graph_dataset_gen/core/io.h create mode 100644 reasoners/data_preprocess/logic/graph_dataset_gen/core/lex.h create mode 100644 reasoners/data_preprocess/logic/graph_dataset_gen/core/map.h create mode 100644 reasoners/data_preprocess/logic/graph_dataset_gen/core/random.h create mode 100644 reasoners/data_preprocess/logic/graph_dataset_gen/core/timer.h create mode 100644 reasoners/data_preprocess/logic/graph_dataset_gen/core/utility.h create mode 100644 reasoners/data_preprocess/logic/graph_dataset_gen/core/xxhash.h create mode 100644 reasoners/data_preprocess/logic/graph_dataset_gen/generator.cpp create mode 100644 reasoners/data_preprocess/logic/graph_dataset_gen/jobfile.sub create mode 100644 reasoners/data_preprocess/logic/graph_dataset_gen/logic.py create mode 100644 reasoners/data_preprocess/logic/process_graph_dataset.py create mode 100644 reasoners/data_preprocess/logic/process_puzzles_dataset.py create mode 100644 reasoners/data_preprocess/logic/process_zebrapuzzle_dataset.py create mode 100644 reasoners/data_preprocess/logic/puzzle_gen.py create mode 100644 reasoners/data_preprocess/logic/zebrapuzzle_gen/generator.py create mode 100644 reasoners/data_preprocess/logic/zebrapuzzle_gen/puzzle_generator.py create mode 100644 reasoners/data_preprocess/math/bigmath_preview_filtered_mar21.py create mode 100644 reasoners/data_preprocess/math/dapo_or1_merge_dedup_apr30.py create mode 100644 reasoners/data_preprocess/math/deepscaler_preview.py create mode 100644 reasoners/data_preprocess/ood/ifeval.py create mode 100644 reasoners/data_preprocess/ood/livebench.py create mode 100644 reasoners/data_preprocess/sft_rlhf/full_hh_rlhf.py create mode 100644 reasoners/data_preprocess/sft_rlhf/limo.py create mode 100644 reasoners/data_preprocess/simulation/arcagi.py create mode 100644 reasoners/data_preprocess/simulation/barc.py create mode 100644 reasoners/data_preprocess/simulation/codeio.py create mode 100644 reasoners/data_preprocess/simulation/cruxeval-i.py create mode 100644 reasoners/data_preprocess/simulation/cruxeval-o.py create mode 100644 reasoners/data_preprocess/stem/bbeh.py create mode 100644 reasoners/data_preprocess/stem/gpqa.py create mode 100644 reasoners/data_preprocess/stem/gpqa_diamond.py create mode 100644 reasoners/data_preprocess/stem/supergpqa.py create mode 100644 reasoners/data_preprocess/stem_web/README.md create mode 100644 reasoners/data_preprocess/stem_web/json2_parquet.py create mode 100644 reasoners/data_preprocess/stem_web/simple_filter.py create mode 100644 reasoners/data_preprocess/stem_web/stem.py create mode 100644 reasoners/data_preprocess/table/finqa.py create mode 100644 reasoners/data_preprocess/table/hitab.py create mode 100644 reasoners/data_preprocess/table/multihier.py create mode 100644 reasoners/reward/__init__.py create mode 100644 reasoners/reward/arcagi.py create mode 100644 reasoners/reward/codeio.py create mode 100644 reasoners/reward/coder1/__init__.py create mode 100644 reasoners/reward/coder1/bwrap_exec.py create mode 100644 reasoners/reward/coder1/ces_exec.py create mode 100644 reasoners/reward/coder1/docker_exec.py create mode 100644 reasoners/reward/coder1/firejail_exec.py create mode 100644 reasoners/reward/coder1/kira_exec.py create mode 100644 reasoners/reward/coder1/unsafe_local_exec.py create mode 100644 reasoners/reward/coder1/utils.py create mode 100644 reasoners/reward/cruxeval/__init__.py create mode 100644 reasoners/reward/cruxeval/cruxeval.py create mode 100644 reasoners/reward/cruxeval/utils.py create mode 100644 reasoners/reward/geo3k.py create mode 100644 reasoners/reward/gpqa.py create mode 100644 reasoners/reward/graph_dataset.py create mode 100644 reasoners/reward/gsm8k.py create mode 100644 reasoners/reward/ifeval/__init__.py create mode 100644 reasoners/reward/ifeval/instructions.py create mode 100644 reasoners/reward/ifeval/instructions_registry.py create mode 100644 reasoners/reward/ifeval/instructions_util.py create mode 100644 reasoners/reward/livebench/__init__.py create mode 100644 reasoners/reward/livebench/data_analysis/cta/utils.py create mode 100644 reasoners/reward/livebench/data_analysis/tablejoin/utils.py create mode 100644 reasoners/reward/livebench/data_analysis/tablereformat/utils.py create mode 100644 reasoners/reward/livebench/reasoning/house_traversal/utils.py create mode 100644 reasoners/reward/livebench/reasoning/spatial/utils.py create mode 100644 reasoners/reward/livebench/reasoning/web_of_lies_v2/utils.py create mode 100644 reasoners/reward/livebench/reasoning/web_of_lies_v3/utils.py create mode 100644 reasoners/reward/livebench/reasoning/zebra_puzzle/utils.py create mode 100644 reasoners/reward/livebench/util.py create mode 100644 reasoners/reward/livebench/writing/connections/utils.py create mode 100644 reasoners/reward/livebench/writing/plot_unscrambling/utils.py create mode 100644 reasoners/reward/livebench/writing/typos/utils.py create mode 100644 reasoners/reward/math.py create mode 100644 reasoners/reward/math_batch.py create mode 100644 reasoners/reward/math_dapo.py create mode 100644 reasoners/reward/math_llm_judge/__init__.py create mode 100644 reasoners/reward/math_llm_judge/grader.py create mode 100644 reasoners/reward/math_llm_judge/math_normalize.py create mode 100644 reasoners/reward/math_verify.py create mode 100644 reasoners/reward/naive_dapo.py create mode 100644 reasoners/reward/orz/__init__.py create mode 100644 reasoners/reward/orz/math_utils.py create mode 100644 reasoners/reward/orz/math_utils_sync.py create mode 100644 reasoners/reward/prime_code/__init__.py create mode 100644 reasoners/reward/prime_code/testing_util.py create mode 100644 reasoners/reward/prime_code/utils.py create mode 100644 reasoners/reward/prime_math/__init__.py create mode 100644 reasoners/reward/prime_math/grader.py create mode 100644 reasoners/reward/prime_math/math_normalize.py create mode 100644 reasoners/reward/puzzles_dataset.py create mode 100644 reasoners/reward/stem_llm_judge/__init__.py create mode 100644 reasoners/reward/supergpqa.py create mode 100644 reasoners/reward/tablereason.py create mode 100644 reasoners/reward/utils.py create mode 100644 reasoners/reward/zebra_puzzle.py create mode 100644 reasoners/tools/model_filtering/README.md create mode 100644 reasoners/tools/model_filtering/__init__.py create mode 100644 reasoners/tools/model_filtering/pipeline.py create mode 100644 reasoners/tools/model_filtering/run_add_pr.py create mode 100644 reasoners/tools/model_filtering/run_filter.py create mode 100644 reasoners/tools/model_filtering/run_inference.py create mode 100644 reasoners/tools/model_filtering/run_reward.py create mode 100644 reasoners/tools/model_filtering/run_sample_and_postprocess.py create mode 100644 reasoners/tools/model_filtering/test_parquet.py create mode 100644 reasoners/tools/model_filtering/utils.py diff --git a/reasoners/data_preprocess/README.md b/reasoners/data_preprocess/README.md new file mode 100644 index 000000000..c249a1559 --- /dev/null +++ b/reasoners/data_preprocess/README.md @@ -0,0 +1,104 @@ +# Guideline + +The processed data will be saved in the `data/train` and `data/test` directories. +The naming convention is `___`, where `` is one of `math`, `codegen`, `logic`, `simulation`, `table` for now. + +# Usage +## Math + +**DAPO+OR1** +```bash +python data_preprocess/math/dapo_or1_merge_dedup_apr30.py +``` +(Please see the comment if want to apply llm as judge) +> Note: +> +> Original: OR1 (105055) + DAPO (17917) +> +> After [dedup](../scripts/tools/cross_dedup.py): 117192 +> +> After merging and removing instance with too long answer >100 char ([script](../scripts/tools/merge_and_upload_dataset.ipynb)): 116632 + +**BigMath** +```bash +python data_preprocess/math/bigmath_preview_filtered_mar21.py --train-sample-size +``` + +**DeepScaler** +```bash +python data_preprocess/math/deepscaler_preview.py --train-sample-size +``` + + +## Code +**leetcode2k** +```bash +python data_preprocess/codegen/leetcode2k.py --train-sample-size +``` +**taco** +```bash +python data_preprocess/codegen/taco.py --train-sample-size +``` +**primeintellect** +```bash +python data_preprocess/codegen/primeintellect.py --train-sample-size +``` +**humaneval** +```bash +python data_preprocess/codegen/humaneval.py +``` +**mbpp** +```bash +python data_preprocess/codegen/mbpp.py +``` +**livecodebench** +```bash +python data_preprocess/codegen/livecodebench.py +``` + +## Logic +**zebra_puzzle_dataset** +```bash +python data_preprocess/logic/zebrapuzzle_gen/puzzle_generator.py --output_dir data/raw --num_puzzles --num_processes +cd .. +python data_preprocess/logic/process_zebrapuzzle_dataset.py +``` + +**graph_logical_dataset** +```bash +uv pip install pybind11 +uv pip install Faker==37.1.0 +cd data_preprocess/logic/graph_dataset_gen/ +python logic.py --num_samples +cd ../../.. # return to Reasoning360 +python data_preprocess/logic/process_graph_dataset.py +``` + +**ordering_puzzle_dataset** +```bash +uv pip install Faker==37.1.0 +python data_preprocess/logic/puzzle_gen.py --test True --num_puzzles +python data_preprocess/logic/process_puzzles_dataset.py +``` + +## Simulation +```bash +python data_preprocess/simulation/codeio.py --train-sample-size --test-sample-size +``` + +## Table +```bash +uv pip install gdown +python data_preprocess/table/multihier.py +``` + +## ARC-AGI +```bash +python data_preprocess/simulation/arcagi.py --name arcagi1 +python data_preprocess/simulation/arcagi.py --name arcagi2 +``` + +# Add a new dataset +1. Add a new script in `data_preprocess//.py` +2. Add a new entry in `tests/data_process/test_data_preprocess.py`. +3. Run `pytest tests/data_process` to check the functionality of the data preprocessing scripts. diff --git a/reasoners/data_preprocess/__init__.py b/reasoners/data_preprocess/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/reasoners/data_preprocess/archived_for_reference/bigmath_preview_filtered_mar21_r1_instruction_style.py b/reasoners/data_preprocess/archived_for_reference/bigmath_preview_filtered_mar21_r1_instruction_style.py new file mode 100644 index 000000000..ad0cb66a5 --- /dev/null +++ b/reasoners/data_preprocess/archived_for_reference/bigmath_preview_filtered_mar21_r1_instruction_style.py @@ -0,0 +1,187 @@ +# 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. +""" +Preprocess the https://huggingface.co/datasets/agentica-org/DeepScaleR-Preview-Dataset to parquet format +Code referring to https://github.com/agentica-project/deepscaler/blob/main/scripts/data/deepscaler_dataset.py +""" + +import os +import datasets +from typing import Dict, List, Optional, Any, Union +import enum +import argparse +import pandas as pd +import json + +from verl.utils.hdfs_io import copy, makedirs +from verl.utils.reward_score.math import remove_boxed, last_boxed_only_string + + +class TrainDataset(enum.Enum): + """Enum for training datasets. + + Contains identifiers for various math problem datasets used during training. + """ + + AIME = "AIME" # American Invitational Mathematics Examination + AMC = "AMC" # American Mathematics Competition + OMNI_MATH = "OMNI_MATH" # Omni Math + NUMINA_OLYMPIAD = "OLYMPIAD" # Unique Olympiad problems from NUMINA + MATH = "MATH" # Dan Hendrycks Math Problems + STILL = "STILL" # STILL dataset + DEEPSCALER = "DEEPSCALER" # DeepScaler (AIME, AMC, OMNI_MATH, MATH, STILL) + + +class TestDataset(enum.Enum): + """Enum for testing/evaluation datasets. + + Contains identifiers for datasets used to evaluate model performance. + """ + + AIME = "AIME" # American Invitational Mathematics Examination + AMC = "AMC" # American Mathematics Competition + MATH = "MATH" # Math 500 problems + MINERVA = "MINERVA" # Minerva dataset + OLYMPIAD_BENCH = "OLYMPIAD_BENCH" # Olympiad benchmark problems + + +"""Type alias for either training or testing dataset types.""" +Dataset = Union[TrainDataset, TestDataset] + + +def load_dataset(dataset: Dataset) -> List[Dict[str, Any]]: + """Load a dataset from a JSON file. + + Loads and parses a JSON dataset file based on the provided dataset enum. + The file path is constructed based on whether it's a training or testing dataset. + + Args: + dataset: A Dataset enum value specifying which dataset to load. + + Returns: + List[Dict[str, Any]]: A list of dictionaries containing the dataset records. + Each dictionary represents one example in the dataset. + + Raises: + ValueError: If the dataset file cannot be found, contains invalid JSON, + or encounters other file access errors. + + Example: + >>> load_dataset(TrainDataset.AIME) + [{'problem': 'Find x...', 'solution': '42', ...}, ...] + """ + dataset_name = dataset.value.lower() + data_dir = "train" if isinstance(dataset, TrainDataset) else "test" + + current_dir = os.path.dirname(os.path.realpath(__file__)) + file_path = os.path.join(data_dir, f"{dataset_name}.json") + file_path = os.path.join(current_dir, file_path) + + if not os.path.exists(file_path): + raise ValueError(f"Dataset file not found: {file_path}") + + try: + with open(file_path, "r", encoding="utf-8") as file: + data = json.load(file) + return data + except json.JSONDecodeError: + raise ValueError(f"Invalid JSON format in {file_path}") + except Exception as exc: # pylint: disable=broad-except + raise ValueError(f"Error loading dataset: {exc}") from exc + + +def extract_solution(solution_str): + return remove_boxed(last_boxed_only_string(solution_str)) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--local_dir", default="data/bigmath_preview_filtered_mar21_r1_instruction_style") + parser.add_argument("--hdfs_dir", default=None) + + args = parser.parse_args() + + train_data_source = "SDSB/big_math_partial_mar21_filtered_basic" + test_data_sources = [ + "nanoverl/minerva", + "SDSB/aime_repeated_8x", + "SDSB/amc_repeated_4x", + "nanoverl/olympiad_bench", + "nanoverl/math", + ] + train_dataset = datasets.load_dataset( + train_data_source, trust_remote_code=True, split="train" + ) + print(f"Loading the {train_data_source} dataset from huggingface...", flush=True) + test_datasets = [ + datasets.load_dataset(test_data_source, trust_remote_code=True, split="test") + for test_data_source in test_data_sources + ] + + print(f"Loading the {test_data_sources} dataset from huggingface...", flush=True) + instruction_following = ( + "Please output the final answer within \\boxed{}." + ) + + + # add a row to each data item that represents a unique id + def make_map_fn(split, data_source): + + def process_fn(example, idx): + question = example.pop("problem") + + question = question + " " + instruction_following + + answer = example.pop("answer") + data = { + "data_source": data_source, + "prompt": [{ + "role": "user", + "content": question + }], + "ability": "math", + "apply_chat_template": True, + "reward_model": {"style": "rule", "ground_truth": answer}, + "extra_info": {"split": split, + "index": idx, + "question": question}, + } + if idx == 0: + print("=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) + print(data) + return data + + return process_fn + + local_dir = args.local_dir + hdfs_dir = args.hdfs_dir + + train_data = train_dataset.map( + function=make_map_fn("train", train_data_source), with_indices=True + ) + train_data.to_parquet(os.path.join(local_dir, "train.parquet")) + print(f"train data size:", len(train_data)) + print(train_data[0]) + for test_data_source, test_data in zip(test_data_sources, test_datasets): + process_fn = make_map_fn("test", test_data_source) + test_data = test_data.map(process_fn, with_indices=True) + dataset_name = os.path.basename(test_data_source.lower()) + test_df = pd.DataFrame(test_data) + test_df.to_parquet(os.path.join(local_dir, f"{dataset_name}.parquet")) + print(f"test data size: ({dataset_name})", len(test_df)) + + if hdfs_dir is not None: + makedirs(hdfs_dir) + + copy(src=local_dir, dst=hdfs_dir) diff --git a/reasoners/data_preprocess/archived_for_reference/bigmath_preview_filtered_mar21_r1_zero_style.py b/reasoners/data_preprocess/archived_for_reference/bigmath_preview_filtered_mar21_r1_zero_style.py new file mode 100644 index 000000000..535b405e5 --- /dev/null +++ b/reasoners/data_preprocess/archived_for_reference/bigmath_preview_filtered_mar21_r1_zero_style.py @@ -0,0 +1,190 @@ +# 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. +""" +Preprocess the https://huggingface.co/datasets/agentica-org/DeepScaleR-Preview-Dataset to parquet format +Code referring to https://github.com/agentica-project/deepscaler/blob/main/scripts/data/deepscaler_dataset.py +""" + +import os +import datasets +from typing import Dict, List, Optional, Any, Union +import enum +import argparse +import pandas as pd +import json + +from verl.utils.hdfs_io import copy, makedirs +from verl.utils.reward_score.math import remove_boxed, last_boxed_only_string + + +class TrainDataset(enum.Enum): + """Enum for training datasets. + + Contains identifiers for various math problem datasets used during training. + """ + + AIME = "AIME" # American Invitational Mathematics Examination + AMC = "AMC" # American Mathematics Competition + OMNI_MATH = "OMNI_MATH" # Omni Math + NUMINA_OLYMPIAD = "OLYMPIAD" # Unique Olympiad problems from NUMINA + MATH = "MATH" # Dan Hendrycks Math Problems + STILL = "STILL" # STILL dataset + DEEPSCALER = "DEEPSCALER" # DeepScaler (AIME, AMC, OMNI_MATH, MATH, STILL) + + +class TestDataset(enum.Enum): + """Enum for testing/evaluation datasets. + + Contains identifiers for datasets used to evaluate model performance. + """ + + AIME = "AIME" # American Invitational Mathematics Examination + AMC = "AMC" # American Mathematics Competition + MATH = "MATH" # Math 500 problems + MINERVA = "MINERVA" # Minerva dataset + OLYMPIAD_BENCH = "OLYMPIAD_BENCH" # Olympiad benchmark problems + + +"""Type alias for either training or testing dataset types.""" +Dataset = Union[TrainDataset, TestDataset] + + +def load_dataset(dataset: Dataset) -> List[Dict[str, Any]]: + """Load a dataset from a JSON file. + + Loads and parses a JSON dataset file based on the provided dataset enum. + The file path is constructed based on whether it's a training or testing dataset. + + Args: + dataset: A Dataset enum value specifying which dataset to load. + + Returns: + List[Dict[str, Any]]: A list of dictionaries containing the dataset records. + Each dictionary represents one example in the dataset. + + Raises: + ValueError: If the dataset file cannot be found, contains invalid JSON, + or encounters other file access errors. + + Example: + >>> load_dataset(TrainDataset.AIME) + [{'problem': 'Find x...', 'solution': '42', ...}, ...] + """ + dataset_name = dataset.value.lower() + data_dir = "train" if isinstance(dataset, TrainDataset) else "test" + + current_dir = os.path.dirname(os.path.realpath(__file__)) + file_path = os.path.join(data_dir, f"{dataset_name}.json") + file_path = os.path.join(current_dir, file_path) + + if not os.path.exists(file_path): + raise ValueError(f"Dataset file not found: {file_path}") + + try: + with open(file_path, "r", encoding="utf-8") as file: + data = json.load(file) + return data + except json.JSONDecodeError: + raise ValueError(f"Invalid JSON format in {file_path}") + except Exception as exc: # pylint: disable=broad-except + raise ValueError(f"Error loading dataset: {exc}") from exc + + +def extract_solution(solution_str): + return remove_boxed(last_boxed_only_string(solution_str)) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--local_dir", default="data/bigmath_preview_filtered_mar21_r1_zero_style" + ) + parser.add_argument("--hdfs_dir", default=None) + + args = parser.parse_args() + + train_data_source = "SDSB/big_math_partial_mar21_filtered_basic" + test_data_sources = [ + "nanoverl/minerva", + "SDSB/aime_repeated_8x", + "SDSB/amc_repeated_4x", + "nanoverl/olympiad_bench", + "nanoverl/math", + ] + train_dataset = datasets.load_dataset( + train_data_source, trust_remote_code=True, split="train" + ) + print(f"Loading the {train_data_source} dataset from huggingface...", flush=True) + test_datasets = [ + datasets.load_dataset(test_data_source, trust_remote_code=True, split="test") + for test_data_source in test_data_sources + ] + + print(f"Loading the {test_data_sources} dataset from huggingface...", flush=True) + instruction_following = ( + "Let's think step by step and output the final answer within \\boxed{}." + ) + + prompt = """A conversation between User and Assistant. The user asks a question, and the Assistant solves it. The assistant first thinks about the reasoning process in the mind and then provides the user with the response. The reasoning process is enclosed within i.e., reasoning process here respond to the user's question here. + +User: {{prompt}} Please put your answer in \\boxed{} tags. +Assistant: +""" + + # add a row to each data item that represents a unique id + def make_map_fn(split, data_source): + + def process_fn(example, idx): + question = example.pop("problem") + + # question = question + " " + instruction_following + + answer = example.pop("answer") + data = { + "data_source": data_source, + "prompt": [], + "raw_prompt": prompt.replace("{{prompt}}", question), + "ability": "math", + "apply_chat_template": False, + "reward_model": {"style": "rule", "ground_truth": answer}, + "extra_info": {"split": split, "index": idx, "question": question}, + } + if idx == 0: + print("=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) + print(data) + return data + + return process_fn + + local_dir = args.local_dir + hdfs_dir = args.hdfs_dir + + train_data = train_dataset.map( + function=make_map_fn("train", train_data_source), with_indices=True + ) + train_data.to_parquet(os.path.join(local_dir, "train.parquet")) + print(f"train data size:", len(train_data)) + print(train_data[0]) + for test_data_source, test_data in zip(test_data_sources, test_datasets): + process_fn = make_map_fn("test", test_data_source) + test_data = test_data.map(process_fn, with_indices=True) + dataset_name = os.path.basename(test_data_source.lower()) + test_df = pd.DataFrame(test_data) + test_df.to_parquet(os.path.join(local_dir, f"{dataset_name}.parquet")) + print(f"test data size: ({dataset_name})", len(test_df)) + + if hdfs_dir is not None: + makedirs(hdfs_dir) + + copy(src=local_dir, dst=hdfs_dir) diff --git a/reasoners/data_preprocess/archived_for_reference/code_utils.py b/reasoners/data_preprocess/archived_for_reference/code_utils.py new file mode 100644 index 000000000..8774e39e2 --- /dev/null +++ b/reasoners/data_preprocess/archived_for_reference/code_utils.py @@ -0,0 +1,241 @@ +""" +Code utils for data preprocessing +""" + +import sys +import pandas as pd +import matplotlib.pyplot as plt + +N_TESTSET_PER_DATASET = 512 # per dataset +MAX_N_TRAINSET_PER_DATASET = 50000 +EMPTY_RETURN = { + "data_source": None, + "prompt": None, + "ability": None, + "reward_model": None, + "extra_info": None +} +SYSTEM_PROMPT = """You are a helpful programming assistant. \ +The user will ask you a question and you as the assistant solve it. \ +The assistant first thinks how to solve the task through reasoning and then provides the user with the final answer. \ +The reasoning process and answer are enclosed within ... and ... tags, respectively.""" +PY_IMPORTS = "import heapq\nfrom math import floor, gcd\nimport random\nimport sys\nfrom typing import *\nfrom functools import *\nimport collections\nfrom collections import *\nfrom itertools import *\nfrom heapq import *\nfrom bisect import *\nfrom string import *\nimport math\nimport datetime\ninf = float('inf')\n" +BLOCK_LIBS = [ + # "requests", + # "subprocess", + # "multiprocessing", + # "threading", + # "os", + # "sys", + "fake-useragent", + "keras", + "socket", + "torch", + "scipy", + "sklearn", + "cv2", + "scipy", + "imageio", + "sphinx-pyproject", + "xgboost", + "tweepy", + "flask", + "matplotlib", + "pillow", + "seaborn", + "smtpd", + "docker", + "psutil", + "paramiko", + "pickle", + "shutil", + "telnetlib", + "ftplib", + "pyautogui", + "pexpect", + "asyncssh", + "cryptography", + "pycrypto", + "scapy", + "nmap", + "pwntools", + "webbrowser", + "ctypes", + "tempfile", +] +BLOCK_USAGES = [ + "open(", + "exec(", + "eval(", + "system(", + "popen(", + "importlib", + "globals(", + "locals(", + "breakpoint(", + "__builtins__", + "compile(", + "fork(", + "chmod(", +] +MAX_PROMPT_LENGTH = 2048 + + +def minimize_stdio(inputs, outputs, max_n_tests=8): + stdin_list = [] + stdout_list = [] + for stdin, stdout in zip(inputs, outputs): + if isinstance(stdin, list): + stdin = "\n".join(stdin) + if isinstance(stdout, list): + stdout = "\n".join(stdout) + if sys.getsizeof(stdin) > 4 * 1024: + continue + stdout.replace("\r\n", "\n") + stdin_list.append(stdin) + stdout_list.append(stdout) + + zipped = sorted(zip(stdin_list, stdout_list), key=lambda x: sys.getsizeof(x[0])) + + if not zipped: + print("No tests found!") + return [], [] + + sorted_stdin, sorted_stdout = zip(*zipped) + return list(sorted_stdin[:max_n_tests]), list(sorted_stdout[:max_n_tests]) + + +# Plot the filter counts +def plot_hist( + data_dict, + file_path, + title, + xlabel, + ylabel, + figsize=(10, 6), + plot_type="bar", + color=None, + sort_by=None, + max_categories=None, + rotation=45, + dpi=300, + log_scale=False, + grid=True, + font_size=10, +): + """ + Plot the histogram of filter counts from a dictionary. + + Args: + data_dict (dict): Dictionary mapping categories to counts + file_path (str): Path to save the output figure + title (str): Title of the plot + xlabel (str): Label for x-axis + ylabel (str): Label for y-axis + figsize (tuple): Figure size in inches (width, height) + plot_type (str): Type of plot ('bar', 'barh', 'pie') + color (str or list): Color or color map for the plot + sort_by (str): How to sort categories ('key', 'value', 'value_desc', None) + max_categories (int): Maximum number of categories to display (None for all) + rotation (int): Rotation angle for x-tick labels + dpi (int): DPI for the output figure + log_scale (bool): Use log scale for y-axis + grid (bool): Show grid lines + font_size (int): Base font size for the plot + + Returns: + None: The figure is saved to file_path + """ + try: + plt.figure(figsize=figsize) + + # Handle empty dictionary + if not data_dict: + plt.text( + 0.5, + 0.5, + "No data available", + ha="center", + va="center", + fontsize=font_size + 4, + ) + plt.savefig(file_path, dpi=dpi, bbox_inches="tight") + plt.close() + return + + # Process data + items = list(data_dict.items()) + + # Sort data if requested + if sort_by == "key": + items.sort(key=lambda x: x[0]) + elif sort_by == "value": + items.sort(key=lambda x: x[1]) + elif sort_by == "value_desc": + items.sort(key=lambda x: x[1], reverse=True) + + # Limit number of categories if requested + if max_categories and len(items) > max_categories: + items = items[: max_categories - 1] + # Add "Other" category if needed + if max_categories < len(data_dict): + other_sum = sum(v for k, v in data_dict.items() if k not in dict(items)) + items.append(("Other", other_sum)) + + labels, counts = zip(*items) if items else ([], []) + + # Create the appropriate plot + if plot_type == "bar": + plt.bar(labels, counts, color=color) + if log_scale and any(counts): + plt.yscale("log") + plt.xticks(rotation=rotation, ha="right", fontsize=font_size) + plt.yticks(fontsize=font_size) + elif plot_type == "barh": + plt.barh(labels, counts, color=color) + if log_scale and any(counts): + plt.xscale("log") + plt.yticks(fontsize=font_size) + plt.xticks(fontsize=font_size) + elif plot_type == "pie": + plt.pie( + counts, labels=labels, autopct="%1.1f%%", startangle=90, colors=color + ) + plt.axis("equal") + + # Set labels and title + plt.xlabel(xlabel, fontsize=font_size + 2) + plt.ylabel(ylabel, fontsize=font_size + 2) + plt.title(title, fontsize=font_size + 4) + + if grid and plot_type != "pie": + plt.grid(alpha=0.3) + + plt.tight_layout() + + # Ensure directory exists + import os + + os.makedirs(os.path.dirname(os.path.abspath(file_path)), exist_ok=True) + + # Save figure + plt.savefig(file_path, dpi=dpi, bbox_inches="tight") + plt.close() + + except Exception as e: + import logging + + logging.error(f"Error plotting histogram: {str(e)}") + # Create a minimal error plot + plt.figure(figsize=(8, 6)) + plt.text( + 0.5, + 0.5, + f"Error creating plot: {str(e)}", + ha="center", + va="center", + color="red", + ) + plt.tight_layout() + plt.savefig(file_path, dpi=dpi) + plt.close() diff --git a/reasoners/data_preprocess/archived_for_reference/codegen.py b/reasoners/data_preprocess/archived_for_reference/codegen.py new file mode 100644 index 000000000..0e6fe4c50 --- /dev/null +++ b/reasoners/data_preprocess/archived_for_reference/codegen.py @@ -0,0 +1,977 @@ +""" +Preprocess LeetCode problems (newfacade/LeetCodeDataset) to parquet format. +Thanks to https://github.com/ganler/code-r1/blob/main/examples/data_preprocess/coder1.py +""" + +import os +import json +import sys +from concurrent.futures import ThreadPoolExecutor, as_completed +from transformers import AutoTokenizer +from datasets import load_dataset, concatenate_datasets +from rich.rule import Rule + +import rich +import matplotlib.pyplot as plt +import datasets + +from verl.utils.hdfs_io import copy, makedirs +from verl.utils.reward_score.coder1 import ( + code_exec, + remote_check_stdio, + _ERROR_MSG_PREFIX, + extract_code_from_string, + fuzzy_equal +) + +# from examples.data_preprocess.code.code_utils import * + +WORKDING_DIR = os.path.join(os.environ.get("HOME"), "Reasoning360") + +SYSTEM_PROMPT = """You are a helpful programming assistant. \ +The user will ask you a question and you as the assistant solve it. \ +The assistant first thinks how to solve the task through reasoning and then provides the user with the final answer. \ +The reasoning process and answer are enclosed within ... and ... tags, respectively.""" +EMPTY_RETURN = { + "data_source": None, + "prompt": None, + "ability": None, + "reward_model": None, + "extra_info": None +} +N_TESTSET_PER_DATASET = 0 + +def minimize_stdio(inputs, outputs, max_n_tests=8): + stdin_list = [] + stdout_list = [] + for stdin, stdout in zip(inputs, outputs): + if isinstance(stdin, list): + stdin = "\n".join(stdin) + if isinstance(stdout, list): + stdout = "\n".join(stdout) + if sys.getsizeof(stdin) > 4 * 1024: + continue + stdout.replace("\r\n", "\n") + stdin_list.append(stdin) + stdout_list.append(stdout) + + zipped = sorted(zip(stdin_list, stdout_list), key=lambda x: sys.getsizeof(x[0])) + + if not zipped: + print("No tests found!") + return [], [] + + sorted_stdin, sorted_stdout = zip(*zipped) + return list(sorted_stdin[:max_n_tests]), list(sorted_stdout[:max_n_tests]) + + +def kodcode(): # Thanks!!! to Zhangchen and Yueqin + # library requirements? + rich.print(Rule("Loading KodCode/KodCode-Light-RL-10K...")) + dataset = load_dataset("KodCode/KodCode-Light-RL-10K") + tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-32B-Instruct") + + def make_map_fn(split): + + def process_fn(example, idx): + reference_solution = example["solution"] + test_code = "from solution import *\n" + example["test"].strip() + # Filtering... + # + block libs are used in reference solution and test code + # + block usages are used in reference solution or test code + # + filter out too long prompts + # + filter out easy problems + # + filter out failed unittests + filter_info = {} + for lib in BLOCK_LIBS: + if ( + f"import {lib}" + in reference_solution # naive import detection; ast then detect would be better + or f"from {lib}" in reference_solution + ): + print("===========Blocked lib in solution===========") + print(f"reference_solution:") + print(reference_solution) + print(f"lib: {lib}") + print(f"question_id: {example['question_id']}") + return { + **EMPTY_RETURN, + "filter_info": { + "type": f"blocked_lib", + "detail": lib, + }, + } + for lib in BLOCK_LIBS: + if f"import {lib}" in test_code or f"from {lib}" in test_code: + print("===========Blocked lib in test===========") + print(f"test_code:") + print(test_code) + print(f"lib: {lib}") + print(f"question_id: {example['question_id']}") + return { + **EMPTY_RETURN, + "filter_info": { + "type": f"blocked_lib", + "detail": lib, + }, + } + for usage in BLOCK_USAGES: + if usage in reference_solution: + return { + **EMPTY_RETURN, + "filter_info": { + "type": f"blocked_usage", + "detail": usage, + }, + } + for usage in BLOCK_USAGES: + if usage in test_code: + return { + **EMPTY_RETURN, + "filter_info": { + "type": f"blocked_usage", + "detail": usage, + }, + } + if ( + len(tokenizer.encode(example["question"])) > MAX_PROMPT_LENGTH - 200 + ): # -200 for (approximately) the prompt template extra tokens + return { + **EMPTY_RETURN, + "filter_info": {"type": "prompt_too_long", "detail": None}, + } + if example["gpt_difficulty"] == "easy": + return { + **EMPTY_RETURN, + "filter_info": {"type": "easy_problem", "detail": None}, + } + succ, err = code_exec(code=reference_solution, pytest=test_code) + if not succ: + # The above code is using the `rich` library in Python to print a formatted message in the console. + # The message is in red color and includes the value of `example['conversation_id']`. + # rich.print( + # f"[bold red]Test code failed for {example['question_id']}" + # ) + print("===========Unittest failed===========") + print(f"reference_solution:") + print(reference_solution) + print(f"test_code:") + print(test_code) + print(f"err:") + print(err) + return { + **EMPTY_RETURN, + "filter_info": {"type": "failed_unittests", "detail": None}, + } + + prompt = f"Please solve the programming task below in Python. Code should wrapped in a markdown code block.\n\n{example['question'].strip()}" + if example["test_info"]: + prompt += f"\n\nNote that the output function should be {str(example['test_info']).strip()}." + + return { + "data_source": "codegen-kodcode", + "prompt": [ + {"role": "system", "content": SYSTEM_PROMPT}, + {"role": "user", "content": prompt}, + ], + "ability": "codegen", + "reward_model": { + "style": "rule", + "ground_truth": json.dumps({"pytest": test_code}), + }, + "extra_info": { + "split": split, + "index": idx, + "reference": reference_solution, + "prompt": prompt, + "dataset": "KodCode/KodCode-Light-RL-10K", + "question_subset": example["subset"], + "question_id": example["question_id"], + "gpt_difficulty": example["gpt_difficulty"], + }, + "filter_info": None, + } + + return process_fn + + dataset = dataset["train"].shuffle(seed=666) + + # Preprocess the dataset + print("Executing tests to ensure correctness...") + dataset = dataset.map( + function=make_map_fn("train"), + with_indices=True, + num_proc=64, + load_from_cache_file=False, + ) + # Analyze the filter reasons + filter_counts, filter_counts_fine_block_libs = {}, {} + for entry in dataset: + if entry["filter_info"] is None: + continue + filter_type = entry["filter_info"]["type"] + if filter_type is not None: + # Filter type distribution + if filter_type not in filter_counts: + filter_counts[filter_type] = 0 + filter_counts[filter_type] += 1 + # Filter detail distribution + filter_detail = entry["filter_info"].get("detail", None) + if filter_detail is not None: + if filter_detail not in filter_counts_fine_block_libs: + filter_counts_fine_block_libs[filter_detail] = 0 + filter_counts_fine_block_libs[filter_detail] += 1 + # entry["filter_info"] = None + + print(f"Filtered samples from KodCode: {filter_counts}") + + plot_hist( + filter_counts, + file_path=os.path.join(WORKDING_DIR, "artifacts", "filter_counts.png"), + title="Filter Sample Distribution from KodCode", + xlabel="Filter Reason", + ylabel="Count", + ) + plot_hist( + filter_counts_fine_block_libs, + file_path=os.path.join( + WORKDING_DIR, "artifacts", "filter_counts_fine_block_libs.png" + ), + title="Blocked Library Distribution from KodCode", + xlabel="Blocked Library", + ylabel="Count", + ) + + print(f"Before filtering, KodCode dataset size: {len(dataset)}") + + dataset = dataset.filter(lambda x: x["data_source"] is not None) + print(f"Remaining samples from KodCode: {len(dataset)}") + + # pick random 50k examples for RL, otherwise it's too large + # dataset = dataset.select(range(50000 + N_TESTSET_PER_DATASET)) + + # Split into train and test + # splits = dataset["train"].train_test_split( + # test_size=N_TESTSET_PER_DATASET, seed=666 + # ) + splits = dataset.train_test_split(test_size=N_TESTSET_PER_DATASET, seed=666) + train_dataset = splits["train"].shuffle(seed=666) + test_dataset = splits["test"] + return train_dataset, test_dataset + + +# this dataset is super noisy and needs code execution to verify the tasks +def taco(): + rich.print(Rule("Loading likaixin/TACO-verified...")) + dataset = load_dataset("likaixin/TACO-verified")["train"] + + + # add a row to each data item that represents a unique id + def make_map_fn(split): + + def process_fn(example, idx): + oracle = json.loads(example["input_output"]) + source = example["source"] + + # skip poorly formatted examples + if source in ["geeksforgeeks", "leetcode"]: + return EMPTY_RETURN + + # too short description + if len("".join([c for c in example["question"] if c.isalnum()])) < 100: + return EMPTY_RETURN + + # no image + if "image" in example["question"].lower() or "\n![" in example["question"]: + return EMPTY_RETURN + + prompt_pieces = [ + "Solve the programming task below in a Python markdown code block.", + example["question"].strip(), + ] + if example["starter_code"].strip(): + prompt_pieces.append( + "Also feel free to reuse/extend the following starter code:" + ) + prompt_pieces.append( + f"```python\n{example['starter_code'].strip()}\n```" + ) + + ## + ## Customization + ## + if "fn_name" in oracle: # the dataset is too noisy + fn_name = oracle["fn_name"] + if source == "leetcode": + fn_name = "Solution()." + fn_name + + test_code = f"""\ +_inputs = {oracle["inputs"]} +_outputs = {oracle["outputs"]} +import math +def _deep_eq(a, b, tol=1e-5): + if isinstance(a, float) or isinstance(b, float): + return math.isclose(a, b, rel_tol=tol, abs_tol=tol) + if isinstance(a, (list, tuple)): + if len(a) != len(b): return False + return all(_deep_eq(x, y, tol) for x, y in zip(a, b)) + return a == b + +for i, o in zip(_inputs, _outputs): +""" + + if source in ["leetcode", "hackerrank"]: + test_code += f" assert _deep_eq({fn_name}(*i), o)" + elif source == "codewars": + test_code += f" assert _deep_eq({fn_name}(*i), o[0])" + else: + raise ValueError(f"Unknown source: {source}") + + _check_test = example["solutions"][-1] + "\n" + test_code + if source in ["leetcode"]: + _check_test = PY_IMPORTS + _check_test + + succ, err = code_exec(_check_test) + if not succ: + rich.print(f"[bold red]Test code failed for {source}") + print(_check_test) + print(err) + return EMPTY_RETURN + oracle = json.dumps({"functional": test_code}) + assert example["starter_code"].strip() != "" + elif "inputs" in oracle and "outputs" in oracle: + stdin_list, stdout_list = minimize_stdio( + oracle["inputs"], oracle["outputs"] + ) + if len(stdin_list) == 0: + return EMPTY_RETURN + + with ThreadPoolExecutor( + max_workers=min(len(stdin_list), 8) + ) as executor: + futures = [] + for stdin, stdout in zip(stdin_list, stdout_list): + futures.append( + executor.submit( + remote_check_stdio, + example["solutions"][-1], + stdin, + stdout, + ) + ) + for future in as_completed(futures): + exec_succ, output, stdin, stdout = future.result() + pass_test = exec_succ and fuzzy_equal(output.strip(), stdout.strip()) + if not pass_test: + rich.print(f"[bold red]Test code failed for {source}") + print(example["solutions"][-1]) + print(f"{exec_succ = }") + print(f"{stdin = }", f"{stdout = }") + if output.startswith(_ERROR_MSG_PREFIX): + print("output = \n", output) + else: + print(f"{output = }") + return EMPTY_RETURN + + oracle = json.dumps({"inputs": stdin_list, "outputs": stdout_list}) + else: + raise ValueError(f"Unknown ground truth format: {oracle}") + + prompt = "\n".join(prompt_pieces) + return { + "data_source": "codegen-taco", + "prompt": [ + {"role": "system", "content": SYSTEM_PROMPT}, + {"role": "user", "content": prompt}, + ], + "ability": "codegen", + "reward_model": { + "style": "rule", + "ground_truth": oracle, + }, + "extra_info": { + "split": split, + "index": idx, + "prompt": prompt, + "reference": ( + example["solutions"][0] if example["solutions"] else "" + ), + "dataset": "likaixin/TACO-verified", + }, + } + + return process_fn + + dataset = dataset.map( + function=make_map_fn("train"), + with_indices=True, + num_proc=64, + remove_columns=dataset.column_names, + load_from_cache_file=False, + ).filter(lambda x: x['reward_model'] is not None) + + if N_TESTSET_PER_DATASET == 0: + # Return no test samples if N_TESTSET_PER_DATASET is 0 + test_dataset = datasets.Dataset.from_dict({ + "data_source": [], + "prompt": [], + "ability": [], + "reward_model": [], + "extra_info": [] + }) + train_dataset = dataset + else: + splits = dataset.train_test_split( + test_size=min(N_TESTSET_PER_DATASET, len(dataset) * 0.1), seed=666 + ) + train_dataset = splits["train"] + test_dataset = splits["test"] + + print(f"Taco train set: {train_dataset}") + print(f"Taco test set: {test_dataset}") + + return train_dataset, test_dataset + + +def leetcode2k(): + rich.print(Rule("Loading LeetCodeDataset...")) + commit_hash = "34803eb64eab1979069ba1f80e7ea474282e28f3" + train_dataset = load_dataset("newfacade/LeetCodeDataset", + revision=commit_hash, + cache_dir=cache_dir)["train"] + test_dataset = load_dataset("newfacade/LeetCodeDataset", + revision=commit_hash, + cache_dir=cache_dir)["test"] + print("Train set:", train_dataset) + print("Test set:", test_dataset) + + + # add a row to each data item that represents a unique id + def make_map_fn(split): + + def process_fn(example, idx): + prefix = example["prompt"] + + prompt = example["query"] + # remove the "### Answer: (use the provided format with backticks)" part + prompt = prompt.replace("### Answer: (use the provided format with backticks)", "").strip() + # adjust the "### Format: " part to be more readable + prompt = prompt.replace("### Format: ", "### Format:\n") + + # Build test code (as before) + test_code = f"{example['test']}\n\ncheck({example['entry_point'].strip()})" + # Extract the candidate solution (the original completion) + solution = example["completion"] + + # Combine all code pieces into a single file to execute. + full_code = f"{prefix}\n{solution}\n{test_code}" + + # Validate that the candidate solution passes the tests + # 20s timeout as some leetcode tests are slow + succ, err = code_exec(full_code, timeout=20) + + if not succ: + print("===========Test code failed for LeetCodeDataset===========") + print("Question:", example["meta"]["question_title"]) + print("Error:", err) + # Skip the example if the test code fails + return EMPTY_RETURN + + return { + "data_source": "codegen-leetcode2k", + "prompt": [ + {"role": "system", "content": SYSTEM_PROMPT}, + {"role": "user", "content": prompt}, + ], + "ability": "codegen", + "reward_model": { + "style": "rule", + "ground_truth": json.dumps({ + "functional": test_code + }), + }, + "extra_info": { + "split": split, + "index": idx, + "reference": solution, + "prompt": prompt, + "prefix": prefix, + "dataset": "LeetCodeDataset", + }, + } + + return process_fn + + # filter out empty examples ("reward_model" is None) + train_dataset = train_dataset.map( + function=make_map_fn("train"), + with_indices=True, + load_from_cache_file=False, + ).filter(lambda x: x["reward_model"] is not None) + test_dataset = test_dataset.map( + function=make_map_fn("test"), + with_indices=True, + load_from_cache_file=False, + ).filter(lambda x: x["reward_model"] is not None) + print(f"Leetcode2k train set: {train_dataset}") + print(f"Leetcode2k test set: {test_dataset}") + return train_dataset, test_dataset + + +def humaneval(): + rich.print(Rule("Loading OpenAI HumanEval...")) + dataset = load_dataset("openai_humaneval")["test"] + print("HumanEval dataset:", dataset) + + def process_fn(example, idx): + # HumanEval's prompt already contains the function signature and docstring + prompt = ( + "Write a complete, self-contained Python solution to the following problem. " + "Your solution must include all necessary imports and the full function definition including " + "the signature exactly as specified. Do not modify the function signature or docstring.\n\n" + f"```python\n{example['prompt'].strip()}\n```" + ) + + # Extract test code + test_code = example['test'] + entry_point = example['entry_point'] + + # Validate that the canonical solution passes the tests + solution = example['canonical_solution'] + + # Combine the prompt code + solution + test code to verify it works + full_code = f"{example['prompt']}\n{solution}\n{test_code}\n\ncheck({entry_point})" + + succ, err = code_exec(full_code) + if not succ: + print(f"Error in canonical solution for task {example['task_id']}: {err}") + return EMPTY_RETURN + + return { + "data_source": "codegen-humaneval", + "prompt": [ + {"role": "system", "content": SYSTEM_PROMPT}, + {"role": "user", "content": prompt}, + ], + "ability": "codegen", + "reward_model": { + "style": "rule", + "ground_truth": json.dumps( + {"functional": f"{test_code}\n\ncheck({entry_point})"} + ), + }, + "extra_info": { + "split": "test", + "index": idx, + "reference": solution, + "prompt": prompt, + "dataset": "openai_humaneval", + "task_id": str(example["task_id"]), + }, + } + + test_dataset = dataset.map( + function=process_fn, + with_indices=True, + load_from_cache_file=False, + ).filter(lambda x: x["reward_model"] is not None) + + # Return empty train dataset and test dataset + empty_train = datasets.Dataset.from_dict({ + "data_source": [], + "prompt": [], + "ability": [], + "reward_model": [], + "extra_info": [] + }) if len(test_dataset) > 0 else datasets.Dataset.from_dict({}) + + print(f"HumanEval test set: {test_dataset}") + return empty_train, test_dataset + +def mbpp(): + rich.print(Rule("Loading MBPP dataset...")) + dataset = load_dataset("google-research-datasets/mbpp") + + def make_map_fn(split): + def process_fn(example, idx): + # rewrite the task_id as it is int + example["task_id"] = "MBPP/" + str(example["task_id"]) + + # Create prompt + prompt = ( + f"{example['text']}\n\n" + f"Your solution should be a complete, self-contained function in a markdown code block. " + f"Make sure your solution passes the following test cases:\n" + ) + + # Construct test code + test_code = "" + if example.get('test_setup_code'): + test_code += example['test_setup_code'] + "\n\n" + + # Add all test assertions + for assertion in example['test_list'] + example.get('challenge_test_list', []): + test_code += assertion + "\n" + + # Add test cases to prompt + prompt += f"```python\n{test_code}```" + prompt += "\n\nPlease do not include the test cases in your solution." + + # Validate that the canonical solution passes the tests + solution = example['code'] + full_code = f"{solution}\n\n{test_code}" + + succ, err = code_exec(full_code) + if not succ: + print(f"Error in canonical solution for task {example['task_id']}: {err}") + return EMPTY_RETURN + + return { + "data_source": "codegen-mbpp", + "prompt": [ + {"role": "system", "content": SYSTEM_PROMPT}, + {"role": "user", "content": prompt}, + ], + "ability": "codegen", + "reward_model": { + "style": "rule", + "ground_truth": json.dumps( + {"functional": test_code} + ), + }, + "extra_info": { + "split": split, + "index": idx, + "reference": solution, + "prompt": prompt, + "dataset": "mbpp", + "task_id": str(example["task_id"]), + }, + } + + return process_fn + + # Process train and test splits + train_dataset = dataset["train"].map( + function=make_map_fn("train"), + with_indices=True, + load_from_cache_file=False, + ).filter(lambda x: x['reward_model'] is not None) + + test_dataset = dataset["test"].map( + function=make_map_fn("test"), + with_indices=True, + load_from_cache_file=False, + ).filter(lambda x: x['reward_model'] is not None) + + print(f"MBPP train set: {train_dataset}") + print(f"MBPP test set: {test_dataset}") + return train_dataset, test_dataset + +def primeintellect(): + rich.print(Rule("Loading PrimeIntellect dataset...")) + dataset = load_dataset("agentica-org/DeepCoder-Preview-Dataset", "primeintellect", split="train") + + def make_map_fn(split): + def process_fn(example, idx): + # Get the problem description + prompt = example["problem"] + + # Get the solution code + solution = example["solutions"][0] if example["solutions"] else "" + solution = extract_code_from_string(solution) + + # Process tests + tests = json.loads(example["tests"]) + if not tests: + return EMPTY_RETURN + + # Handle different test types + if tests[0]["type"] == "function_call": + # Function call tests + fn_name = tests[0]["fn_name"] + test_code = f"""\ +def check_{fn_name}(): +""" + for test in tests: + input_args = ", ".join([ + repr(arg) if isinstance(arg, (str, list, tuple, dict)) else str(arg) + for arg in test["input"] + ]) + expected_output = repr(test["output"][0]) if isinstance(test["output"][0], (str, list, tuple, dict)) else test["output"][0] + test_code += f""" assert {fn_name}({input_args}) == {expected_output} +""" + test_code += f""" +check_{fn_name}() +""" + + # Validate the solution + full_code = f"{solution}\n{test_code}" + succ, err = code_exec(full_code) + if not succ: + rich.print(f"[bold red]Test code failed for PrimeIntellect example {idx}") + print(f"===========Full code===========") + print(full_code) + print(f"===========Error===========") + print(err) + return EMPTY_RETURN + + oracle = json.dumps({"functional": test_code}) + + elif tests[0]["type"] == "stdin_stdout": + # STDIN/STDOUT tests + stdin_list = [] + stdout_list = [] + for test in tests: + stdin_list.append(test["input"]) + stdout_list.append(test["output"]) + + # Validate the solution + with ThreadPoolExecutor(max_workers=min(len(stdin_list), 8)) as executor: + futures = [] + for stdin, stdout in zip(stdin_list, stdout_list): + futures.append( + executor.submit( + remote_check_stdio, + solution, + stdin, + stdout, + ) + ) + for future in as_completed(futures): + exec_succ, output, stdin, stdout = future.result() + pass_test = exec_succ and fuzzy_equal(output.strip(), stdout.strip()) + if not pass_test: + rich.print(f"[bold red]Test code failed for PrimeIntellect example {idx}") + print(f"===========Solution===========") + print(solution) + print(f"===========Input===========") + print(stdin) + print(f"===========Expected output===========") + print(stdout) + print(f"===========Actual output===========") + print(output) + return EMPTY_RETURN + + oracle = json.dumps({"inputs": stdin_list, "outputs": stdout_list}) + else: + raise ValueError(f"Unknown test type: {tests[0]['type']}") + + return { + "data_source": "codegen-primeintellect", + "prompt": [ + {"role": "system", "content": SYSTEM_PROMPT}, + {"role": "user", "content": prompt}, + ], + "ability": "codegen", + "reward_model": { + "style": "rule", + "ground_truth": oracle, + }, + "extra_info": { + "split": split, + "index": idx, + "prompt": prompt, + "reference": solution, + "dataset": "PrimeIntellect", + "function_name": fn_name if tests[0]["type"] == "function_call" else None, + }, + } + + return process_fn + + # Process train and test splits + train_dataset = dataset.map( + function=make_map_fn("train"), + with_indices=True, + num_proc=64, + remove_columns=dataset.column_names, + load_from_cache_file=False, + ).filter(lambda x: x['reward_model'] is not None) + + # Split into train and test sets + if N_TESTSET_PER_DATASET == 0: + # Return no test samples if N_TESTSET_PER_DATASET is 0 + test_dataset = datasets.Dataset.from_dict({ + "data_source": [], + "prompt": [], + "ability": [], + "reward_model": [], + "extra_info": [] + }) + else: + splits = train_dataset.train_test_split( + test_size=min(N_TESTSET_PER_DATASET, len(train_dataset) * 0.1), seed=666 + ) + train_dataset = splits["train"].shuffle(seed=666) + test_dataset = splits["test"] + + print(f"PrimeIntellect train set: {train_dataset}") + print(f"PrimeIntellect test set: {test_dataset}") + return train_dataset, test_dataset + +def livecodebench(): + rich.print(Rule("Loading LiveCodeBench dataset...")) + # Load both train and test splits directly + train_dataset = load_dataset("agentica-org/DeepCoder-Preview-Dataset", "lcbv5", split="train") + test_dataset = load_dataset("agentica-org/DeepCoder-Preview-Dataset", "lcbv5", split="test") + + def make_map_fn(split): + def process_fn(example, idx): + # Get the problem description + problem_desc = example["problem"] + starter_code = example["starter_code"] + + # Create the prompt with the starter code + prompt = f"{problem_desc}\n\nComplete the implementation using the provided starter code:\n```python\n{starter_code}\n```\n\nYour solution should implement the method(s) in the Solution class." + + # Process tests + tests = json.loads(example["tests"]) + if not tests: + return EMPTY_RETURN + + # Process metadata to get function name + metadata = example["metadata"] + function_name = metadata.get("func_name") + + # Handle different test types + if tests[0]["testtype"] == "functional": + if not function_name: + return EMPTY_RETURN + + # Function call tests + test_code = f"""\ +def check_{function_name}(): +""" + for test in tests: + # Parse input string by splitting on '\n' + input_parts = test["input"].split('\n') + # Create proper comma-separated arguments for the function call + input_args = ', '.join(input_parts) + # Get the output value + output_val = test["output"] + + test_code += f""" assert Solution().{function_name}({input_args}) == {output_val} +""" + + test_code += f""" +check_{function_name}() +""" + + # For debugging - print a few examples + if idx < 20: + print(f"Generated test code for example {idx}:") + print(test_code) + + oracle = json.dumps({"functional": test_code}) + + elif tests[0]["testtype"] == "stdin": + # STDIN/STDOUT tests + stdin_list = [] + stdout_list = [] + for test in tests: + stdin_list.append(test["input"]) + stdout_list.append(test["output"]) + + # For debugging - print a few examples + if idx < 20: + print(f"Generated test code for example {idx}:") + for i in range(len(stdin_list)): + print(f"Test {i+1}:") + print(f"Input: {stdin_list[i]}") + print(f"Output: {stdout_list[i]}") + + oracle = json.dumps({"inputs": stdin_list, "outputs": stdout_list}) + else: + raise ValueError(f"Unknown test type: {tests[0]['testtype']}") + + return { + "data_source": "codegen-livecodebench", + "prompt": [ + {"role": "system", "content": SYSTEM_PROMPT}, + {"role": "user", "content": prompt}, + ], + "ability": "codegen", + "reward_model": { + "style": "rule", + "ground_truth": oracle, + }, + "extra_info": { + "split": split, + "index": idx, + "prompt": prompt, + "reference": "", # No solution data + "dataset": "LiveCodeBench", + "function_name": function_name, + }, + } + + return process_fn + + # Process train and test datasets + train_dataset = train_dataset.map( + function=make_map_fn("train"), + with_indices=True, + num_proc=64, + load_from_cache_file=False, + ).filter(lambda x: x['reward_model'] is not None) + + test_dataset = test_dataset.map( + function=make_map_fn("test"), + with_indices=True, + num_proc=64, + load_from_cache_file=False, + ).filter(lambda x: x['reward_model'] is not None) + + print(f"LiveCodeBench train set: {train_dataset}") + print(f"LiveCodeBench test set: {test_dataset}") + return train_dataset, test_dataset + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument("--root_dir", default=os.path.join(WORKDING_DIR, "data")) + parser.add_argument( + "--dataset_names", default="kodcode", help="comma separated dataset names" + ) + + args = parser.parse_args() + + root_dir = args.root_dir + dataset_names = args.dataset_names.split(",") + + train_datasets = [] + test_datasets = [] + + dataset_map = { + "kodcode": kodcode, + "taco": taco, + "leetcode2k": leetcode2k, + "humaneval": humaneval, + "mbpp": mbpp, + "primeintellect": primeintellect, + "livecodebench": livecodebench, + } + dataset_makes = [dataset_map[name] for name in dataset_names] + names = "-".join([make.__name__ for make in dataset_makes]) + + for make in dataset_makes: + train, test = make() + train_datasets.append(train) + test_datasets.append(test) + + train_dataset = concatenate_datasets(train_datasets).shuffle(seed=666) + test_dataset = concatenate_datasets(test_datasets) + + rich.print(Rule("Saving the final dataset")) + print("Train set:", train_dataset) + print("Test set:", test_dataset) + + local_dir = os.path.join( + root_dir, f"codegen-{round(len(train_dataset) / 1000)}k-{names}" + ) + rich.print(f"[bold green]Saving to {local_dir}...") + train_dataset.to_parquet(os.path.join(local_dir, "train.parquet")) + test_dataset.to_parquet(os.path.join(local_dir, "test.parquet")) \ No newline at end of file diff --git a/reasoners/data_preprocess/archived_for_reference/deepscaler_preview_r1_zero_style_filter.py b/reasoners/data_preprocess/archived_for_reference/deepscaler_preview_r1_zero_style_filter.py new file mode 100644 index 000000000..942540697 --- /dev/null +++ b/reasoners/data_preprocess/archived_for_reference/deepscaler_preview_r1_zero_style_filter.py @@ -0,0 +1,192 @@ +# 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. +""" +Preprocess the https://huggingface.co/datasets/agentica-org/DeepScaleR-Preview-Dataset to parquet format +Code referring to https://github.com/agentica-project/deepscaler/blob/main/scripts/data/deepscaler_dataset.py +""" + +import os +import datasets +from typing import Dict, List, Optional, Any, Union +import enum +import argparse +import pandas as pd +import json + +from verl.utils.hdfs_io import copy, makedirs +from verl.utils.reward_score.math import remove_boxed, last_boxed_only_string + + +class TrainDataset(enum.Enum): + """Enum for training datasets. + + Contains identifiers for various math problem datasets used during training. + """ + + AIME = "AIME" # American Invitational Mathematics Examination + AMC = "AMC" # American Mathematics Competition + OMNI_MATH = "OMNI_MATH" # Omni Math + NUMINA_OLYMPIAD = "OLYMPIAD" # Unique Olympiad problems from NUMINA + MATH = "MATH" # Dan Hendrycks Math Problems + STILL = "STILL" # STILL dataset + DEEPSCALER = "DEEPSCALER" # DeepScaler (AIME, AMC, OMNI_MATH, MATH, STILL) + + +class TestDataset(enum.Enum): + """Enum for testing/evaluation datasets. + + Contains identifiers for datasets used to evaluate model performance. + """ + + AIME = "AIME" # American Invitational Mathematics Examination + AMC = "AMC" # American Mathematics Competition + MATH = "MATH" # Math 500 problems + MINERVA = "MINERVA" # Minerva dataset + OLYMPIAD_BENCH = "OLYMPIAD_BENCH" # Olympiad benchmark problems + + +"""Type alias for either training or testing dataset types.""" +Dataset = Union[TrainDataset, TestDataset] + + +def load_dataset(dataset: Dataset) -> List[Dict[str, Any]]: + """Load a dataset from a JSON file. + + Loads and parses a JSON dataset file based on the provided dataset enum. + The file path is constructed based on whether it's a training or testing dataset. + + Args: + dataset: A Dataset enum value specifying which dataset to load. + + Returns: + List[Dict[str, Any]]: A list of dictionaries containing the dataset records. + Each dictionary represents one example in the dataset. + + Raises: + ValueError: If the dataset file cannot be found, contains invalid JSON, + or encounters other file access errors. + + Example: + >>> load_dataset(TrainDataset.AIME) + [{'problem': 'Find x...', 'solution': '42', ...}, ...] + """ + dataset_name = dataset.value.lower() + data_dir = "train" if isinstance(dataset, TrainDataset) else "test" + + current_dir = os.path.dirname(os.path.realpath(__file__)) + file_path = os.path.join(data_dir, f"{dataset_name}.json") + file_path = os.path.join(current_dir, file_path) + + if not os.path.exists(file_path): + raise ValueError(f"Dataset file not found: {file_path}") + + try: + with open(file_path, "r", encoding="utf-8") as file: + data = json.load(file) + return data + except json.JSONDecodeError: + raise ValueError(f"Invalid JSON format in {file_path}") + except Exception as exc: # pylint: disable=broad-except + raise ValueError(f"Error loading dataset: {exc}") from exc + + +def extract_solution(solution_str): + return remove_boxed(last_boxed_only_string(solution_str)) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--local_dir", default="data/deepscaler_preview_zero_style_mar21_filter") + parser.add_argument("--hdfs_dir", default=None) + + args = parser.parse_args() + + train_data_source = "SDSB/deepscale_partial_mar21_filtered_basic" + + print(f"Loading the {train_data_source} dataset from huggingface...", flush=True) + test_data_sources = [ + "nanoverl/minerva", + "SDSB/aime_repeated_8x", + "SDSB/amc_repeated_4x", + "nanoverl/olympiad_bench", + "nanoverl/math", + ] + print(f"Loading the {test_data_sources} dataset from huggingface...", flush=True) + train_dataset = datasets.load_dataset( + train_data_source, trust_remote_code=True, split="train" + ) + test_datasets = [ + datasets.load_dataset(test_data_source, trust_remote_code=True, split="test") + for test_data_source in test_data_sources + ] + + instruction_following = ( + "Let's think step by step and output the final answer within \\boxed{}." + ) + + + prompt = """A conversation between User and Assistant. The user asks a question, and the Assistant solves it. The assistant first thinks about the reasoning process in the mind and then provides the user with the response. The reasoning process is enclosed within i.e., reasoning process here respond to the user's question here. + +User: {{prompt}} Please put your answer in \\boxed{} tags. +Assistant: +""" + + # add a row to each data item that represents a unique id + def make_map_fn(split, data_source): + + def process_fn(example, idx): + question = example.pop("problem") + + # question = question + " " + instruction_following + + answer = example.pop("answer") + data = { + "data_source": data_source, + "prompt": [], + "raw_prompt": prompt.replace("{{prompt}}", question), + "ability": "math", + "apply_chat_template": False, + "reward_model": {"style": "rule", "ground_truth": answer}, + "extra_info": {"split": split, + "index": idx, + "question": question}, + } + if idx == 0: + print("=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) + print(data) + return data + + return process_fn + + local_dir = args.local_dir + hdfs_dir = args.hdfs_dir + + train_data = train_dataset.map( + function=make_map_fn("train", train_data_source), with_indices=True + ) + train_data.to_parquet(os.path.join(local_dir, "train.parquet")) + print(f"train data size:", len(train_data)) + print(train_data[0]) + for test_data_source, test_data in zip(test_data_sources, test_datasets): + process_fn = make_map_fn("test", test_data_source) + test_data = test_data.map(process_fn, with_indices=True) + dataset_name = os.path.basename(test_data_source.lower()) + test_df = pd.DataFrame(test_data) + test_df.to_parquet(os.path.join(local_dir, f"{dataset_name}.parquet")) + print(f"test data size: ({dataset_name})", len(test_df)) + + if hdfs_dir is not None: + makedirs(hdfs_dir) + + copy(src=local_dir, dst=hdfs_dir) diff --git a/reasoners/data_preprocess/archived_for_reference/deepscaler_preview_r1_zero_style_sampled.py b/reasoners/data_preprocess/archived_for_reference/deepscaler_preview_r1_zero_style_sampled.py new file mode 100644 index 000000000..f278701f1 --- /dev/null +++ b/reasoners/data_preprocess/archived_for_reference/deepscaler_preview_r1_zero_style_sampled.py @@ -0,0 +1,201 @@ +# 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. +""" +Preprocess the https://huggingface.co/datasets/agentica-org/DeepScaleR-Preview-Dataset to parquet format +Code referring to https://github.com/agentica-project/deepscaler/blob/main/scripts/data/deepscaler_dataset.py +""" + +import os +import datasets +from typing import Dict, List, Optional, Any, Union +import enum +import argparse +import pandas as pd +import json + +from verl.utils.hdfs_io import copy, makedirs +from verl.utils.reward_score.math import remove_boxed, last_boxed_only_string + + +class TrainDataset(enum.Enum): + """Enum for training datasets. + + Contains identifiers for various math problem datasets used during training. + """ + + AIME = "AIME" # American Invitational Mathematics Examination + AMC = "AMC" # American Mathematics Competition + OMNI_MATH = "OMNI_MATH" # Omni Math + NUMINA_OLYMPIAD = "OLYMPIAD" # Unique Olympiad problems from NUMINA + MATH = "MATH" # Dan Hendrycks Math Problems + STILL = "STILL" # STILL dataset + DEEPSCALER = "DEEPSCALER" # DeepScaler (AIME, AMC, OMNI_MATH, MATH, STILL) + + +class TestDataset(enum.Enum): + """Enum for testing/evaluation datasets. + + Contains identifiers for datasets used to evaluate model performance. + """ + + AIME = "AIME" # American Invitational Mathematics Examination + AMC = "AMC" # American Mathematics Competition + MATH = "MATH" # Math 500 problems + MINERVA = "MINERVA" # Minerva dataset + OLYMPIAD_BENCH = "OLYMPIAD_BENCH" # Olympiad benchmark problems + + +"""Type alias for either training or testing dataset types.""" +Dataset = Union[TrainDataset, TestDataset] + + +def load_dataset(dataset: Dataset) -> List[Dict[str, Any]]: + """Load a dataset from a JSON file. + + Loads and parses a JSON dataset file based on the provided dataset enum. + The file path is constructed based on whether it's a training or testing dataset. + + Args: + dataset: A Dataset enum value specifying which dataset to load. + + Returns: + List[Dict[str, Any]]: A list of dictionaries containing the dataset records. + Each dictionary represents one example in the dataset. + + Raises: + ValueError: If the dataset file cannot be found, contains invalid JSON, + or encounters other file access errors. + + Example: + >>> load_dataset(TrainDataset.AIME) + [{'problem': 'Find x...', 'solution': '42', ...}, ...] + """ + dataset_name = dataset.value.lower() + data_dir = "train" if isinstance(dataset, TrainDataset) else "test" + + current_dir = os.path.dirname(os.path.realpath(__file__)) + file_path = os.path.join(data_dir, f"{dataset_name}.json") + file_path = os.path.join(current_dir, file_path) + + if not os.path.exists(file_path): + raise ValueError(f"Dataset file not found: {file_path}") + + try: + with open(file_path, "r", encoding="utf-8") as file: + data = json.load(file) + return data + except json.JSONDecodeError: + raise ValueError(f"Invalid JSON format in {file_path}") + except Exception as exc: # pylint: disable=broad-except + raise ValueError(f"Error loading dataset: {exc}") from exc + + +def extract_solution(solution_str): + return remove_boxed(last_boxed_only_string(solution_str)) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--local_dir", default="data/deepscaler_default") + parser.add_argument("--hdfs_dir", default=None) + parser.add_argument( + "--train-sample-size", + type=int, + default=None, + help="Number of samples to use from the training dataset. If None, use all samples.", + ) + + args = parser.parse_args() + + train_data_source = "agentica-org/DeepScaleR-Preview-Dataset" + print(f"Loading the {train_data_source} dataset from huggingface...", flush=True) + test_data_sources = [ + "nanoverl/minerva", + "nanoverl/aime", + "nanoverl/amc", + "nanoverl/olympiad_bench", + "nanoverl/math", + ] + print(f"Loading the {test_data_sources} dataset from huggingface...", flush=True) + train_dataset = datasets.load_dataset( + train_data_source, trust_remote_code=True, split="train" + ) + test_datasets = [ + datasets.load_dataset(test_data_source, trust_remote_code=True, split="test") + for test_data_source in test_data_sources + ] + + instruction_following = ( + "Let's think step by step and output the final answer within \\boxed{}." + ) + + prompt = """A conversation between User and Assistant. The user asks a question, and the Assistant solves it. The assistant first thinks about the reasoning process in the mind and then provides the user with the response. The reasoning process is enclosed within i.e., reasoning process here respond to the user's question here. + +User: {{prompt}} Please put your answer in \\boxed{} tags. +Assistant: +""" + + # add a row to each data item that represents a unique id + def make_map_fn(split, data_source): + + def process_fn(example, idx): + question = example.pop("problem") + + # question = question + " " + instruction_following + + answer = example.pop("answer") + data = { + "data_source": data_source, + "prompt": [], + "raw_prompt": prompt.replace("{{prompt}}", question), + "ability": "math", + "apply_chat_template": False, + "reward_model": {"style": "rule", "ground_truth": answer}, + "extra_info": {"split": split, "index": idx}, + } + if idx == 0: + print("=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) + print(data) + return data + + return process_fn + + local_dir = args.local_dir + hdfs_dir = args.hdfs_dir + + train_data = train_dataset.map( + function=make_map_fn("train", train_data_source), with_indices=True + ) + + # Sample the training data if sample size is specified + if args.train_sample_size is not None: + train_data = train_data.select( + range(min(args.train_sample_size, len(train_data))) + ) + + train_data.to_parquet(os.path.join(local_dir, "train.parquet")) + print(f"train data size:", len(train_data)) + print(train_data[0]) + for test_data_source, test_data in zip(test_data_sources, test_datasets): + process_fn = make_map_fn("test", test_data_source) + test_data = test_data.map(process_fn, with_indices=True) + dataset_name = os.path.basename(test_data_source.lower()) + test_df = pd.DataFrame(test_data) + test_df.to_parquet(os.path.join(local_dir, f"{dataset_name}.parquet")) + print(f"test data size: ({dataset_name})", len(test_df)) + + if hdfs_dir is not None: + makedirs(hdfs_dir) + + copy(src=local_dir, dst=hdfs_dir) diff --git a/reasoners/data_preprocess/archived_for_reference/geo3k.py b/reasoners/data_preprocess/archived_for_reference/geo3k.py new file mode 100644 index 000000000..3a0c77673 --- /dev/null +++ b/reasoners/data_preprocess/archived_for_reference/geo3k.py @@ -0,0 +1,86 @@ +# 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. +""" +Preprocess the Geometry3k dataset to parquet format +""" + +import os +import datasets + +from verl.utils.hdfs_io import copy, makedirs +import argparse + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--local_dir', default='~/data/geo3k') + parser.add_argument('--hdfs_dir', default=None) + + args = parser.parse_args() + + data_source = 'hiyouga/geometry3k' + + dataset = datasets.load_dataset(data_source) + + train_dataset = dataset['train'] + test_dataset = dataset['test'] + + instruction_following = ( + r'You FIRST think about the reasoning process as an internal monologue and then provide the final answer. ' + r'The reasoning process MUST BE enclosed within tags. The final answer MUST BE put in \boxed{}.' + ) + + # add a row to each data item that represents a unique id + def make_map_fn(split): + + def process_fn(example, idx): + problem = example.pop('problem') + prompt = problem + ' ' + instruction_following + answer = example.pop('answer') + images = example.pop('images') + + data = { + "data_source": data_source, + "prompt": [{ + "role": "user", + "content": prompt, + }], + "images": images, + "ability": "math", + "reward_model": { + "style": "rule", + "ground_truth": answer + }, + "extra_info": { + 'split': split, + 'index': idx, + 'answer': answer, + "question": problem, + } + } + return data + + return process_fn + + train_dataset = train_dataset.map(function=make_map_fn('train'), with_indices=True, num_proc=8) + test_dataset = test_dataset.map(function=make_map_fn('test'), with_indices=True, num_proc=8) + + local_dir = args.local_dir + hdfs_dir = args.hdfs_dir + + train_dataset.to_parquet(os.path.join(local_dir, 'train.parquet')) + test_dataset.to_parquet(os.path.join(local_dir, 'test.parquet')) + + if hdfs_dir is not None: + makedirs(hdfs_dir) + copy(src=local_dir, dst=hdfs_dir) diff --git a/reasoners/data_preprocess/archived_for_reference/orz_data.py b/reasoners/data_preprocess/archived_for_reference/orz_data.py new file mode 100644 index 000000000..98b1b5f3c --- /dev/null +++ b/reasoners/data_preprocess/archived_for_reference/orz_data.py @@ -0,0 +1,163 @@ +# 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 os +import datasets +from typing import Dict, List, Optional, Any, Union +import enum +import argparse +import pandas as pd +import json + +from verl.utils.hdfs_io import copy, makedirs +from verl.utils.reward_score.math import remove_boxed, last_boxed_only_string + + +class TrainDataset(enum.Enum): + """Enum for training datasets. + + Contains identifiers for various math problem datasets used during training. + """ + + AIME = "AIME" # American Invitational Mathematics Examination + AMC = "AMC" # American Mathematics Competition + OMNI_MATH = "OMNI_MATH" # Omni Math + NUMINA_OLYMPIAD = "OLYMPIAD" # Unique Olympiad problems from NUMINA + MATH = "MATH" # Dan Hendrycks Math Problems + STILL = "STILL" # STILL dataset + DEEPSCALER = "DEEPSCALER" # DeepScaler (AIME, AMC, OMNI_MATH, MATH, STILL) + + +class TestDataset(enum.Enum): + """Enum for testing/evaluation datasets. + + Contains identifiers for datasets used to evaluate model performance. + """ + + AIME = "AIME" # American Invitational Mathematics Examination + AMC = "AMC" # American Mathematics Competition + MATH = "MATH" # Math 500 problems + MINERVA = "MINERVA" # Minerva dataset + OLYMPIAD_BENCH = "OLYMPIAD_BENCH" # Olympiad benchmark problems + + +"""Type alias for either training or testing dataset types.""" +Dataset = Union[TrainDataset, TestDataset] + + +def load_dataset(dataset: Dataset) -> List[Dict[str, Any]]: + """Load a dataset from a JSON file. + + Loads and parses a JSON dataset file based on the provided dataset enum. + The file path is constructed based on whether it's a training or testing dataset. + + Args: + dataset: A Dataset enum value specifying which dataset to load. + + Returns: + List[Dict[str, Any]]: A list of dictionaries containing the dataset records. + Each dictionary represents one example in the dataset. + + Raises: + ValueError: If the dataset file cannot be found, contains invalid JSON, + or encounters other file access errors. + + Example: + >>> load_dataset(TrainDataset.AIME) + [{'problem': 'Find x...', 'solution': '42', ...}, ...] + """ + dataset_name = dataset.value.lower() + data_dir = "train" if isinstance(dataset, TrainDataset) else "test" + + current_dir = os.path.dirname(os.path.realpath(__file__)) + file_path = os.path.join(data_dir, f"{dataset_name}.json") + file_path = os.path.join(current_dir, file_path) + + if not os.path.exists(file_path): + raise ValueError(f"Dataset file not found: {file_path}") + + try: + with open(file_path, "r", encoding="utf-8") as file: + data = json.load(file) + return data + except json.JSONDecodeError: + raise ValueError(f"Invalid JSON format in {file_path}") + except Exception as exc: # pylint: disable=broad-except + raise ValueError(f"Error loading dataset: {exc}") from exc + + +def extract_solution(solution_str): + return remove_boxed(last_boxed_only_string(solution_str)) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--local_dir", default="data/orz") + parser.add_argument("--hdfs_dir", default=None) + args = parser.parse_args() + + # Load train data from local JSON file + + # https://raw.githubusercontent.com/Open-Reasoner-Zero/Open-Reasoner-Zero/refs/heads/main/data/orz_math_57k_collected.json + train_data_source = "orz_math_57k_collected.json" + print(f"Loading training data from {train_data_source}...", flush=True) + + train_dataset = json.load(open(train_data_source, "r")) + + # Convert to Dataset format that's compatible with the rest of the code + train_data = datasets.Dataset.from_list([{ + "problem": item[0]["value"], + "answer": item[1]["ground_truth"]["value"] + } for item in train_dataset]) + + print(train_data[0]) + + instruction_following = ( + "Let's think step by step and output the final answer within \\boxed{}." + ) + + # Rest of the processing remains the same + def make_map_fn(split, data_source): + def process_fn(example, idx): + question = example.pop("problem") + question = question + " " + instruction_following + answer = example.pop("answer") + data = { + "data_source": data_source, + "prompt": [{"role": "user", "content": question}], + "ability": "math", + "reward_model": {"style": "rule", "ground_truth": answer}, + "extra_info": {"split": split, "index": idx}, + } + if idx == 0: + print("=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) + print(data) + return data + return process_fn + + local_dir = args.local_dir + hdfs_dir = args.hdfs_dir + + # Process train data + train_data = train_data.map( + function=make_map_fn("train", train_data_source), with_indices=True + ) + print(train_data[0]) + train_data.to_parquet(os.path.join(local_dir, "train.parquet")) + print(f"train data size:", len(train_data)) + + # Remove test data processing since we're only working with the training data + if hdfs_dir is not None: + makedirs(hdfs_dir) + copy(src=local_dir, dst=hdfs_dir) diff --git a/reasoners/data_preprocess/archived_for_reference/orz_data_r1_zero_style.py b/reasoners/data_preprocess/archived_for_reference/orz_data_r1_zero_style.py new file mode 100644 index 000000000..95a8d972c --- /dev/null +++ b/reasoners/data_preprocess/archived_for_reference/orz_data_r1_zero_style.py @@ -0,0 +1,192 @@ +# 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 os +import datasets +from typing import Dict, List, Optional, Any, Union +import enum +import argparse +import pandas as pd +import json + +from verl.utils.hdfs_io import copy, makedirs +from verl.utils.reward_score.math import remove_boxed, last_boxed_only_string + + +class TrainDataset(enum.Enum): + """Enum for training datasets. + + Contains identifiers for various math problem datasets used during training. + """ + + AIME = "AIME" # American Invitational Mathematics Examination + AMC = "AMC" # American Mathematics Competition + OMNI_MATH = "OMNI_MATH" # Omni Math + NUMINA_OLYMPIAD = "OLYMPIAD" # Unique Olympiad problems from NUMINA + MATH = "MATH" # Dan Hendrycks Math Problems + STILL = "STILL" # STILL dataset + DEEPSCALER = "DEEPSCALER" # DeepScaler (AIME, AMC, OMNI_MATH, MATH, STILL) + + +class TestDataset(enum.Enum): + """Enum for testing/evaluation datasets. + + Contains identifiers for datasets used to evaluate model performance. + """ + + AIME = "AIME" # American Invitational Mathematics Examination + AMC = "AMC" # American Mathematics Competition + MATH = "MATH" # Math 500 problems + MINERVA = "MINERVA" # Minerva dataset + OLYMPIAD_BENCH = "OLYMPIAD_BENCH" # Olympiad benchmark problems + + +"""Type alias for either training or testing dataset types.""" +Dataset = Union[TrainDataset, TestDataset] + + +def load_dataset(dataset: Dataset) -> List[Dict[str, Any]]: + """Load a dataset from a JSON file. + + Loads and parses a JSON dataset file based on the provided dataset enum. + The file path is constructed based on whether it's a training or testing dataset. + + Args: + dataset: A Dataset enum value specifying which dataset to load. + + Returns: + List[Dict[str, Any]]: A list of dictionaries containing the dataset records. + Each dictionary represents one example in the dataset. + + Raises: + ValueError: If the dataset file cannot be found, contains invalid JSON, + or encounters other file access errors. + + Example: + >>> load_dataset(TrainDataset.AIME) + [{'problem': 'Find x...', 'solution': '42', ...}, ...] + """ + dataset_name = dataset.value.lower() + data_dir = "train" if isinstance(dataset, TrainDataset) else "test" + + current_dir = os.path.dirname(os.path.realpath(__file__)) + file_path = os.path.join(data_dir, f"{dataset_name}.json") + file_path = os.path.join(current_dir, file_path) + + if not os.path.exists(file_path): + raise ValueError(f"Dataset file not found: {file_path}") + + try: + with open(file_path, "r", encoding="utf-8") as file: + data = json.load(file) + return data + except json.JSONDecodeError: + raise ValueError(f"Invalid JSON format in {file_path}") + except Exception as exc: # pylint: disable=broad-except + raise ValueError(f"Error loading dataset: {exc}") from exc + + +def extract_solution(solution_str): + return remove_boxed(last_boxed_only_string(solution_str)) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--local_dir", default="data/orz_zero_style_v2") + parser.add_argument("--hdfs_dir", default=None) + args = parser.parse_args() + + # Load train data from local JSON file + + # https://raw.githubusercontent.com/Open-Reasoner-Zero/Open-Reasoner-Zero/refs/heads/main/data/orz_math_57k_collected.json + train_data_source = "examples/data_preprocess/orz_math_57k_collected.json" + print(f"Loading training data from {train_data_source}...", flush=True) + + train_dataset = json.load(open(train_data_source, "r")) + + # Convert to Dataset format that's compatible with the rest of the code + train_data = datasets.Dataset.from_list([{ + "problem": item[0]["value"], + "answer": item[1]["ground_truth"]["value"] + } for item in train_dataset]) + + print(train_data[0]) + + test_data_sources = [ + "nanoverl/minerva", + "nanoverl/aime", + "nanoverl/amc", + "nanoverl/olympiad_bench", + "nanoverl/math", + ] + print(f"Loading the {test_data_sources} dataset from huggingface...", flush=True) + + test_datasets = [ + datasets.load_dataset(test_data_source, trust_remote_code=True, split="test") + for test_data_source in test_data_sources + ] + + prompt = """A conversation between User and Assistant. The user asks a question, and the Assistant solves it. The assistant first thinks about the reasoning process in the mind and then provides the user with the response. The reasoning process is enclosed within i.e., reasoning process here respond to the user's question here. + +User: {{prompt}} Please put your answer in \\boxed{} tags. +Assistant: +""" + # Rest of the processing remains the same + def make_map_fn(split, data_source): + def process_fn(example, idx): + question = example.pop("problem") + answer = example.pop("answer") + data = { + "data_source": data_source, + "question": question, + "prompt": [], + "raw_prompt": prompt.replace("{{prompt}}", question), + "ability": "math", + "apply_chat_template": False, + "reward_model": {"style": "rule", "ground_truth": answer}, + "extra_info": { + "split": split, + "index": idx, + "question": question, + }, + } + if idx == 0: + print("=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) + print(data) + return data + return process_fn + + local_dir = args.local_dir + hdfs_dir = args.hdfs_dir + + # Process train data + train_data = train_data.map( + function=make_map_fn("train", train_data_source), with_indices=True + ) + print(train_data[0]) + train_data.to_parquet(os.path.join(local_dir, "train.parquet")) + print(f"train data size:", len(train_data)) + + for test_data_source, test_data in zip(test_data_sources, test_datasets): + process_fn = make_map_fn("test", test_data_source) + test_data = test_data.map(process_fn, with_indices=True) + dataset_name = os.path.basename(test_data_source.lower()) + test_df = pd.DataFrame(test_data) + test_df.to_parquet(os.path.join(local_dir, f"{dataset_name}.parquet")) + print(f"test data size: ({dataset_name})", len(test_df)) + + # Remove test data processing since we're only working with the training data + if hdfs_dir is not None: + makedirs(hdfs_dir) + copy(src=local_dir, dst=hdfs_dir) diff --git a/reasoners/data_preprocess/codegen/humaneval.py b/reasoners/data_preprocess/codegen/humaneval.py new file mode 100644 index 000000000..6557b843a --- /dev/null +++ b/reasoners/data_preprocess/codegen/humaneval.py @@ -0,0 +1,156 @@ +"""Downloads, processes, and saves HumanEval dataset.""" + +import os +import argparse +import json +import transformers +import datasets +from datasets import load_dataset, Dataset + +from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset +from verl.utils.data_process.filter import LengthFilter +from verl.utils.reward_score.coder1 import code_exec + + +def get_datasets(cache_dir: str): + """ + Loads the HumanEval dataset. + """ + try: + dataset = load_dataset("openai_humaneval", cache_dir=cache_dir)["test"] + print(f"HumanEval dataset: {len(dataset)} examples") + return None, dataset + except Exception as e: + print(f"Error loading dataset: {e}") + return None, None + + +def make_map_fn(split: str, data_source: str) -> callable: + def process_fn(example, idx): + # Create a default "skip" response with all required fields + skip_response = { + "data_source": None, + "prompt": None, + "ability": None, + "reward_model": None, + "extra_info": None + } + + # Extract task ID and prompt + task_id = example["task_id"] + + prompt = ( + "Write a complete, self-contained Python solution to the following problem. " + "Your solution must include all necessary imports and the full function definition including " + "the signature exactly as specified. Do not modify the function signature or docstring.\n\n" + f"```python\n{example['prompt'].strip()}\n```" + ) + + # Extract test function, entry point, and canonical solution + test_code = example["test"] + entry_point = example["entry_point"] + solution = example["canonical_solution"] + + # Build test code that calls the entry point + test_code_with_check = f"{test_code}\n\ncheck({entry_point})" + + # Verify the canonical solution passes the tests + # full_code = f"{prompt}\n{solution}\n{test_code}\n\ncheck({entry_point})" + full_code = f"{example['prompt']}\n{solution}\n{test_code}\n\ncheck({entry_point})" + succ, err = code_exec(full_code, timeout=5) + print(f"[DEBUG] succ: {succ}, err: {err}") + if not succ: + print(f"Error in canonical solution for task {task_id}: {err}") + return skip_response + + + data = { + "data_source": data_source, + "prompt": [ + {"role": "user", "content": prompt} + ], + "ability": "codegen", + "apply_chat_template": True, + "reward_model": { + "style": "rule", + "ground_truth": json.dumps({ + "functional": test_code_with_check + }), + }, + "extra_info": { + "split": split, + "index": idx, + "reference": solution, # Include the canonical solution as reference + "original_prompt": prompt, + "dataset": "openai_humaneval", + "task_id": task_id, + }, + } + + if idx == 0 or idx == 1: + print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) + print(data) + + return data + + return process_fn + + +if __name__ == '__main__': + """Main script execution: parse args, load, process, and save datasets.""" + parser = argparse.ArgumentParser(description="Process and save HumanEval dataset.") + parser.add_argument('--data-dir', default='data', + help='Base directory to save the processed data files.') + parser.add_argument('--domain', default="codegen", + help='Domain of the dataset.') + parser.add_argument('--name', default="humaneval", + help='Name of the dataset.') + parser.add_argument('--sample-size', type=int, default=None, + help='Number of samples to use from dataset. If None, use all samples.') + parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility') + + args = parser.parse_args() + + set_seed(args.seed) + + data_source = f"{args.domain}__{args.name}" + test_output_dir = os.path.join(args.data_dir, 'test') + os.makedirs(test_output_dir, exist_ok=True) + + # Load the dataset + cache_dir = datasets.config.HF_DATASETS_CACHE + _, dataset = get_datasets(cache_dir=cache_dir) + print(f"[DEBUG] dataset: {len(dataset)}") + + # Process the dataset + process_fn = make_map_fn('test', data_source) + + dataset = dataset.map(function=process_fn, with_indices=True) + print(f"[DEBUG] processed dataset: {len(dataset)}") + + # Filter out examples where processing failed + dataset = dataset.filter(lambda x: x["data_source"] == data_source) + print(f"[DEBUG] filtered dataset: {len(dataset)}") + # Length filter + try: + tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B") + length_filter = LengthFilter(tokenizer=tokenizer, max_length=4096) + dataset = dataset.filter(lambda x: length_filter.check(x)) + except Exception as e: + print(f"Warning: Could not perform length filtering. Error: {e}") + print("Proceeding without length filtering.") + + # Sample the dataset + dataset = sample_dataset(dataset, args.sample_size) + + # Save the dataset to test directory + test_output_path = save_dataset( + dataset=dataset, + output_dir=test_output_dir, + filename_prefix=data_source, + sample_size=len(dataset) + ) + + print(f"\nDone! \n" + f"Data source: {data_source}\n" + f"Test data saved to {test_output_path} ({len(dataset)} samples)") \ No newline at end of file diff --git a/reasoners/data_preprocess/codegen/humanevalplus.py b/reasoners/data_preprocess/codegen/humanevalplus.py new file mode 100644 index 000000000..a03d13938 --- /dev/null +++ b/reasoners/data_preprocess/codegen/humanevalplus.py @@ -0,0 +1,157 @@ +"""Downloads, processes, and saves HumanEval dataset.""" + +import os +import argparse +import json +import transformers +import datasets +from datasets import load_dataset, Dataset + +from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset +from verl.utils.data_process.filter import LengthFilter +from verl.utils.reward_score.coder1 import code_exec + + +def get_datasets(cache_dir: str): + """ + Loads the HumanEvalPlus dataset. + """ + try: + dataset = load_dataset("autoprogrammer/humanevalplus_corrected", cache_dir=cache_dir)["test"] + print(f"HumanEvalPlus dataset: {len(dataset)} examples") + return None, dataset + except Exception as e: + print(f"Error loading dataset: {e}") + return None, None + + +def make_map_fn(split: str, data_source: str) -> callable: + def process_fn(example, idx): + # Create a default "skip" response with all required fields + skip_response = { + "data_source": None, + "prompt": None, + "ability": None, + "apply_chat_template": None, + "reward_model": None, + "extra_info": None + } + + # Extract task ID and prompt + task_id = example["task_id"] + + prompt = ( + "Write a complete, self-contained Python solution to the following problem. " + "Your solution must include all necessary imports and the full function definition including " + "the signature exactly as specified. Do not modify the function signature or docstring.\n\n" + f"```python\n{example['prompt'].strip()}\n```" + ) + + # Extract test function, entry point, and canonical solution + test_code = example["test"] + entry_point = example["entry_point"] + solution = example["canonical_solution"] + + # Build test code that calls the entry point + test_code_with_check = f"{test_code}\n\ncheck({entry_point})" + + # Verify the canonical solution passes the tests + # full_code = f"{prompt}\n{solution}\n{test_code}\n\ncheck({entry_point})" + full_code = f"{example['prompt']}\n{solution}\n{test_code}\n\ncheck({entry_point})" + succ, err = code_exec(full_code, timeout=30) + print(f"[DEBUG] succ: {succ}, err: {err}") + if not succ: + print(f"Error in canonical solution for task {task_id}: {err}") + return skip_response + + + data = { + "data_source": data_source, + "prompt": [ + {"role": "user", "content": prompt} + ], + "ability": "codegen", + "apply_chat_template": True, + "reward_model": { + "style": "rule", + "ground_truth": json.dumps({ + "functional": test_code_with_check + }), + }, + "extra_info": { + "split": split, + "index": idx, + "reference": solution, # Include the canonical solution as reference + "original_prompt": prompt, + "dataset": "evalplus_humanevalplus", + "task_id": task_id, + }, + } + + if idx == 0 or idx == 1: + print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) + print(data) + + return data + + return process_fn + + +if __name__ == '__main__': + """Main script execution: parse args, load, process, and save datasets.""" + parser = argparse.ArgumentParser(description="Process and save HumanEvalPlus dataset.") + parser.add_argument('--data-dir', default='data', + help='Base directory to save the processed data files.') + parser.add_argument('--domain', default="codegen", + help='Domain of the dataset.') + parser.add_argument('--name', default="humanevalplus", + help='Name of the dataset.') + parser.add_argument('--sample-size', type=int, default=None, + help='Number of samples to use from dataset. If None, use all samples.') + parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility') + + args = parser.parse_args() + + set_seed(args.seed) + + data_source = f"{args.domain}__{args.name}" + test_output_dir = os.path.join(args.data_dir, 'test') + os.makedirs(test_output_dir, exist_ok=True) + + # Load the dataset + cache_dir = datasets.config.HF_DATASETS_CACHE + _, dataset = get_datasets(cache_dir=cache_dir) + print(f"[DEBUG] dataset: {len(dataset)}") + + # Process the dataset + process_fn = make_map_fn('test', data_source) + + dataset = dataset.map(function=process_fn, with_indices=True,load_from_cache_file=False) + print(f"[DEBUG] processed dataset: {len(dataset)}") + + # Filter out examples where processing failed + dataset = dataset.filter(lambda x: x["data_source"] == data_source) + print(f"[DEBUG] filtered dataset: {len(dataset)}") + # Length filter + try: + tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B") + length_filter = LengthFilter(tokenizer=tokenizer, max_length=4096) + dataset = dataset.filter(lambda x: length_filter.check(x)) + except Exception as e: + print(f"Warning: Could not perform length filtering. Error: {e}") + print("Proceeding without length filtering.") + + # Sample the dataset + dataset = sample_dataset(dataset, args.sample_size) + + # Save the dataset to test directory + test_output_path = save_dataset( + dataset=dataset, + output_dir=test_output_dir, + filename_prefix=data_source, + sample_size=len(dataset) + ) + + print(f"\nDone! \n" + f"Data source: {data_source}\n" + f"Test data saved to {test_output_path} ({len(dataset)} samples)") \ No newline at end of file diff --git a/reasoners/data_preprocess/codegen/leetcode2k.py b/reasoners/data_preprocess/codegen/leetcode2k.py new file mode 100644 index 000000000..22d48fc24 --- /dev/null +++ b/reasoners/data_preprocess/codegen/leetcode2k.py @@ -0,0 +1,174 @@ +"""Downloads, processes, and saves LeetCode2K datasets.""" + +import os +import argparse +import json +import transformers +import datasets +from datasets import load_dataset, Dataset + +from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset +from verl.utils.data_process.filter import LengthFilter +from verl.utils.reward_score.coder1 import code_exec + +EMPTY_EXAMPLE = { + "data_source": None, + "prompt": None, + "apply_chat_template": False, + "ability": None, + "reward_model": None, + "extra_info": None +} + + +def get_datasets(cache_dir: str): + """ + Loads the LeetCode2K dataset. + """ + try: + commit_hash = "34803eb64eab1979069ba1f80e7ea474282e28f3" + train_dataset = load_dataset("newfacade/LeetCodeDataset", + revision=commit_hash, + cache_dir=cache_dir)["train"] + test_dataset = load_dataset("newfacade/LeetCodeDataset", + revision=commit_hash, + cache_dir=cache_dir)["test"] + print(f"Train set: {len(train_dataset)} examples") + print(f"Test set: {len(test_dataset)} examples") + return train_dataset, test_dataset + except Exception as e: + print(f"Error loading dataset: {e}") + return None, None + + +def make_map_fn(split: str, data_source: str) -> callable: + def process_fn(example, idx): + prefix = example["prompt"] + + # Clean up the prompt + prompt = example["query"] + prompt = prompt.replace("### Answer: (use the provided format with backticks)", "").strip() + prompt = prompt.replace("### Format: ", "### Format:\n") + + # Build test code + test_code = f"{example['test']}\n\ncheck({example['entry_point'].strip()})" + + # Extract the candidate solution + solution = example["completion"] + + # Combine all code pieces into a single file to execute + full_code = f"{prefix}\n{solution}\n{test_code}" + + # Validate that the candidate solution passes the tests + # Skip examples where the test code fails + succ, err = code_exec(full_code, timeout=20) + if not succ: + print(f"Test code failed for example {idx}: {example['meta']['question_title']}") + return EMPTY_EXAMPLE + + data = { + "data_source": data_source, + "prompt": [ + {"role": "user", "content": prompt} + ], + "ability": "codegen", + "apply_chat_template": True, + "reward_model": { + "style": "rule", + "ground_truth": json.dumps({ + "functional": test_code + }), + }, + "extra_info": { + "split": split, + "index": idx, + "reference": solution, + "original_prompt": prompt, + "prefix": prefix, + "dataset": "LeetCodeDataset", + "question_title": example["meta"]["question_title"], + "difficulty": example["meta"]["difficulty"], + }, + } + + if idx == 0 or idx == 1: + print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) + print(data) + + return data + + return process_fn + + +if __name__ == '__main__': + """Main script execution: parse args, load, process, and save datasets.""" + parser = argparse.ArgumentParser(description="Process and save LeetCode2K datasets.") + parser.add_argument('--data-dir', default='data', + help='Base directory to save the processed data files.') + parser.add_argument('--domain', default="codegen", + help='Domain of the dataset.') + parser.add_argument('--name', default="leetcode2k", + help='Name of the dataset.') + parser.add_argument('--train-sample-size', type=int, default=None, + help='Number of samples to use from training dataset. If None, use all samples.') + parser.add_argument('--test-sample-size', type=int, default=None, + help='Number of samples to use from test dataset. If None, use all samples.') + parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility') + + args = parser.parse_args() + + set_seed(args.seed) + + data_source = f"{args.domain}__{args.name}" + train_output_dir = os.path.join(args.data_dir, 'train') + test_output_dir = os.path.join(args.data_dir, 'test') + os.makedirs(train_output_dir, exist_ok=True) + os.makedirs(test_output_dir, exist_ok=True) + + # Load the dataset + cache_dir = datasets.config.HF_DATASETS_CACHE + train_dataset, test_dataset = get_datasets(cache_dir=cache_dir) + + # Process the dataset + process_train_fn = make_map_fn('train', data_source) + process_test_fn = make_map_fn('test', data_source) + + train_dataset = train_dataset.map(function=process_train_fn, with_indices=True, num_proc=64) + test_dataset = test_dataset.map(function=process_test_fn, with_indices=True, num_proc=64) + + # Filter out examples where processing failed + train_dataset = train_dataset.filter(lambda x: x["data_source"] == data_source) + test_dataset = test_dataset.filter(lambda x: x["data_source"] == data_source) + + # Length filter + try: + tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B") + length_filter = LengthFilter(tokenizer=tokenizer, max_length=4096) + train_dataset = train_dataset.filter(lambda x: length_filter.check(x)) + test_dataset = test_dataset.filter(lambda x: length_filter.check(x)) + except Exception as e: + print(f"Warning: Could not perform length filtering. Error: {e}") + print("Proceeding without length filtering.") + + # Sample the datasets using utility function + train_dataset = sample_dataset(train_dataset, args.train_sample_size) + test_dataset = sample_dataset(test_dataset, args.test_sample_size) + + # Save the datasets using utility function + train_output_path = save_dataset( + dataset=train_dataset, + output_dir=train_output_dir, + filename_prefix=data_source, + sample_size=len(train_dataset) + ) + test_output_path = save_dataset( + dataset=test_dataset, + output_dir=test_output_dir, + filename_prefix=data_source, + sample_size=len(test_dataset) + ) + + print(f"\nDone! \n" + f"Data source: {data_source}\n" + f"Train data saved to {train_output_path} ({len(train_dataset)} samples)\n" + f"Test data saved to {test_output_path} ({len(test_dataset)} samples)") \ No newline at end of file diff --git a/reasoners/data_preprocess/codegen/livecodebench.py b/reasoners/data_preprocess/codegen/livecodebench.py new file mode 100644 index 000000000..090ea9345 --- /dev/null +++ b/reasoners/data_preprocess/codegen/livecodebench.py @@ -0,0 +1,230 @@ +"""Downloads, processes, and saves MBPP datasets.""" + +import os +import argparse +import json +import transformers +import datasets +from datasets import load_dataset, Dataset + +from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset +from verl.utils.data_process.filter import LengthFilter + + +def get_datasets(cache_dir: str): + """ + Loads the LiveCodeBench dataset. + """ + try: + dataset = load_dataset("agentica-org/DeepCoder-Preview-Dataset", "lcbv5", cache_dir=cache_dir) + print(f"Train set: {len(dataset['train'])} examples") + print(f"Test set: {len(dataset['test'])} examples") + return dataset["train"], dataset["test"] + except Exception as e: + print(f"Error loading dataset: {e}") + return None, None + + +def make_map_fn(split: str, data_source: str) -> callable: + def process_fn(example, idx): + # Get the problem description + problem_desc = example["problem"] + starter_code = example["starter_code"] + + # Process tests + try: + tests = json.loads(example["tests"]) + if not tests: + return { + "data_source": None, + "prompt": None, + "ability": None, + "reward_model": None, + "extra_info": None + } + + # Process metadata to get function name + metadata = example["metadata"] + function_name = metadata.get("func_name") + + # Handle different test types + if tests[0]["testtype"] == "functional": + if not function_name: + return { + "data_source": None, + "prompt": None, + "raw_prompt": None, + "ability": None, + "reward_model": None, + "extra_info": None + } + + # create the prompt with the starter code + # prompt adopted from livecodebench official implementation + # https://github.com/LiveCodeBench/LiveCodeBench/blob/main/lcb_runner/prompts/code_generation.py + prompt = f"""You are an expert Python programmer. You will be given a question (problem specification) and will generate a correct Python program that matches the specification and passes all tests. + +Below is the question: + +{problem_desc} + +You will use the following starter code to write the solution to the problem and enclose your code within ```python delimiters. +```python +{starter_code} +```""" + + # Function call tests + test_code = f"""\ +def check_{function_name}(): +""" + for test in tests: + # Parse input string by splitting on '\n' + input_parts = test["input"].split('\n') + # Create proper comma-separated arguments for the function call + input_args = ', '.join(input_parts) + # Get the output value + output_val = test["output"] + + test_code += f""" assert Solution().{function_name}({input_args}) == {output_val} +""" + + test_code += f""" +check_{function_name}() +""" + + oracle = json.dumps({"functional": test_code}) + + elif tests[0]["testtype"] == "stdin": + # create the prompt, and there will be no starter code + # prompt adopted from livecodebench official implementation + # https://github.com/LiveCodeBench/LiveCodeBench/blob/main/lcb_runner/prompts/code_generation.py + prompt = f"""You are an expert Python programmer. You will be given a question (problem specification) and will generate a correct Python program that matches the specification and passes all tests. + +Below is the question: + +{problem_desc} + +Read the inputs from stdin solve the problem and write the answer to stdout (do not directly test on the sample inputs). Enclose your code within delimiters. Ensure that when the python program runs, it reads the inputs, runs the algorithm and writes output to STDOUT.""" + + # STDIN/STDOUT tests + stdin_list = [] + stdout_list = [] + for test in tests: + stdin_list.append(test["input"]) + stdout_list.append(test["output"]) + + oracle = json.dumps({"inputs": stdin_list, "outputs": stdout_list}) + else: + raise ValueError(f"Unknown test type: {tests[0]['testtype']}") + except Exception as e: + print(f"Error processing tests for example {idx}: {e}") + return { + "data_source": None, + "prompt": None, + "ability": None, + "reward_model": None, + "extra_info": None + } + + data = { + "data_source": data_source, + "prompt": [ + {"role": "user", "content": prompt} + ], + "ability": "codegen", + "apply_chat_template": True, + "reward_model": { + "style": "rule", + "ground_truth": oracle, + }, + "extra_info": { + "split": split, + "index": idx, + "reference": "", # No solution data in LiveCodeBench + "dataset": "LiveCodeBench", + "function_name": function_name if function_name else "", + "original_prompt": prompt, + }, + } + + if idx == 0 or idx == 1: + print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) + print(data) + + return data + + return process_fn + + +if __name__ == '__main__': + """Main script execution: parse args, load, process, and save datasets.""" + parser = argparse.ArgumentParser(description="Process and save LiveCodeBench datasets.") + parser.add_argument('--data-dir', default='data', + help='Base directory to save the processed data files.') + parser.add_argument('--domain', type=str, default='codegen', + help='Domain of the dataset.') + parser.add_argument('--name', type=str, default='livecodebench', + help='Name of the dataset.') + parser.add_argument('--train-sample-size', type=int, default=None, + help='Number of samples to use from training dataset. If None, use all samples.') + parser.add_argument('--test-sample-size', type=int, default=None, + help='Number of samples to use from test dataset. If None, use all samples.') + parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility') + + args = parser.parse_args() + + set_seed(args.seed) + + data_source = f"{args.domain}__{args.name}" + train_output_dir = os.path.join(args.data_dir, 'train') + test_output_dir = os.path.join(args.data_dir, 'test') + os.makedirs(train_output_dir, exist_ok=True) + os.makedirs(test_output_dir, exist_ok=True) + + # Load the dataset + cache_dir = datasets.config.HF_DATASETS_CACHE + train_dataset, test_dataset = get_datasets(cache_dir) + + # Process the dataset + process_train_fn = make_map_fn('train', data_source) + process_test_fn = make_map_fn('test', data_source) + + train_dataset = train_dataset.map(function=process_train_fn, with_indices=True, num_proc=64) + test_dataset = test_dataset.map(function=process_test_fn, with_indices=True, num_proc=64) + + # Filter out examples where processing failed + train_dataset = train_dataset.filter(lambda x: x["data_source"] == data_source) + test_dataset = test_dataset.filter(lambda x: x["data_source"] == data_source) + + # Length filter + try: + tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B") + length_filter = LengthFilter(tokenizer=tokenizer, max_length=4096) + train_dataset = train_dataset.filter(lambda x: length_filter.check(x), num_proc=64) + test_dataset = test_dataset.filter(lambda x: length_filter.check(x), num_proc=64) + except Exception as e: + print(f"Warning: Could not perform length filtering. Error: {e}") + print("Proceeding without length filtering.") + + # Sample the datasets using utility function + train_dataset = sample_dataset(train_dataset, args.train_sample_size) + test_dataset = sample_dataset(test_dataset, args.test_sample_size) + + # Save the datasets using utility function + train_output_path = save_dataset( + dataset=train_dataset, + output_dir=train_output_dir, + filename_prefix=data_source, + sample_size=len(train_dataset) + ) + test_output_path = save_dataset( + dataset=test_dataset, + output_dir=test_output_dir, + filename_prefix=data_source, + sample_size=len(test_dataset) + ) + + print(f"\nDone! \n" + f"Data source: {data_source}\n" + f"Train data saved to {train_output_path} ({len(train_dataset)} samples)\n" + f"Test data saved to {test_output_path} ({len(test_dataset)} samples)") \ No newline at end of file diff --git a/reasoners/data_preprocess/codegen/mbpp.py b/reasoners/data_preprocess/codegen/mbpp.py new file mode 100644 index 000000000..d77d5efac --- /dev/null +++ b/reasoners/data_preprocess/codegen/mbpp.py @@ -0,0 +1,169 @@ +"""Downloads, processes, and saves MBPP datasets.""" + +import os +import argparse +import json +import transformers +import datasets +from datasets import load_dataset, Dataset + +from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset +from verl.utils.data_process.filter import LengthFilter +from verl.utils.reward_score.coder1 import code_exec + + +def get_datasets(cache_dir: str): + """ + Loads the MBPP dataset. + """ + try: + dataset = load_dataset("google-research-datasets/mbpp", cache_dir=cache_dir) + print(f"Train set: {len(dataset['train'])} examples") + print(f"Test set: {len(dataset['test'])} examples") + return dataset["train"], dataset["test"] + except Exception as e: + print(f"Error loading dataset: {e}") + return None, None + + +def make_map_fn(split: str, data_source: str) -> callable: + def process_fn(example, idx): + # Create prompt + prompt = ( + f"{example['text']}\n\n" + f"Your solution should be a complete, self-contained function in a markdown code block. " + f"Make sure your solution passes the following test cases:\n" + ) + + # Construct test code + test_code = "" + if example.get('test_setup_code'): + test_code += example['test_setup_code'] + "\n\n" + + # Add all test assertions + for assertion in example['test_list'] + example.get('challenge_test_list', []): + test_code += assertion + "\n" + + # Add test cases to prompt + prompt += f"```python\n{test_code}```" + prompt += "\n\nPlease do not include the test cases in your solution." + + # Validate that the canonical solution passes the tests + solution = example['code'] + full_code = f"{solution}\n\n{test_code}" + + succ, err = code_exec(full_code) + if not succ: + print(f"Test code failed for example {idx}: {example.get('task_id', 'unknown')}") + return { + "data_source": None, + "prompt": None, + "ability": None, + "reward_model": None, + "extra_info": None + } + + data = { + "data_source": data_source, + "prompt": [ + {"role": "user", "content": prompt} + ], + "ability": "codegen", + "apply_chat_template": True, + "reward_model": { + "style": "rule", + "ground_truth": json.dumps({ + "functional": test_code + }), + }, + "extra_info": { + "split": split, + "index": idx, + "reference": solution, + "original_prompt": prompt, + "dataset": "mbpp", + "task_id": str(example.get("task_id", "")), + }, + } + + if idx == 0 or idx == 1: + print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) + print(data) + + return data + + return process_fn + + +if __name__ == '__main__': + """Main script execution: parse args, load, process, and save datasets.""" + parser = argparse.ArgumentParser(description="Process and save MBPP datasets.") + parser.add_argument('--data-dir', default='data', + help='Base directory to save the processed data files.') + parser.add_argument('--domain', default="codegen", + help='Domain of the dataset.') + parser.add_argument('--name', default="mbpp", + help='Name of the dataset.') + parser.add_argument('--train-sample-size', type=int, default=None, + help='Number of samples to use from training dataset. If None, use all samples.') + parser.add_argument('--test-sample-size', type=int, default=None, + help='Number of samples to use from test dataset. If None, use all samples.') + parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility') + + args = parser.parse_args() + + set_seed(args.seed) + + data_source = f"{args.domain}__{args.name}" + train_output_dir = os.path.join(args.data_dir, 'train') + test_output_dir = os.path.join(args.data_dir, 'test') + os.makedirs(train_output_dir, exist_ok=True) + os.makedirs(test_output_dir, exist_ok=True) + + # Load the dataset + cache_dir = datasets.config.HF_DATASETS_CACHE + train_dataset, test_dataset = get_datasets(cache_dir=cache_dir) + + # Process the dataset + process_train_fn = make_map_fn('train', data_source) + process_test_fn = make_map_fn('test', data_source) + + train_dataset = train_dataset.map(function=process_train_fn, with_indices=True, num_proc=64) + test_dataset = test_dataset.map(function=process_test_fn, with_indices=True, num_proc=64) + + # Filter out examples where processing failed + train_dataset = train_dataset.filter(lambda x: x["data_source"] == data_source) + test_dataset = test_dataset.filter(lambda x: x["data_source"] == data_source) + + # Length filter + try: + tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B") + length_filter = LengthFilter(tokenizer=tokenizer, max_length=4096) + train_dataset = train_dataset.filter(lambda x: length_filter.check(x)) + test_dataset = test_dataset.filter(lambda x: length_filter.check(x)) + except Exception as e: + print(f"Warning: Could not perform length filtering. Error: {e}") + print("Proceeding without length filtering.") + + # Sample the datasets using utility function + train_dataset = sample_dataset(train_dataset, args.train_sample_size) + test_dataset = sample_dataset(test_dataset, args.test_sample_size) + + # Save the datasets using utility function + train_output_path = save_dataset( + dataset=train_dataset, + output_dir=train_output_dir, + filename_prefix=data_source, + sample_size=len(train_dataset) + ) + test_output_path = save_dataset( + dataset=test_dataset, + output_dir=test_output_dir, + filename_prefix=data_source, + sample_size=len(test_dataset) + ) + + print(f"\nDone! \n" + f"Data source: {data_source}\n" + f"Train data saved to {train_output_path} ({len(train_dataset)} samples)\n" + f"Test data saved to {test_output_path} ({len(test_dataset)} samples)") \ No newline at end of file diff --git a/reasoners/data_preprocess/codegen/mbppplus.py b/reasoners/data_preprocess/codegen/mbppplus.py new file mode 100644 index 000000000..603e7e367 --- /dev/null +++ b/reasoners/data_preprocess/codegen/mbppplus.py @@ -0,0 +1,158 @@ +"""Downloads, processes, and saves MBPP+ datasets.""" + +import os +import argparse +import json +import transformers +import datasets +from datasets import load_dataset, Dataset + +from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset +from verl.utils.data_process.filter import LengthFilter +from verl.utils.reward_score.coder1 import code_exec + + +def get_datasets(cache_dir: str): + """ + Loads the MBPP+ dataset. + """ + try: + dataset = load_dataset("evalplus/mbppplus", cache_dir=cache_dir) + print(f"Test set: {len(dataset['test'])} examples") + return None, dataset["test"] + except Exception as e: + print(f"Error loading dataset: {e}") + return None, None + + +def make_map_fn(split: str, data_source: str) -> callable: + def process_fn(example, idx): + # Create prompt + prompt = ( + f"{example['prompt']}\n\n" + f"Your solution should be a complete, self-contained function in a markdown code block. " + f"Make sure your solution passes the following test cases:\n" + ) + # Construct test code: imports + assertions + test_code = "" + if example.get('test_imports'): + for imp in example['test_imports']: + test_code += imp + "\n" + test_code += "\n" + # Add all test assertions + test_code += example['test'] + example_code="" + for assertion in example['test_list'] + example.get('challenge_test_list', []): + example_code += assertion + "\n" + + # Add test cases to prompt + prompt += f"```python\n{example_code}```" + prompt += "\n\nPlease do not include the test cases in your solution." + + # Validate that the canonical solution passes the tests + solution = example['code'] + full_code = f"{solution}\n\n{test_code}" + succ, err = code_exec(full_code) + if not succ: + print(f"Test code failed for example {idx}: {example.get('task_id', 'unknown')}") + return { + "data_source": None, + "prompt": None, + "ability": None, + "reward_model": None, + "extra_info": None + } + + data = { + "data_source": data_source, + "prompt": [ + {"role": "user", "content": prompt} + ], + "ability": "codegen", + "apply_chat_template": True, + "reward_model": { + "style": "rule", + "ground_truth": json.dumps({ + "functional": test_code + }), + }, + "extra_info": { + "split": split, + "index": idx, + "reference": solution, + "original_prompt": prompt, + "dataset": "mbpp", + "task_id": str(example.get("task_id", "")), + }, + } + + if idx == 0 or idx == 1: + print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) + print(data) + + return data + + return process_fn + + +if __name__ == '__main__': + """Main script execution: parse args, load, process, and save datasets.""" + parser = argparse.ArgumentParser(description="Process and save MBPP datasets.") + parser.add_argument('--data-dir', default='data', + help='Base directory to save the processed data files.') + parser.add_argument('--domain', default="codegen", + help='Domain of the dataset.') + parser.add_argument('--name', default="mbppplus", + help='Name of the dataset.') + parser.add_argument('--train-sample-size', type=int, default=None, + help='Number of samples to use from training dataset. If None, use all samples.') + parser.add_argument('--test-sample-size', type=int, default=None, + help='Number of samples to use from test dataset. If None, use all samples.') + parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility') + + args = parser.parse_args() + + set_seed(args.seed) + + data_source = f"{args.domain}__{args.name}" + train_output_dir = os.path.join(args.data_dir, 'train') + test_output_dir = os.path.join(args.data_dir, 'test') + os.makedirs(train_output_dir, exist_ok=True) + os.makedirs(test_output_dir, exist_ok=True) + + # Load the dataset + cache_dir = datasets.config.HF_DATASETS_CACHE + + _, test_dataset = get_datasets(cache_dir=cache_dir) + + # Process the dataset + process_test_fn = make_map_fn('test', data_source) + + test_dataset = test_dataset.map(function=process_test_fn, with_indices=True, num_proc=64, load_from_cache_file=False) + + # Filter out examples where processing failed + test_dataset = test_dataset.filter(lambda x: x["data_source"] == data_source) + + # # Length filter + # try: + # tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B") + # length_filter = LengthFilter(tokenizer=tokenizer, max_length=4096) + # test_dataset = test_dataset.filter(lambda x: length_filter.check(x)) + # except Exception as e: + # print(f"Warning: Could not perform length filtering. Error: {e}") + # print("Proceeding without length filtering.") + + # Sample the datasets using utility function + test_dataset = sample_dataset(test_dataset, args.test_sample_size) + + # Save the datasets using utility function + test_output_path = save_dataset( + dataset=test_dataset, + output_dir=test_output_dir, + filename_prefix=data_source, + sample_size=len(test_dataset) + ) + + print(f"\nDone! \n" + f"Data source: {data_source}\n" + f"Test data saved to {test_output_path} ({len(test_dataset)} samples)") \ No newline at end of file diff --git a/reasoners/data_preprocess/codegen/primeintellect.py b/reasoners/data_preprocess/codegen/primeintellect.py new file mode 100644 index 000000000..f122a4e31 --- /dev/null +++ b/reasoners/data_preprocess/codegen/primeintellect.py @@ -0,0 +1,236 @@ +"""Downloads, processes, and saves PrimeIntellect datasets.""" + +import os +import argparse +import json +import transformers +import datasets +from datasets import load_dataset, Dataset +from concurrent.futures import ThreadPoolExecutor, as_completed + +from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset +from verl.utils.data_process.filter import LengthFilter +from verl.utils.reward_score.coder1 import code_exec, remote_check_stdio, fuzzy_equal, extract_code_from_string + +# Define a constant for the empty/failed example +EMPTY_EXAMPLE = { + "data_source": None, + "prompt": None, + "apply_chat_template": False, + "ability": None, + "reward_model": None, + "extra_info": None +} + +def get_datasets(cache_dir: str): + """ + Loads the PrimeIntellect dataset. + """ + try: + dataset = load_dataset("agentica-org/DeepCoder-Preview-Dataset", "primeintellect", split="train", cache_dir=cache_dir) + + # TODO: Remove this line before production - only selecting 100 examples for debugging + # dataset = dataset.select(range(100)) + + print(f"Dataset: {len(dataset)} examples") + return dataset, None + except Exception as e: + print(f"Error loading dataset: {e}") + return None, None + + +def make_map_fn(split: str, data_source: str, verbose: bool) -> callable: + def process_fn(example, idx): + # Get the problem description + prompt = example["problem"] + + # Get the solution code + solution = example["solutions"][0] if example["solutions"] else "" + solution = extract_code_from_string(solution) + + # Process tests + tests = json.loads(example["tests"]) + + # Now let's do some filtering + # 1. Remove examples with no tests + if not tests: + print(f"No tests found for example {idx}") + return EMPTY_EXAMPLE + + # 2. Remove examples with "image" in prompt, image typically will be in the following format: + # or [image] + if "" in prompt.lower() or "[image]" in prompt.lower(): + print(f"Image found in prompt for example {idx}") + return EMPTY_EXAMPLE + + # 3. Remove examples with no problem description + # Check if prompt starts with unwanted patterns after removing common prefix + check_prompt = prompt + if "Solve the following coding problem using the programming language python:" in check_prompt: + check_prompt = check_prompt.split("Solve the following coding problem using the programming language python:")[1].lstrip() + + if (check_prompt.lower().startswith("example") or + check_prompt.lower().startswith("input") or + check_prompt.lower().startswith("-----input-----")): + print(f"Example starts with unwanted pattern for example {idx}") + return EMPTY_EXAMPLE + + # Handle different test types + if tests[0]["type"] == "function_call": + # Function call tests + fn_name = tests[0]["fn_name"] + + inputs = [] + outputs = [] + for test in tests: + inputs.append(test["input"]) + outputs.append(test["output"]) + + test_code = f"""\ +_inputs = {inputs} +_outputs = {outputs} +import math +def _deep_eq(a, b, tol=1e-5): + if isinstance(a, float) or isinstance(b, float): + return math.isclose(a, b, rel_tol=tol, abs_tol=tol) + if isinstance(a, (list, tuple)): + if len(a) != len(b): return False + return all(_deep_eq(x, y, tol) for x, y in zip(a, b)) + return a == b + +for i, o in zip(_inputs, _outputs): + assert _deep_eq({fn_name}(*i), o[0] if len(o) == 1 else tuple(o)) +""" + + # Validate the solution + full_code = f"{solution}\n{test_code}" + succ, err = code_exec(full_code) + if not succ: + print(f"Test code failed for example {idx}") + print(f"Error: {err}") + return EMPTY_EXAMPLE + + oracle = json.dumps({"functional": test_code}) + + elif tests[0]["type"] == "stdin_stdout": + # STDIN/STDOUT tests + stdin_list = [] + stdout_list = [] + for test in tests: + stdin_list.append(test["input"]) + stdout_list.append(test["output"]) + + # Validate the solution + with ThreadPoolExecutor(max_workers=min(len(stdin_list), 8)) as executor: + futures = [] + for stdin, stdout in zip(stdin_list, stdout_list): + futures.append( + executor.submit( + remote_check_stdio, + solution, + stdin, + stdout, + ) + ) + for future in as_completed(futures): + exec_succ, output, stdin, stdout = future.result() + pass_test = exec_succ and fuzzy_equal(output.strip(), stdout.strip(), verbose=verbose) + if not pass_test: + print(f"Test code failed for example {idx}") + if verbose: + print(f"Input: {stdin}") + print(f"Expected output: {stdout}") + print(f"Actual output: {output}") + return EMPTY_EXAMPLE + + oracle = json.dumps({"inputs": stdin_list, "outputs": stdout_list}) + else: + print(f"Unknown test type: {tests[0]['type']} for example {idx}") + return EMPTY_EXAMPLE + + data = { + "data_source": data_source, + "prompt": [ + {"role": "user", "content": prompt} + ], + "ability": "codegen", + "apply_chat_template": True, + "reward_model": { + "style": "rule", + "ground_truth": oracle, + }, + "extra_info": { + "split": split, + "index": idx, + "reference": solution, + "original_prompt": prompt, + "dataset": "PrimeIntellect", + "function_name": fn_name if tests[0]["type"] == "function_call" else "" + }, + } + + if idx == 0 or idx == 1: + print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) + print(data) + + return data + + return process_fn + + +if __name__ == '__main__': + """Main script execution: parse args, load, process, and save datasets.""" + parser = argparse.ArgumentParser(description="Process and save PrimeIntellect datasets.") + parser.add_argument('--data-dir', type=str, default='data', help='Directory to save the processed datasets.') + parser.add_argument('--domain', default="codegen", + help='Domain of the dataset.') + parser.add_argument('--name', default="primeintellect", + help='Name of the dataset.') + parser.add_argument('--train-sample-size', type=int, default=None, + help='Number of samples to use from training dataset. If None, use all samples.') + parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility') + parser.add_argument('--verbose', type=bool, default=False, help='Whether to print verbose output.') + + args = parser.parse_args() + + set_seed(args.seed) + + data_source = f"{args.domain}__{args.name}" + train_output_dir = os.path.join(args.data_dir, 'train') + os.makedirs(train_output_dir, exist_ok=True) + + # Load the dataset + cache_dir = datasets.config.HF_DATASETS_CACHE + dataset, _ = get_datasets(cache_dir) + + # Process the dataset + process_fn = make_map_fn('train', data_source, args.verbose) + + dataset = dataset.map(function=process_fn, with_indices=True, num_proc=64) + + # Filter out examples where processing failed + dataset = dataset.filter(lambda x: x["data_source"] == data_source) + + # Length filter + try: + tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B") + length_filter = LengthFilter(tokenizer=tokenizer, max_length=4096) + dataset = dataset.filter(lambda x: length_filter.check(x),) + except Exception as e: + print(f"Warning: Could not perform length filtering. Error: {e}") + print("Proceeding without length filtering.") + + # Sample the dataset using utility function + train_dataset = sample_dataset(dataset, args.train_sample_size) + + # Save the dataset using utility function + train_output_path = save_dataset( + dataset=train_dataset, + output_dir=train_output_dir, + filename_prefix=data_source, + sample_size=len(train_dataset) + ) + + print(f"\nDone! \n" + f"Data source: {data_source}\n" + f"Train data saved to {train_output_path} ({len(train_dataset)} samples)") \ No newline at end of file diff --git a/reasoners/data_preprocess/codegen/taco.py b/reasoners/data_preprocess/codegen/taco.py new file mode 100644 index 000000000..31eb8be70 --- /dev/null +++ b/reasoners/data_preprocess/codegen/taco.py @@ -0,0 +1,246 @@ +"""Downloads, processes, and saves TACO dataset.""" + +import os +import argparse +import json +import sys +from concurrent.futures import ThreadPoolExecutor, as_completed +import datasets +import transformers +from datasets import load_dataset, Dataset + +from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset +from verl.utils.data_process.filter import LengthFilter +from verl.utils.reward_score.coder1 import ( + code_exec, + remote_check_stdio, + fuzzy_equal +) + +EMPTY_EXAMPLE = { + "data_source": None, + "prompt": None, + "apply_chat_template": False, + "ability": None, + "reward_model": None, + "extra_info": None +} + + + +def get_datasets(cache_dir: str): + """ + Loads the TACO dataset. + """ + try: + dataset = load_dataset("likaixin/TACO-verified", trust_remote_code=True, split="train", cache_dir=cache_dir) + print(f"TACO dataset: {len(dataset)} examples") + return dataset, None + except Exception as e: + print(f"Error loading dataset: {e}") + return Dataset.from_list([]), None + + +def make_map_fn(split: str, data_source: str) -> callable: + def process_fn(example, idx): + # Create a default "skip" response with all required fields + oracle = json.loads(example["input_output"]) + source = example["source"] + + # Skip poorly formatted examples + if source in ["geeksforgeeks", "leetcode"]: + return EMPTY_EXAMPLE + + # Skip examples with too short descriptions + if len("".join([c for c in example["question"] if c.isalnum()])) < 100: + return EMPTY_EXAMPLE + + # Skip examples with images + if "image" in example["question"].lower() or "\n![" in example["question"]: + return EMPTY_EXAMPLE + + # Build prompt + prompt_pieces = [ + "Solve the programming task below in a Python markdown code block.", + example["question"].strip(), + ] + if example["starter_code"].strip(): + prompt_pieces.append( + "You will use the following starter code to write the solution to the problem and enclose your code within ```python delimiters." + ) + prompt_pieces.append( + f"```python\n{example['starter_code'].strip()}\n```" + ) + else: + prompt_pieces.append( + "Read the inputs from stdin solve the problem and write the answer to stdout (do not directly test on the sample inputs). Enclose your code within ```python delimiters." + ) + + # Process oracle based on format + if "fn_name" in oracle: # Function-based tests + fn_name = oracle["fn_name"] + if source == "leetcode": + fn_name = "Solution()." + fn_name + + test_code = f"""\ +_inputs = {oracle["inputs"]} +_outputs = {oracle["outputs"]} +import math +def _deep_eq(a, b, tol=1e-5): + if isinstance(a, float) or isinstance(b, float): + return math.isclose(a, b, rel_tol=tol, abs_tol=tol) + if isinstance(a, (list, tuple)): + if len(a) != len(b): return False + return all(_deep_eq(x, y, tol) for x, y in zip(a, b)) + return a == b + +for i, o in zip(_inputs, _outputs): +""" + + if source in ["leetcode", "hackerrank"]: + test_code += f" assert _deep_eq({fn_name}(*i), o)" + elif source == "codewars": + test_code += f" assert _deep_eq({fn_name}(*i), o[0])" + else: + print(f"Unknown source: {source}") + return EMPTY_EXAMPLE + + # Verify the solution passes tests + _check_test = example["solutions"][-1] + "\n" + test_code + succ, err = code_exec(_check_test) + if not succ: + print(f"Test code failed for {source}") + return EMPTY_EXAMPLE + + oracle_json = json.dumps({"functional": test_code}) + + elif "inputs" in oracle and "outputs" in oracle: # STDIN/STDOUT tests + stdin_list, stdout_list = oracle["inputs"], oracle["outputs"] + if len(stdin_list) == 0: + return EMPTY_EXAMPLE + + # handle list inputs and normalize line endings + stdin_list = [ + "\n".join(stdin) if isinstance(stdin, list) else stdin + for stdin in stdin_list + ] + stdout_list = [ + ("\n".join(stdout) if isinstance(stdout, list) else stdout).replace("\r\n", "\n") + for stdout in stdout_list + ] + + # Verify the solution passes tests + with ThreadPoolExecutor(max_workers=min(len(stdin_list), 8)) as executor: + futures = [] + for stdin, stdout in zip(stdin_list, stdout_list): + futures.append( + executor.submit( + remote_check_stdio, + example["solutions"][-1], + stdin, + stdout, + ) + ) + for future in as_completed(futures): + exec_succ, output, stdin, stdout = future.result() + pass_test = exec_succ and fuzzy_equal(output.strip(), stdout.strip()) + if not pass_test: + print(f"Test code failed for {source}") + return EMPTY_EXAMPLE + + oracle_json = json.dumps({"inputs": stdin_list, "outputs": stdout_list}) + else: + print(f"Unknown ground truth format: {oracle}") + return EMPTY_EXAMPLE + + # Format the final prompt + prompt = "\n".join(prompt_pieces) + + data = { + "data_source": data_source, + "prompt": [ + {"role": "user", "content": prompt} + ], + "ability": "codegen", + "apply_chat_template": True, + "reward_model": { + "style": "rule", + "ground_truth": oracle_json, + }, + "extra_info": { + "split": split, + "index": idx, + "reference": example["solutions"][0] if example["solutions"] else "", + "original_prompt": prompt, + "dataset": "likaixin/TACO-verified", + "source": source, + }, + } + + if idx == 0 or idx == 1: + print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) + print(data) + + return data + + return process_fn + + +if __name__ == '__main__': + """Main script execution: parse args, load, process, and save datasets.""" + parser = argparse.ArgumentParser(description="Process and save TACO dataset.") + parser.add_argument('--data-dir', default='data', + help='Base directory to save the processed data files.') + parser.add_argument('--name', default="taco", + help='Name of the dataset.') + parser.add_argument('--domain', default="codegen", + help='Domain of the dataset.') + parser.add_argument('--train-sample-size', type=int, default=None, + help='Number of samples to use for training. If None, use all samples.') + parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility') + + args = parser.parse_args() + + set_seed(args.seed) + + data_source = f"{args.domain}__{args.name}" + train_output_dir = os.path.join(args.data_dir, 'train') + test_output_dir = os.path.join(args.data_dir, 'test') + os.makedirs(train_output_dir, exist_ok=True) + os.makedirs(test_output_dir, exist_ok=True) + + # Load the dataset + cache_dir = datasets.config.HF_DATASETS_CACHE + dataset, _ = get_datasets(cache_dir=cache_dir) + + # Process the dataset + process_fn = make_map_fn('train', data_source) + + dataset = dataset.map(function=process_fn, with_indices=True, num_proc=64) + + # Filter out examples where processing failed + dataset = dataset.filter(lambda x: x["data_source"] == data_source) + + # Length filter + try: + tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B") + length_filter = LengthFilter(tokenizer=tokenizer, max_length=4096) + dataset = dataset.filter(lambda x: length_filter.check(x), num_proc=64) + except Exception as e: + print(f"Warning: Could not perform length filtering. Error: {e}") + print("Proceeding without length filtering.") + + # Sample the dataset + dataset = sample_dataset(dataset, args.train_sample_size) + + # Save the datasets + train_output_path = save_dataset( + dataset=dataset, + output_dir=train_output_dir, + filename_prefix=data_source, + sample_size=args.train_sample_size + ) + + print(f"\nDone! \n" + f"Data source: {data_source}\n" + f"Train data saved to {train_output_path} ({len(dataset)} samples)\n") \ No newline at end of file diff --git a/reasoners/data_preprocess/logic/gen_proof.py b/reasoners/data_preprocess/logic/gen_proof.py new file mode 100644 index 000000000..96782ba36 --- /dev/null +++ b/reasoners/data_preprocess/logic/gen_proof.py @@ -0,0 +1,781 @@ + +import os +from typing import List, Set, Dict, Tuple, Optional +import re +from dataclasses import dataclass, field +from copy import deepcopy +import random +import sys +import pdb +import signal +import time +from contextlib import contextmanager + + +class Constraint: + text: str # Human-readable constraint + object1: Optional[int] # First object's encoded value (if applicable) + object2: Optional[int] # Second object's encoded value (if applicable) + constraint_type: str # Type of mathematical constraint + value: Optional[int] # Additional value needed for constraint (e.g., number of birds between) + + def evaluate(self, positions: List[int]) -> bool: + """Evaluate if the constraint is satisfied for given positions.""" + if self.constraint_type == "left_of": + return positions.index(self.object1) < positions.index(self.object2) + elif self.constraint_type == "right_of": + return positions.index(self.object1) > positions.index(self.object2) + elif self.constraint_type == "adjacent_left": + return positions.index(self.object1) + 1 == positions.index(self.object2) + elif self.constraint_type == "adjacent_right": + return positions.index(self.object1) - 1 == positions.index(self.object2) + elif self.constraint_type == "birds_between": + return abs(positions.index(self.object1) - positions.index(self.object2)) - 1 == self.value + elif self.constraint_type == "absolute_position_left": + return positions.index(self.object1) == self.value - 1 + elif self.constraint_type == "absolute_position_right": + return positions.index(self.object1) == len(positions) - self.value + return False + +DEBUG = False + +@dataclass +class SearchNode: + """Data class to represent a search node in the tree.""" + depth: int + arrangement: Dict[str, int] = field(default_factory=dict) + satisfied_constraints: Set[int] = field(default_factory=set) + children: List['SearchNode'] = field(default_factory=list) + is_leaf: bool = False + parent: Optional['SearchNode'] = None + + + def __str__(self): + positions = [-1] * len(self.arrangement) + for bird, pos in self.arrangement.items(): + if pos != -1: + positions[pos] = bird + return f"Arrangement: {positions}, Satisfied: {self.satisfied_constraints}" + +class PuzzleSolver: + def __init__(self, birds: List[str], puzzle_constraints: List[str]=None): + self.birds = birds + self.num_positions = len(birds) + self.raw_constraints = puzzle_constraints + self.constraints: List[Constraint] = [] + self.absolute_constraints: List[Constraint] = [] + self.relative_constraints: List[Constraint] = [] + self.steps: List[str] = [] + self.ground_truth: List[str] = [] + + + def add_step(self, step: str): + """Add a step to the list of reasoning steps.""" + self.steps.append(step) + # print(step) + + + + def parse_constraints(self): + """Parse text constraints into structured Constraint objects.""" + for text in self.raw_constraints: + constraint = self._parse_single_constraint(text) + self.constraints.append(constraint) + + # Separate absolute and relative constraints + if constraint.constraint_type in ['absolute_position_left', 'absolute_position_right']: + self.absolute_constraints.append(constraint) + else: + self.relative_constraints.append(constraint) + + def _parse_single_constraint(self, text:str)->Constraint: + """Parse a single text constraint into a Constraint object.""" + # Pattern for extracting bird names + bird_pattern = '|'.join(self.birds) + # Absolute position patterns + if "from the left" in text or "from the right" in text: + match = re.search(rf"The ({bird_pattern}) is the (\d+)[a-z]+ from the (left|right)", text) + if match: + bird, position, direction = match.groups() + return Constraint( + text=text, + bird1=bird, + bird2=None, + constraint_type=f"absolute_position_{direction}", + value=int(position) + ) + + # Birds between pattern + if "bird" in text and "between" in text: + match = re.search(rf"There (?:is|are) (\d+) birds? between the ({bird_pattern}) and the ({bird_pattern})", text) + if match: + num_birds, bird1, bird2 = match.groups() + return Constraint( + text=text, + bird1=bird1, + bird2=bird2, + constraint_type="birds_between", + value=int(num_birds) + ) + + # Immediately adjacent patterns + if "immediately" in text: + if "to the left of" in text: + match = re.search(f"The ({bird_pattern}) is immediately to the left of the ({bird_pattern})", text) + bird1, bird2 = match.groups() + return Constraint( + text=text, + bird1=bird1, + bird2=bird2, + constraint_type="adjacent_left", + value=None + ) + elif "to the right of" in text: + match = re.search(f"The ({bird_pattern}) is immediately to the right of the ({bird_pattern})", text) + bird1, bird2 = match.groups() + return Constraint( + text=text, + bird1=bird1, + bird2=bird2, + constraint_type="adjacent_right", + value=None + ) + + # General left/right patterns + if "to the left of" in text: + match = re.search(f"The ({bird_pattern}) is to the left of the ({bird_pattern})", text) + bird1, bird2 = match.groups() + return Constraint( + text=text, + bird1=bird1, + bird2=bird2, + constraint_type="left_of", + value=None + ) + elif "to the right of" in text: + match = re.search(f"The ({bird_pattern}) is to the right of the ({bird_pattern})", text) + bird1, bird2 = match.groups() + return Constraint( + text=text, + bird1=bird1, + bird2=bird2, + constraint_type="right_of", + value=None + ) + + raise ValueError(f"Could not parse constraint: {text}") + + + + + def init_helper(self, file_path: str): + """Initialize the solver with the puzzle file.""" + self.ground_truth = self.read_file(file_path) + self.birds = random.shuffle(self.ground_truth)# Shuffle the birds + self.parse_constraints() + + def verify_constraint(self, arrangement: Dict[str, int], constraint: Constraint) -> bool: + """Check if a single constraint is satisfied by current arrangement""" + """Return true only if the constraint is satisfied else false""" + if constraint.constraint_type == "absolute_position_left": + pos = arrangement.get(constraint.bird1, -1) + return pos == constraint.value - 1 + + elif constraint.constraint_type == "absolute_position_right": + pos = arrangement.get(constraint.bird1, -1) + return pos == self.num_positions - constraint.value if pos != -1 else False + + elif constraint.constraint_type == "adjacent_left": + pos1 = arrangement.get(constraint.bird1, -1) + pos2 = arrangement.get(constraint.bird2, -1) + return pos1 != -1 and pos2 != -1 and pos1 + 1 == pos2 + + elif constraint.constraint_type == "adjacent_right": + pos1 = arrangement.get(constraint.bird1, -1) + pos2 = arrangement.get(constraint.bird2, -1) + return pos1 != -1 and pos2 != -1 and pos1 - 1 == pos2 + + elif constraint.constraint_type == "birds_between": + pos1 = arrangement.get(constraint.bird1, -1) + pos2 = arrangement.get(constraint.bird2, -1) + if pos1 != -1 and pos2 != -1: + return abs(pos1 - pos2) - 1 == constraint.value + return False # Can't verify yet + + elif constraint.constraint_type == "left_of": + pos1 = arrangement.get(constraint.bird1, -1) + pos2 = arrangement.get(constraint.bird2, -1) + return pos1 < pos2 if pos1 != -1 and pos2 != -1 else False + + elif constraint.constraint_type == "right_of": + pos1 = arrangement.get(constraint.bird1, -1) + pos2 = arrangement.get(constraint.bird2, -1) + return pos1 > pos2 if pos1 != -1 and pos2 != -1 else False + + return False + + def solve(self) -> Optional[Dict[str, int]]: + """Main solving function""" + # Parse constraints if not done already + if not self.constraints: + self.parse_constraints() + + # Start with empty arrangement + initial_arrangement = {bird: -1 for bird in self.birds} + # First apply absolute constraints + self.add_step("Starting with absolute position constraints:") + for i, constraint in enumerate(self.absolute_constraints): + bird = constraint.bird1 + if constraint.constraint_type == "absolute_position_left": + initial_arrangement[bird] = constraint.value - 1 + self.add_step(f"Placing {bird} at position {constraint.value} from left") + else: # absolute_position_right + pos = self.num_positions - constraint.value + initial_arrangement[bird] = pos + self.add_step(f"Placing {bird} at position {pos + 1} from left ({constraint.value} from right)") + + # Create initial node + initial_node = SearchNode( + arrangement=initial_arrangement, + satisfied_constraints=set(), + depth=0 + ) + + # Start Tree Search + self.add_step("Starting search to find complete arrangement...") + result = self.dfs(initial_node) + + if result: + positions = [-1] * self.num_positions + for bird, pos in result.arrangement.items(): + positions[pos] = bird + if -1 in positions: + index = positions.index(-1) + positions[index] = [bird for bird in self.birds if bird not in positions][0] + self.add_step(f"{positions[index]} is placed at position {index} since only one spot was left") + self.add_step("Found solution!") + self.add_step(f"Final arrangement: {positions}") + return result.arrangement + else: + self.add_step("\nNo solution found!") + return None + + def dfs(self, node: SearchNode) -> Optional[SearchNode]: + """Depth-first search implementation""" + # First verify all constraints are satisfied for current arrangement + if DEBUG: + pdb.set_trace() + print(f"\nDEBUG: Current depth: {node.depth}") + print(f"DEBUG: Current arrangement: {node.arrangement}") + print(f"DEBUG: Satisfied constraints: {node.satisfied_constraints}") + violated_constraint = None + current_constraint = None + current_constraint_index = None + for i, constraint in enumerate(self.relative_constraints): + if i not in node.satisfied_constraints: + consistent = self.verify_constraint(node.arrangement, constraint) + if consistent: + node.satisfied_constraints.add(i) + else: + current_constraint = constraint + current_constraint_index = i + break + + # return if solution found + if not current_constraint: + if DEBUG: print("DEBUG: No more constraints to satisfy!") + node.is_leaf = True + return node + + filled = set(pos for pos in node.arrangement.values() if pos != -1) + empty_positions = sorted(set(range(self.num_positions)) - filled) + if DEBUG: print(f"DEBUG: Empty positions: {empty_positions}") + # empty_positions = [i for bird, i in node.arrangement.items() if i == -1] + # import pdb; pdb.set_trace() + k = len(empty_positions) + if k==0 and current_constraint: + if DEBUG: print("DEBUG: No empty positions left but constraints remain!") + return None + current_arrangement = ["_"] * self.num_positions + for bird, pos in node.arrangement.items(): + if pos != -1: + current_arrangement[pos] = bird + + if current_constraint: + if current_constraint.constraint_type == "left_of": + if DEBUG: pdb.set_trace() + bird1 = current_constraint.bird1 + bird2 = current_constraint.bird2 + pos1 = node.arrangement.get(bird1, -1) + pos2 = node.arrangement.get(bird2, -1) + if pos1 != -1 and pos2 != -1 and pos1 >= pos2: + violated_constraint = current_constraint + self.add_step(f"Current arrangement: {current_arrangement}; Constraint violated: {current_constraint.text}, cannot proceed on this path") + return None + elif pos1 ==-1 and pos2 == -1: + for i in range(k-1): + for j in range(i+1, k): + new_node = deepcopy(node) + new_node.arrangement[bird1] = empty_positions[i] + new_node.arrangement[bird2] = empty_positions[j] + new_node.satisfied_constraints.add(current_constraint_index) + new_node.depth += 1 + new_node.parent = node + node.children.append(new_node) + self.add_step(f"Current arrangement: {current_arrangement}; Trying to place {bird1} to the left of {bird2} at positions {empty_positions[i]} and {empty_positions[j]}") + result = self.dfs(new_node) + if result: + return result + else: + self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird1} to the left of {bird2} at positions {empty_positions[i]} and {empty_positions[j]}") + else: + if pos1 == -1: + for i in range(k): + if empty_positions[i] < pos2: + new_node = deepcopy(node) + new_node.arrangement[bird1] = empty_positions[i] + new_node.satisfied_constraints.add(current_constraint_index) + new_node.depth += 1 + new_node.parent = node + node.children.append(new_node) + self.add_step(f"Current arrangement: {current_arrangement}; Trying to place {bird1} to the left of {bird2} at positions {empty_positions[i]}") + result = self.dfs(new_node) + if result: + return result + else: + self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird1} to the left of {bird2} at positions {empty_positions[i]}") + elif pos2 == -1: + for i in range(k): + if empty_positions[i] > pos1: + new_node = deepcopy(node) + new_node.arrangement[bird2] = empty_positions[i] + new_node.satisfied_constraints.add(current_constraint_index) + new_node.depth += 1 + new_node.parent = node + node.children.append(new_node) + self.add_step(f"Current arrangement: {current_arrangement}; Trying to place {bird2} to the right of {bird1} at positions {empty_positions[i]}") + result = self.dfs(new_node) + if result: + return result + else: + self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird2} to the right of {bird1} at positions {empty_positions[i]}") + elif current_constraint.constraint_type == "right_of": + if DEBUG: pdb.set_trace() + bird1 = current_constraint.bird1 + bird2 = current_constraint.bird2 + pos1 = node.arrangement.get(bird1, -1) + pos2 = node.arrangement.get(bird2, -1) + if pos1 != -1 and pos2 != -1 and pos1 <= pos2: + violated_constraint = current_constraint + self.add_step(f"Current arrangement: {current_arrangement}; Constraint violated: {current_constraint.text}, cannot proceed on this path") + return None + elif pos1 == -1 and pos2 == -1: + for i in range(k-1): + for j in range(i+1, k): + new_node = deepcopy(node) + new_node.arrangement[bird1] = empty_positions[j] + new_node.arrangement[bird2] = empty_positions[i] + new_node.satisfied_constraints.add(current_constraint_index) + new_node.depth += 1 + new_node.parent = node + node.children.append(new_node) + self.add_step(f"Current arrangement: {current_arrangement}; Trying to place {bird1} to the right of {bird2} at positions {empty_positions[j]} and {empty_positions[i]}") + result = self.dfs(new_node) + if result: + return result + else: + self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird1} to the right of {bird2} at positions {empty_positions[j]} and {empty_positions[i]}") + else: + if pos1 == -1: + for i in range(k): + if empty_positions[i] > pos2: + new_node = deepcopy(node) + new_node.arrangement[bird1] = empty_positions[i] + new_node.satisfied_constraints.add(current_constraint_index) + new_node.depth += 1 + new_node.parent = node + node.children.append(new_node) + self.add_step(f"Current arrangement: {current_arrangement}; Trying to place {bird1} to the right of {bird2} at positions {empty_positions[i]}") + result = self.dfs(new_node) + if result: + return result + else: + self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird1} to the right of {bird2} at positions {empty_positions[i]}") + elif pos2 == -1: + for i in range(k): + if empty_positions[i] < pos1: + new_node = deepcopy(node) + new_node.arrangement[bird2] = empty_positions[i] + new_node.satisfied_constraints.add(current_constraint_index) + new_node.depth += 1 + new_node.parent = node + node.children.append(new_node) + self.add_step(f"Current arrangement: {current_arrangement}; Trying to place {bird2} to the left of {bird1} at positions {empty_positions[i]}") + result = self.dfs(new_node) + if result: + return result + else: + self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird2} to the left of {bird1} at positions {empty_positions[i]}") + elif current_constraint.constraint_type == "adjacent_left": + if DEBUG: pdb.set_trace() + bird1 = current_constraint.bird1 + bird2 = current_constraint.bird2 + pos1 = node.arrangement.get(bird1, -1) + pos2 = node.arrangement.get(bird2, -1) + if pos1 != -1 and pos2 != -1 and pos1 + 1 != pos2: + violated_constraint = current_constraint + self.add_step(f"Current arrangement: {current_arrangement}; Constraint violated: {current_constraint.text}, cannot proceed on this path") + return None + elif pos1 == -1 and pos2 == -1: + for i in range(k-1): + for j in range(i+1, k): + if empty_positions[i] + 1 == empty_positions[j]: + new_node = deepcopy(node) + new_node.arrangement[bird1] = empty_positions[i] + new_node.arrangement[bird2] = empty_positions[j] + new_node.satisfied_constraints.add(current_constraint_index) + new_node.depth += 1 + new_node.parent = node + node.children.append(new_node) + self.add_step(f"Current arrangement: {current_arrangement}; Trying to place {bird1} immediately to the left of {bird2} at positions {empty_positions[i]} and {empty_positions[j]}") + result = self.dfs(new_node) + if result: + return result + else: + self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird1} immediately to the left of {bird2} at positions {empty_positions[i]} and {empty_positions[j]}") + else: + if pos1 == -1: + if pos2-1 not in empty_positions: + self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird1} immediately to the left of {bird2} at positions {pos2-1}") + return None + new_node = deepcopy(node) + new_node.arrangement[bird1] = pos2 - 1 + new_node.satisfied_constraints.add(current_constraint_index) + new_node.depth += 1 + new_node.parent = node + node.children.append(new_node) + self.add_step(f"Current arrangement: {current_arrangement}; Trying to place {bird1} immediately to the left of {bird2} at positions {pos2-1}") + result = self.dfs(new_node) + if result: + return result + else: + self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird1} immediately to the left of {bird2} at positions {pos2-1}") + elif pos2 == -1: + if pos1+1 not in empty_positions: + self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird1} immediately to the left of {bird2} at positions {pos1+1}") + return None + new_node = deepcopy(node) + new_node.arrangement[bird2] = pos1 + 1 + new_node.satisfied_constraints.add(current_constraint_index) + new_node.depth += 1 + new_node.parent = node + node.children.append(new_node) + self.add_step(f"Current arrangement: {current_arrangement}; Trying to place {bird2} immediately to the right of {bird1} at positions {pos1+1}") + result = self.dfs(new_node) + if result: + return result + else: + self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird2} immediately to the right of {bird1} at positions {pos1+1}") + elif current_constraint.constraint_type == "adjacent_right": + if DEBUG: pdb.set_trace() + bird1 = current_constraint.bird1 + bird2 = current_constraint.bird2 + pos1 = node.arrangement.get(bird1, -1) + pos2 = node.arrangement.get(bird2, -1) + if pos1 != -1 and pos2 != -1 and pos1 - 1 != pos2: + violated_constraint = current_constraint + self.add_step(f"Current arrangement: {current_arrangement}; Constraint violated: {current_constraint.text}, cannot proceed on this path ") + return None + elif pos1 == -1 and pos2 == -1: + for i in range(k-1): + for j in range(i+1, k): + if empty_positions[i] + 1 == empty_positions[j]: + new_node = deepcopy(node) + new_node.arrangement[bird1] = empty_positions[j] + new_node.arrangement[bird2] = empty_positions[i] + new_node.satisfied_constraints.add(current_constraint_index) + new_node.depth += 1 + new_node.parent = node + node.children.append(new_node) + self.add_step(f"Current arrangement: {current_arrangement}; Trying to place {bird1} immediately to the right of {bird2} at positions {empty_positions[j]} and {empty_positions[i]}") + result = self.dfs(new_node) + if result: + return result + else: + self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird1} immediately to the right of {bird2} at positions {empty_positions[j]} and {empty_positions[i]}") + else: + if pos1 == -1: + if pos2+1 not in empty_positions: + self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird1} immediately to the right of {bird2} at positions {pos2+1}") + return None + new_node = deepcopy(node) + new_node.arrangement[bird1] = pos2 + 1 + new_node.satisfied_constraints.add(current_constraint_index) + new_node.depth += 1 + new_node.parent = node + node.children.append(new_node) + self.add_step(f"Current arrangement: {current_arrangement}; Trying to place {bird1} immediately to the right of {bird2} at positions {pos2+1}") + result = self.dfs(new_node) + if result: + return result + else: + self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird1} immediately to the right of {bird2} at positions {pos2+1}") + elif pos2 == -1: + if pos1-1 not in empty_positions: + self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird1} immediately to the right of {bird2} at positions {pos1-1}") + return None + new_node = deepcopy(node) + new_node.arrangement[bird2] = pos1 - 1 + new_node.satisfied_constraints.add(current_constraint_index) + new_node.depth += 1 + new_node.parent = node + node.children.append(new_node) + self.add_step(f"Current arrangement: {current_arrangement}; Trying to place {bird2} immediately to the left of {bird1} at positions {pos1-1}") + result = self.dfs(new_node) + if result: + return result + else: + self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird2} immediately to the left of {bird1} at positions {pos1-1}") + elif current_constraint.constraint_type == "birds_between": + if DEBUG: pdb.set_trace() + bird1 = current_constraint.bird1 + bird2 = current_constraint.bird2 + pos1 = node.arrangement.get(bird1, -1) + pos2 = node.arrangement.get(bird2, -1) + value = current_constraint.value + if pos1 != -1 and pos2 != -1 and abs(pos1 - pos2) - 1 != value: + violated_constraint = current_constraint + self.add_step(f"Current arrangement: {current_arrangement}; Constraint violated: {current_constraint.text}, cannot proceed on this path") + return None + elif pos1 == -1 and pos2 == -1: + for p1 in empty_positions: + for p2 in empty_positions: + if p1 != p2 and abs(p1-p2) - 1 == value: + # 2 paths bird 1 left of bird 2 and bird 1 right of bird 2 + new_node = deepcopy(node) + new_node.arrangement[bird1] = p1 + new_node.arrangement[bird2] = p2 + new_node.satisfied_constraints.add(current_constraint_index) + new_node.depth += 1 + new_node.parent = node + node.children.append(new_node) + self.add_step(f"Current arrangement: {current_arrangement}; Trying to place {bird1} to the left of {bird2} at positions {p1} and {p2}") + result = self.dfs(new_node) + if result: + return result + else: + self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird1} to the left of {bird2} at positions {p1} and {p2}") + new_node = deepcopy(node) + new_node.arrangement[bird1] = p2 + new_node.arrangement[bird2] = p1 + new_node.satisfied_constraints.add(current_constraint_index) + new_node.depth += 1 + new_node.parent = node + node.children.append(new_node) + self.add_step(f"Current arrangement: {current_arrangement}; Trying to place {bird1} to the right of {bird2} at positions {p2} and {p1}") + result = self.dfs(new_node) + if result: + return result + else: + self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird1} to the right of {bird2} at positions {p2} and {p1}") + else: + if pos1 == -1: + if pos2 - value - 1 in empty_positions: + new_node = deepcopy(node) + new_node.arrangement[bird1] = pos2 - value - 1 + new_node.satisfied_constraints.add(current_constraint_index) + new_node.depth += 1 + new_node.parent = node + node.children.append(new_node) + self.add_step(f"Current arrangement: {current_arrangement}; Trying to place {bird1} to the left of {bird2} at positions {pos2 - value - 1}") + result = self.dfs(new_node) + if result: + return result + else: + self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird1} to the left of {bird2} at positions {pos2 - value - 1}") + if pos2 + value +1 in empty_positions: + new_node = deepcopy(node) + new_node.arrangement[bird1] = pos2 + value + 1 + new_node.satisfied_constraints.add(current_constraint_index) + new_node.depth += 1 + new_node.parent = node + node.children.append(new_node) + self.add_step(f"Current arrangement: {current_arrangement}; Trying to place {bird1} to the right of {bird2} at positions {pos2 + value + 1}") + result = self.dfs(new_node) + if result: + return result + else: + self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird1} to the right of {bird2} at positions {pos2 + value + 1}") + else: + return None + elif pos2 == -1: + if DEBUG: pdb.set_trace() + if pos1 - value - 1 in empty_positions: + new_node = deepcopy(node) + new_node.arrangement[bird2] = pos1 - value - 1 + new_node.satisfied_constraints.add(current_constraint_index) + new_node.depth += 1 + new_node.parent = node + node.children.append(new_node) + self.add_step(f"Current arrangement: {current_arrangement}; Trying to place {bird2} to the left of {bird1} at positions {pos1 - value - 1}") + result = self.dfs(new_node) + if result: + return result + else: + self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird2} to the left of {bird1} at positions {pos1 - value - 1}") + if pos1 + value + 1 in empty_positions: + new_node = deepcopy(node) + new_node.arrangement[bird2] = pos1 + value + 1 + new_node.satisfied_constraints.add(current_constraint_index) + new_node.depth += 1 + new_node.parent = node + node.children.append(new_node) + self.add_step(f"Current arrangement: {current_arrangement}; Trying to place {bird2} to the right of {bird1} at positions {pos1 + value + 1}") + result = self.dfs(new_node) + if result: + return result + else: + self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird2} to the right of {bird1} at positions {pos1 + value + 1}") + else: + return None + return None + + + +def read_file(file_path: str) -> List[str]: + """Read the puzzle file and extract the ground truth and statements.""" + with open(file_path, 'r') as file: + lines = file.readlines() + + constraints = [] + + for line in lines: + if line[0].isdigit(): + match = re.match(r"\d+\.\s*(.+)", line.strip()) + constraints.append(match.group(1)) + # Extract ground truth and statements + start = lines[0].index("[") + ground_truth = lines[0][start:].strip().strip("[]").split(", ") + ground_truth = [item.strip("'") for item in ground_truth] + raw_constraints = constraints + return ground_truth, raw_constraints + +class TimeoutException(Exception): + pass + +@contextmanager +def time_limit(seconds): + """Context manager for timeout""" + def signal_handler(signum, frame): + raise TimeoutException("Solution took too long!") + + # Register signal handler + signal.signal(signal.SIGALRM, signal_handler) + signal.alarm(seconds) # Start timer + + try: + yield + finally: + signal.alarm(0) # Disable timer + +def solve_with_timeout(solver: PuzzleSolver, timeout_seconds: int) -> Optional[Dict[str, int]]: + """Wrapper to call solve() with timeout""" + try: + with time_limit(timeout_seconds): + return solver.solve() + except TimeoutException: + print(f"Solver timed out after {timeout_seconds} seconds!") + return None + +def solve_all_puzzles(num_puzzles: int): + random.seed(0) + timeout_seconds = 30 # Set timeout to 30 seconds + + base_dir = os.path.dirname(os.path.abspath(__file__)) + output_folder = os.path.join(base_dir, "QA_sols") + if not os.path.exists(output_folder): + os.makedirs(output_folder) + + num_sols = 0 + num_timeouts = 0 + + for i in range(num_puzzles): + # filepath = f"generated_QAs/birds_{i}.txt" + + try: + # ground_truth, raw_constraints = read_file(filepath) + + # if len(ground_truth) <= 7: + print(f"Solving puzzle {i}, num_sols: {num_sols}") + birds = deepcopy(ground_truth) + random.shuffle(birds) + sol_file = os.path.join(output_folder, f"sol_{i}.txt") + + try: + with open(sol_file, 'w') as file: + sys.stdout = file + solver = PuzzleSolver(birds=birds, puzzle_constraints=raw_constraints) + solver.parse_constraints() + + start_time = time.time() + final_arrangement = solve_with_timeout(solver, timeout_seconds) + solve_time = time.time() - start_time + + if final_arrangement: + # print(f"Solution found in {solve_time:.2f} seconds.") + num_sols += 1 + else: + # print(f"No solution found or timed out after {solve_time:.2f} seconds.") + if solve_time >= timeout_seconds: + num_timeouts += 1 + # Delete the solution file for timed out puzzles + sys.stdout = sys.__stdout__ + os.remove(sol_file) + continue + + except Exception as e: + print(f"Error solving puzzle {i}: {str(e)}") + finally: + sys.stdout = sys.__stdout__ + + except Exception as e: + print(f"Error reading puzzle {i}: {str(e)}") + continue + + # print(f"Number of solvable puzzles: {num_sols}") + # print(f"Number of timeouts: {num_timeouts}") + + + + + # filepath = f"generated_QAs/birds_923.txt" + # ground_truth, raw_constraints = read_file(filepath) + # birds = deepcopy(ground_truth) + # random.shuffle(birds) + # solver = PuzzleSolver(birds=birds, puzzle_constraints= raw_constraints) + # solver.parse_constraints() + # final_arrangement = solver.solve() + # if not final_arrangement: + # print("No solution found.") + # else: + # print("Solution found.") + + +def solve_puzzle(ground_truth, raw_constraints): + timeout_seconds = 30 + birds = deepcopy(ground_truth) + random.shuffle(birds) + solver = PuzzleSolver(birds=birds, puzzle_constraints=raw_constraints) + solver.parse_constraints() + start_time = time.time() + final_arrangement = solve_with_timeout(solver, timeout_seconds) + solve_time = time.time() - start_time + if final_arrangement: + return solver.steps + else: + return None + + + +if __name__ == "__main__": + solve_all_puzzles(1000) \ No newline at end of file diff --git a/reasoners/data_preprocess/logic/graph_dataset_gen/core/LICENSE b/reasoners/data_preprocess/logic/graph_dataset_gen/core/LICENSE new file mode 100644 index 000000000..924bb45c8 --- /dev/null +++ b/reasoners/data_preprocess/logic/graph_dataset_gen/core/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "{}" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright 2017 Abulhair Saparov + + 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. diff --git a/reasoners/data_preprocess/logic/graph_dataset_gen/core/README.md b/reasoners/data_preprocess/logic/graph_dataset_gen/core/README.md new file mode 100644 index 000000000..7a253f0f4 --- /dev/null +++ b/reasoners/data_preprocess/logic/graph_dataset_gen/core/README.md @@ -0,0 +1,21 @@ +Documentation and code examples are available at . The repository is located at . + +This repository is a small collection of general-purpose data structures and algorithms. Originally, I developed this library to explore and tinker with different programming paradigms and design patterns. Since then, the code has evolved into a collection of generic data structures, using [templates](http://en.cppreference.com/w/cpp/language/templates) for polymorphism. + +### Usage and Dependencies + +To use the code, simply download the files into a folder named "core". Add this folder to the include path for the compiler, for example by using the `-I` flag. + +The code makes use of the [xxhash](https://github.com/Cyan4973/xxHash) library as its default hash function implementation, but the required files are included in the repository. Otherwise, there are no dependencies on external libraries. The code makes use of `C++11` and is regularly tested with `gcc 11` but I have previously compiled it with `gcc 4.8`, `clang 4.0`, and `Microsoft Visual C++ 14.0 (2015)`. The code is intended to be platform-independent, so please create an issue if there are any compilation bugs. + +### Overview + +The library currently implements the following data structures: + - [array](https://asaparov.org/docs/core/array.h.html#struct%20array) is a self-expanding sequential array, + - [hash_set](https://asaparov.org/docs/core/map.h.html#struct%20hash_set) is an unordered associative set of unique elements, implemented using a hashtable, + - [hash_map](https://asaparov.org/docs/core/map.h.html#struct%20hash_map) is an unordered associative map of key-value pairs, where the keys are unique, implemented using a hashtable, + - [array_map](https://asaparov.org/docs/core/map.h.html#struct%20array_map) is a sequential map of key-value pairs, where the keys are unique. + +In addition, the library implements a handful of sorting algorithms in [array.h](https://asaparov.org/docs/core/array.h.html). Set operations, such as union, intersection, and subtraction, can also be found in [array.h](https://asaparov.org/docs/core/array.h.html), which operate on sets represented as sorted arrays. The library implements serialization/deserialization with the functions `read`, `write`, and `print` for all data structures using a very regular pattern involving ["scribes"](https://asaparov.org/docs/core/io.h.html#scribes) as described in [io.h](https://asaparov.org/docs/core/io.h.html). Functions and structures useful for lexical analysis are provided in [lex.h](https://asaparov.org/docs/core/lex.h.html). [random.h](https://asaparov.org/docs/core/random.h.html) provides pseudo-random number generation and sampling from a handful of built-in distributions. + +The implementations of every data structure are transparent, with `public` visibility for the underlying fields, which gives the user more control over data structure behavior. All procedures and structures do *not* use exceptions. Rather, errors are returned, typically as a `bool` type, where `true` indicates success and `false` indicates failure. I observed, in practice, that the overhead of this kind of error handling is negligible when no errors are thrown, and the overhead is much smaller than exception handling when errors do occur. The library does not provide automatic memory management, and expects the user to manage memory appropriately. The `free` function is implemented for all data structures, but it does not automatically deallocate child objects in the container classes. The library does not use any thread synchronization methods to guarantee thread safety, and the user must implement their own synchronization or avoid simultaneous writes to data structures. diff --git a/reasoners/data_preprocess/logic/graph_dataset_gen/core/array.h b/reasoners/data_preprocess/logic/graph_dataset_gen/core/array.h new file mode 100644 index 000000000..d361f6faa --- /dev/null +++ b/reasoners/data_preprocess/logic/graph_dataset_gen/core/array.h @@ -0,0 +1,2429 @@ +/** + * \file array.h + * + * This file contains the implementation of the `array` data structure, which + * is a resizeable ordered collection of elements with generic type `T`. In + * addition, the file contains global helper functions for working with arrays, + * such as resizing, linear and binary search, insertion sort, and quicksort. + * The `pair` structure is also defined here. + * + * This file also contains functions that perform set operations on sorted + * arrays and native arrays, such as union, intersection, subtraction, and + * subset. + * + * + */ + +#ifndef ARRAY_H_ +#define ARRAY_H_ + +#include +#include +#include +#include + +#include + +/** + * The multiplicative factor by which array capacity is changed. + */ +#define RESIZE_FACTOR 2 + + +namespace core { + + +namespace detail { + template static auto test_resizeable(int32_t) -> + decltype(void(std::declval().on_resize()), std::true_type{}); + template static auto test_resizeable(int64_t) -> std::false_type; +} + +/** + * This type trait is [true_type](http://en.cppreference.com/w/cpp/types/integral_constant) + * if and only if `T` is class with a public function `void on_resize()`. + */ +template struct is_resizeable : decltype(core::detail::test_resizeable(0)){}; + +/** + * Resizes the given native array `data` with the requested capacity `new_capacity`. + * \tparam SizeType a type that satisfies [is_integral](http://en.cppreference.com/w/cpp/types/is_integral). + * \tparam T the generic type of the elements in `data`. + * \param data the array to resize. + * \param new_capacity the requested size. + * \return `true` on success; `data` may point to a different memory address. + * \return `false` on failure; `data` is unchanged. + */ +template::value>::type* = nullptr> +inline bool resize(T*& data, const SizeType& new_capacity) { + T* new_data = (T*) realloc(static_cast(data), new_capacity * sizeof(T)); + if (new_data == NULL) { + fprintf(stderr, "resize ERROR: Out of memory.\n"); + return false; + } + data = new_data; + return true; +} + +/** + * This function multiplies `capacity` by #RESIZE_FACTOR. It then repeats this + * until `capacity >= new_length`. + */ +template::value>::type* = nullptr> +inline void expand_capacity(SizeType& capacity, size_t new_length) { + do { + /* increase the size of the underlying array */ + capacity *= RESIZE_FACTOR; + } while (new_length > capacity); +} + +/** + * For a given requested length `new_length`, this function calls + * expand_capacity() to determine the new `capacity` of the native array + * `data`. The function then attempts to resize the array with this capacity. + * Note this function does not check whether `new_length <= capacity`. + */ +template::value>::type* = nullptr> +inline bool expand(T*& data, SizeType& capacity, size_t new_length) { + expand_capacity(capacity, new_length); + return resize(data, capacity); +} + +/** + * For a given requested length `new_length`, this function expands the native + * array `data` by factors of #RESIZE_FACTOR until `capacity >= new_length`. + * If initially `new_length <= capacity`, this function does nothing. + */ +template::value>::type* = nullptr> +inline bool ensure_capacity(T*& data, SizeType& capacity, size_t new_length) +{ + if (new_length <= capacity) + return true; + SizeType new_capacity = capacity; + if (!expand(data, new_capacity, new_length)) + return false; + capacity = new_capacity; + return true; +} + +/** + * Performs a linear search through the array `data` to find the smallest index + * `i >= start` such that `element == data[i]`. + * \tparam Key a generic type for which operator `==` is defined for arguments of type `Key` and `T`. + * \tparam T a generic type for which operator `==` is defined for arguments of type `Key` and `T`. + * \return an index in `start, start + 1, ..., length - 1` if the element was found. + * \return `length` if the element was not found. + */ +template::value>::type* = nullptr> +inline SizeType index_of(const Key& element, const T* data, + const SizeType& length, const SizeType& start = 0) +{ + for (SizeType i = start; i < length; i++) + if (element == data[i]) + return i; + return length; +} + +/** + * Performs a linear search through the array `data` to find the largest index + * `i` such that `element == data[i]`. + * \tparam Key a generic type for which operator `==` is defined for arguments of type `Key` and `T`. + * \tparam T a generic type for which operator `==` is defined for arguments of type `Key` and `T`. + * \return an index in `0, 1, ..., length - 1` if the element was found. + * \return `static_cast(-1)` if the element was not found. + */ +template::value>::type* = nullptr> +inline unsigned int last_index_of(const Key& element, const T* data, const SizeType& length) +{ + unsigned int i = length; + while (i != 0) { + i--; + if (element == data[i]) + return i; + } + return static_cast(-1); +} + +/** + * A resizeable sequence of objects, stored contiguously, each with generic + * type `T`. This structure does not automatically initialize or free its + * elements, and so the user must appropriately free each element before the + * array is destroyed. + * + * In the following example, we demonstrate a simple use-case of array. Here, + * `a` is automatically freed by the destructor since it was initialized on the + * stack. The expected output is `-1 -1 0 3 `. + * + * ```{.cpp} + * #include + * #include + * using namespace core; + * + * int main() { + * array a = array(8); + * a.add(-1); a.add(-4); + * a.add(3); a.add(0); + * a.remove(1); + * + * printf("%d ", a[0]); + * for (int element : a) + * printf("%d ", element); + * } + * ``` + * + * + * However, if `a` is not allocated on the stack, the destructor will not be + * automatically called, and so it must be freed manually using `core::free` or + * `array::free`. + * + * ```{.cpp} + * #include + * #include + * using namespace core; + * + * int main() { + * array& a = *((array*) alloca(sizeof(array))); + * array_init(a, 8); + * a.add(-1); a.add(-4); + * a.add(3); a.add(0); + * a.remove(1); + * + * printf("%d ", a[0]); + * for (int element : a) + * printf("%d ", element); + * free(a); + * } + * ``` + * + * + * Also note that a number of member functions require that `T` be + * [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). + * In other cases, elements should be added manually to the underlying native + * array array::data. This structure also defines array::begin and array::end, + * similar to Standard Template Library iterators, which enables the use of the + * range-based for loop in the example below. In this example, the expected + * output is `first second `. + * + * ```{.cpp} + * #include + * #include + * #include + * using namespace core; + * + * struct custom_string { + * char* buffer; + * + * static void free(custom_string& s) { + * core::free(s.buffer); + * } + * }; + * + * bool init(custom_string& s, const char* src) { + * s.buffer = (char*) malloc(sizeof(char) * (strlen(src) + 1)); + * if (s.buffer == NULL) + * return false; + * memcpy(s.buffer, src, sizeof(char) * (strlen(src) + 1)); + * return true; + * } + * + * int main() { + * array a = array(8); + * init(a[0], "first"); + * init(a[1], "second"); + * a.length = 2; + * + * for (const custom_string& s : a) + * printf("%s ", s.buffer); + * for (custom_string& s : a) + * free(s); + * } + * ``` + * Note in the above example that since the array struct does not automatically + * free its elements, they must be freed manually. + */ +template +struct array { + /** + * The underlying native array of elements. + */ + T* data; + + /** + * The length of the array. + */ + size_t length; + + /** + * The capacity of array::data. + */ + size_t capacity; + + /** + * Constructs an array with zero size and the given `initial_capacity`. + */ + array(size_t initial_capacity) + { + if (!initialize(initial_capacity)) + exit(EXIT_FAILURE); + } + + ~array() { free(); } + + /** + * Returns a reference to the element at the given `index`. No bounds-checking is performed. + */ + inline T& operator[] (size_t index) { + return data[index]; + } + + /** + * Returns a const reference to the element at the given `index`. No bounds-checking is performed. + */ + inline const T& operator[] (size_t index) const { + return data[index]; + } + + /** + * Sets the length of the array to `0`. Any elements are not freed. + */ + inline void clear() + { + length = 0; + } + + /** + * Moves the last element in the array to the position given by `index` and + * decrements array::length by `1`. The element initially at `index` is not + * freed. + */ + void remove(size_t index) + { + core::move(data[length - 1], data[index]); + length--; + } + + /** + * For a given requested length `new_length`, this function expands the + * array by factors of #RESIZE_FACTOR until `array::capacity >= new_length`. + * If initially `new_length <= array::capacity`, this function does + * nothing. If the resize operation moved array::data to a new memory + * address, and `T` satisfies is_resizeable, then `x.on_resize()` is called + * for every element `x` in the array. + */ + template::value && is_resizeable::value>::type* = nullptr> + bool ensure_capacity(size_t new_length) { + const T* old_data = data; + if (!core::ensure_capacity(data, capacity, new_length)) return false; + if (data != old_data) { + for (unsigned int i = 0; i < length; i++) + data[i].on_resize(); + } + return true; + } + + template::value && !is_resizeable::value>::type* = nullptr> + bool ensure_capacity(size_t new_length) { + return core::ensure_capacity(data, capacity, new_length); + } + + /** + * Adds the given native array of elements to this structure. This function + * uses [memcpy](http://en.cppreference.com/w/cpp/string/byte/memcpy), + * and so it should not be used if the elements are not + * [TriviallyCopyable](https://en.cppreference.com/w/cpp/named_req/TriviallyCopyable). + * In such a case, addition should be performed manually using the public fields. + * \tparam T is [TriviallyCopyable](https://en.cppreference.com/w/cpp/named_req/TriviallyCopyable). + */ + bool append(const T* elements, size_t size) + { + if (!ensure_capacity(length + size)) + return false; + memcpy(&data[length], elements, sizeof(T) * size); + length += size; + return true; + } + + /** + * Calls array::index_of to determine whether `element` exists in this array. + * \tparam Key a generic type for which operator `==` is defined for arguments of type `Key` and `T`. + */ + template + inline bool contains(const Key& element) const { + return index_of(element) < length; + } + + /** + * Performs a linear search of the array to find the smallest index `i` + * such that `element == array::data[i]`. + * \tparam Key a generic type for which operator `==` is defined for arguments of type `Key` and `T`. + * \return an index in `0, 1, ..., array::length - 1` if the element was found. + * \return `array::length` if the element was not found. + */ + template + inline unsigned int index_of(const Key& element) const { + return core::index_of(element, data, (unsigned int) length); + } + + /** + * Performs a linear search through the array to find the smallest index + * `i >= start` such that `element == array::data[i]`. + * \tparam Key a generic type for which operator `==` is defined for arguments of type `Key` and `T`. + * \return an index in `start, start + 1, ..., length - 1` if the element was found. + * \return `length` if the element was not found. + */ + template + inline unsigned int index_of(const Key& element, SizeType start) const { + return core::index_of(element, data, (unsigned int) length, start); + } + + /** + * Returns a reference to `array::data[0]`, ignoring any bounds. + */ + T& first() + { + return data[0]; + } + + /** + * Returns a const reference to `array::data[0]`, ignoring any bounds. + */ + const T& first() const + { + return data[0]; + } + + /** + * Returns a reference to `array::data[array::length - 1]`, ignoring any bounds. + */ + T& last() + { + return data[length - 1]; + } + + /** + * Returns a const reference to `array::data[array::length - 1]`, ignoring any bounds. + */ + const T& last() const + { + return data[length - 1]; + } + + /** + * Adds the given element to this array, increasing its capacity if + * necessary. The assignment operator performs the addition, and so this + * function should not be used if `T` is not + * [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). + * In such a case, addition should be performed manually using the public fields. + * \tparam T is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). + */ + bool add(const T& element) + { + if (!ensure_capacity(length + 1)) + return false; + data[length] = element; + length++; + return true; + } + + /** + * Returns the element at `array::length - 1` and decrements array::length + * by `1`, ignoring any bounds. + */ + T pop() + { + length--; + return data[length]; + } + + /** + * Returns an iterator to the beginning of the array. + */ + inline T* begin() { + return data; + } + + /** + * Returns an iterator to the end of the array. + */ + inline T* end() { + return data + length; + } + + /** + * Returns a const iterator to the beginning of the array. + */ + inline const T* begin() const { + return data; + } + + /** + * Returns a const iterator to the end of the array. + */ + inline const T* end() const { + return data + length; + } + + /** + * Copies the underlying fields from `src` into `dst`, effectively moving + * the array from `src` into `dst`. + */ + static inline void move(const array& src, array& dst) { + dst.length = src.length; + dst.capacity = src.capacity; + dst.data = src.data; + } + + template + static inline long unsigned int size_of(const array& a, const Metric& metric) { + long unsigned int sum = core::size_of(a.capacity) + core::size_of(a.length); + for (unsigned int i = 0; i < a.length; i++) + sum += core::size_of(a.data[i], metric); + return sum + (a.capacity - a.length) * sizeof(T); + } + + /** + * Frees array::data. This should not be used if `a` was constructed on the + * stack, as the destructor will automatically free array::data. The + * elements of `a` are not freed. + */ + static inline void free(array& a) { a.free(); } + +private: + inline bool initialize(size_t initial_capacity) + { +#if !defined(NDEBUG) + if (initial_capacity == 0) + fprintf(stderr, "array.initialize WARNING: Initial capacity is zero.\n"); +#endif + + capacity = initial_capacity; + length = 0; + data = (T*) malloc(sizeof(T) * capacity); + if (data == NULL) { + fprintf(stderr, "array.initialize ERROR: Out of memory.\n"); + return false; + } + return true; + } + + inline void free() { + if (data != NULL) { + core::free(data); + data = NULL; + } + } + + template + friend bool array_init(array& m, size_t initial_capacity); +}; + +/** + * Initializes the array `m` with the given `initial_capacity`. + */ +template +bool array_init(array& m, size_t initial_capacity) { + return m.initialize(initial_capacity); +} + +/** + * Returns array::length of `m`. + */ +template +inline size_t size(const array& m) { + return m.length; +} + +/** + * Swaps the underlying buffers of the given arrays. + */ +template +inline void swap(array& a, array& b) +{ + T* temp = a.data; + a.data = b.data; + b.data = temp; + + size_t swap = a.length; + a.length = b.length; + b.length = swap; + + swap = a.capacity; + a.capacity = b.capacity; + b.capacity = swap; +} + +template +inline long unsigned int size(const array& a, const Metric& metric) { + long unsigned int sum = size(a.capacity) + size(a.length); + for (unsigned int i = 0; i < a.length; i++) + sum += size(a.data[i], metric); + return sum + (a.capacity - a.length) * sizeof(T); +} + +/** + * Compares the two given arrays and returns true if and only if `a.length == b.length` + * and there is no index `i` such that `a.data[i] != b.data[i]`. + */ +template +inline bool operator == (const array& a, const array& b) { + if (a.length != b.length) + return false; + for (unsigned int i = 0; i < a.length; i++) + if (a.data[i] != b.data[i]) + return false; + return true; +} + +/** + * Compares the two given arrays and returns true if and only if `a.length != b.length` + * or there is some index `i` such that `a.data[i] != b.data[i]`. + */ +template +inline bool operator != (const array& a, const array& b) { + if (a.length != b.length) + return true; + for (unsigned int i = 0; i < a.length; i++) + if (a.data[i] != b.data[i]) + return true; + return false; +} + +/** + * Performs insertion sort on the given native array `keys` with given `length`. + * \tparam T is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable), + * [CopyConstructible](https://en.cppreference.com/w/cpp/named_req/CopyConstructible), + * and [LessThanComparable](https://en.cppreference.com/w/cpp/named_req/LessThanComparable). + */ +template +void insertion_sort(T* keys, unsigned int length) +{ + for (unsigned int i = 1; i < length; i++) { + T item = keys[i]; + unsigned int hole = i; + + while (hole > 0 && item < keys[hole - 1]) { + keys[hole] = keys[hole - 1]; + hole--; + } + + keys[hole] = item; + } +} + +/** + * Performs insertion sort on the given native array `keys` with given `length` + * and `sorter`. + * \tparam T satisfies is_moveable, and for which a function + * `bool less_than(const T&, const T&, const Sorter&)` is defined. + */ +template::value>::type* = nullptr> +void insertion_sort(T* keys, unsigned int length, const Sorter& sorter) +{ + T& item = *((T*) malloc(sizeof(T))); + for (unsigned int i = 1; i < length; i++) { + move(keys[i], item); + unsigned int hole = i; + + while (hole > 0 && less_than(item, keys[hole - 1], sorter)) { + move(keys[hole - 1], keys[hole]); + hole--; + } + + move(item, keys[hole]); + } + free(&item); +} + +/** + * Performs insertion sort on the given native arrays `keys` and `values` with given `length`. + * \tparam K is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable), + * [CopyConstructible](https://en.cppreference.com/w/cpp/named_req/CopyConstructible), + * and [LessThanComparable](https://en.cppreference.com/w/cpp/named_req/LessThanComparable). + * \tparam V is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable), + * [CopyConstructible](https://en.cppreference.com/w/cpp/named_req/CopyConstructible). + */ +template +void insertion_sort(K* keys, V* values, unsigned int length) +{ + for (unsigned int i = 1; i < length; i++) { + K item = keys[i]; + V value = values[i]; + unsigned int hole = i; + + while (hole > 0 && item < keys[hole - 1]) { + keys[hole] = keys[hole - 1]; + values[hole] = values[hole - 1]; + hole--; + } + + keys[hole] = item; + values[hole] = value; + } +} + +/** + * Performs insertion sort on the given native arrays `keys` and `values` with + * given `length` and `sorter`. + * \tparam K satisfies is_moveable, and for which a function + * `bool less_than(const K&, const K&, const Sorter&)` is defined. + * \tparam V satisfies is_moveable. + */ +template::value>::type* = nullptr> +void insertion_sort(K* keys, V* values, unsigned int length, const Sorter& sorter) +{ + K& item = *((K*) malloc(sizeof(K))); + V& value = *((V*) malloc(sizeof(V))); + for (unsigned int i = 1; i < length; i++) { + move(keys[i], item); + move(values[i], value); + unsigned int hole = i; + + while (hole > 0 && less_than(item, keys[hole - 1], sorter)) { + move(keys[hole - 1], keys[hole]); + move(values[hole - 1], values[hole]); + hole--; + } + + move(item, keys[hole]); + move(value, values[hole]); + } + free(&item); free(&value); +} + +/** + * Performs insertion sort on the given array `keys`. + * \tparam T is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable), + * [CopyConstructible](https://en.cppreference.com/w/cpp/named_req/CopyConstructible), + * and [LessThanComparable](https://en.cppreference.com/w/cpp/named_req/LessThanComparable). + */ +template +inline void insertion_sort(array& keys) { + insertion_sort(keys.data, (unsigned int) keys.length); +} + +/** + * Performs insertion sort on the given array `keys` with the given `sorter`. + * \tparam T satisfies is_moveable, and for which a function + * `bool less_than(const T&, const T&, const Sorter&)` is defined. + */ +template::value>::type* = nullptr> +inline void insertion_sort(array& keys, const Sorter& sorter) +{ + insertion_sort(keys.data, (unsigned int) keys.length, sorter); +} + +/** + * Performs insertion sort on the given arrays `keys` and `values`. + * \tparam K is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable), + * [CopyConstructible](https://en.cppreference.com/w/cpp/named_req/CopyConstructible), + * and [LessThanComparable](https://en.cppreference.com/w/cpp/named_req/LessThanComparable). + * \tparam V is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable), + * [CopyConstructible](https://en.cppreference.com/w/cpp/named_req/CopyConstructible). + */ +template +inline void insertion_sort(array& keys, array& values) { + insertion_sort(keys.data, values.data, (unsigned int) keys.length); +} + +/** + * Performs insertion sort on the given arrays `keys` and `values` with the + * given `sorter`. + * \tparam K satisfies is_moveable, and for which a function + * `bool less_than(const K&, const K&, const Sorter&)` is defined. + * \tparam V satisfies is_moveable. + */ +template::value>::type* = nullptr> +inline void insertion_sort(array& keys, array& values, const Sorter& sorter) +{ + insertion_sort(keys.data, values.data, (unsigned int) keys.length, sorter); +} + +/** + * Reverses the order of the elements in the given native `array` with given `length`. + * \tparam T satisfies is_swappable. + */ +template +void reverse(T* array, unsigned int length) { + for (unsigned int i = 0; i < length / 2; i++) { + unsigned int other = length - i - 1; + swap(array[i], array[other]); + } +} + +/** + * Reverses the order of the elements in the given `array`. + * \tparam T satisfies is_swappable. + */ +template +inline void reverse(array& array) { + reverse(array.data, (unsigned int) array.length); +} + +template +inline const T& get_pivot(T* array, unsigned int start, unsigned int end) { + return array[(end + start) / 2]; +} + +template +inline void quick_sort_partition(T* array, + unsigned int start, unsigned int end, + unsigned int& l, unsigned int& r) +{ + const T p = get_pivot(array, start, end); + while (true) { + while (array[l] < p) + l++; + while (p < array[r]) + r--; + if (l == r) { + l++; + if (r > 0) r--; + return; + } else if (l > r) return; + swap(array[l++], array[r--]); + } +} + +template::value>::type* = nullptr> +inline void quick_sort_partition( + T* array, unsigned int start, unsigned int end, + unsigned int& l, unsigned int& r, const Sorter& sorter) +{ + T& p = *((T*) malloc(sizeof(T))); + move(get_pivot(array, start, end), p); + while (true) { + while (less_than(array[l], p, sorter)) + l++; + while (less_than(p, array[r], sorter)) + r--; + if (l == r) { + l++; + if (r > 0) r--; + break; + } else if (l > r) break; + swap(array[l++], array[r--]); + } + free(&p); +} + +template +inline void quick_sort_partition(K* keys, V* values, + unsigned int start, unsigned int end, unsigned int& l, unsigned int& r) +{ + const K p = get_pivot(keys, start, end); + while (true) { + while (keys[l] < p) + l++; + while (p < keys[r]) + r--; + if (l == r) { + l++; + if (r > 0) r--; + return; + } else if (l > r) return; + swap(values[l], values[r]); + swap(keys[l++], keys[r--]); + } +} + +template::value>::type* = nullptr> +inline void quick_sort_partition( + K* keys, V* values, unsigned int start, unsigned int end, + unsigned int& l, unsigned int& r, const Sorter& sorter) +{ + K& p = *((K*) malloc(sizeof(K))); + move(get_pivot(keys, start, end), p); + while(true) { + while (less_than(keys[l], p, sorter)) + l++; + while (less_than(p, keys[r], sorter)) + r--; + if (l == r) { + l++; + if (r > 0) r--; + break; + } else if (l > r) break; + swap(values[l], values[r]); + swap(keys[l++], keys[r--]); + } + free(&p); +} + +template +void quick_sort(T* array, unsigned int start, unsigned int end) +{ + if (start >= end) + return; + unsigned int l = start, r = end; + quick_sort_partition(array, start, end, l, r); + quick_sort(array, start, r); + quick_sort(array, l, end); +} + +template::value>::type* = nullptr> +void quick_sort(T* array, unsigned int start, unsigned int end, const Sorter& sorter) +{ + if (start >= end) + return; + unsigned int l = start, r = end; + quick_sort_partition(array, start, end, l, r, sorter); + quick_sort(array, start, r, sorter); + quick_sort(array, l, end, sorter); +} + +template +void quick_sort(K* keys, V* values, unsigned int start, unsigned int end) +{ + if (start >= end) + return; + unsigned int l = start, r = end; + quick_sort_partition(keys, values, start, end, l, r); + quick_sort(keys, values, start, r); + quick_sort(keys, values, l, end); +} + +template::value>::type* = nullptr> +void quick_sort(K* keys, V* values, + unsigned int start, unsigned int end, const Sorter& sorter) +{ + if (start >= end) + return; + unsigned int l = start, r = end; + quick_sort_partition(keys, values, start, end, l, r, sorter); + quick_sort(keys, values, start, r, sorter); + quick_sort(keys, values, l, end, sorter); +} + +/** + * Performs Quicksort on the given native array `keys` with given `length`. + * This function assumes `length > 0`. If the preprocessor `NDEBUG` is not + * defined, this function outputs a warning when `length == 0`. + * \tparam T is [CopyConstructible](https://en.cppreference.com/w/cpp/named_req/CopyConstructible), + * [LessThanComparable](https://en.cppreference.com/w/cpp/named_req/LessThanComparable), + * and is_swappable. + */ +template +inline void quick_sort(T* keys, unsigned int length) { +#if !defined(NDEBUG) + if (length == 0) { + fprintf(stderr, "quick_sort WARNING: Length is zero.\n"); + return; + } +#endif + quick_sort(keys, 0, length - 1); +} + +/** + * Performs Quicksort on the given native array `keys` with given `length` and + * `sorter`. This function assumes `length > 0`. If the preprocessor `NDEBUG` + * is not defined, this function outputs a warning when `length == 0`. + * \tparam T a generic type that satisfies is_swappable, and for which a + * function `bool less_than(const T&, const T&, const Sorter&)` is + * defined. + */ +template::value>::type* = nullptr> +inline void quick_sort(T* keys, unsigned int length, const Sorter& sorter) +{ +#if !defined(NDEBUG) + if (length == 0) { + fprintf(stderr, "quick_sort WARNING: Length is zero.\n"); + return; + } +#endif + quick_sort(keys, 0, length - 1, sorter); +} + +/** + * Performs Quicksort on the given native arrays `keys` and `values` with given + * `length`. This function assumes `length > 0`. If the preprocessor `NDEBUG` + * is not defined, this function outputs a warning when `length == 0`. + * \tparam K is [CopyConstructible](https://en.cppreference.com/w/cpp/named_req/CopyConstructible), + * [LessThanComparable](https://en.cppreference.com/w/cpp/named_req/LessThanComparable), + * and is_swappable. + * \tparam V satisfies is_swappable. + */ +template +inline void quick_sort(K* keys, V* values, unsigned int length) { +#if !defined(NDEBUG) + if (length == 0) { + fprintf(stderr, "quick_sort WARNING: Length is zero.\n"); + return; + } +#endif + quick_sort(keys, values, 0, length - 1); +} + +/** + * Performs Quicksort on the given native arrays `keys` and `values` with given + * `length` and `sorter`. This function assumes `length > 0`. If the preprocessor + * `NDEBUG` is not defined, this function outputs a warning when `length == 0`. + * \tparam K a generic type that satisfies is_swappable, and for which a + * function `bool less_than(const K&, const K&, const Sorter&)` is + * defined. + * \tparam V satisfies is_swappable. + */ +template::value>::type* = nullptr> +inline void quick_sort(K* keys, V* values, unsigned int length, const Sorter& sorter) +{ +#if !defined(NDEBUG) + if (length == 0) { + fprintf(stderr, "quick_sort WARNING: Length is zero.\n"); + return; + } +#endif + quick_sort(keys, values, 0, length - 1, sorter); +} + +/** + * Performs Quicksort on the given array `keys`. This function assumes + * `length > 0`. If the preprocessor `NDEBUG` is not defined, this function + * outputs a warning when `length == 0`. + * \tparam T is [CopyConstructible](https://en.cppreference.com/w/cpp/named_req/CopyConstructible), + * [LessThanComparable](https://en.cppreference.com/w/cpp/named_req/LessThanComparable), + * and is_swappable. + */ +template +inline void quick_sort(array& keys) { +#if !defined(NDEBUG) + if (keys.length == 0) { + fprintf(stderr, "quick_sort WARNING: Length is zero.\n"); + return; + } +#endif + quick_sort(keys.data, 0, (unsigned int) keys.length - 1); +} + +/** + * Performs Quicksort on the given native array `keys` with the given `sorter`. + * This function assumes `length > 0`. If the preprocessor `NDEBUG` is not + * defined, this function outputs a warning when `length == 0`. + * \tparam T a generic type that satisfies is_swappable, and for which a + * function `bool less_than(const T&, const T&, const Sorter&)` is + * defined. + */ +template::value>::type* = nullptr> +inline void quick_sort(array& keys, const Sorter& sorter) +{ +#if !defined(NDEBUG) + if (keys.length == 0) { + fprintf(stderr, "quick_sort WARNING: Length is zero.\n"); + return; + } +#endif + quick_sort(keys.data, 0, (unsigned int) keys.length - 1, sorter); +} + +/** + * Performs Quicksort on the given arrays `keys` and `values`. This function + * assumes `length > 0`. If the preprocessor `NDEBUG` is not defined, this + * function outputs a warning when `length == 0`. + * \tparam K is [CopyConstructible](https://en.cppreference.com/w/cpp/named_req/CopyConstructible), + * [LessThanComparable](https://en.cppreference.com/w/cpp/named_req/LessThanComparable), + * and is_swappable. + * \tparam V satisfies is_swappable. + */ +template +inline void quick_sort(array& keys, array& values) { +#if !defined(NDEBUG) + if (keys.length == 0) { + fprintf(stderr, "quick_sort WARNING: Length is zero.\n"); + return; + } +#endif + quick_sort(keys.data, values.data, 0, (unsigned int) keys.length - 1); +} + +/** + * Performs Quicksort on the given arrays `keys` and `values` with the given + * `sorter`. This function assumes `length > 0`. If the preprocessor `NDEBUG` + * is not defined, this function outputs a warning when `length == 0`. + * \tparam K a generic type that satisfies is_swappable, and for which a + * function `bool less_than(const K&, const K&, const Sorter&)` is + * defined. + * \tparam V satisfies is_swappable. + */ +template::value>::type* = nullptr> +inline void quick_sort(array& keys, array& values, const Sorter& sorter) +{ +#if !defined(NDEBUG) + if (keys.length == 0) { + fprintf(stderr, "quick_sort WARNING: Length is zero.\n"); + return; + } +#endif + quick_sort(keys.data, values.data, 0, (unsigned int) keys.length - 1, sorter); +} + + +/** + * + */ + +template +void sort(T* array, unsigned int start, unsigned int end) +{ + if (start >= end) + return; + else if (start + 16 >= end) { + insertion_sort(&array[start], end - start + 1); + return; + } + unsigned int l = start, r = end; + quick_sort_partition(array, start, end, l, r); + sort(array, start, r); + sort(array, l, end); +} + +template::value>::type* = nullptr> +void sort(T* array, unsigned int start, unsigned int end, const Sorter& sorter) +{ + if (start >= end) + return; + else if (start + 16 >= end) { + insertion_sort(&array[start], end - start + 1, sorter); + return; + } + unsigned int l = start, r = end; + quick_sort_partition(array, start, end, l, r, sorter); + sort(array, start, r, sorter); + sort(array, l, end, sorter); +} + +template +void sort(K* keys, V* values, unsigned int start, unsigned int end) +{ + if (start >= end) + return; + else if (start + 16 >= end) { + insertion_sort(&keys[start], &values[start], end - start + 1); + return; + } + unsigned int l = start, r = end; + quick_sort_partition(keys, values, start, end, l, r); + sort(keys, values, start, r); + sort(keys, values, l, end); +} + +template::value>::type* = nullptr> +void sort(K* keys, V* values, unsigned int start, unsigned int end, const Sorter& sorter) +{ + if (start >= end) + return; + else if (start + 16 >= end) { + insertion_sort(&keys[start], &values[start], end - start + 1, sorter); + return; + } + unsigned int l = start, r = end; + quick_sort_partition(keys, values, start, end, l, r, sorter); + sort(keys, values, start, r, sorter); + sort(keys, values, l, end, sorter); +} + +/** + * Performs hybrid Quicksort-insertion sort on the given native array `keys` + * with given `length`. To improve performance, the Quicksort switches to + * insertion sort for small partitions. This function assumes `length > 0`. If + * the preprocessor `NDEBUG` is not defined, this function outputs a warning + * when `length == 0`. + * \tparam T is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable), + * [CopyConstructible](https://en.cppreference.com/w/cpp/named_req/CopyConstructible) + * [LessThanComparable](https://en.cppreference.com/w/cpp/named_req/LessThanComparable), + * and is_swappable. + */ +template +inline void sort(T* keys, unsigned int length) { +#if !defined(NDEBUG) + if (length == 0) { + fprintf(stderr, "sort WARNING: Length is zero.\n"); + return; + } +#endif + sort(keys, 0, length - 1); +} + +/** + * Performs hybrid Quicksort-insertion sort on the given native array `keys` + * with given `length` and `sorter`. To improve performance, the Quicksort + * switches to insertion sort for small partitions. This function assumes + * `length > 0`. If the preprocessor `NDEBUG` is not defined, this function + * outputs a warning when `length == 0`. + * \tparam T a generic type that satisfies is_swappable and is_moveable, and + * for which a function `bool less_than(const T&, const T&, const Sorter&)` + * is defined. + */ +template::value>::type* = nullptr> +inline void sort(T* keys, unsigned int length, const Sorter& sorter) +{ +#if !defined(NDEBUG) + if (length == 0) { + fprintf(stderr, "sort WARNING: Length is zero.\n"); + return; + } +#endif + sort(keys, 0, length - 1, sorter); +} + +/** + * Performs hybrid Quicksort-insertion sort on the given native arrays `keys` + * and `values` with given `length`. To improve performance, the Quicksort + * switches to insertion sort for small partitions. This function assumes + * `length > 0`. If the preprocessor `NDEBUG` is not defined, this function + * outputs a warning when `length == 0`. + * \tparam K is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable), + * [CopyConstructible](https://en.cppreference.com/w/cpp/named_req/CopyConstructible) + * [LessThanComparable](https://en.cppreference.com/w/cpp/named_req/LessThanComparable), + * and is_swappable. + * \tparam V satisfies is_swappable and is_moveable. + */ +template +inline void sort(K* keys, V* values, unsigned int length) { +#if !defined(NDEBUG) + if (length == 0) { + fprintf(stderr, "sort WARNING: Length is zero.\n"); + return; + } +#endif + sort(keys, values, 0, length - 1); +} + +/** + * Performs hybrid Quicksort-insertion sort on the given native arrays `keys` + * and `values` with given `length` and `sorter`. To improve performance, the + * Quicksort switches to insertion sort for small partitions. This function + * assumes `length > 0`. If the preprocessor `NDEBUG` is not defined, this + * function outputs a warning when `length == 0`. + * \tparam K a generic type that satisfies is_swappable and is_moveable, and + * for which a function `bool less_than(const K&, const K&, const Sorter&)` + * is defined. + * \tparam V satisfies is_swappable and is_moveable. + */ +template::value>::type* = nullptr> +inline void sort(K* keys, V* values, unsigned int length, const Sorter& sorter) +{ +#if !defined(NDEBUG) + if (length == 0) { + fprintf(stderr, "sort WARNING: Length is zero.\n"); + return; + } +#endif + sort(keys, values, 0, length - 1, sorter); +} + +/** + * Performs hybrid Quicksort-insertion sort on the given array `keys`. To + * improve performance, the Quicksort switches to insertion sort for small + * partitions. This function assumes `length > 0`. If the preprocessor `NDEBUG` + * is not defined, this function outputs a warning when `length == 0`. + * \tparam T is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable), + * [CopyConstructible](https://en.cppreference.com/w/cpp/named_req/CopyConstructible) + * [LessThanComparable](https://en.cppreference.com/w/cpp/named_req/LessThanComparable), + * and is_swappable. + */ +template +inline void sort(array& keys) { +#if !defined(NDEBUG) + if (keys.length == 0) { + fprintf(stderr, "sort WARNING: Length is zero.\n"); + return; + } +#endif + sort(keys.data, 0, (unsigned int) keys.length - 1); +} + +/** + * Performs hybrid Quicksort-insertion sort on the given array `keys` with the + * given `sorter`. To improve performance, the Quicksort switches to insertion + * sort for small partitions. This function assumes `length > 0`. If the + * preprocessor `NDEBUG` is not defined, this function outputs a warning when + * `length == 0`. + * \tparam T a generic type that satisfies is_swappable and is_moveable, and + * for which a function `bool less_than(const T&, const T&, const Sorter&)` + * is defined. + */ +template::value>::type* = nullptr> +inline void sort(array& keys, const Sorter& sorter) +{ +#if !defined(NDEBUG) + if (keys.length == 0) { + fprintf(stderr, "sort WARNING: Length is zero.\n"); + return; + } +#endif + sort(keys.data, 0, (unsigned int) keys.length - 1, sorter); +} + +/** + * Performs hybrid Quicksort-insertion sort on the given arrays `keys` and + * `values`. To improve performance, the Quicksort switches to insertion sort + * for small partitions. This function assumes `length > 0`. If the + * preprocessor `NDEBUG` is not defined, this function outputs a warning when + * `length == 0`. + * \tparam K is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable), + * [CopyConstructible](https://en.cppreference.com/w/cpp/named_req/CopyConstructible) + * [LessThanComparable](https://en.cppreference.com/w/cpp/named_req/LessThanComparable), + * and is_swappable. + * \tparam V satisfies is_swappable and is_moveable. + */ +template +inline void sort(array& keys, array& values) { +#if !defined(NDEBUG) + if (keys.length == 0) { + fprintf(stderr, "sort WARNING: Length is zero.\n"); + return; + } +#endif + sort(keys.data, values.data, 0, (unsigned int) keys.length - 1); +} + +/** + * Performs hybrid Quicksort-insertion sort on the given arrays `keys` and + * `values` with the given `sorter`. To improve performance, the Quicksort + * switches to insertion sort for small partitions. This function assumes + * `length > 0`. If the preprocessor `NDEBUG` is not defined, this function + * outputs a warning when `length == 0`. + * \tparam K a generic type that satisfies is_swappable and is_moveable, and + * for which a function `bool less_than(const K&, const K&, const Sorter&)` + * is defined. + * \tparam V satisfies is_swappable and is_moveable. + */ +template::value>::type* = nullptr> +inline void sort(array& keys, array& values, const Sorter& sorter) +{ +#if !defined(NDEBUG) + if (keys.length == 0) { + fprintf(stderr, "sort WARNING: Length is zero.\n"); + return; + } +#endif + sort(keys.data, values.data, 0, (unsigned int) keys.length - 1, sorter); +} + +/** + * Returns `true` if and only if the given native array `items` with the given + * `length` is sorted in non-decreasing order. + * \tparam T a generic type for which a function `bool less_than(const T&, const T&, const Sorter&)` + * is defined. + */ +template +inline bool is_sorted(const T* items, size_t length, const Sorter& sorter) { + if (length == 0) + return true; + size_t first = 0, next = 0; + while (++next != length) { + if (less_than(items[next], items[first], sorter)) + return false; + first = next; + } + return true; +} + +/** + * Returns `true` if and only if the given array `items` is sorted in + * non-decreasing order. + * \tparam T a generic type for which a function `bool less_than(const T&, const T&, const Sorter&)` + * is defined. + */ +template +inline bool is_sorted(const array& items, const Sorter& sorter) { + return is_sorted(items.data, items.length, sorter); +} + +/** + * The default_sorter provides a default implementation of + * `bool less_than(const T&, const T&, const default_sorter&)` + * where the comparison is done using the operator `<`. + */ +struct default_sorter { }; + +template +inline bool less_than(const T& first, const T& second, const default_sorter& sorter) { + return (first < second); +} + +/** + * The pointer_sorter compares items using the `<` operator on the dereferenced + * arguments. + */ +struct pointer_sorter { }; + +template +inline bool less_than(const T* first, const T* second, const pointer_sorter& sorter) { + return (*first < *second); +} + +/** + * Deletes consecutive duplicates in the given native `array` with given + * `length` and returns the new length. Note the deleted elements are not freed. + * \tparam T is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). + */ +template +unsigned int unique(T* array, size_t length) +{ + unsigned int result = 0; + for (unsigned int i = 1; i < length; i++) { + if (array[result] != array[i]) + array[++result] = array[i]; + } + return result + 1; +} + +/** + * Deletes consecutive duplicates in the given `array` with given and returns + * the new length. Note the deleted elements are not freed. + * \tparam T is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). + */ +template +inline void unique(array& a) { + a.length = unique(a.data, a.length); +} + +/* forward declarations */ + +unsigned int sample_uniform(unsigned int); + +/** + * Performs a Knuth shuffle on the given native `array` with given `length`. + * \tparam T satisfies is_swappable. + */ +template +void shuffle(T* array, unsigned int length) { +#if !defined(NDEBUG) + if (length == 0) { + fprintf(stderr, "shuffle WARNING: Length is zero.\n"); + return; + } +#endif + for (unsigned int i = length - 1; i > 0; i--) { + unsigned int next = sample_uniform(i + 1); + if (next != i) + core::swap(array[next], array[i]); + } +} + +/** + * Performs a Knuth shuffle on the given native arrays `keys` and `values` with given `length`. + * \tparam T satisfies is_swappable. + */ +template +void shuffle(K* keys, V* values, unsigned int length) { +#if !defined(NDEBUG) + if (length == 0) { + fprintf(stderr, "shuffle WARNING: Length is zero.\n"); + return; + } +#endif + for (unsigned int i = length - 1; i > 0; i--) { + unsigned int next = sample_uniform(i + 1); + if (next != i) { + core::swap(keys[next], keys[i]); + core::swap(values[next], values[i]); + } + } +} + +/** + * Performs a Knuth shuffle on the given `array`. + * \tparam T satisfies is_swappable. + */ +template +inline void shuffle(array& items) { + shuffle(items.data, items.length); +} + +/** + * Given sorted array `a`, this function finds the smallest index `i` such that + * `a[i] >= b` and `i >= start` and `i < end`. + */ +template +unsigned int linear_search( + const T* a, const T& b, + unsigned int start, + unsigned int end) +{ + for (unsigned int i = start; i < end; i++) + if (a[i] >= b) return i; + return end; +} + +/** + * Given sorted array `a`, this function finds the smallest index `i` such that + * `a[i] > b` and `i >= start` and `i < end`. + */ +template +unsigned int strict_linear_search( + const T* a, const T& b, + unsigned int start, + unsigned int end) +{ + for (unsigned int i = start; i < end; i++) + if (a[i] > b) return i; + return end; +} + +/** + * Given sorted array `a`, this function finds the smallest index `i` such that + * `a[i] > b` and `i >= start` and `i < end`. + */ +template +unsigned int reverse_strict_linear_search( + const T* a, const T& b, + unsigned int start, + unsigned int end) +{ + for (unsigned int i = end; i > start; i--) + if (a[i - 1] <= b) return i; + return start; +} + +/** + * Given sorted array `a`, this function finds the smallest index `i` such that + * `a[i] >= b` and `i >= min` and `i <= max`. + */ +/* TODO: implement a strict variant */ +template +unsigned int binary_search( + const T* a, const T& b, + unsigned int min, + unsigned int max) +{ + if (a[max] < b) + return max + 1; + + while (min < max) { + unsigned int mid = (max + min) / 2; + if (a[mid] < b) + min = mid + 1; + else max = mid; + } + + return min; +} + + +/** + * A simple pair data structure. + */ + +template +struct pair { + /** + * The key object in the field. + */ + K key; + + /** + * The value object in the field. + */ + V value; + + pair(const K& key, const V& value) : key(key), value(value) { } + + inline bool operator == (const pair& other) const { + return key == other.key && value == other.value; + } + + inline bool operator != (const pair& other) const { + return key != other.key || value != other.value; + } + + static inline void move(const pair& src, pair& dst) { + core::move(src.key, dst.key); + core::move(src.value, dst.value); + } + + static inline void swap(pair& first, pair& second) { + core::swap(first.key, second.key); + core::swap(first.value, second.value); + } + + static inline unsigned int hash(const pair& pair) { + return hasher::hash(pair.key) + hasher::hash(pair.value); + } + + static inline bool is_empty(const pair& pair) { + return hasher::is_empty(pair.key); + } + + static inline void set_empty(pair& pair) { + hasher::set_empty(pair.key); + } + + static inline void free(pair& pair) { + core::free(pair.key); + core::free(pair.value); + } +}; + +template +inline bool operator < (const pair& first, const pair& second) { + return first.key < second.key; +} + +template +inline bool operator <= (const pair& first, const pair& second) { + return first.key <= second.key; +} + +template +inline bool operator > (const pair& first, const pair& second) { + return first.key > second.key; +} + +template +inline bool operator >= (const pair& first, const pair& second) { + return first.key >= second.key; +} + +template +constexpr pair make_pair(const K& key, const V& value) { + return pair(key, value); +} + +template +inline void swap(pair& first, pair& second) { + swap(first.key, second.key); + swap(first.value, second.value); +} + + +/** + * + */ + +/** + * For every index `i >= index`, this function moves each element at `i` to index `i + 1`. + * \tparam T satisfies is_moveable. + */ +template +void shift_right(T* list, unsigned int length, unsigned int index) +{ + for (unsigned int i = length; i > index; i--) + move(list[i - 1], list[i]); +} + +/** + * For every index `i < index`, this function moves each element at `i + 1` to index `i`. + * \tparam T satisfies is_moveable. + */ +template +void shift_left(T* list, unsigned int index) +{ + for (unsigned int i = 0; i < index; i++) + move(list[i + 1], list[i]); +} + +template +inline void set_union_helper(UnionFunc do_union, const T& item, + unsigned int i, unsigned int j, const T*& prev) +{ + if (RemoveDuplicates) { + if (prev == NULL || *prev != item) { + do_union(item, i, j); + prev = &item; + } + } else { + do_union(item, i, j); + } +} + +/** + * Given ordered native arrays `first` and `second`, this function visits their + * elements sequentially, and: + * 1. For every element `x` that appears in both arrays, `union_both(x, i, j)` + * is called where `first[i] == x` and `second[j] == x`. + * 2. For every element `x` that appears in `first` but not `second`, + * `union_first(x, i, j)` is called where `first[i] == x`, and `j` is the + * smallest such index where `second[j] > x`. + * 3. For every element `x` that appears in `second` but not `first`, + * `union_second(x, i, j)` is called where `second[j] == x`, and `i` is + * the smallest such index where `first[i] > x`. + * + * The visited elements are also ordered. + * + * \tparam T a generic type for which the operators `==`, `!=` and `<` are implemented. + * \tparam RemoveDuplicates if `true`, this function will avoid calling the + * union functions more than once for each element. + */ +template +void set_union(UnionBoth union_both, + UnionFirst union_first, UnionSecond union_second, + const T* first, unsigned int first_length, + const T* second, unsigned int second_length) +{ + unsigned int i = 0, j = 0; + const T* prev = NULL; + while (i < first_length && j < second_length) + { + if (first[i] == second[j]) { + set_union_helper(union_both, first[i], i, j, prev); + i++; j++; + } else if (first[i] < second[j]) { + set_union_helper(union_first, first[i], i, j, prev); + i++; + } else { + set_union_helper(union_second, second[j], i, j, prev); + j++; + } + } + + while (i < first_length) { + set_union_helper(union_first, first[i], i, j, prev); + i++; + } while (j < second_length) { + set_union_helper(union_second, second[j], i, j, prev); + j++; + } +} + +template::value>::type* = nullptr> +inline void set_union_helper(T* dst, SizeType& dst_length, const T& item) { + if (!RemoveDuplicates || dst_length == 0 || dst[dst_length - 1] != item) { + dst[dst_length] = item; + dst_length++; + } +} + +/** + * Given ordered native arrays `first` and `second`, compute their union and + * appends the result to `dst`. The union is also ordered. This function + * assumes `dst` has sufficient capacity to store the union. + * \tparam T satisfies [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable) + * and implements the operators `==` and `<`. + * \tparam RemoveDuplicates if `true`, this function ignores duplicate elements. + */ +template::value>::type* = nullptr> +void set_union(T* dst, SizeType& dst_length, + const T* first, unsigned int first_length, + const T* second, unsigned int second_length) +{ + unsigned int i = 0, j = 0; + while (i < first_length && j < second_length) + { + if (first[i] == second[j]) { + set_union_helper(dst, dst_length, first[i]); + i++; j++; + } else if (first[i] < second[j]) { + set_union_helper(dst, dst_length, first[i]); + i++; + } else { + set_union_helper(dst, dst_length, second[j]); + j++; + } + } + + while (i < first_length) { + set_union_helper(dst, dst_length, first[i]); + i++; + } while (j < second_length) { + set_union_helper(dst, dst_length, second[j]); + j++; + } +} + +/** + * Given ordered native arrays `first` and `second`, compute their union and + * appends the result to `dst`. The union is also ordered. + * \tparam T satisfies [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable) + * and implements the operators `==` and `<`. + * \tparam RemoveDuplicates if `true`, this function ignore duplicate elements. + */ +template +inline bool set_union(array& dst, + const T* first, unsigned int first_length, + const T* second, unsigned int second_length) +{ + if (!dst.ensure_capacity(dst.length + first_length + second_length)) + return false; + set_union(dst.data, dst.length, first, first_length, second, second_length); + return true; +} + +/** + * Given ordered arrays `first` and `second`, compute their union and appends + * the result to `dst`. The union is also ordered. + * \tparam T satisfies [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable) + * and implements the operators `==` and `<`. + * \tparam RemoveDuplicates if `true`, this function ignore duplicate elements. + */ +template +inline bool set_union(array& dst, const array& first, const array& second) { + return set_union(dst, + first.data, (unsigned int) first.length, + second.data, (unsigned int) second.length); +} + +template +struct array_position { + unsigned int array_id; + unsigned int position; + T* element; +}; + +/** + * Given a collection of ordered arrays `arrays`, compute their union and + * appends the result to `dst`. The union is also ordered. + * NOTE: this function assumes the given arrays are all non-empty. + * \tparam T satisfies [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable) + * and implements the operators `==` and `<`. + * \tparam ArraySetCollection a collection type where each element is accessed + * using `arrays[i]` and the size of each array can be obtained using + * `size(arrays[i])`. The element at index `j` for array `i` is accessed + * using `arrays[i][j]`. + */ +template +bool set_union(array& dst, const ArraySetCollection& arrays, unsigned int array_count) +{ + /* first ensure the destination array has enough space */ + unsigned int total_size = dst.length; + for (unsigned int i = 0; i < array_count; i++) { +#if !defined(NDEBUG) + if (size(arrays[i]) == 0) + fprintf(stderr, "set_union WARNING: Input array %u is empty.\n", i); +#endif + total_size += size(arrays[i]); + } + if (!dst.ensure_capacity(total_size)) + return false; + + /* TODO: we can probably use a faster heap structure */ + array_position* heap = (array_position*) malloc(array_count * sizeof(array_position)); + if (heap == NULL) { + fprintf(stderr, "set_union ERROR: Out of memory.\n"); + return false; + } + for (unsigned int i = 0; i < array_count; i++) + heap[i] = { i, 0, &arrays[i][0] }; + std::make_heap(heap, heap + array_count); + + /* add the first item to the destination set */ + unsigned int heap_size = array_count; + std::pop_heap(heap, heap + heap_size); + const array_position& first = heap[heap_size]; + dst.data[dst.length] = first.key; + dst.length++; + if (size(arrays[first.value]) > 1) { + heap[heap_size] = { first.array_id, 1, &arrays[first.array_id][1] }; + std::push_heap(heap, heap + heap_size); + } else { heap_size--; } + + while (heap_size > 0) + { + std::pop_heap(heap, heap + heap_size); + const array_position& next = heap[heap_size]; + if (next.key != dst.last()) { + dst.data[dst.length] = next.key; + dst.length++; + } + if (next.value + 1 < size(arrays[next.value])) { + heap[heap_size] = { next.array_id, next.position + 1, arrays[next.array_id][next.position + 1] }; + std::push_heap(heap, heap + heap_size); + } else { heap_size--; } + } + free(heap); + return true; +} + +/** + * Given ordered native arrays `first` and `second`, this function visits their + * elements sequentially, and for every element `x` that exists in both arrays, + * `intersect(i, j)` is called where `first[i] == x` and `second[j] == x`. The + * visited elements are also ordered. + * \tparam T a generic type for which the operators `==` and `<` are implemented. + * \tparam BinarySearch if `true`, binary search is used to find indices of + * identical elements rather than linear search. + */ +template +bool set_intersect(Intersect intersect, + const T* first, unsigned int first_length, + const T* second, unsigned int second_length) +{ + unsigned int i = 0, j = 0; + while (i < first_length && j < second_length) + { + if (first[i] == second[j]) { + intersect(i, j); + i++; j++; + } else if (first[i] < second[j]) { + if (BinarySearch) { + /* use binary search to find the value of i + such that first.data[i] >= second.data[j] */ + i = binary_search(first, second[j], i, first_length - 1); + } else { + i++; + } + } else { + if (BinarySearch) { + /* use binary search to find the value of j + such that second.data[j] >= first.data[i] */ + j = binary_search(second, first[i], j, second_length - 1); + } else { + j++; + } + } + } + return true; +} + +/** + * Given ordered native arrays `first` and `second`, compute the intersection + * and append it to the native array `intersection`. The computed intersection + * is also ordered. This function assumes `intersection` has sufficient capacity. + * \tparam T satisfies [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable) + * and implements the operators `==` and `<`. + * \tparam BinarySearch if `true`, binary search is used to find indices of + * identical elements rather than linear search. + */ +template::value>::type* = nullptr> +bool set_intersect( + T* intersection, SizeType& intersection_length, + const T* first, unsigned int first_length, + const T* second, unsigned int second_length) +{ + unsigned int i = 0, j = 0; + while (i < first_length && j < second_length) + { + if (first[i] == second[j]) { + intersection[intersection_length] = first[i]; + intersection_length++; + i++; j++; + } else if (first[i] < second[j]) { + if (BinarySearch) { + /* use binary search to find the value of i + such that first.data[i] >= second.data[j] */ + i = binary_search(first, second[j], i, first_length - 1); + } else { + i++; + } + } else { + if (BinarySearch) { + /* use binary search to find the value of j + such that second.data[j] >= first.data[i] */ + j = binary_search(second, first[i], j, second_length - 1); + } else { + j++; + } + } + } + return true; +} + +/** + * Given ordered native arrays `first` and `second`, compute the intersection + * and append it to the array `intersection`. The computed intersection is also + * ordered. + * \tparam T satisfies [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable) + * and implements the operators `==` and `<`. + * \tparam BinarySearch if `true`, binary search is used to find indices of + * identical elements rather than linear search. + */ +template +inline bool set_intersect( + array& intersection, + const T* first, unsigned int first_length, + const T* second, unsigned int second_length) +{ + if (!intersection.ensure_capacity(intersection.length + max(first_length, second_length))) + return false; + + return set_intersect( + intersection.data, intersection.length, + first, first_length, second, second_length); +} + +/** + * Given ordered arrays `first` and `second`, compute the intersection and + * append it to the array `intersection`. The computed intersection is also + * ordered. + * \tparam T satisfies [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable) + * and implements the operators `==` and `<`. + * \tparam BinarySearch if `true`, binary search is used to find indices of + * identical elements rather than linear search. + */ +template +inline bool set_intersect( + array& intersection, + const array& first, + const array& second) +{ + return set_intersect(intersection, + first.data, (unsigned int) first.length, + second.data, (unsigned int) second.length); +} + +/** + * Given ordered native arrays `first` and `second`, compute the intersection + * in-place and store it in `first`. The computed intersection is also ordered. + * \tparam T satisfies [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable) + * and implements the operators `==` and `<`. + * \tparam BinarySearch if `true`, binary search is used to find indices of + * identical elements rather than linear search. + */ +template::value>::type* = nullptr> +void set_intersect( + T* first, SizeType& first_length, + const T* second, unsigned int second_length) +{ + unsigned int index = 0; + unsigned int i = 0, j = 0; + while (i < first_length && j < second_length) + { + if (first[i] == second[j]) { + first[index] = first[i]; + index++; i++; j++; + } else if (first[i] < second[j]) { + if (BinarySearch) { + /* use binary search to find the value of i + such that first.data[i] >= second.data[j] */ + i = binary_search(first, second[j], i, first_length - 1); + } else { + i++; + } + } else { + if (BinarySearch) { + /* use binary search to find the value of j + such that second.data[j] >= first.data[i] */ + j = binary_search(second, first[i], j, second_length - 1); + } else { + j++; + } + } + } + first_length = index; +} + +/** + * Given ordered array `first` and ordered native array `second`, compute the + * intersection in-place and store it in `first`. The computed intersection is + * also ordered. + * \tparam T satisfies [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable) + * and implements the operators `==` and `<`. + * \tparam BinarySearch if `true`, binary search is used to find indices of + * identical elements rather than linear search. + */ +template +inline void set_intersect(array& first, + const T* second, unsigned int second_length) +{ + return set_intersect( + first.data, first.length, second, second_length); +} + +/** + * Given ordered arrays `first` and `second`, compute the intersection in-place + * and store it in `first`. The computed intersection is also ordered. + * \tparam T satisfies [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable) + * and implements the operators `==` and `<`. + * \tparam BinarySearch if `true`, binary search is used to find indices of + * identical elements rather than linear search. + */ +template +inline void set_intersect(array& first, const array& second) { + return set_intersect(first, second.data, second.length); +} + +/** + * Returns true if the intersection of `first` and `second` is non-empty, + * where `first` and `second` are ordered native arrays. + * \tparam T a generic type that implements the operators `==` and `<`. + * \tparam BinarySearch if `true`, binary search is used to find indices of + * identical elements rather than linear search. + */ +template +bool has_intersection( + const T* first, unsigned int first_length, + const T* second, unsigned int second_length) +{ + unsigned int i = 0, j = 0; + while (i < first_length && j < second_length) + { + if (first[i] == second[j]) { + return true; + i++; j++; + } else if (first[i] < second[j]) { + if (BinarySearch) { + /* use binary search to find the value of i + such that first.data[i] >= second.data[j] */ + i = binary_search(first, second[j], i, first_length - 1); + } else { + i++; + } + } else { + if (BinarySearch) { + /* use binary search to find the value of j + such that second.data[j] >= first.data[i] */ + j = binary_search(second, first[i], j, second_length - 1); + } else { + j++; + } + } + } + return false; +} + +/** + * Returns true if the intersection of `first` and `second` is non-empty, + * where `first` and `second` are ordered arrays. + * \tparam T a generic type that implements the operators `==` and `<`. + * \tparam BinarySearch if `true`, binary search is used to find indices of + * identical elements rather than linear search. + */ +template +inline bool has_intersection(const array& first, const array& second) { + return has_intersection( + first.data, (unsigned int) first.length, + second.data, (unsigned int) second.length); +} + +/** + * Returns true if `first` is a subset of `second`, where `first` and `second` + * are ordered native arrays. + * \tparam T a generic type that implements the operators `==` and `<`. + * \tparam BinarySearch if `true`, binary search is used to find indices of + * identical elements rather than linear search. + */ +template +bool is_subset( + const T* first, unsigned int first_length, + const T* second, unsigned int second_length) +{ + unsigned int i = 0, j = 0; + while (i < first_length && j < second_length) + { + if (first[i] == second[j]) { + i++; j++; + } else if (first[i] < second[j]) { + return false; + } else { + if (BinarySearch) { + /* use binary search to find the value of j + such that second.data[j] >= first.data[i] */ + j = binary_search(second, first[i], j, second_length - 1); + } else { + j++; + } + } + } + return (i == first_length); +} + +/** + * Given ordered native arrays `first` and `second`, this function visits their + * elements sequentially, and for every element `x` that exists in `first` but + * not in `second`, `emit(i)` is called where `first[i] == x`. The visited + * elements are also ordered. + * \tparam T a generic type that implements the operators `==` and `<`. + * \tparam BinarySearch if `true`, binary search is used to find indices of + * identical elements rather than linear search. + */ +template +void set_subtract(EmitFunction emit, + const T* first, unsigned int first_length, + const T* second, unsigned int second_length) +{ + unsigned int i = 0, j = 0; + while (i < first_length && j < second_length) + { + if (first[i] == second[j]) { + i++; j++; + } else if (first[i] < second[j]) { + if (BinarySearch) { + /* use binary search to find the value of i + such that first.data[i] >= second.data[j] */ + unsigned int next_i = binary_search(first, second[j], i, first_length - 1); + for (; i < next_i; i++) + emit(i); + i = next_i; + } else { + emit(i); + i++; + } + } else { + if (BinarySearch) { + /* use binary search to find the value of j + such that second.data[j] >= first.data[i] */ + j = binary_search(second, first[i], j, second_length - 1); + } else { + j++; + } + } + } + + while (i < first_length) { + emit(i); + i++; + } +} + +/** + * Given ordered native arrays `first` and `second`, this function computes the + * set difference between `first` and `second` and stores the result in `dst`. + * This function assumes `dst` has sufficient capacity. The set difference is + * also ordered. + * \tparam T satisfies [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable) + * and implements the operators `==` and `<`. + * \tparam BinarySearch if `true`, binary search is used to find indices of + * identical elements rather than linear search. + */ +template::value>::type* = nullptr> +void set_subtract(T* dst, SizeType& dst_length, + const T* first, unsigned int first_length, + const T* second, unsigned int second_length) +{ + unsigned int i = 0, j = 0; + while (i < first_length && j < second_length) + { + if (first[i] == second[j]) { + i++; j++; + } else if (first[i] < second[j]) { + if (BinarySearch) { + /* use binary search to find the value of i + such that first.data[i] >= second.data[j] */ + unsigned int next_i = binary_search(first, second[j], i, first_length - 1); + for (; i < next_i; i++) { + dst[dst_length] = first[i]; + dst_length++; + } + i = next_i; + } else { + dst[dst_length] = first[i]; + dst_length++; + i++; + } + } else { + if (BinarySearch) { + /* use binary search to find the value of j + such that second.data[j] >= first.data[i] */ + j = binary_search(second, first[i], j, second_length - 1); + } else { + j++; + } + } + } + + memcpy(dst + dst_length, first + i, (first_length - i) * sizeof(T)); + dst_length += first_length - i; +} + +/** + * Given ordered native arrays `first` and `second`, this function computes the + * set difference between `first` and `second` and stores the result in `dst`. + * The set difference is also ordered. + * \tparam T satisfies [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable) + * and implements the operators `==` and `<`. + * \tparam BinarySearch if `true`, binary search is used to find indices of + * identical elements rather than linear search. + */ +template +bool set_subtract(array& dst, + const T* first, unsigned int first_length, + const T* second, unsigned int second_length) +{ + if (!dst.ensure_capacity(dst.length + max(first_length, second_length))) + return false; + set_subtract(dst.data, dst.length, first, first_length, second, second_length); + return true; +} + +/** + * Given ordered arrays `first` and `second`, this function computes the set + * difference between `first` and `second` and stores the result in `dst`. The + * set difference is also ordered. + * \tparam T satisfies [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable) + * and implements the operators `==` and `<`. + * \tparam BinarySearch if `true`, binary search is used to find indices of + * identical elements rather than linear search. + */ +template +inline bool set_subtract(array& dst, + const array& first, + const array& second) +{ + return set_subtract(dst, + first.data, (unsigned int) first.length, + second.data, (unsigned int) second.length); +} + +/** + * Given ordered native arrays `first` and `second`, this function computes the + * set difference in-place between `first` and `second` and stores the result + * in `first`. The set difference is also ordered. + * \tparam T satisfies [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable) + * and implements the operators `==` and `<`. + * \tparam BinarySearch if `true`, binary search is used to find indices of + * identical elements rather than linear search. + */ +template::value>::type* = nullptr> +void set_subtract( + T* first, SizeType& first_length, + const T* second, unsigned int second_length) +{ + unsigned int index = 0; + unsigned int i = 0, j = 0; + while (i < first_length && j < second_length) + { + if (first[i] == second[j]) { + i++; j++; + } else if (first[i] < second[j]) { + if (BinarySearch) { + /* use binary search to find the value of i + such that first.data[i] >= second.data[j] */ + unsigned int next_i = binary_search(first, second[j], i, first_length - 1); + for (; i < next_i; i++) { + first[index] = first[i]; + index++; + } + i = next_i; + } else { + first[index] = first[i]; + index++; i++; + } + } else { + if (BinarySearch) { + /* use binary search to find the value of j + such that second.data[j] >= first.data[i] */ + j = binary_search(second, first[i], j, second_length - 1); + } else { + j++; + } + } + } + + while (i < first_length) { + first[index] = first[i]; + index++; i++; + } + first_length = index; +} + +/** + * Given ordered arrays `first` and `second`, this function computes the set + * difference in-place between `first` and `second` and stores the result in + * `first`. The set difference is also ordered. + * \tparam T satisfies [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable) + * and implements the operators `==` and `<`. + * \tparam BinarySearch if `true`, binary search is used to find indices of + * identical elements rather than linear search. + */ +template +inline void set_subtract(array& first, const array& second) +{ + return set_subtract(first.data, first.length, second.data, second.length); +} + +inline void array_test(void) +{ + array buf(1); + buf.append("0123456789 ", 11); + + buf[(unsigned int) buf.length] = '\0'; + if (buf.length != 11) + fprintf(stderr, "array test ERROR: First array length test failed.\n"); + if (buf.capacity != 16) + fprintf(stderr, "array test ERROR: First array capacity test failed.\n"); + if (strcmp(buf.data, "0123456789 ") != 0) + fprintf(stderr, "array test ERROR: First string comparison test failed.\n"); + + buf.append("abcdefghijklmnopqrstuvwxyz ", 27); + + buf[(unsigned int) buf.length] = '\0'; + if (buf.length != 11 + 27) + fprintf(stderr, "array test ERROR: Second array length test failed.\n"); + if (buf.capacity != 64) + fprintf(stderr, "array test ERROR: Second array capacity test failed.\n"); + if (strcmp(buf.data, "0123456789 abcdefghijklmnopqrstuvwxyz ") != 0) + fprintf(stderr, "array test ERROR: Second string comparison test failed.\n"); + + buf.append("9876543210 ", 11); + + buf[(unsigned int) buf.length] = '\0'; + if (buf.length != 11 + 27 + 11) + fprintf(stderr, "array test ERROR: Third array length test failed.\n"); + if (buf.capacity != 64) + fprintf(stderr, "array test ERROR: Third array capacity test failed.\n"); + if (strcmp(buf.data, "0123456789 abcdefghijklmnopqrstuvwxyz 9876543210 ") != 0) + fprintf(stderr, "array test ERROR: Third string comparison test failed.\n"); + + /* test some of the helper functions */ + array numbers = array(10); + numbers.add(4); + numbers.add(-6); + numbers.add(4); + numbers.add(2); + numbers.add(0); + numbers.add(-6); + numbers.add(1); + numbers.add(4); + numbers.add(2); + + array numbers_copy = array(10); + numbers_copy.append(numbers.data, numbers.length); + insertion_sort(numbers); + quick_sort(numbers_copy); + + int expected[] = {-6, -6, 0, 1, 2, 2, 4, 4, 4}; + for (unsigned int i = 0; i < 9; i++) { + if (numbers[i] != expected[i]) { + fprintf(stderr, "array test ERROR: insertion_sort failed.\n"); + return; + } + if (numbers_copy[i] != expected[i]) { + fprintf(stderr, "array test ERROR: quick_sort failed.\n"); + return; + } + } + + int expected_unique[] = {-6, 0, 1, 2, 4}; + unique(numbers); + if (numbers.length != 5) { + fprintf(stderr, "array test ERROR: unique failed.\n"); + return; + } + for (unsigned int i = 0; i < 5; i++) { + if (numbers[i] != expected_unique[i]) { + fprintf(stderr, "array test ERROR: unique failed.\n"); + return; + } + } + + printf("array test completed.\n"); +} + +} /* namespace core */ + +#endif /* ARRAY_H_ */ diff --git a/reasoners/data_preprocess/logic/graph_dataset_gen/core/core.h b/reasoners/data_preprocess/logic/graph_dataset_gen/core/core.h new file mode 100644 index 000000000..18986681f --- /dev/null +++ b/reasoners/data_preprocess/logic/graph_dataset_gen/core/core.h @@ -0,0 +1,479 @@ +/** + * \file core.h + * + * This file contains functions widely used by the core library. It contains + * functions for object movement, swapping, copying, and hashing, as well as + * the `free` function, which deinitializes objects and releases its + * memory resources. + * + * Memory management + * ----------------- + * + * This library does not manage memory, and leaves it to the user to ensure + * memory is allocated and freed correctly. The core data structures, such as + * core::array, core::hash_set, core::hash_map, core::array_map do not + * automatically free their elements upon deallocation. The library uses the + * core::free() function to free memory resources. When developing a new type, + * the class/struct should implement the public static function `free`, as in + * the example below. + * + * ```{.cpp} + * #include + * #include + * #include + * using namespace core; + * + * struct custom_string { + * char* buffer; + * + * custom_string() { } + * + * custom_string(const char* src) { + * buffer = (char*) malloc(sizeof(char) * (strlen(src) + 1)); + * memcpy(buffer, src, sizeof(char) * (strlen(src) + 1)); + * } + * + * ~custom_string() { + * core::free(buffer); + * } + * + * static void free(custom_string& s) { + * core::free(s.buffer); + * } + * }; + * + * bool init(custom_string& s, const char* src) { + * s.buffer = (char*) malloc(sizeof(char) * (strlen(src) + 1)); + * if (s.buffer == NULL) + * return false; + * memcpy(s.buffer, src, sizeof(char) * (strlen(src) + 1)); + * return true; + * } + * + * int main() { + * custom_string first = custom_string("first"); + * custom_string& second = *((custom_string*) alloca(sizeof(custom_string))); + * init(second, "second"); + * + * printf("%s ", first.buffer); + * printf("%s", second.buffer); + * + * core::free(second); + * } + * ``` + * In this example, core::free() is not called on `first` since it was + * initialized on the stack, and so its destructor will automatically free + * `buffer`. However, this was not the case for `second`, and so we must call + * core::free() directly. The expected output of the program is `first second`. + * + * + */ + +#ifndef CORE_H_ +#define CORE_H_ + +#include +#include +#include +#include +#include + +#if defined(_WIN32) +#include +#endif + + +namespace core { + + +/** + * This function returns the compile-time length of the static array `array`. + */ +template +inline constexpr std::size_t array_length(const T (&array)[N]) { + return N; +} + +/** + * Returns `first` if `first < second`, and returns `second` otherwise. + */ +template +inline typename std::common_type::type min(const A& first, const B& second) { + typedef typename std::common_type::type C; + return (static_cast(first) < static_cast(second)) ? first : second; +} + +/** + * Returns `second` if `first < second`, and returns `first` otherwise. + */ +template +inline typename std::common_type::type max(const A& first, const B& second) { + typedef typename std::common_type::type C; + return (static_cast(first) < static_cast(second)) ? second : first; +} + +/** + * If `T` satisfies [is_fundamental](http://en.cppreference.com/w/cpp/types/is_fundamental), + * [is_enum](http://en.cppreference.com/w/cpp/types/is_enum), or + * [is_pointer](http://en.cppreference.com/w/cpp/types/is_pointer), this + * function assigns `a = b`. + */ +template::value || std::is_enum::value || std::is_pointer::value>::type* = nullptr> +inline bool init(T& a, const T& b) { + a = b; + return true; +} + +/** + * If `T` satisfies [is_fundamental](http://en.cppreference.com/w/cpp/types/is_fundamental), + * [is_enum](http://en.cppreference.com/w/cpp/types/is_enum), + * or [is_pointer](http://en.cppreference.com/w/cpp/types/is_pointer), + * `dst` is assigned to `src`. Otherwise, this function calls + * `T::move(src, dst)`. This function is intended to be used to effectively + * move the object stored in `src` into the location specified by `dst`. + */ +template::value && !std::is_enum::value && !std::is_pointer::value>::type* = nullptr> +inline void move(const T& src, T& dst) { + T::move(src, dst); +} + +template::value || std::is_enum::value || std::is_pointer::value>::type* = nullptr> +inline void move(const T& a, T& b) { + b = a; +} + +/** + * If `T` satisfies [is_fundamental](http://en.cppreference.com/w/cpp/types/is_fundamental), + * [is_enum](http://en.cppreference.com/w/cpp/types/is_enum), + * or [is_pointer](http://en.cppreference.com/w/cpp/types/is_pointer), + * `dst` is assigned to `src`. Otherwise, this function calls and returns the + * result of `T::copy(src, dst)`. This function is intended to be used to copy + * the object stored in `src` into the location specified by `dst`. + */ +template::value && !std::is_enum::value && !std::is_pointer::value>::type* = nullptr> +inline bool copy(const T& src, T& dst) { + return T::copy(src, dst); +} + +template::value || std::is_enum::value || std::is_pointer::value>::type* = nullptr> +inline bool copy(const T& a, T& b) { + b = a; + return true; +} + +/** + * If `T` satisfies [is_fundamental](http://en.cppreference.com/w/cpp/types/is_fundamental), + * [is_enum](http://en.cppreference.com/w/cpp/types/is_enum), + * or [is_pointer](http://en.cppreference.com/w/cpp/types/is_pointer), + * the values of `a` and `b` are swapped. Otherwise, this function calls + * `T::swap(a, b)`. This function is intended to be used to swap the object + * stored in `a` with the object stored in `b`. + */ +template::value && !std::is_enum::value && !std::is_pointer::value>::type* = nullptr> +inline void swap(T& a, T& b) { + T::swap(a, b); +} + +template::value || std::is_enum::value || std::is_pointer::value>::type* = nullptr> +inline void swap(T& a, T& b) { + T temp = a; + a = b; + b = temp; +} + +template::value || std::is_enum::value>::type* = nullptr> +inline constexpr T size_of(const T& a) { + return sizeof(T); +} + +template::value>::type* = nullptr> +inline constexpr long unsigned int size_of(const E& a) { + return sizeof(a); +} + +struct default_metric { }; + +template::value && !std::is_enum::value && !std::is_pointer::value>::type* = nullptr> +inline long unsigned int size_of(const T& a) { + return T::size_of(a, default_metric()); +} + +template +inline constexpr long unsigned int size_of(const T* const& a) { + return sizeof(T*); +} + +template +inline long unsigned int size_of(const T (&array)[N]) { + long unsigned int size = 0; + for (size_t i = 0; i < N; i++) + size += size_of(array[i]); + return size; +} + +template::value || std::is_enum::value>::type* = nullptr> +inline long unsigned int size_of(const T& a, const default_metric& metric) { + return size_of(a); +} + +template +inline long unsigned int size_of(const T& a, const Metric& metric) { + return T::size_of(a, metric); +} + +template::value || std::is_enum::value>::type* = nullptr> +inline void free(T& a) { } + +/** + * If `T` satisfies [is_fundamental](http://en.cppreference.com/w/cpp/types/is_fundamental) + * or [is_enum](http://en.cppreference.com/w/cpp/types/is_enum), + * this function does nothing. Otherwise, this function calls + * `T::free(a, std::forward(args)...)`. This function is intended to be + * used to free the object stored in `a`, passing along any additional + * arguments. Note that since `args` is variadic, it may also be empty. + */ +template::value && !std::is_enum::value>::type* = nullptr> +inline void free(T& a, Args&&... args) { + T::free(a, std::forward(args)...); +} + +/** + * This function calls [free](http://en.cppreference.com/w/c/memory/free) + * on `a`. + */ +template +inline void free(T* a) { + ::free(a); +} + + +/** + * + */ + +namespace detail { + template static auto test_swappable(int32_t) -> + decltype(void(core::swap(std::declval(), std::declval())), std::true_type{}); + template static auto test_swappable(int64_t) -> std::false_type; + + template static auto test_moveable(int32_t) -> + decltype(void(core::move(std::declval(), std::declval())), std::true_type{}); + template static auto test_moveable(int64_t) -> std::false_type; + + template static auto test_copyable(int32_t) -> + decltype(void(core::copy(std::declval(), std::declval())), std::true_type{}); + template static auto test_copyable(int64_t) -> std::false_type; +} + +/** + * This type trait is [true_type](http://en.cppreference.com/w/cpp/types/integral_constant) + * if and only if `T` satisfies any of the following: + * 1. [is_fundamental](http://en.cppreference.com/w/cpp/types/is_fundamental), + * 2. [is_enum](http://en.cppreference.com/w/cpp/types/is_enum), + * 3. [is_pointer](http://en.cppreference.com/w/cpp/types/is_pointer), + * 4. or implements the public static method `void T::move(const T&, T&)`. + */ +template struct is_moveable : decltype(core::detail::test_moveable(0)){}; + +/** + * This type trait is [true_type](http://en.cppreference.com/w/cpp/types/integral_constant) + * if and only if `T` satisfies any of the following: + * 1. [is_fundamental](http://en.cppreference.com/w/cpp/types/is_fundamental), + * 2. [is_enum](http://en.cppreference.com/w/cpp/types/is_enum), + * 3. [is_pointer](http://en.cppreference.com/w/cpp/types/is_pointer), + * 4. or implements the public static method `void T::swap(T&, T&)`. + */ +template struct is_swappable : decltype(core::detail::test_swappable(0)){}; + +/** + * This type trait is [true_type](http://en.cppreference.com/w/cpp/types/integral_constant) + * if and only if `T` satisfies any of the following: + * 1. [is_fundamental](http://en.cppreference.com/w/cpp/types/is_fundamental), + * 2. [is_enum](http://en.cppreference.com/w/cpp/types/is_enum), + * 3. [is_pointer](http://en.cppreference.com/w/cpp/types/is_pointer), + * 4. or implements the public static method `bool T::copy(const T&, T&)`. + */ +template struct is_copyable : decltype(core::detail::test_copyable(0)){}; + +/** + * This type trait is [true_type](http://en.cppreference.com/w/cpp/types/integral_constant) + * if and only if every template argument is also `true_type`. + */ +template struct and_type; +template<> struct and_type<> { + typedef std::true_type type; + static constexpr bool value = true; +}; +template struct and_type { + typedef typename and_type::type type; + static constexpr bool value = and_type::value; +}; +template struct and_type { + typedef std::false_type type; + static constexpr bool value = false; +}; + + +/** + * + */ + +#define DEFAULT_HASH_SEED 0 + +/* these are defined in map.h */ +#if defined(__LP64__) || defined(_WIN64) || defined(__x86_64__) || defined(__ppc64__) + +template +inline uint_fast32_t default_hash(const K& key); + +template +inline uint_fast32_t default_hash(const K* keys, unsigned int length); + +#else +template +inline unsigned int default_hash(const K& key); + +template +inline unsigned int default_hash(const K* keys, unsigned int length); +#endif + +template +inline void set_all_empty(K* keys, unsigned int length) { + memset(keys, 0, sizeof(K) * length); +} + +template +struct hasher { + static inline bool is_empty(const K& key) { + return K::is_empty(key); + } + + static inline void set_empty(K& key) { + K::set_empty(key); + } + + static inline void set_empty(K* keys, unsigned int length) { + K::set_empty(keys, length); + } + + static inline unsigned int hash(const K& key) { + return K::hash(key); + } +}; + +template +struct hasher::value || std::is_enum::value || std::is_pointer::value>::type> { + static inline bool is_empty(const K& key) { + return (key == static_cast(0)); + } + + static inline void set_empty(K& key) { + key = static_cast(0); + } + + static inline void set_empty(K* keys, unsigned int length) { + set_all_empty(keys, length); + } + + static inline unsigned int hash(const K& key) { + return default_hash(key); + } +}; + +/** + * Hashtables in this library require the type `K` to define a special + * "empty value" to indicate that a bucket is empty. For `K` that satisfies + * [is_fundamental](http://en.cppreference.com/w/cpp/types/is_fundamental), + * [is_enum](http://en.cppreference.com/w/cpp/types/is_enum), + * or [is_pointer](http://en.cppreference.com/w/cpp/types/is_pointer), + * this function returns whether `key == static_cast(0)`. Otherwise, this + * function returns `K::is_empty(key)`. Thus, to enable the use of a custom + * struct/class as a hashtable key, it must implement the public static + * function `is_empty`. + */ +template +inline bool is_empty(const K& key) { + return hasher::is_empty(key); +} + +/** + * Hashtables in this library require the type `K` to define a special + * "empty value" to indicate that a bucket is empty. For `K` that satisfies + * [is_fundamental](http://en.cppreference.com/w/cpp/types/is_fundamental), + * [is_enum](http://en.cppreference.com/w/cpp/types/is_enum), + * or [is_pointer](http://en.cppreference.com/w/cpp/types/is_pointer), + * this function sets `key` to `static_cast(0)`. Otherwise, this function + * calls `K::set_empty(key)`. Some hashtable operations use this operation, and + * therefore, if a custom struct/class is used as a hashtable key, it must + * implement the public static function `set_empty` in order to use the + * aforementioned hashtable operations. + */ +template +inline void set_empty(K& key) { + hasher::set_empty(key); +} + + +/** + * Returns the base-2 logarithm of the given 32-bit unsigned integer argument. + * This function assumes the argument is not `0`. + */ +template::value && std::is_integral::value && sizeof(T) == sizeof(uint32_t)>::type* = nullptr> +inline unsigned int log2(T x) { +#if defined(_WIN32) + unsigned long index; + _BitScanReverse(&index, x); + return index; +#else + return (unsigned int) sizeof(uint32_t) * CHAR_BIT - __builtin_clz(x) - 1; +#endif +} + +/** + * Returns the base-2 logarithm of the given 64-bit unsigned integer argument. + * This function assumes the argument is not `0`. + */ +template::value && std::is_integral::value && sizeof(T) == sizeof(uint64_t)>::type* = nullptr> +inline unsigned int log2(T x) { +#if defined(_WIN64) + unsigned long index; + _BitScanReverse64(&index, x); + return index; +#elif defined(_WIN32) + uint32_t left = (x >> 32); + if (left == 0) return log2((uint32_t) x); + else return 32 + log2(left); +#else + return (unsigned int) sizeof(uint64_t) * CHAR_BIT - __builtin_clzll(x) - 1; +#endif +} + +/** + * Returns the base-2 logarithm of the given unsigned integer argument. This + * function assumes the argument is not `0`. + */ +template::value && std::is_integral::value>::type* = nullptr> +constexpr unsigned int static_log2(T x) { + return (x < 2) ? 1 : (1 + static_log2(x)); +} + + +} /* namespace core */ + +#endif /* CORE_H_ */ diff --git a/reasoners/data_preprocess/logic/graph_dataset_gen/core/io.h b/reasoners/data_preprocess/logic/graph_dataset_gen/core/io.h new file mode 100644 index 000000000..0ddad8cac --- /dev/null +++ b/reasoners/data_preprocess/logic/graph_dataset_gen/core/io.h @@ -0,0 +1,1496 @@ +/** + * \file io.h + * + * This file defines serialization and deserialization functions `read` and + * `write`, as well as the `print` function for all fundamental types and core + * data structures. + * + * The file also contains the definition and implementation of the + * `core::memory_stream` class, which may be used to read/write to an in-memory + * buffer. + * + * Scribes + * ------- + * + * The `read`, `write`, and `print` functions in this library follow a very + * regular argument structure: The first argument is the object to + * read/write/print. The second argument is the stream to read from/write + * to/print to. Most functions also require a third (optional) argument called + * the `scribe`, which controls how the subobjects are read/written/printed. + * The scribe is typically passed to the read/write/print functions when + * operating on subobjects of the given object. + * + * Calling read/write/print with core::default_scribe will call the same + * function without the third argument. This largely corresponds to the + * "default" behavior of those functions. + * + * For example, `write(const core::array& a, Stream& out, Writer&&... writer)` + * will call the function `write(a[i], out, writer)` for every element `a[i]` + * in `a`. This enables users to define their own scribes, and define new ways + * to read/write/print objects without having to re-implement the + * read/write/print functions for container structures such as core::array. The + * following example demonstrates how the behavior of the `print` function for + * an array of integers can be altered using a custom scribe. + * + * ```{.cpp} + * #include + * using namespace core; + * + * struct my_string_scribe { + * const char* strings[3]; + * }; + * + * template + * bool print(int i, Stream& out, const my_string_scribe& printer) { + * return print(printer.strings[i], out); + * } + * + * int main() { + * array a = array(8); + * a.add(1); a.add(2); a.add(0); + * + * print(a, stdout); print(' ', stdout); + * + * default_scribe def; + * print(a, stdout, def); print(' ', stdout); + * + * my_string_scribe printer; + * printer.strings[0] = "vici"; + * printer.strings[1] = "veni"; + * printer.strings[2] = "vidi"; + * + * print(a, stdout, printer); + * } + * ``` + * + * This example has expected output `[1, 2, 0] [1, 2, 0] [veni, vidi, vici]`. + * + * + */ + +#ifndef IO_H_ +#define IO_H_ + +#include "array.h" +#include "map.h" + +#include +#include +#include + +#if defined(__APPLE__) +#include + +inline size_t mbrtoc32(char32_t* wc, const char* s, size_t n, mbstate_t* ps) { + wchar_t c; + size_t result = mbrtowc(&c, s, n, ps); + if (result < (size_t) -2 && wc != nullptr) + *wc = c; + return result; +} + +#else /* __APPLE__ */ +#include +#endif /* __APPLE__ */ + +namespace core { + +/** + * Reads `sizeof(T)` bytes from `in` and writes them to the memory referenced + * by `value`. This function does not perform endianness transformations. + * \param in the stream given by a [FILE](https://en.cppreference.com/w/c/io) pointer. + * \tparam T satisfies [is_fundamental](https://en.cppreference.com/w/cpp/types/is_fundamental). + */ +template::value>::type* = nullptr> +inline bool read(T& value, FILE* in) { + return (fread(&value, sizeof(T), 1, in) == 1); +} + +/** + * Reads `length` elements from `in` and writes them to the native array + * `values`. This function does not perform endianness transformations. + * \param in the stream given by a [FILE](https://en.cppreference.com/w/c/io) pointer. + * \tparam T satisfies [is_fundamental](https://en.cppreference.com/w/cpp/types/is_fundamental). + */ +template::value>::type* = nullptr> +inline bool read(T* values, FILE* in, SizeType length) { + return (fread(values, sizeof(T), length, in) == length); +} + +/** + * Writes `sizeof(T)` bytes to `out` from the memory referenced by `value`. + * This function does not perform endianness transformations. + * \param out the stream given by a [FILE](https://en.cppreference.com/w/c/io) pointer. + * \tparam T satisfies [is_fundamental](https://en.cppreference.com/w/cpp/types/is_fundamental). + */ +template::value>::type* = nullptr> +inline bool write(const T& value, FILE* out) { + return (fwrite(&value, sizeof(T), 1, out) == 1); +} + +/** + * Writes `length` elements to `out` from the native array `values`. This + * function does not perform endianness transformations. + * \param out the stream given by a [FILE](https://en.cppreference.com/w/c/io) pointer. + * \tparam T satisfies [is_fundamental](https://en.cppreference.com/w/cpp/types/is_fundamental). + */ +template::value>::type* = nullptr> +inline bool write(const T* values, FILE* out, SizeType length) { + return (fwrite(values, sizeof(T), length, out) == length); +} + +/** + * Prints the character `value` to the stream given by the + * [FILE](https://en.cppreference.com/w/c/io) pointer `out`. + */ +inline bool print(const char& value, FILE* out) { + return (fputc(value, out) != EOF); +} + +/** + * Prints the int `value` to the stream given by the + * [FILE](https://en.cppreference.com/w/c/io) pointer `out`. + */ +inline bool print(const int& value, FILE* out) { + return (fprintf(out, "%d", value) > 0); +} + +/** + * Prints the long `value` to the stream given by the + * [FILE](https://en.cppreference.com/w/c/io) pointer `out`. + */ +inline bool print(const long& value, FILE* out) { + return (fprintf(out, "%ld", value) > 0); +} + +/** + * Prints the long long `value` to the stream given by the + * [FILE](https://en.cppreference.com/w/c/io) pointer `out`. + */ +inline bool print(const long long& value, FILE* out) { + return (fprintf(out, "%lld", value) > 0); +} + +/** + * Prints the unsigned int `value` to the stream given by the + * [FILE](https://en.cppreference.com/w/c/io) pointer `out`. + */ +inline bool print(const unsigned int& value, FILE* out) { + return (fprintf(out, "%u", value) > 0); +} + +/** + * Prints the unsigned long `value` to the stream given by the + * [FILE](https://en.cppreference.com/w/c/io) pointer `out`. + */ +inline bool print(const unsigned long& value, FILE* out) { + return (fprintf(out, "%lu", value) > 0); +} + +/** + * Prints the unsigned long long `value` to the stream given by the + * [FILE](https://en.cppreference.com/w/c/io) pointer `out`. + */ +inline bool print(const unsigned long long& value, FILE* out) { + return (fprintf(out, "%llu", value) > 0); +} + +/** + * Prints the float `value` to the stream given by the + * [FILE](https://en.cppreference.com/w/c/io) pointer `out`. + */ +inline bool print(const float& value, FILE* out) { + return (fprintf(out, "%f", (double) value) > 0); +} + +/** + * Prints the double `value` to the stream given by the + * [FILE](https://en.cppreference.com/w/c/io) pointer `out`. + */ +inline bool print(const double& value, FILE* out) { + return (fprintf(out, "%lf", value) > 0); +} + +/** + * Prints the float `value` to the stream given by the + * [FILE](https://en.cppreference.com/w/c/io) pointer `out` + * with the given `precision`. + */ +inline bool print(const float& value, FILE* out, unsigned int precision) { + return (fprintf(out, "%.*f", precision, (double) value) > 0); +} + +/** + * Prints the double `value` to the stream given by the + * [FILE](https://en.cppreference.com/w/c/io) pointer `out` + * with the given `precision`. + */ +inline bool print(const double& value, FILE* out, unsigned int precision) { + return (fprintf(out, "%.*lf", precision, value) > 0); +} + +/** + * Prints the null-terminated C string `value` to the stream given by the + * [FILE](https://en.cppreference.com/w/c/io) pointer `out`. + */ +inline bool print(const char* values, FILE* out) { + return (fprintf(out, "%s", values) >= 0); +} + +namespace detail { + template static auto test_readable(int32_t) -> + decltype(bool(read(std::declval(), std::declval())), std::true_type{}); + template static auto test_readable(int64_t) -> std::false_type; + + template static auto test_writeable(int32_t) -> + decltype(bool(write(std::declval(), std::declval())), std::true_type{}); + template static auto test_writeable(int64_t) -> std::false_type; + + template static auto test_printable(int32_t) -> + decltype(bool(print(std::declval(), std::declval())), std::true_type{}); + template static auto test_printable(int64_t) -> std::false_type; +} + +/** + * This type trait is [true_type](https://en.cppreference.com/w/cpp/types/integral_constant) + * if and only if the function `bool read(integral&, T&)` is defined where + * `integral` is any integral type. + */ +template struct is_readable : and_type< + decltype(core::detail::test_readable(0)), + decltype(core::detail::test_readable(0)), + decltype(core::detail::test_readable(0)), + decltype(core::detail::test_readable(0)), + decltype(core::detail::test_readable(0)), + decltype(core::detail::test_readable(0)), + decltype(core::detail::test_readable(0)), + decltype(core::detail::test_readable(0)), + decltype(core::detail::test_readable(0)), + decltype(core::detail::test_readable(0)), + decltype(core::detail::test_readable(0)), + decltype(core::detail::test_readable(0)), + decltype(core::detail::test_readable(0)), + decltype(core::detail::test_readable(0))>::type {}; + +/** + * This type trait is [true_type](https://en.cppreference.com/w/cpp/types/integral_constant) + * if and only if the function `bool fwrite(const integral&, T&)` is defined where + * `integral` is any integral type. + */ +template struct is_writeable : and_type< + decltype(core::detail::test_writeable(0)), + decltype(core::detail::test_writeable(0)), + decltype(core::detail::test_writeable(0)), + decltype(core::detail::test_writeable(0)), + decltype(core::detail::test_writeable(0)), + decltype(core::detail::test_writeable(0)), + decltype(core::detail::test_writeable(0)), + decltype(core::detail::test_writeable(0)), + decltype(core::detail::test_writeable(0)), + decltype(core::detail::test_writeable(0)), + decltype(core::detail::test_writeable(0)), + decltype(core::detail::test_writeable(0)), + decltype(core::detail::test_writeable(0)), + decltype(core::detail::test_writeable(0))>::type {}; + +/** + * This type trait is [true_type](https://en.cppreference.com/w/cpp/types/integral_constant) + * if and only if the function `bool print(value, T&)` is defined. + */ +template struct is_printable : and_type< + decltype(core::detail::test_printable(0)), + decltype(core::detail::test_printable(0)), + decltype(core::detail::test_printable(0)), + decltype(core::detail::test_printable(0)), + decltype(core::detail::test_printable(0)), + decltype(core::detail::test_printable(0)), + decltype(core::detail::test_printable(0))>::type {}; + +/** + * Represents a stream to read/write from an in-memory buffer. + */ +struct memory_stream { + /** + * The size of the stream. + */ + unsigned int length; + + /** + * The current position of the stream in the buffer. + */ + unsigned int position; + + /** + * The underlying buffer. + */ + char* buffer; + + /** + * The default constructor does not initialize any fields. + */ + memory_stream() { } + + /** + * Initializes the stream with memory_stream::length given by + * `initial_capacity` and memory_stream::position set to `0`. + * memory_stream::buffer is allocated but not initialized to any value. + */ + memory_stream(unsigned int initial_capacity) : length(initial_capacity), position(0) { + buffer = (char*) malloc(sizeof(char) * length); + if (buffer == NULL) { + fprintf(stderr, "memory_stream ERROR: Unable to initialize buffer.\n"); + exit(EXIT_FAILURE); + } + } + + /** + * Initializes the stream with the memory_stream::buffer given by `buf`, + * memory_stream::length given by `length`, and memory_stream::position set + * to `0`. + */ + memory_stream(const char* buf, unsigned int length) : length(length), position(0) { + buffer = (char*) malloc(sizeof(char) * length); + if (buffer == NULL) { + fprintf(stderr, "memory_stream ERROR: Unable to initialize buffer.\n"); + exit(EXIT_FAILURE); + } + memcpy(buffer, buf, sizeof(char) * length); + } + + ~memory_stream() { + free(buffer); + } + + /** + * Reads a number of bytes given by `bytes` from the memory_stream and + * writes them to `dst`. This function assumes `dst` has sufficient capacity. + */ + inline bool read(void* dst, unsigned int bytes) { + if (position + bytes >= length) + return false; + memcpy(dst, buffer + position, bytes); + position += bytes; + return true; + } + + /** + * Checks whether the stream has sufficient size for an additional number + * of bytes given by `bytes` at its current memory_stream::position. If + * not, this function attempts to expand the buffer to a new size computed + * as `memory_stream::position + bytes`. + */ + inline bool ensure_capacity(unsigned int bytes) { + if (position + bytes <= length) + return true; + + unsigned int new_length = length; + if (!core::expand(buffer, new_length, position + bytes)) + return false; + length = new_length; + return true; + } + + /** + * Writes a number of bytes given by `bytes` from the given native array + * `src` to the current position in this stream. memory_stream::ensure_capacity + * is called to ensure the underlying buffer has sufficient size. + */ + inline bool write(const void* src, unsigned int bytes) { + if (!ensure_capacity(bytes)) + return false; + memcpy(buffer + position, src, bytes); + position += bytes; + return true; + } +}; + +/** + * Reads `sizeof(T)` bytes from `in` and writes them to the memory referenced + * by `value`. This function does not perform endianness transformations. + * \param in a memory_stream. + * \tparam T satisfies [is_fundamental](https://en.cppreference.com/w/cpp/types/is_fundamental). + */ +template::value>::type* = nullptr> +inline bool read(T& value, memory_stream& in) { + return in.read(&value, sizeof(T)); +} + +/** + * Reads `length` elements from `in` and writes them to the native array + * `values`. This function does not perform endianness transformations. + * \param in a memory_stream. + * \tparam T satisfies [is_fundamental](https://en.cppreference.com/w/cpp/types/is_fundamental). + */ +template::value>::type* = nullptr> +inline bool read(T* values, memory_stream& in, unsigned int length) { + return in.read(values, (unsigned int) sizeof(T) * length); +} + +/** + * Writes `sizeof(T)` bytes to `out` from the memory referenced by `value`. + * This function does not perform endianness transformations. + * \param out a memory_stream. + * \tparam T satisfies [is_fundamental](https://en.cppreference.com/w/cpp/types/is_fundamental). + */ +template::value>::type* = nullptr> +inline bool write(const T& value, memory_stream& out) { + return out.write(&value, sizeof(T)); +} + +/** + * Writes `length` elements to `out` from the native array `values`. This + * function does not perform endianness transformations. + * \param out a memory_stream. + * \tparam T satisfies [is_fundamental](https://en.cppreference.com/w/cpp/types/is_fundamental). + */ +template::value>::type* = nullptr> +inline bool write(const T* values, memory_stream& out, unsigned int length) { + return out.write(values, (unsigned int) sizeof(T) * length); +} + +/** + * Reads an array of `n` elements, each with a size of `size` bytes, from the + * memory_stream `in`, to the memory address referenced by `dst`. + * \see This function mirrors the equivalent [fread](http://en.cppreference.com/w/cpp/io/c/fread) + * for [FILE](https://en.cppreference.com/w/c/io) pointer streams. + * \returns the number of elements read. + */ +inline size_t fread(void* dst, size_t size, size_t count, memory_stream& in) { + size_t num_bytes = size * count; + if (in.position + num_bytes > in.length) { + count = (in.length - in.position) / size; + num_bytes = size * count; + } + + memcpy(dst, in.buffer + in.position, num_bytes); + in.position += num_bytes; + return count; +} + +/** + * Writes the array of `n` elements, each with a size of `size` bytes, from the + * memory address referenced by `src` to the memory_stream `out`. + * \see This function mirrors the equivalent [fwrite](http://en.cppreference.com/w/cpp/io/c/fwrite) + * for [FILE](https://en.cppreference.com/w/c/io) pointer streams. + * \returns either `n` if the write is successful, or `0` upon failure. + */ +inline size_t fwrite(const void* src, size_t size, size_t n, memory_stream& out) { + if (out.write(src, (unsigned int) (size * n))) + return n; + else return 0; +} + +/** + * Retrieves the current position in the given memory_stream. + * \see This function mirrors the equivalent [fgetpos](https://en.cppreference.com/w/c/io/fgetpos) + * for [FILE](https://en.cppreference.com/w/c/io) pointer streams. + * \returns 0 on success; nonzero value otherwise. + */ +inline int fgetpos(const memory_stream& stream, fpos_t* pos) { +#if defined(_WIN32) || defined(__APPLE__) /* on Windows or Mac */ + *pos = (fpos_t) stream.position; +#else /* on Windows or Linux */ + pos->__pos = stream.position; +#endif + return 0; +} + +/** + * Sets the current position in the given memory_stream. + * \see This function mirrors the equivalent [fsetpos](https://en.cppreference.com/w/c/io/fsetpos) + * for [FILE](https://en.cppreference.com/w/c/io) pointer streams. + * \returns 0 on success; nonzero value otherwise. + */ +inline int fsetpos(memory_stream& stream, const fpos_t* pos) { +#if defined(_WIN32) || defined(__APPLE__) /* on Windows or Mac */ + stream.position = (unsigned int) *pos; +#else /* on Windows or Linux */ + stream.position = pos->__pos; +#endif + return 0; +} + +/** + * Writes the given character `c` to the memory_stream `out`. + * \see This function mirrors the equivalent [fputc](http://en.cppreference.com/w/cpp/io/c/fputc) + * for [FILE](https://en.cppreference.com/w/c/io) pointer streams. + */ +inline int fputc(int c, memory_stream& out) { + char ch = (char) c; + if (out.write(&ch, sizeof(char))) + return c; + else return EOF; +} + +/** + * Writes the given null-terminated C string `s` to the memory_stream `out`. + * \see This function mirrors the equivalent [fputs](http://en.cppreference.com/w/cpp/io/c/fputs) + * for [FILE](https://en.cppreference.com/w/c/io) pointer streams. + */ +inline int fputs(const char* s, memory_stream& out) { + if (out.write(s, (unsigned int) strlen(s))) + return 1; + else return EOF; +} + +/** + * Writes the given arguments according to the format string `format` to the + * memory_stream `out`. + * \see This function mirrors the equivalent [fprintf](http://en.cppreference.com/w/cpp/io/c/fprintf) + * for [FILE](https://en.cppreference.com/w/c/io) pointer streams. + * \returns the number of bytes written to the stream, or `-1` upon error. + */ +inline int fprintf(memory_stream& out, const char* format, ...) { + va_list argptr; + va_start(argptr, format); + +#if defined(_WIN32) + int required = _vscprintf(format, argptr); + if (!out.ensure_capacity((unsigned int) required + 1)) { + fprintf(stderr, "fprintf ERROR: Unable to expand memory stream buffer.\n"); + va_end(argptr); + return -1; + } + + int written = vsnprintf_s(out.buffer + out.position, out.length - out.position, (size_t) required, format, argptr); +#else + int written = vsnprintf(out.buffer + out.position, out.length - out.position, format, argptr); + if (written < 0) { + va_end(argptr); + return -1; + } else if ((unsigned) written < out.length - out.position) { + va_end(argptr); + out.position += written; + return written; + } + + if (!out.ensure_capacity(written + 1)) { + fprintf(stderr, "fprintf ERROR: Unable to expand memory stream buffer.\n"); + va_end(argptr); + return -1; + } + written = vsnprintf(out.buffer + out.position, out.length - out.position, format, argptr); +#endif + + va_end(argptr); + if (written < 0) return -1; + out.position += written; + return written; +} + +/** + * Prints the character `value` to the stream given by the memory_stream `out`. + */ +inline bool print(const char& value, memory_stream& out) { + return (fputc(value, out) != EOF); +} + +/** + * Prints the int `value` to the stream given by the memory_stream `out`. + */ +inline bool print(const int& value, memory_stream& out) { + return (fprintf(out, "%d", value) > 0); +} + +/** + * Prints the long `value` to the stream given by the memory_stream `out`. + */ +inline bool print(const long& value, memory_stream& out) { + return (fprintf(out, "%ld", value) > 0); +} + +/** + * Prints the long long `value` to the stream given by the memory_stream `out`. + */ +inline bool print(const long long& value, memory_stream& out) { + return (fprintf(out, "%lld", value) > 0); +} + +/** + * Prints the unsigned int `value` to the stream given by the memory_stream + * `out`. + */ +inline bool print(const unsigned int& value, memory_stream& out) { + return (fprintf(out, "%u", value) > 0); +} + +/** + * Prints the unsigned long `value` to the stream given by the memory_stream + * `out`. + */ +inline bool print(const unsigned long& value, memory_stream& out) { + return (fprintf(out, "%lu", value) > 0); +} + +/** + * Prints the unsigned long long `value` to the stream given by the + * memory_stream `out`. + */ +inline bool print(const unsigned long long& value, memory_stream& out) { + return (fprintf(out, "%llu", value) > 0); +} + +/** + * Prints the float `value` to the stream given by the memory_stream `out`. + */ +inline bool print(const float& value, memory_stream& out) { + return (fprintf(out, "%f", (double) value) > 0); +} + +/** + * Prints the double `value` to the stream given by the memory_stream `out`. + */ +inline bool print(const double& value, memory_stream& out) { + return (fprintf(out, "%lf", value) > 0); +} + +/** + * Prints the float `value` to the stream given by the memory_stream `out` with + * the given `precision`. + */ +inline bool print(const float& value, memory_stream& out, unsigned int precision) { + return (fprintf(out, "%.*f", precision, (double) value) > 0); +} + +/** + * Prints the double `value` to the stream given by the memory_stream `out` + * with the given `precision`. + */ +inline bool print(const double& value, memory_stream& out, unsigned int precision) { + return (fprintf(out, "%.*lf", precision, value) > 0); +} + +/** + * Prints the null-terminated C string `value` to the stream given by the + * memory_stream `out`. + */ +inline bool print(const char* values, memory_stream& out) { + return (fprintf(out, "%s", values) >= 0); +} + +/** + * A stream wrapper for reading UTF-32 characters from an underlying multibyte + * stream (such as UTF-8). + */ +template +struct buffered_stream { + Stream& underlying_stream; + char buffer[BufferSize]; + unsigned int position; + unsigned int length; + mbstate_t shift; + + buffered_stream(Stream& underlying_stream) : underlying_stream(underlying_stream), position(0) { + shift = {0}; + fill_buffer(); + } + + inline void fill_buffer() { + length = fread(buffer, sizeof(char), BufferSize, underlying_stream); + } + + /** + * Returns the next UTF-32 character (as a `char32_t`) from the stream. If + * there are no further bytes in the underlying stream or an error occurred, + * `static_cast(-1)` is returned. + */ + char32_t fgetc32() { + static_assert(BufferSize >= MB_LEN_MAX, "BufferSize must be at least MB_LEN_MAX"); + + while (true) + { + if (length == 0) + return static_cast(-1); + + char32_t c; + size_t status = mbrtoc32(&c, buffer + position, sizeof(char) * (length - position), &shift); + if (status == static_cast(-1)) { + /* encoding error occurred */ + return static_cast(-1); + } else if (status == static_cast(-2)) { + /* the character is valid but incomplete */ + position = 0; + fill_buffer(); + continue; + } else { + if (status == 0) + position += 1; + else position += status; + if (position == length) { + position = 0; + fill_buffer(); + } + return c; + } + } + } +}; + +template +struct buffered_stream { + memory_stream& underlying_stream; + mbstate_t shift; + + buffered_stream(memory_stream& underlying_stream) : underlying_stream(underlying_stream) { + shift = {0}; + } + + /** + * Returns the next UTF-32 character (as a `char32_t`) from the stream. If + * there are no further bytes in the underlying stream or an error occurred, + * `static_cast(-1)` is returned. + */ + char32_t fgetc32() { + if (underlying_stream.position == underlying_stream.length) + return static_cast(-1); + + char32_t c; + size_t status = mbrtoc32(&c, underlying_stream.buffer + underlying_stream.position, sizeof(char) * (underlying_stream.length - underlying_stream.position), &shift); + if (status == static_cast(-1) || status == static_cast(-2)) { + /* encoding error occurred or the character is incomplete */ + return static_cast(-1); + } + + if (status == 0) + underlying_stream.position += 1; + else underlying_stream.position += status; + return c; + } +}; + +/** + * Returns the next UTF-32 character (as a `char32_t`) from the buffered_stream + * `input`. If there are no further bytes in the underlying stream or an error + * occurred, `static_cast(-1)` is returned. + */ +template +inline char32_t fgetc32(buffered_stream& input) +{ + return input.fgetc32(); +} + +/** + * A stream wrapper for reading/writing integral types as fixed-width integral + * values. This is useful for cross-platform readability and writeability. + */ +template +struct fixed_width_stream { + Stream& stream; + + fixed_width_stream(Stream& stream) : stream(stream) { } + + template struct type { }; + template struct type { typedef BoolType value; }; + template struct type { typedef CharType value; }; + template struct type { typedef UCharType value; }; + template struct type { typedef ShortType value; }; + template struct type { typedef UShortType value; }; + template struct type { typedef IntType value; }; + template struct type { typedef UIntType value; }; + template struct type { typedef LongType value; }; + template struct type { typedef ULongType value; }; + template struct type { typedef LongLongType value; }; + template struct type { typedef ULongLongType value; }; + template struct type { typedef FloatType value; }; + template struct type { typedef DoubleType value; }; +}; + +/** + * Reads `size(K)` bytes from `in` where `K` is the appropriate template + * argument in the fixed_width_stream and writes them to the memory referenced + * by `value`. This function does not perform endianness transformations. + * \param in a fixed_width_stream. + * \tparam T satisfies [is_fundamental](https://en.cppreference.com/w/cpp/types/is_fundamental). + */ +template::value>::type* = nullptr> +inline bool read(T& value, fixed_width_stream& in) { + typedef typename fixed_width_stream::template type::value value_type; + value_type c; + if (!read(c, in.stream)) return false; + value = (T) c; + return true; +} + +/** + * Reads `length` elements from `in` and writes them to the native array + * `values`. This function does not perform endianness transformations. + * \param in a fixed_width_stream. + * \tparam T satisfies [is_fundamental](https://en.cppreference.com/w/cpp/types/is_fundamental). + */ +template::value>::type* = nullptr> +inline bool read(T* values, fixed_width_stream& in, SizeType length) { + for (SizeType i = 0; i < length; i++) + if (!read(values[i], in)) return false; + return true; +} + +/** + * Writes `sizeof(K)` bytes to `out` from the memory referenced by `value` + * where `K` is the appropriate template argument in the fixed_width_stream. + * This function does not perform endianness transformations. + * \param out a fixed_width_stream. + * \tparam T satisfies [is_fundamental](https://en.cppreference.com/w/cpp/types/is_fundamental). + */ +template::value>::type* = nullptr> +inline bool write(const T& value, fixed_width_stream& out) { + typedef typename fixed_width_stream::template type::value value_type; + return write((value_type) value, out.stream); +} + +/** + * Writes `length` elements to `out` from the native array `values`. This + * function does not perform endianness transformations. + * \param out a fixed_width_stream. + * \tparam T satisfies [is_fundamental](https://en.cppreference.com/w/cpp/types/is_fundamental). + */ +template::value>::type* = nullptr> +inline bool write(const T* values, fixed_width_stream& out, SizeType length) { + for (SizeType i = 0; i < length; i++) + if (!write(values[i], out)) return false; + return true; +} + +/** + * Writes the given null-terminated C string `values` to the stream `out`. + * \tparam Stream satisfies is_writeable. + */ +template::value>::type* = nullptr> +inline bool write(const char* values, Stream& out) { + return write(values, out, strlen(values)); +} + +/** + * The default scribe implementation that provides the default behavior for + * read/write/print functions. + * \see [Section on scribes](#scribes). + */ +struct default_scribe { }; + +/* a type trait for detecting whether the function 'print' is defined with a default_scribe argument */ +namespace detail { + template static auto test_default_read(int) -> + decltype(bool(read(std::declval(), std::declval(), std::declval())), std::true_type{}); + template static auto test_default_read(long) -> std::false_type; + template static auto test_default_write(int) -> + decltype(bool(write(std::declval(), std::declval(), std::declval())), std::true_type{}); + template static auto test_default_write(long) -> std::false_type; + template static auto test_default_print(int) -> + decltype(bool(print(std::declval(), std::declval(), std::declval())), std::true_type{}); + template static auto test_default_print(long) -> std::false_type; +} + +template struct has_default_read : decltype(core::detail::test_default_read(0)){}; +template struct has_default_write : decltype(core::detail::test_default_write(0)){}; +template struct has_default_print : decltype(core::detail::test_default_print(0)){}; + +/** + * Calls and returns `read(value, in)`, dropping the default_scribe argument. + * \tparam Stream satisfies is_readable. + */ +template::value && !has_default_read::value>::type** = nullptr> +inline bool read(T& value, Stream& in, default_scribe& scribe) { + return read(value, in); +} + +/** + * Calls and returns `write(value, out)`, dropping the default_scribe argument. + * \tparam Stream satisfies is_writeable. + */ +template::value && !has_default_write::value>::type* = nullptr> +inline bool write(const T& value, Stream& out, default_scribe& scribe) { + return write(value, out); +} + +/** + * Calls and returns `print(value, out)`, dropping the default_scribe argument. + * \tparam Stream satisfies is_printable. + */ +template::value && !has_default_print::value>::type* = nullptr> +inline auto print(const T& value, Stream& out, default_scribe& scribe) -> decltype(print(value, out)) { + return print(value, out); +} + +/** + * A scribe that prints pointers by dereferencing the pointer and calling the + * appropriate read/write/print function. + * \see [Section on scribes](#scribes). + */ +struct pointer_scribe { }; + +/** + * Allocates memory and stores the address in `value`, and then calls + * `read(*value, in, reader)`, dropping the pointer_scribe argument. Note that + * since `reader` is a variadic argument, it may be empty. + * \tparam Stream satisfies is_readable. + */ +template::value>::type* = nullptr> +inline bool read(T*& value, Stream& in, const pointer_scribe& scribe, Reader&&... reader) { + value = (T*) malloc(sizeof(T)); + if (value == NULL) { + fprintf(stderr, "read ERROR: Out of memory.\n"); + return false; + } else if (!read(*value, in, std::forward(reader)...)) { + free(value); + return false; + } + return true; +} + +/** + * Calls and returns `write(*value, out, writer)`, dropping the pointer_scribe + * argument. Note that since `writer` is a variadic argument, it may be empty. + * \tparam Stream satisfies is_writeable. + */ +template::value>::type* = nullptr> +inline bool write(const T* const value, Stream& out, const pointer_scribe& scribe, Writer&&... writer) { + return write(*value, out, std::forward(writer)...); +} + +/** + * Calls and returns `print(*value, out, printer)`, dropping the pointer_scribe + * argument. Note that since `printer` is a variadic argument, it may be empty. + * \tparam Stream satisfies is_printable. + */ +template::value>::type* = nullptr> +inline bool print(const T* const value, Stream& out, const pointer_scribe& scribe, Printer&&... printer) { + return print(*value, out, std::forward(printer)...); +} + +/** + * The default left delimitter "[" for the array print functions. + */ +char default_left_bracket[] = "["; + +/** + * The default right delimitter "]" for the array print functions. + */ +char default_right_bracket[] = "]"; + +/** + * The default separator between elements ", " for the array print functions. + */ +char default_array_separator[] = ", "; + +/** + * Prints the given native array of `values` each of type `T`, where `length` + * is the number of elements in the array. The output stream is `out`. + * \param printer a scribe for which the function `bool print(const T&, Stream&, Printer&)` + * is defined. Note that since this is a variadic argument, it may be empty. + * \tparam Stream satisfies is_printable. + */ +template::value>::type* = nullptr> +bool print(const T* values, SizeType length, Stream& out, Printer&&... printer) { + if (!print(LeftBracket, out)) return false; + if (length == 0) + return print(RightBracket, out); + if (!print(values[0], out, std::forward(printer)...)) return false; + for (SizeType i = 1; i < length; i++) { + if (!print(Separator, out) || !print(values[i], out, std::forward(printer)...)) + return false; + } + return print(RightBracket, out); +} + +/** + * Prints the given native static array of `values` each of type `T`, where `N` + * is the number of elements in the array. The output stream is `out`. + * \param printer a scribe for which the function `bool print(const T&, Stream&, Printer&)` + * is defined. Note that since this is a variadic argument, it may be empty. + * \tparam Stream satisfies is_printable. + */ +template::value>::type* = nullptr> +bool print(const T (&values)[N], Stream& out, Printer&&... printer) { + if (!print(LeftBracket, out)) return false; + if (N == 0) + return print(RightBracket, out); + if (!print(values[0], out, std::forward(printer)...)) return false; + for (size_t i = 1; i < N; i++) { + if (!print(Separator, out) || !print(values[i], out, std::forward(printer)...)) + return false; + } + return print(RightBracket, out); +} + +/** + * Reads an array of `length` elements from `in` and stores the result in the + * given native array `a`. This function assumes `a` has sufficient capacity. + * \param reader a scribe for which the function `bool read(T&, Stream&, Reader&&...)` + * is defined. Note that since this is a variadic argument, it may be empty. + * \tparam Stream satisfies is_readable. + */ +template::value>::type* = nullptr> +inline bool read(T* a, Stream& in, SizeType length, Reader&&... reader) { + for (SizeType i = 0; i < length; i++) + if (!read(a[i], in, std::forward(reader)...)) return false; + return true; +} + +/** + * Reads an array of `N` elements from `in` and stores the result in the given + * native array `a`. + * \param reader a scribe for which the function `bool read(T&, Stream&, Reader&&...)` + * is defined. Note that since this is a variadic argument, it may be empty. + * \tparam Stream satisfies is_readable. + */ +template::value>::type* = nullptr> +inline bool read(T (&a)[N], Stream& in, Reader&&... reader) { + for (size_t i = 0; i < N; i++) + if (!read(a[i], in, std::forward(reader)...)) return false; + return true; +} + +/** + * Reads a core::array structure from `in` and stores the result in `a`. + * \param a an uninitialized core::array structure. This function initializes + * `a`, and the caller is responsible for its memory and must call free + * to release its memory resources. + * \param reader a scribe for which the function `bool read(T&, Stream&, Reader&&...)` + * is defined. Note that since this is a variadic argument, it may be empty. + * \tparam Stream satisfies is_readable. + */ +template::value>::type* = nullptr> +bool read(array& a, Stream& in, Reader&&... reader) { + size_t length; + if (!read(length, in)) + return false; + size_t capacity = ((size_t) 1) << (core::log2(length == 0 ? 1 : length) + 1); + a.data = (T*) malloc(sizeof(T) * capacity); + if (a.data == NULL) return false; + if (!read(a.data, in, length, std::forward(reader)...)) { + free(a.data); + return false; + } + a.length = length; + a.capacity = capacity; + return true; +} + +/** + * Writes the given native array `a` of elements to `out`, each of type `T`, + * where the number of elements is given by `length`. + * \param writer a scribe for which the function `bool write(const T&, Stream&, Writer&&...)` + * is defined. Note that since this is a variadic argument, it may be empty. + * \tparam Stream satisfies is_writeable. + */ +template::value>::type* = nullptr> +inline bool write(const T* a, Stream& out, SizeType length, Writer&&... writer) { + for (SizeType i = 0; i < length; i++) + if (!write(a[i], out, std::forward(writer)...)) return false; + return true; +} + +/** + * Writes the given native array `a` of elements to `out`, each of type `T`, + * where the number of elements is given by `N`. + * \param writer a scribe for which the function `bool write(const T&, Stream&, Writer&&...)` + * is defined. Note that since this is a variadic argument, it may be empty. + * \tparam Stream satisfies is_writeable. + */ +template::value>::type* = nullptr> +inline bool write(const T (&a)[N], Stream& out, Writer&&... writer) { + for (size_t i = 0; i < N; i++) + if (!write(a[i], out, std::forward(writer)...)) return false; + return true; +} + +/** + * Writes the given core::array structure `a` of elements to `out`, each of type `T`. + * \param writer a scribe for which the function `bool write(const T&, Stream&, Writer&&...)` + * is defined. Note that since this is a variadic argument, it may be empty. + * \tparam Stream satisfies is_writeable. + */ +template::value>::type* = nullptr> +bool write(const array& a, Stream& out, Writer&&... writer) { + return write(a.length, out) + && write(a.data, out, a.length, std::forward(writer)...); +} + +/** + * Prints the given core::array structure `a` of elements to `out`, each of type `T`. + * \param printer a scribe for which the function `bool print(const T&, Stream&, Printer&&...)` + * is defined. Note that since this is a variadic argument, it may be empty. + * \tparam Stream satisfies is_printable. + */ +template::value>::type* = nullptr> +inline bool print(const array& a, Stream&& out, Printer&&... printer) { + return print(a.data, a.length, out, std::forward(printer)...); +} + +/** + * Reads a core::hash_set structure `set` from `in`. + * \param set an uninitialized core::hash_set structure. This function + * initializes `set`, and the caller is responsible for its memory and + * must call free to release its memory resources. + * \param alloc_keys a memory allocation function with prototype + * `void* alloc_keys(size_t count, size_t size)` that allocates space for + * `count` items, each with size `size`, and initializes them such that + * core::is_empty() returns `true` for each element. + * \param reader a scribe for which the function `bool read(T&, Stream&, Reader&&...)` + * is defined. Note that since this is a variadic argument, it may be empty. + * \tparam Stream satisfies is_readable. + */ +template::value>::type* = nullptr> +bool read(hash_set& set, Stream& in, alloc_keys_func alloc_keys, Reader&&... reader) { + unsigned int length; + if (!read(length, in)) return false; + + set.size = 0; + set.capacity = 1 << (core::log2(RESIZE_THRESHOLD_INVERSE * (length == 0 ? 1 : length)) + 1); + set.keys = (T*) alloc_keys(set.capacity, sizeof(T)); + if (set.keys == NULL) return false; + + for (unsigned int i = 0; i < length; i++) { + T& key = *((T*) alloca(sizeof(T))); + if (!read(key, in, std::forward(reader)...)) return false; + move(key, set.keys[set.index_to_insert(key)]); + set.size++; + } + return true; +} + +/** + * Reads a core::hash_set structure `set` from `in`. The keys in the hash_set + * are allocated using [calloc](http://en.cppreference.com/w/c/memory/calloc). + * \param set an uninitialized core::hash_set structure. This function + * initializes `set`, and the caller is responsible for its memory and + * must call free to release its memory resources. + * \param reader a scribe for which the function `bool read(T&, Stream&, Reader&&...)` + * is defined. Note that since this is a variadic argument, it may be empty. + * \tparam Stream satisfies is_readable. + */ +template::value>::type* = nullptr> +inline bool read(hash_set& set, Stream& in, Reader&&... reader) { + return read(set, in, calloc, std::forward(reader)...); +} + +/** + * Writes the given core::hash_set structure `set` to `out`. + * \param writer a scribe for which the function `bool write(const T&, Stream&, Writer&&...)` + * is defined. Note that since this is a variadic argument, it may be empty. + * \tparam Stream satisfies is_writeable. + */ +template::value>::type* = nullptr> +bool write(const hash_set& set, Stream& out, Writer&&... writer) { + if (!write(set.size, out)) return false; + for (unsigned int i = 0; i < set.capacity; i++) { + if (is_empty(set.keys[i])) continue; + if (!write(set.keys[i], out, std::forward(writer)...)) return false; + } + return true; +} + +/** + * Reads a core::hash_map structure `map` from `in`. + * \param map an uninitialized core::hash_map structure. This function + * initializes `map`, and the caller is responsible for its memory and + * must call free to release its memory resources. + * \param alloc_keys a memory allocation function with prototype + * `void* alloc_keys(size_t count, size_t size)` that allocates space for + * `count` items, each with size `size`, and initializes them such that + * core::is_empty() returns `true` for each element of type `K`. + * \param key_reader a scribe for which the function `bool read(K&, Stream&, KeyReader&)` is defined. + * \param value_reader a scribe for which the function `bool read(V&, Stream&, ValueReader&)` is defined. + * \tparam Stream satisfies is_readable. + */ +template::value>::type* = nullptr> +bool read(hash_map& map, + Stream& in, alloc_keys_func alloc_keys, + KeyReader& key_reader, ValueReader& value_reader) +{ + unsigned int length; + if (!read(length, in)) return false; + + map.table.size = 0; + map.table.capacity = 1 << (core::log2(RESIZE_THRESHOLD_INVERSE * (length == 0 ? 1 : length)) + 1); + map.table.keys = (K*) alloc_keys(map.table.capacity, sizeof(K)); + if (map.table.keys == NULL) return false; + map.values = (V*) malloc(sizeof(V) * map.table.capacity); + if (map.values == NULL) { + free(map.table.keys); + return false; + } + + for (unsigned int i = 0; i < length; i++) { + K& key = *((K*) alloca(sizeof(K))); + if (!read(key, in, key_reader)) return false; + + bool contains; + unsigned int bucket; + map.get(key, contains, bucket); + if (!read(map.values[bucket], in, value_reader)) + return false; + if (!contains) { + move(key, map.table.keys[bucket]); + map.table.size++; + } else { + fprintf(stderr, "read WARNING: Serialized hash_map contains duplicates.\n"); + } + } + return true; +} + +/** + * Reads a core::hash_map structure `map` from `in`. + * \param map an uninitialized core::hash_map structure. This function + * initializes `map`, and the caller is responsible for its memory and + * must call free to release its memory resources. + * \param alloc_keys a memory allocation function with prototype + * `void* alloc_keys(size_t count, size_t size)` that allocates space for + * `count` items, each with size `size`, and initializes them such that + * core::is_empty() returns `true` for each element of type `K`. + * \param key_reader a scribe for which the function `bool read(K&, Stream&, KeyReader&)` is defined. + * \tparam Stream satisfies is_readable. + */ +template::value>::type* = nullptr> +inline bool read(hash_map& map, Stream& in, + KeyReader& key_reader, + alloc_keys_func alloc_keys = calloc) +{ + default_scribe scribe; + return read(map, in, alloc_keys, key_reader, scribe); +} + +/** + * Reads a core::hash_map structure `map` from `in`. + * \param map an uninitialized core::hash_map structure. This function + * initializes `map`, and the caller is responsible for its memory and + * must call free to release its memory resources. + * \param alloc_keys a memory allocation function with prototype + * `void* alloc_keys(size_t count, size_t size)` that allocates space for + * `count` items, each with size `size`, and initializes them such that + * core::is_empty() returns `true` for each element of type `K`. + * \tparam Stream satisfies is_readable. + */ +template::value>::type* = nullptr> +inline bool read(hash_map& map, Stream& in, + alloc_keys_func alloc_keys = calloc) +{ + default_scribe scribe; + return read(map, in, alloc_keys, scribe, scribe); +} + +/** + * Writes the core::hash_map structure `map` to `out`. + * \param key_writer a scribe for which the function `bool write(const K&, Stream&, KeyWriter&)` is defined. + * \param value_writer a scribe for which the function `bool write(const V&, Stream&, ValueWriter&)` is defined. + * \tparam Stream satisfies is_writeable. + */ +template::value>::type* = nullptr> +bool write(const hash_map& map, Stream& out, + KeyWriter& key_writer, ValueWriter& value_writer) +{ + if (!write(map.table.size, out)) return false; + for (unsigned int i = 0; i < map.table.capacity; i++) { + if (is_empty(map.table.keys[i])) continue; + if (!write(map.table.keys[i], out, key_writer) + || !write(map.values[i], out, value_writer)) + return false; + } + return true; +} + +/** + * Writes the core::hash_map structure `map` to `out`. + * \param key_writer a scribe for which the function `bool write(const K&, Stream&, KeyWriter&)` is defined. + * \tparam Stream satisfies is_writeable. + */ +template::value>::type* = nullptr> +inline bool write(const hash_map& map, Stream& out, KeyWriter& key_writer) { + default_scribe scribe; + return write(map, out, key_writer, scribe); +} + +/** + * Writes the core::hash_map structure `map` to `out`. + * \tparam Stream satisfies is_writeable. + */ +template::value>::type* = nullptr> +inline bool write(const hash_map& map, Stream& out) { + default_scribe scribe; + return write(map, out, scribe, scribe); +} + +/** + * Reads a core::array_map structure `map` from `in`. + * \param map an uninitialized core::array_map structure. This function + * initializes `map`, and the caller is responsible for its memory and + * must call free to release its memory resources. + * \param key_reader a scribe for which the function `bool read(K&, Stream&, KeyReader&)` is defined. + * \param value_reader a scribe for which the function `bool read(V&, Stream&, ValueReader&)` is defined. + * \tparam Stream satisfies is_readable. + */ +template::value>::type* = nullptr> +bool read(array_map& map, Stream& in, + KeyReader& key_reader, ValueReader& value_reader) +{ + size_t length; + if (!read(length, in)) return false; + + map.size = 0; + map.capacity = 1 << (core::log2(length == 0 ? 1 : length) + 1); + map.keys = (K*) malloc(sizeof(K) * map.capacity); + if (map.keys == NULL) return false; + map.values = (V*) malloc(sizeof(V) * map.capacity); + if (map.values == NULL) { + free(map.keys); + return false; + } + + for (unsigned int i = 0; i < length; i++) { + if (!read(map.keys[i], in, key_reader)) return false; + if (!read(map.values[i], in, value_reader)) return false; + map.size++; + } + return true; +} + +/** + * Reads a core::array_map structure `map` from `in`. + * \param map an uninitialized core::array_map structure. This function + * initializes `map`, and the caller is responsible for its memory and + * must call free to release its memory resources. + * \param key_reader a scribe for which the function `bool read(K&, Stream&, KeyReader&)` is defined. + * \tparam Stream satisfies is_readable. + */ +template::value>::type* = nullptr> +inline bool read(array_map& map, Stream& in, KeyReader& key_reader) { + default_scribe scribe; + return read(map, in, key_reader, scribe); +} + +/** + * Reads a core::array_map structure `map` from `in`. + * \param map an uninitialized core::array_map structure. This function + * initializes `map`, and the caller is responsible for its memory and + * must call free to release its memory resources. + * \tparam Stream satisfies is_readable. + */ +template::value>::type* = nullptr> +inline bool read(array_map& map, Stream& in) { + default_scribe scribe; + return read(map, in, scribe, scribe); +} + +/** + * Writes the given core::array_map structure `map` to `out`. + * \param key_writer a scribe for which the function `bool write(const K&, Stream&, KeyWriter&)` is defined. + * \param value_writer a scribe for which the function `bool write(const V&, Stream&, ValueWriter&)` is defined. + * \tparam Stream satisfies is_writeable. + */ +template::value>::type* = nullptr> +bool write(const array_map& map, Stream& out, + KeyWriter& key_writer, ValueWriter& value_writer) +{ + if (!write(map.size, out)) return false; + for (unsigned int i = 0; i < map.size; i++) { + if (!write(map.keys[i], out, key_writer)) return false; + if (!write(map.values[i], out, value_writer)) return false; + } + return true; +} + +/** + * Writes the given core::array_map structure `map` to `out`. + * \param key_writer a scribe for which the function `bool write(const K&, Stream&, KeyWriter&)` is defined. + * \tparam Stream satisfies is_writeable. + */ +template::value>::type* = nullptr> +inline bool write(const array_map& map, Stream& out, KeyWriter& key_writer) { + default_scribe scribe; + return write(map, out, key_writer, scribe); +} + +/** + * Writes the given core::array_map structure `map` to `out`. + * \tparam Stream satisfies is_writeable. + */ +template::value>::type* = nullptr> +inline bool write(const array_map& map, Stream& out) { + default_scribe scribe; + return write(map, out, scribe, scribe); +} + +/** + * Reads a core::pair structure `p` from `stream` by calling `read(p.key, stream)` and `read(p.value, stream)`. + * \tparam Stream satisfies is_readable. + */ +template::value>::type* = nullptr> +inline bool read(pair& p, Stream& stream) +{ + return read(p.key, stream) && read(p.value, stream); +} + +/** + * Writes the given core::pair structure `p` to `stream` by calling + * `write(p.key, stream)` and `write(p.value, stream)`. + * \tparam Stream satisfies is_writeable. + */ +template::value>::type* = nullptr> +inline bool write(const pair& p, Stream& stream) +{ + return write(p.key, stream) && write(p.value, stream); +} + +} /* namespace core */ + +#endif /* IO_H_ */ diff --git a/reasoners/data_preprocess/logic/graph_dataset_gen/core/lex.h b/reasoners/data_preprocess/logic/graph_dataset_gen/core/lex.h new file mode 100644 index 000000000..1d38055b1 --- /dev/null +++ b/reasoners/data_preprocess/logic/graph_dataset_gen/core/lex.h @@ -0,0 +1,504 @@ +/** + * \file lex.h + * + * This file implements common functionality for lexical analysis, such as + * string comparison, tokenization (splitting strings by whitespace), and + * parsing arithmetic types from strings. + * + * + */ + +#ifndef LEX_H_ +#define LEX_H_ + +#include "map.h" +#include "utility.h" + +#include + +#if defined(__APPLE__) +#include + +inline size_t c32rtomb(char* s, char32_t wc, mbstate_t* ps) { + return wcrtomb(s, (wchar_t) wc, ps); +} + +#else /* __APPLE__ */ +#include +#endif /* __APPLE__ */ + + +namespace core { + +/** + * Compares the strings given by the array structure `first` and the + * null-terminated C string `second`. + * \returns `true` if the strings are equivalent, and `false` otherwise. + */ +/* TODO: test with benchmarks whether we should inline these functions */ +inline bool compare_strings(const array& first, const char* second) { + for (unsigned int i = 0; i < first.length; i++) { + if (first[i] != second[i]) + return false; + } + if (second[first.length] != '\0') + return false; + return true; +} + +/** + * Compares the strings given by the core::string `first` and the native char + * array `second` whose length is given by `second_length`. + * \returns `true` if the strings are equivalent, and `false` otherwise. + */ +inline bool compare_strings(const string& first, const char* second, unsigned int second_length) { + if (first.length != second_length) + return false; + for (unsigned int i = 0; i < first.length; i++) { + if (first[i] != second[i]) + return false; + } + return true; +} + +/** + * Tokenizes the given native char array `str` with length `length`, assigning + * to each unique token an `unsigned int` identifier. These identifiers are + * stored in the core::hash_map `names`. The tokenized identifiers are added to + * the core::array `tokens`. + */ +bool tokenize(const char* str, unsigned int length, + array& tokens, hash_map& names) +{ + bool whitespace = true; + unsigned int token_start = 0; + for (unsigned int i = 0; i < length; i++) { + if (whitespace) { + if (!isspace(str[i])) { + token_start = i; + whitespace = false; + } + } else { + if (isspace(str[i])) { + unsigned int id; + if (!get_token(string(str + token_start, i - token_start), id, names) + || !tokens.add(id)) + return false; + whitespace = true; + } + } + } + + if (!whitespace) { + unsigned int id; + if (!get_token(string(str + token_start, length - token_start), id, names) + || !tokens.add(id)) + return false; + } + return true; +} + +/** + * Attempts to parse the string given by `token` as a `double`. + * \tparam CharArray a string type that implements two fields: (1) `data` which + * returns the underlying `char*` array, and (2) `length` which returns + * the length of the string. + * \returns `true` if successful, or `false` if there is insufficient memory or + * `token` is not an appropriate string representation of a floating-point + * number. + */ +template +inline bool parse_float(const CharArray& token, double& value) { + char* buffer = (char*) malloc(sizeof(char) * (token.length + 1)); + if (buffer == NULL) { + fprintf(stderr, "parse_float ERROR: Unable to allocate temporary string buffer.\n"); + return false; + } + memcpy(buffer, token.data, sizeof(char) * token.length); + buffer[token.length] = '\0'; + + char* end_ptr; + value = strtod(buffer, &end_ptr); + if (*end_ptr != '\0') { + free(buffer); + return false; + } + free(buffer); + return true; +} + +/** + * Attempts to parse the string given by `token` as an `unsigned int`. + * \tparam CharArray a string type that implements two fields: (1) `data` which + * returns the underlying `char*` array, and (2) `length` which returns + * the length of the string. + * \returns `true` if successful, or `false` if there is insufficient memory or + * `token` is not an appropriate string representation of a unsigned + * integer. + */ +template +inline bool parse_uint(const CharArray& token, unsigned int& value, unsigned int base = 0) { + char* buffer = (char*) malloc(sizeof(char) * (token.length + 1)); + if (buffer == NULL) { + fprintf(stderr, "parse_uint ERROR: Unable to allocate temporary string buffer.\n"); + return false; + } + memcpy(buffer, token.data, sizeof(char) * token.length); + buffer[token.length] = '\0'; + + char* end_ptr; + value = strtoul(buffer, &end_ptr, base); + if (*end_ptr != '\0') { + free(buffer); + return false; + } + free(buffer); + return true; +} + +/** + * Attempts to parse the string given by `token` as an `unsigned int`. + * \tparam CharArray a string type that implements two fields: (1) `data` which + * returns the underlying `char*` array, and (2) `length` which returns + * the length of the string. + * \returns `true` if successful, or `false` if there is insufficient memory or + * `token` is not an appropriate string representation of a unsigned + * integer. + */ +template +inline bool parse_ulonglong(const CharArray& token, unsigned long long& value) { + char* buffer = (char*) malloc(sizeof(char) * (token.length + 1)); + if (buffer == NULL) { + fprintf(stderr, "parse_ulonglong ERROR: Unable to allocate temporary string buffer.\n"); + return false; + } + memcpy(buffer, token.data, sizeof(char) * token.length); + buffer[token.length] = '\0'; + + char* end_ptr; + value = strtoull(buffer, &end_ptr, 0); + if (*end_ptr != '\0') { + free(buffer); + return false; + } + free(buffer); + return true; +} + +/** + * Attempts to parse the string given by `token` as an `int`. + * \tparam CharArray a string type that implements two fields: (1) `data` which + * returns the underlying `char*` array, and (2) `length` which returns + * the length of the string. + * \returns `true` if successful, or `false` if there is insufficient memory or + * `token` is not an appropriate string representation of a integer. + */ +template +inline bool parse_int(const CharArray& token, int& value) { + char* buffer = (char*) malloc(sizeof(char) * (token.length + 1)); + if (buffer == NULL) { + fprintf(stderr, "parse_int ERROR: Unable to allocate temporary string buffer.\n"); + return false; + } + memcpy(buffer, token.data, sizeof(char) * token.length); + buffer[token.length] = '\0'; + + char* end_ptr; + value = strtol(buffer, &end_ptr, 0); + if (*end_ptr != '\0') { + free(buffer); + return false; + } + free(buffer); + return true; +} + +/** + * Attempts to parse the string given by `token` as a `long`. + * \tparam CharArray a string type that implements two fields: (1) `data` which + * returns the underlying `char*` array, and (2) `length` which returns + * the length of the string. + * \returns `true` if successful, or `false` if there is insufficient memory or + * `token` is not an appropriate string representation of a long. + */ +template +inline bool parse_long(const CharArray& token, long& value) { + char* buffer = (char*) malloc(sizeof(char) * (token.length + 1)); + if (buffer == NULL) { + fprintf(stderr, "parse_longlong ERROR: Unable to allocate temporary string buffer.\n"); + return false; + } + memcpy(buffer, token.data, sizeof(char) * token.length); + buffer[token.length] = '\0'; + + char* end_ptr; + value = strtol(buffer, &end_ptr, 0); + if (*end_ptr != '\0') { + free(buffer); + return false; + } + free(buffer); + return true; +} + +/** + * Attempts to parse the string given by `token` as a `long`. + * \tparam CharArray a string type that implements two fields: (1) `data` which + * returns the underlying `char*` array, and (2) `length` which returns + * the length of the string. + * \returns `true` if successful, or `false` if there is insufficient memory or + * `token` is not an appropriate string representation of a long. + */ +template +inline bool parse_long_long(const CharArray& token, long long& value) { + char* buffer = (char*)malloc(sizeof(char) * (token.length + 1)); + if (buffer == NULL) { + fprintf(stderr, "parse_longlong ERROR: Unable to allocate temporary string buffer.\n"); + return false; + } + memcpy(buffer, token.data, sizeof(char) * token.length); + buffer[token.length] = '\0'; + + char* end_ptr; + value = strtoll(buffer, &end_ptr, 0); + if (*end_ptr != '\0') { + free(buffer); + return false; + } + free(buffer); + return true; +} + +/** + * Attempts to parse the string given by `token` as an `unsigned int`. + * \param base if `0`, the numeric base of the integer is detected + * automatically in the same way as [strtoul](http://en.cppreference.com/w/cpp/string/byte/strtoul). + * Otherwise, the numeric base can be specified explicitly. + * \returns `true` if successful, or `false` if `token` is not an appropriate + * string representation of an unsigned integer. + */ +template +inline bool parse_uint(const char (&token)[N], unsigned int& value, unsigned int base = 0) { + char buffer[N + 1]; + for (unsigned int i = 0; i < N; i++) + buffer[i] = token[i]; + buffer[N] = '\0'; + + char* end_ptr; + value = strtoul(buffer, &end_ptr, base); + return (*end_ptr == '\0'); +} + +/** + * Represents a position in a file. This structure is typically used to provide + * informative errors during lexical analysis of data from a file. + */ +struct position { + /** + * The line number of the position in the file. + */ + unsigned int line; + + /** + * The column number of the position in the file. + */ + unsigned int column; + + /** + * Constructs the position structure with the given `line` and `column`. + */ + position(unsigned int line, unsigned int column) : + line(line), column(column) { } + + /** + * Constructs the position structure by copying from `p`. + */ + position(const position& p) : + line(p.line), column(p.column) { } + + /** + * Returns a position with the column number increased by `i`. + */ + position operator + (unsigned int i) const { + return position(line, column + i); + } + + /** + * Returns a position with the column number decreased by `i`. + */ + position operator - (unsigned int i) const { + return position(line, column - i); + } + + /** + * Copies the position structure from `src` into `dst`. + */ + static inline bool copy(const position& src, position& dst) { + dst.line = src.line; + dst.column = src.column; + return true; + } +}; + +/** + * A structure representing a single token during lexical analysis. This + * structure is generic, intended for use across multiple lexical analyzers. + */ +template +struct lexical_token { + /** + * The generic type of this token. + */ + TokenType type; + + /** + * The start position (inclusive) of the token in the source file. + */ + position start; + + /** + * The end position (exclusive) of the token in the source file. + */ + position end; + + /** + * An (optional) string representing the contents of the token. + */ + string text; +}; + +/** + * Prints the given lexical_token `token` to the output `stream`. + * \tparam Printer a scribe type for which the functions + * `print(const TokenType&, Stream&, Printer&)` and + * `print(const core::string& s, Stream&, Printer&)` are defined. + */ +template +bool print(const lexical_token& token, Stream& stream, Printer& printer) { + bool success = true; + success &= print(token.type, stream, printer); + if (!is_empty(token.text)) { + success &= print('(', stream); + success &= print(token.text, stream, printer); + success &= print(')', stream); + } + return success; +} + +/** + * Reports an error with the given message `error` as a null-terminated C + * string at the given source file position `pos` to [stderr](http://en.cppreference.com/w/cpp/io/c). + */ +inline void read_error(const char* error, const position& pos) { + fprintf(stderr, "ERROR at %d:%d: %s.\n", pos.line, pos.column, error); +} + +/** + * Constructs a lexical_token with the given `start` and `end` positions, and + * TokenType `type`, with an empty lexical_token::text message and appends it + * to the `tokens` array. + */ +template +bool emit_token(array>& tokens, + const position& start, const position& end, TokenType type) +{ + if (!tokens.ensure_capacity(tokens.length + 1)) { + fprintf(stderr, "emit_token ERROR: Unable to create token.\n"); + return false; + } + + lexical_token& new_token = tokens[(unsigned int) tokens.length]; + new_token.text.data = NULL; + new_token.type = type; + new_token.start = start; + new_token.end = end; + tokens.length++; + return true; +} + +/** + * Constructs a lexical_token with the given `start` and `end` positions, and + * TokenType `type`, with lexical_token::text copied from `token` and appends + * it to the `tokens` array. + */ +template +bool emit_token( + array>& tokens, array& token, + const position& start, const position& end, TokenType type) +{ + if (!tokens.ensure_capacity(tokens.length + 1)) { + fprintf(stderr, "emit_token ERROR: Unable to create token.\n"); + return false; + } + + lexical_token& new_token = tokens[(unsigned int) tokens.length]; + if (!init(new_token.text, token.data, (unsigned int) token.length)) { + fprintf(stderr, "emit_token ERROR: Unable to create string.\n"); + return false; + } + new_token.type = type; + new_token.start = start; + new_token.end = end; + tokens.length++; + token.clear(); + return true; +} + +/** + * Frees every element in the given `tokens` array. This function does not free + * the array itself. + */ +template +void free_tokens(array>& tokens) { + for (unsigned int i = 0; i < tokens.length; i++) + if (tokens[i].text.data != NULL) + core::free(tokens[i].text); +} + +/** + * Inspects the element at the given `index` in the `tokens` array. If `index` + * is not out of bounds, and the token at that index has type that matches the + * given `type`, the function returns `true`. Otherwise, an error message is + * printed to [stderr](http://en.cppreference.com/w/cpp/io/c) indicating that + * the expected token was missing, with its `name` as part of the error + * message, and `false` is returned. + */ +template +bool expect_token(const array>& tokens, + const unsigned int& index, TokenType type, const char* name) +{ + FILE* out = stderr; + if (index == tokens.length) { + /* unexpected end of input */ + fprintf(out, "ERROR: Unexpected end of input. Expected %s.\n", name); + return false; + } + else if (tokens[index].type != type) { + fprintf(out, "ERROR at %d:%d: Unexpected token ", tokens[index].start.line, tokens[index].start.column); + print(tokens[index].type, out); + fprintf(out, ". Expected %s.\n", name); + return false; + } + return true; +} + +/** + * Appends the given wide character `next` to the char array `token` which + * represents a multi-byte string. + */ +inline bool append_to_token( + array& token, char32_t next, mbstate_t& shift) +{ + if (!token.ensure_capacity(token.length + MB_CUR_MAX)) + return false; + size_t written = c32rtomb(token.data + token.length, next, &shift); + if (written == static_cast(-1)) + return false; + token.length += written; + return true; +} + +} /* namespace core */ + +#endif /* LEX_H_ */ diff --git a/reasoners/data_preprocess/logic/graph_dataset_gen/core/map.h b/reasoners/data_preprocess/logic/graph_dataset_gen/core/map.h new file mode 100644 index 000000000..9ddba2a02 --- /dev/null +++ b/reasoners/data_preprocess/logic/graph_dataset_gen/core/map.h @@ -0,0 +1,2508 @@ +/** + * \file map.h + * + * This file contains the hash_set, hash_map, and array_map data structures. It + * also defines the default_hash function, which provides a default hash + * algorithm (currently implemented using [xxhash](https://github.com/Cyan4973/xxHash)). + * + * + */ + +#ifndef MAP_H_ +#define MAP_H_ + +#include +#include + +#define XXH_INLINE_ALL +#include "array.h" +#include "xxhash.h" + + +namespace core { + +/** + * For hash_set and hash_map, `capacity` is the size of the underlying array + * (i.e. the number of buckets), whereas `size` is the number of elements (i.e. + * number of non-empty buckets). The functions hash_set::check_size and + * hash_map::check_size compute the load factor `size / capacity` and compare + * it to RESIZE_THRESHOLD. If the load factor is too large, the hashtable is + * resized and the capacity is increased. + */ +#define RESIZE_THRESHOLD 1 / 2 + +/** + * The multiplicative inverse of RESIZE_THRESHOLD. + */ +#define RESIZE_THRESHOLD_INVERSE 2 / 1 + +/** + * The multiplicative factor by which hash_set, hash_map, and array_map capacity is changed. + */ +#define RESIZE_FACTOR 2 + +/** + * A function pointer type describing a function that returns a pointer to + * allocated memory. The first argument is the number of elements to allocate, + * and the second argument is the number of bytes for each element. + * [calloc](http://en.cppreference.com/w/c/memory/calloc) is an example of a + * function with this type. + */ +typedef void*(alloc_keys_func)(size_t, size_t); + + +/* forward declarations */ +#if !defined(DOXYGEN_IGNORE) + +template +struct hash_set; + +template +struct hash_map; + +template +struct array_map; + +#endif /* DOXYGEN_IGNORE */ + + +#if defined(__LP64__) || defined(_WIN64) || defined(__x86_64__) || defined(__ppc64__) +template +inline uint_fast32_t default_hash(const K& key) { + return (uint_fast32_t) XXH64(&key, sizeof(K), Seed); +} + +template +inline uint_fast32_t default_hash(const K* keys, unsigned int length) { + return (uint_fast32_t) XXH64(keys, sizeof(K) * length, Seed); +} + +#else +/** + * Evaluates the hash function of the given value `key` with the given `Seed` using the default implementation. + */ +template +inline unsigned int default_hash(const K& key) { + return XXH32(&key, sizeof(K), Seed); +} + +/** + * Evaluates the hash function of the given native array of values `keys` with the given `Seed` using the default implementation. + */ +template +inline unsigned int default_hash(const K* keys, unsigned int length) { + return XXH32(keys, sizeof(K) * length, Seed); +} +#endif + +template +struct key_value_metric { + const KeyMetric& key_metric; + const ValueMetric& value_metric; + + constexpr key_value_metric(const KeyMetric& key_metric, const ValueMetric& value_metric) : + key_metric(key_metric), value_metric(value_metric) { } +}; + +template +inline constexpr key_value_metric make_key_value_metric( + const KeyMetric& key_metric, const ValueMetric& value_metric) { + return key_value_metric(key_metric, value_metric); +} + +inline constexpr key_value_metric make_key_value_metric() { + return make_key_value_metric(default_metric(), default_metric()); +} + + +/** + * + */ + +/** + * An iterator implementation, similar to those in the Standard Template + * Library, to enable iteration of elements in a hash_set. This iterator is + * typically initialized using hash_set::begin. + * + * This definition enables the use of range-based for loops. + */ +template +struct hash_set_iterator { + typedef typename std::conditional&, hash_set&>::type container_type; + + /** + * The type of the entries returned by this iterator. If this is a const + * iterator, `value_type` is `const T&`. Otherwise, `value_type` is `T&`. + */ + typedef typename std::conditional::type value_type; + + container_type set; + unsigned int position; + + hash_set_iterator(container_type& set, unsigned int position) : set(set), position(position) { } + + /** + * Returns whether this iterator is in the same position as `other`. This + * function assumes the two iterators were created from the same hash_set, + * and that it was not modified. + */ + inline bool operator != (const hash_set_iterator& other) const { + return position != other.position; + } + + /** + * Returns the element in the hash_set at the current iterator position. + * This function assumes the hash_set was not resized and no element was + * removed since the last call to either the operator `++` or the + * constructor of this iterator, whichever came later. + */ + inline value_type operator * () { + return set.keys[position]; + } + + /** + * Advances the position of the iterator to the next element in the hash_set. + */ + inline const hash_set_iterator& operator ++ () { + do { + ++position; + } while (position < set.capacity && hasher::is_empty(set.keys[position])); + return *this; + } +}; + +/** + * An iterator implementation, similar to those in the Standard Template + * Library, to enable iteration of elements in a hash_map. This iterator is + * typically initialized using hash_map::begin. + * + * This definition enables the use of range-based for loops. + */ +template +struct hash_map_iterator { + typedef typename std::conditional&, hash_map&>::type container_type; + + /** + * The type of the entries returned by this iterator. If this is a const + * iterator, `value_type` is `core::pair`. Otherwise, + * `value_type` is `core::pair`. + */ + typedef typename std::conditional, pair>::type value_type; + + container_type map; + unsigned int position; + + hash_map_iterator(container_type& map, unsigned int position) : map(map), position(position) { } + + /** + * Returns whether this iterator is in the same position as `other`. This + * function assumes the two iterators were created from the same hash_map, + * and that it was not modified. + */ + inline bool operator != (const hash_map_iterator& other) const { + return position != other.position; + } + + /** + * Returns the entry in the hash_map at the current iterator position. This + * function assumes the hash_map was not resized and no element was removed + * since the last call to either the operator `++` or the constructor of + * this iterator, whichever came later. + */ + inline value_type operator * () { + return { map.table.keys[position], map.values[position] }; + } + + /** + * Advances the position of the iterator to the next entry in the hash_map. + */ + inline const hash_map_iterator& operator ++ () { + do { + ++position; + } while (position < map.table.capacity && hasher::is_empty(map.table.keys[position])); + return *this; + } +}; + +/** + * An iterator implementation, similar to those in the Standard Template + * Library, to enable iteration of elements in an array_map. This iterator is + * typically initialized using array_map::begin. + * + * This definition enables the use of range-based for loops. + */ +template +struct array_map_iterator { + typedef typename std::conditional&, array_map&>::type container_type; + + /** + * The type of the entries returned by this iterator. If this is a const + * iterator, `value_type` is `core::pair`. Otherwise, + * `value_type` is `core::pair`. + */ + typedef typename std::conditional, pair>::type value_type; + + container_type map; + size_t position; + + array_map_iterator(container_type map, size_t position) : map(map), position(position) { } + + /** + * Returns whether this iterator is in the same position as `other`. This + * function assumes the two iterators were created from the same array_map, + * and that it was not modified. + */ + inline bool operator != (const array_map_iterator& other) const { + return position != other.position; + } + + /** + * Returns the entry in the array_map at the current iterator position. + * This function assumes the array_map was not resized and no element was + * removed since the last call to either the operator `++` or the + * constructor of this iterator, whichever came later. + */ + inline value_type operator * () { + return{ map.keys[position], map.values[position] }; + } + + /** + * Advances the position of the iterator to the next entry in the array_map. + */ + inline const array_map_iterator& operator ++ () { + ++position; + return *this; + } +}; + +/** + * Returns `true` only if `probe > start` and `probe <= end` where `probe`, `start`, + * and `end` are in the additive group of integers modulo the capacity of this set. + */ +inline bool index_between(unsigned int probe, unsigned int start, unsigned int end) { + if (end >= start) { + return (probe > start && probe <= end); + } else { + return (probe <= end || probe > start); + } +} + +/** + * An unordered associative container that contains a set of unique elements, + * each of type `T`. The elements are stored in the native array + * hash_set::keys, which has capacity hash_set::capacity. To compute the index + * of an element in hash_set::keys, we first compute its hash. To do so: + * 1. If `T` [is_pointer](http://en.cppreference.com/w/cpp/types/is_pointer), + * [is_fundamental](http://en.cppreference.com/w/cpp/types/is_fundamental), + * or [is_enum](http://en.cppreference.com/w/cpp/types/is_enum), + * core::default_hash() provides the hash. + * 2. For all other types, the `T::hash` function provides the hash. + * Once the hash is computed, the index into hash_set::keys is computed using + * `hash % hash_set::capacity` (modular division by the capacity). + * + * The above approach could produce the same index for distinct elements. This + * event is known as a *collision*. We use + * [linear probing](https://en.wikipedia.org/wiki/Linear_probing) + * to resolve collisions: When adding an element to the hash_set, we compute + * its `index` using the above procedure, then we inspect + * `hash_set::keys[index]` and use core::is_empty() to determine if another + * element already occupies that position. If so, we try the next position + * `(index + 1) % hash_set::capacity`. We continue until we find an empty + * index, and the element is inserted there. + * + * Thus, the function core::is_empty() is used to determine if a position in + * hash_set::keys is occupied or empty. The total number of occupied positions + * is given by hash_set::size. **WARNING:** If hash_set::keys becomes full, the + * above linear probing mechanism could lead to an infinite loop. The function + * hash_set::check_size should be used whenever adding new elements to avoid + * this scenario. + * + * **Performance:** This data structure provides constant-time access and + * modification, given that the load factor (`size / capacity`) is not too + * large. + * + * + * Below is an example of a simple use-case of hash_set, where the expected + * output is `a.contains(2): false, a.contains(3): true, -4 9 3`. + * + * ```{.cpp} + * #include + * #include + * using namespace core; + * + * int main() { + * hash_set a = hash_set(8); + * a.add(-1); a.add(-4); + * a.add(3); a.add(9); + * a.remove(-1); + * + * printf("a.contains(2): %s, ", a.contains(2) ? "true" : "false"); + * printf("a.contains(3): %s, ", a.contains(3) ? "true" : "false"); + * for (int element : a) + * printf("%d ", element); + * } + * ``` + * + * + * However, if `a` is not allocated on the stack, the destructor will not be + * automatically called, and so it must be freed manually using `core::free` or + * `hash_set::free`. In the example below, the expected output is the same as + * that of the program above: `a.contains(2): false, a.contains(3): true, -4 9 3`. + * + * ```{.cpp} + * #include + * #include + * using namespace core; + * + * int main() { + * hash_set& a = *((hash_set*) alloca(sizeof(hash_set))); + * hash_set_init(a, 8); + * a.add(-1); a.add(-4); + * a.add(3); a.add(9); + * a.remove(-1); + * + * printf("a.contains(2): %s, ", a.contains(2) ? "true" : "false"); + * printf("a.contains(3): %s, ", a.contains(3) ? "true" : "false"); + * for (int element : a) + * printf("%d ", element); + * free(a); + * } + * ``` + * + * + * A number of member functions require `T` to be + * [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). + * If this is not the case, those operations can be performed directly on the + * public fields, as in the following example. In addition, when using a custom + * struct/class with hash_set, it must implement public static functions, like + * `hash`, `is_empty`, and `move`, as well as the operator `==`. The expected + * output of the following example is `first second `. + * + * ```{.cpp} + * #include + * #include + * #include + * using namespace core; + * + * struct custom_string { + * char* buffer; + * + * static unsigned int hash(const custom_string& s) { + * return default_hash(s.buffer, strlen(s.buffer)); + * } + * + * static bool is_empty(const custom_string& s) { + * return s.buffer == NULL; + * } + * + * static void move(const custom_string& src, custom_string& dst) { + * dst.buffer = src.buffer; + * } + * + * static void free(custom_string& s) { + * core::free(s.buffer); + * } + * }; + * + * inline bool operator == (const custom_string& first, const custom_string& second) { + * if (first.buffer == NULL) + * return second.buffer == NULL; + * return strcmp(first.buffer, second.buffer) == 0; + * } + * + * bool init(custom_string& s, const char* src) { + * s.buffer = (char*) malloc(sizeof(char) * (strlen(src) + 1)); + * if (s.buffer == NULL) + * return false; + * memcpy(s.buffer, src, sizeof(char) * (strlen(src) + 1)); + * return true; + * } + * + * int main() { + * custom_string first, second; + * init(first, "first"); + * init(second, "second"); + * + * hash_set a = hash_set(8); + * a.check_size(a.size + 2); + * + * bool contains; unsigned int index; + * index = a.index_of(first, contains); + * if (!contains) { + * core::move(first, a.keys[index]); + * a.size++; + * } + * + * index = a.index_of(second, contains); + * if (!contains) { + * core::move(second, a.keys[index]); + * a.size++; + * } + * + * for (const custom_string& s : a) + * printf("%s ", s.buffer); + * for (custom_string& s : a) + * free(s); + * } + * ``` + * + * \tparam T the generic type of the elements in the set. `T` must satisfy either: + * 1. [is_fundamental](http://en.cppreference.com/w/cpp/types/is_fundamental), + * 2. [is_enum](http://en.cppreference.com/w/cpp/types/is_enum), + * 3. [is_pointer](http://en.cppreference.com/w/cpp/types/is_pointer), + * 4. implements the public static method `unsigned int hash(const T&)`, + * the public static method `void is_empty(const T&)`, implements the + * operators `==`, and satisfies is_moveable. Some operations also + * require the operator `!=`, and public static methods + * `void set_empty(T&)` and `void set_empty(T*, unsigned int)`. + * **NOTE:** The first argument to the `==` and `!=` operators may be + * empty. + */ +/* TODO: consider other collision resolution mechanisms */ +template +struct hash_set +{ + /** + * The native array of keys underlying the hashtable. + */ + T* keys; + + /** + * The capacity of hash_set::keys. + */ + unsigned int capacity; + + /** + * The number of elements in the hashtable (i.e. the number of non-empty buckets). + */ + unsigned int size; + + /** + * Constructs the hash_set with the given `initial_capacity`. + * \param alloc_keys a memory allocation function with prototype + * `void* alloc_keys(size_t count, size_t size)` that allocates space for + * `count` items, each with size `size`, and initializes them such that + * core::is_empty() returns `true` for each element. + */ + hash_set(unsigned int initial_capacity, alloc_keys_func alloc_keys = calloc) { + if (!initialize(initial_capacity, alloc_keys)) { + fprintf(stderr, "hash_set ERROR: Unable to allocate memory.\n"); + exit(EXIT_FAILURE); + } + } + + /** + * Constructs the hash_set and inserts the given native array of elements + * `set` with given `length`. + * \param alloc_keys a memory allocation function with prototype + * `void* alloc_keys(size_t count, size_t size)` that allocates space for + * `count` items, each with size `size`, and initializes them such that + * core::is_empty() returns `true` for each element. + */ + hash_set(const T* set, unsigned int length, + alloc_keys_func alloc_keys = calloc) : + hash_set(length * RESIZE_THRESHOLD_INVERSE + 1, alloc_keys) + { + for (unsigned int i = 0; i < length; i++) + insert(set[i]); + } + + /** + * Constructs the hash_set and inserts the given + * [initializer_list](http://en.cppreference.com/w/cpp/language/list_initialization) + * of elements `list`. + * \param alloc_keys a memory allocation function with prototype + * `void* alloc_keys(size_t count, size_t size)` that allocates space for + * `count` items, each with size `size`, and initializes them such that + * core::is_empty() returns `true` for each element. + */ + hash_set(const std::initializer_list& list, + alloc_keys_func alloc_keys = calloc) : + hash_set(list.size() * RESIZE_THRESHOLD_INVERSE + 1, alloc_keys) + { + typename std::initializer_list::iterator i; + for (i = list.begin(); i != list.end(); i++) + insert(*i); + } + + ~hash_set() { free(); } + + /** + * Forces the underlying hash_set::keys to be resized to the requested + * `capacity`. + * + * **WARNING:** If `new_capacity <= hash_set::size`, the hashtable could + * become full during the resize process, leading to an infinite loop due + * to the linear probing collision resolution mechanism. + * + * \param alloc_keys a memory allocation function with prototype + * `void* alloc_keys(size_t count, size_t size)` that allocates space for + * `count` items, each with size `size`, and initializes them such that + * core::is_empty() returns `true` for each element. + */ + bool resize(unsigned int new_capacity, + alloc_keys_func alloc_keys = calloc) + { + T* old_keys = keys; + + keys = (T*) alloc_keys(new_capacity, sizeof(T)); + if (keys == NULL) { + /* revert changes and return error */ + keys = old_keys; + return false; + } + + /* iterate through keys and re-hash the elements */ + unsigned int old_capacity = capacity; + capacity = new_capacity; + for (unsigned int i = 0; i < old_capacity; i++) { + if (!hasher::is_empty(old_keys[i])) + core::move(old_keys[i], keys[next_empty(old_keys[i])]); + } + core::free(old_keys); + return true; + } + + /** + * This function first determines whether `hash_set::size < hash_set::capacity * RESIZE_THRESHOLD`. + * If not, the capacity of the underlying hash_set::keys is increased by + * RESIZE_FACTOR until the condition is satisfied. This is useful to ensure + * the hashtable is sufficiently large when preparing to add new elements. + * \param alloc_keys a memory allocation function with prototype + * `void* alloc_keys(size_t count, size_t size)` that allocates space for + * `count` items, each with size `size`, and initializes them such that + * core::is_empty() returns `true` for each element. + * \returns `true` if the resize was successful, and `false` if there is insufficient memory. + */ + inline bool check_size(alloc_keys_func alloc_keys = calloc) { + return check_size(size, alloc_keys); + } + + /** + * For a requested number of elements `new_size`, this function first + * determines whether `new_size < hash_set::capacity * RESIZE_THRESHOLD`. + * If not, the capacity of the underlying hashtable is increased by + * RESIZE_FACTOR until the condition is satisfied. This is useful to ensure + * the hashtable is sufficiently large when preparing to add new elements. + * \param alloc_keys a memory allocation function with prototype + * `void* alloc_keys(size_t count, size_t size)` that allocates space for + * `count` items, each with size `size`, and initializes them such that + * core::is_empty() returns `true` for each element. + * \returns `true` if the resize was successful, and `false` if there is insufficient memory. + */ + inline bool check_size(unsigned int new_size, alloc_keys_func alloc_keys = calloc) + { + while (new_size >= capacity * RESIZE_THRESHOLD) { + if (!resize(RESIZE_FACTOR * capacity, alloc_keys)) { + fprintf(stderr, "hash_set.put ERROR: Unable to resize hashtable.\n"); + return false; + } + } + return true; + } + + /** + * Add the given `element` to this set. The assignment operator is used to + * insert each element, and so this function should not be used if `T` is not + * [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). + * In such a case, insertion should be performed manually using + * hash_set::index_of to find the appropriate index and directly modifying + * hash_set::keys and hash_set::size. See the example in the hash_set + * description. + * \param alloc_keys a memory allocation function with prototype + * `void* alloc_keys(size_t count, size_t size)` that allocates space for + * `count` items, each with size `size`, and initializes them such that + * core::is_empty() returns `true` for each element. + * \tparam T is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). + */ + bool add(const T& element, alloc_keys_func alloc_keys = calloc) + { + if (!check_size(size, alloc_keys)) return false; + insert(element); + return true; + } + + /** + * Adds all the elements in the hash_set `elements` to this set. The + * assignment operator is used to insert each element, and so this function + * should not be used if `T` is not + * [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). + * In such a case, insertion should be performed manually using + * hash_set::index_of or hash_set::index_to_insert to find the appropriate + * index and directly modifying hash_set::keys and hash_set::size. See the + * example in the hash_set description. + * \param alloc_keys a memory allocation function with prototype + * `void* alloc_keys(size_t count, size_t size)` that allocates space for + * `count` items, each with size `size`, and initializes them such that + * core::is_empty() returns `true` for each element. + * \tparam T is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). + */ + bool add_all(const hash_set& elements, + alloc_keys_func alloc_keys = calloc) + { + if (!check_size(size + elements.size, alloc_keys)) return false; + for (unsigned int i = 0; i < elements.capacity; i++) + if (!hasher::is_empty(elements.keys[i])) + insert(elements.keys[i]); + return true; + } + + /** + * Adds all the elements in the native array `elements` with length `count` + * to this set. The assignment operator is used to insert each element, and + * so this function should not be used if `T` is not + * [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). + * In such a case, insertion should be performed manually using + * hash_set::index_of or hash_set::index_to_insert and direct modification + * of the public fields. + * \param alloc_keys a memory allocation function with prototype + * `void* alloc_keys(size_t count, size_t size)` that allocates space for + * `count` items, each with size `size`, and initializes them such that + * core::is_empty() returns `true` for each element. + * \tparam T is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). + */ + bool add_all(const T* elements, unsigned int count, + alloc_keys_func alloc_keys = calloc) + { + if (!check_size(size + count, alloc_keys)) return false; + for (unsigned int i = 0; i < count; i++) + insert(elements[i]); + return true; + } + + /** + * This function removes the given `element` from the set. This function + * does not free the removed element. + * \returns `true` if the element is removed, and `false` if the set does not contain `element`. + */ + bool remove(const T& element) + { +#if !defined(NDEBUG) + if (is_empty(element)) + fprintf(stderr, "hash_set.remove WARNING: Specified key is empty.\n"); +#endif + + unsigned int index = hasher::hash(element) % capacity; + while (true) { + if (keys[index] == element) { + break; + } else if (hasher::is_empty(keys[index])) + return false; + index = (index + 1) % capacity; + } + + remove_at(index); + return true; + } + + template + bool remove(const T& element, V* values) + { +#if !defined(NDEBUG) + if (is_empty(element)) + fprintf(stderr, "hash_set.remove WARNING: Specified key is empty.\n"); +#endif + + unsigned int index = hasher::hash(element) % capacity; + while (true) { + if (keys[index] == element) { + break; + } else if (hasher::is_empty(keys[index])) + return false; + index = (index + 1) % capacity; + } + + remove_at(values, index); + return true; + } + + /** + * This function removes the element at the bucket given by `index`. This + * function assumes that an element is located at the given bucket with the + * correct provided hash value. This function does not free the removed + * element. + */ + void remove_at(unsigned int index) + { + unsigned int last = index; + unsigned int search = (index + 1) % capacity; + while (!hasher::is_empty(keys[search])) + { + unsigned int search_hash = hasher::hash(keys[search]) % capacity; + if (!index_between(search_hash, last, search)) { + core::move(keys[search], keys[last]); + last = search; + } + search = (search + 1) % capacity; + } + + hasher::set_empty(keys[last]); + size--; + } + + /** + * Returns `true` if `element` exists in the set, and `false` otherwise. + */ + bool contains(const T& element) const + { +#if !defined(NDEBUG) + if (is_empty(element)) + fprintf(stderr, "hashtable.contains WARNING: Specified key is empty.\n"); + if (size == capacity) + fprintf(stderr, "hashtable.contains WARNING: Hashtable is full!\n"); +#endif + + unsigned int index = hasher::hash(element) % capacity; + while (true) { + if (keys[index] == element) { + return true; + } else if (hasher::is_empty(keys[index])) { + return false; + } + index = (index + 1) % capacity; + } + } + + /** + * If the given `element` exists in this set, this function returns the + * index of the bucket that contains it. If not, this function returns the + * index where the key would be located, if it had existed in the set. + */ + unsigned int index_of(const T& element) const + { +#if !defined(NDEBUG) + if (is_empty(element)) + fprintf(stderr, "hashtable.index_of WARNING: Specified key is empty.\n"); + if (size == capacity) + fprintf(stderr, "hashtable.index_of WARNING: Hashtable is full!\n"); +#endif + + unsigned int index = hasher::hash(element) % capacity; + while (keys[index] != element && !hasher::is_empty(keys[index])) + index = (index + 1) % capacity; + return index; + } + + /** + * If the given `element` exists in this set, this function returns the + * index of the bucket that contains it, and sets `contains` to `true`. + * If `element` is not in the set, `contains` is set to false, and this + * function returns the index where the key would be located, if it had + * existed in the set. + */ + inline unsigned int index_of( + const T& element, bool& contains) const + { + unsigned int hash_value; + return index_of(element, contains, hash_value); + } + + /** + * If the given `element` exists in this set, this function returns the + * index of the bucket that contains it, and sets `contains` to `true`. + * If `element` is not in the set, `contains` is set to false, and this + * function returns the index where the key would be located, if it had + * existed in the set. In any case, the evaluated hash function of + * `element` is stored in `hash_value`. + */ + unsigned int index_of(const T& element, + bool& contains, unsigned int& hash_value) const + { +#if !defined(NDEBUG) + if (is_empty(element)) + fprintf(stderr, "hashtable.index_of WARNING: Specified key is empty.\n"); + if (size == capacity) + fprintf(stderr, "hashtable.index_of WARNING: Hashtable is full!\n"); +#endif + + hash_value = hasher::hash(element); + unsigned int index = hash_value % capacity; + while (true) { + if (keys[index] == element) { + contains = true; + return index; + } else if (hasher::is_empty(keys[index])) { + contains = false; + return index; + } + index = (index + 1) % capacity; + } + } + + /** + * For a given `element`, this function computes and returns the index of + * the bucket where the element would be inserted, for example by a call to + * hash_set::add. `contains` is set to `true` if and only if the given + * `element` already exists in the set. + */ + inline unsigned int index_to_insert(const T& element, bool& contains) + { +#if !defined(NDEBUG) + if (size == capacity) + fprintf(stderr, "hashtable.index_to_insert WARNING: Hashtable is full!\n"); +#endif + unsigned int index = hasher::hash(element) % capacity; + while (true) { + if (hasher::is_empty(keys[index])) { + contains = false; break; + } if (keys[index] == element) { + contains = true; break; + } + index = (index + 1) % capacity; + } + return index; + } + + /** + * For a given `element`, this function computes and returns the index of + * the bucket where the element would be inserted, for example by a call to + * hash_set::add, **assuming** the given element is not in the set. + */ + inline unsigned int index_to_insert(const T& element) + { +#if !defined(NDEBUG) + if (size == capacity) + fprintf(stderr, "hashtable.index_to_insert WARNING: Hashtable is full!\n"); +#endif + unsigned int index = hasher::hash(element) % capacity; + while (true) { + if (hasher::is_empty(keys[index])) break; + index = (index + 1) % capacity; + } + return index; + } + + /** + * Removes all elements from this hash_set. Note that this function does + * not free each element beforehand. + */ + void clear() { + hasher::set_empty(keys, capacity); + size = 0; + } + + /** + * Returns `true` if this hash_set is a subset of `other`. + */ + bool is_subset(const hash_set& other) const + { + for (unsigned int i = 0; i < capacity; i++) + if (!hasher::is_empty(keys[i]) && !other.contains(keys[i])) + return false; + return true; + } + + bool equals(const hash_set& other) const + { + if (size != other.size) return false; + return is_subset(other); + } + + /** + * Returns a hash_set_iterator pointing to the first element in this container. + * + * **NOTE:** Unlike the libstdc++ [unordered_set](http://en.cppreference.com/w/cpp/container/unordered_set) + * and [unordered_map](http://en.cppreference.com/w/cpp/container/unordered_map) + * iterators, we do not keep a linked list among the elements, so if rapid + * iteration over elements is more critical than rapid queries, consider + * using an array_map. + */ + inline hash_set_iterator begin() { + return hash_set_iterator(*this, first_empty()); + } + + /** + * Returns a hash_set_iterator pointing to the end of this container. + * + * **NOTE:** Unlike the libstdc++ [unordered_set](http://en.cppreference.com/w/cpp/container/unordered_set) + * and [unordered_map](http://en.cppreference.com/w/cpp/container/unordered_map) + * iterators, we do not keep a linked list among the elements, so if rapid + * iteration over elements is more critical than rapid queries, consider + * using an array_map. + */ + inline hash_set_iterator end() { + return hash_set_iterator(*this, capacity); + } + + /** + * Returns a const hash_set_iterator pointing to the first element in this container. + * + * **NOTE:** Unlike the libstdc++ [unordered_set](http://en.cppreference.com/w/cpp/container/unordered_set) + * and [unordered_map](http://en.cppreference.com/w/cpp/container/unordered_map) + * iterators, we do not keep a linked list among the elements, so if rapid + * iteration over elements is more critical than rapid queries, consider + * using an array_map. + */ + inline hash_set_iterator begin() const { + return hash_set_iterator(*this, first_empty()); + } + + /** + * Returns a const hash_set_iterator pointing to the end of this container. + * + * **NOTE:** Unlike the libstdc++ [unordered_set](http://en.cppreference.com/w/cpp/container/unordered_set) + * and [unordered_map](http://en.cppreference.com/w/cpp/container/unordered_map) + * iterators, we do not keep a linked list among the elements, so if rapid + * iteration over elements is more critical than rapid queries, consider + * using an array_map. + */ + inline hash_set_iterator end() const { + return hash_set_iterator(*this, capacity); + } + + /** + * Swaps the contents of the hash_set `first` with that of `second`. + */ + static void swap(hash_set& first, hash_set& second) { + core::swap(first.keys, second.keys); + core::swap(first.capacity, second.capacity); + core::swap(first.size, second.size); + } + + /** + * Moves the contents of the hash_set `src` into `dst`. Note this function + * does not copy the contents of the underlying hash_set::keys, it merely + * copies the pointer. + */ + static inline void move(const hash_set& src, hash_set& dst) { + dst.keys = src.keys; + dst.capacity = src.capacity; + dst.size = src.size; + } + + /** + * Copies the contents of the hash_set `src` into `dst`. + * \param alloc_keys a memory allocation function with prototype + * `void* alloc_keys(size_t count, size_t size)` that allocates space for + * `count` items, each with size `size`, and initializes them such that + * core::is_empty() returns `true` for each element. + */ + static inline bool copy(const hash_set& src, hash_set& dst, alloc_keys_func alloc_keys = calloc) { + dst.capacity = src.capacity; + dst.size = src.size; + dst.keys = (T*) alloc_keys(src.capacity, sizeof(T)); + if (dst.keys == NULL) return false; + + for (unsigned int i = 0; i < src.capacity; i++) { + if (is_empty(src.keys[i])) continue; + if (!copy(src.keys[i], dst.keys[i])) { + free(dst); return false; + } + } + return true; + } + + template + static inline long unsigned int size_of(const hash_set& set, const Metric& metric) + { + long unsigned int sum = core::size_of(set.capacity) + core::size_of(set.size); + for (unsigned int i = 0; i < set.capacity; i++) { + if (is_empty(set.keys[i])) + sum += sizeof(T); + else sum += core::size_of(set.keys[i], metric); + } + return sum; + } + + /** + * Frees hash_set::keys. This should not be used if `set` was constructed + * on the stack, as the destructor will automatically free hash_set::data. + * The elements of `set` are not freed. + */ + static inline void free(hash_set& set) { set.free(); } + +private: + bool initialize( + unsigned int initial_capacity, + alloc_keys_func alloc_keys) + { +#if !defined(NDEBUG) + if (initial_capacity == 0) + fprintf(stderr, "hashtable.initialize WARNING: Initial capacity is zero.\n"); +#endif + + size = 0; + capacity = initial_capacity; + + keys = (T*) alloc_keys(capacity, sizeof(T)); + return (keys != NULL); + } + + inline void free() { + core::free(keys); + } + + inline void place( + const T& element, unsigned int index) + { +#if !defined(NDEBUG) + if (is_empty(element)) + fprintf(stderr, "hashtable.place WARNING: Specified key is empty.\n"); +#endif + + keys[index] = element; + } + + inline void insert_unique(const T& element) + { + place(element, next_empty(element)); + } + + inline unsigned int first_empty() const { + unsigned int index = 0; + while (index < capacity && hasher::is_empty(keys[index])) + index++; + return index; + } + + inline unsigned int next_empty(const T& element) + { +#if !defined(NDEBUG) + if (size == capacity) + fprintf(stderr, "hashtable.next_empty WARNING: Hashtable is full!\n"); +#endif + unsigned int index = hasher::hash(element) % capacity; + while (!hasher::is_empty(keys[index])) + index = (index + 1) % capacity; + return index; + } + + inline void insert(const T& element) + { + bool contains; + unsigned int index = index_to_insert(element, contains); + if (!contains) { + place(element, index); + size++; + } + } + + template + void remove_at(V* values, unsigned int index) + { + unsigned int last = index; + unsigned int search = (index + 1) % capacity; + if (!hasher::is_empty(keys[search])) + { + do { + unsigned int search_hash = hasher::hash(keys[search]) % capacity; + if (!index_between(search_hash, last, search)) { + core::move(keys[search], keys[last]); + core::move(values[search], values[last]); + last = search; + } + search = (search + 1) % capacity; + } while (!hasher::is_empty(keys[search])); + } + + hasher::set_empty(keys[last]); + size--; + } + + template + friend struct hash_map; + + template + friend bool hash_set_init(hash_set& set, + unsigned int capacity, alloc_keys_func alloc_keys); + + template + friend bool hash_set_init(hash_set& set, + const std::initializer_list& list, + alloc_keys_func alloc_keys); +}; + +/** + * Initializes the hash_set `set` with the given initial `capacity`. + * \param alloc_keys a memory allocation function with prototype + * `void* alloc_keys(size_t count, size_t size)` that allocates space for + * `count` items, each with size `size`, and initializes them such that + * core::is_empty() returns `true` for each element. + */ +template +bool hash_set_init( + hash_set& set, unsigned int capacity, + alloc_keys_func alloc_keys = calloc) +{ + if (!set.initialize(capacity, alloc_keys)) { + fprintf(stderr, "hash_set_init ERROR: Unable to allocate memory.\n"); + return false; + } + return true; +} + +/** + * Swaps the underlying buffers of the given hash_sets `first` and `second`. + */ +template +void swap(hash_set& first, hash_set& second) { + T* keys_swap = first.keys; + first.keys = second.keys; + second.keys = keys_swap; + + unsigned int ui_swap = first.size; + first.size = second.size; + second.size = ui_swap; + + ui_swap = first.capacity; + first.capacity = second.capacity; + second.capacity = ui_swap; +} + +/** + * An unordered associative container that contains a set of key-value pairs, + * where the keys are unique and have type `K` and the values have type `V`. + * The keys are stored in a hash_set structure, and the values are stored in + * hash_map::values, which is a native array parallel to the key array in + * hash_set::keys. To compute the index of a key-value pair, we compute the + * index of the key using the algorithm described in hash_set. The value has + * the same index in hash_map::values. This structure uses the same linear + * probing approach to resolve collisions (i.e. when two distinct keys are + * computed to have the same index). + * + * **WARNING:** As with hash_set, if the map becomes full, the linear probing + * mechanism could lead to an infinite loop in many hash_map operations. The + * function hash_map::check_size should be used whenever adding new elements + * to avoid this scenario. + * + * **Performance:** This data structure provides constant-time access and + * modification, given that the load factor (`size / capacity`) is not too + * large. + * + * Below is a simple example using a hash_map, where the expected output is + * `a.get(3): c, a.get(4): d, 1:a 4:d 3:c`. + * + * ```{.cpp} + * #include + * #include + * using namespace core; + * + * int main() { + * hash_map a = hash_map(8); + * a.put(1, 'a'); a.put(2, 'b'); + * a.put(3, 'c'); a.put(4, 'd'); + * a.remove(2); + * + * printf("a.get(3): %c, ", a.get(3)); + * printf("a.get(4): %c, ", a.get(4)); + * for (const auto& entry : a) + * printf("%d:%c ", entry.key, entry.value); + * } + * ``` + * + * + * However, if `a` is not allocated on the stack, the destructor will not be + * automatically called, and so it must be freed manually using `core::free` + * or `hash_map::free`. In the example below, the expected output is the same + * as that of the program above: `a.get(3): c, a.get(4): d, 1:a 4:d 3:c`. + * + * ```{.cpp} + * #include + * #include + * using namespace core; + * + * int main() { + * hash_map& a = *((hash_map*) malloc(sizeof(hash_map))); + * hash_map_init(a, 8); + * a.put(1, 'a'); a.put(2, 'b'); + * a.put(3, 'c'); a.put(4, 'd'); + * a.remove(2); + * + * printf("a.get(3): %c, ", a.get(3)); + * printf("a.get(4): %c, ", a.get(4)); + * for (const auto& entry : a) + * printf("%d:%c ", entry.key, entry.value); + * free(a); + * } + * ``` + * + * + * A number of member functions require `K` and `V` to be + * [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). + * If this is not the case, those operations can be performed directly on the + * public fields, as in the following example. In addition, when using a custom + * struct/class as the key type in hash_map, it must implement public static + * functions, like `hash`, `is_empty`, and `move`, as well as the operator + * `==`. The expected output of the following example is `first:1 second:2 `. + * + * ```{.cpp} + * #include + * #include + * #include + * using namespace core; + * + * struct custom_string { + * char* buffer; + * + * static unsigned int hash(const custom_string& s) { + * return default_hash(s.buffer, strlen(s.buffer)); + * } + * + * static bool is_empty(const custom_string& s) { + * return s.buffer == NULL; + * } + * + * static void move(const custom_string& src, custom_string& dst) { + * dst.buffer = src.buffer; + * } + * + * static void free(custom_string& s) { + * core::free(s.buffer); + * } + * }; + * + * inline bool operator == (const custom_string& first, const custom_string& second) { + * if (first.buffer == NULL) + * return second.buffer == NULL; + * return strcmp(first.buffer, second.buffer) == 0; + * } + * + * bool init(custom_string& s, const char* src) { + * s.buffer = (char*) malloc(sizeof(char) * (strlen(src) + 1)); + * if (s.buffer == NULL) + * return false; + * memcpy(s.buffer, src, sizeof(char) * (strlen(src) + 1)); + * return true; + * } + * + * int main() { + * custom_string key_one, key_two; + * int value_one = 1, value_two = 2; + * init(key_one, "first"); + * init(key_two, "second"); + * + * hash_map a = hash_map(8); + * a.check_size(a.table.size + 2); + * + * bool contains; unsigned int index; + * a.get(key_one, contains, index); + * if (!contains) { + * core::move(key_one, a.table.keys[index]); + * a.values[index] = value_one; + * a.table.size++; + * } + * + * a.get(key_two, contains, index); + * if (!contains) { + * core::move(key_two, a.table.keys[index]); + * a.values[index] = value_two; + * a.table.size++; + * } + * + * for (const auto& entry : a) + * printf("%s:%d ", entry.key.buffer, entry.value); + * for (auto entry : a) + * free(entry.key); + * } + * ``` + * + * \tparam K the generic type of the keys in the map. `K` must satisfy either: + * 1. [is_fundamental](http://en.cppreference.com/w/cpp/types/is_fundamental), + * 2. [is_enum](http://en.cppreference.com/w/cpp/types/is_enum), + * 3. [is_pointer](http://en.cppreference.com/w/cpp/types/is_pointer), + * 4. implements the public static method `unsigned int hash(const T&)`, + * the public static method `void is_empty(const T&)`, implements the + * operators `==`, and satisfies is_moveable. Some operations also + * require the operator `!=`, and public static methods + * `void set_empty(T&)` and `void set_empty(T*, unsigned int)`. + * **NOTE:** The first argument to the `==` and `!=` operators may be + * empty. + * \tparam V the generic type of the values in the map. `V` must satisfy is_moveable. + */ +/* TODO: consider other collision resolution mechanisms */ +template +struct hash_map +{ + /** + * The type of the keys in this hash_map. + */ + typedef K key_type; + + /** + * The type of the values in this hash_map. + */ + typedef V value_type; + + /** + * The underlying hashtable containing the keys. + */ + hash_set table; + + /** + * An array parallel to hash_set::keys in hash_map::table, + * containing a value at every non-empty bucket index. + */ + V* values; + + /** + * Constructs the hash_map with the given initial `capacity`. + * \param alloc_keys a memory allocation function with prototype + * `void* alloc_keys(size_t count, size_t size)` that allocates space for + * `count` items, each with size `size`, and initializes them such that + * core::is_empty() returns `true` for each key. + */ + hash_map(unsigned int capacity, alloc_keys_func alloc_keys = calloc) : + table(capacity, alloc_keys) + { + if (!initialize_values()) { + fprintf(stderr, "hash_map ERROR: Unable to allocate memory.\n"); + exit(EXIT_FAILURE); + } + } + + /** + * Constructs the hash_map and inserts the array of keys `map`, where each + * element in `map` is interpreted as a key, and its corresponding value is + * the index of the element. + * \param alloc_keys a memory allocation function with prototype + * `void* alloc_keys(size_t count, size_t size)` that allocates space for + * `count` items, each with size `size`, and initializes them such that + * core::is_empty() returns `true` for each key. + * \tparam V satisfies [is_integral](http://en.cppreference.com/w/cpp/types/is_integral). + */ + hash_map(const K* map, unsigned int length, + alloc_keys_func alloc_keys = calloc) : + table(length * RESIZE_THRESHOLD_INVERSE + 1, alloc_keys) + { + if (!initialize_values()) { + fprintf(stderr, "hash_map ERROR: Unable to allocate memory.\n"); + exit(EXIT_FAILURE); + } + for (unsigned int i = 0; i < length; i++) + insert(map[i], i); + } + + /** + * Constructs the hash_map and inserts the list of key-value pairs given by + * the native arrays `keys` and `values`. + * \param alloc_keys a memory allocation function with prototype + * `void* alloc_keys(size_t count, size_t size)` that allocates space for + * `count` items, each with size `size`, and initializes them such that + * core::is_empty() returns `true` for each key. + */ + hash_map(const K* keys, const V* values, + unsigned int length, alloc_keys_func alloc_keys = calloc) : + table(length * RESIZE_THRESHOLD_INVERSE + 1, alloc_keys) + { + if (!initialize_values()) { + fprintf(stderr, "hash_map ERROR: Unable to allocate memory.\n"); + exit(EXIT_FAILURE); + } + for (unsigned int i = 0; i < length; i++) + insert(keys[i], values[i]); + } + + /** + * Constructs the hash_map and inserts the given + * [initializer_list](http://en.cppreference.com/w/cpp/language/list_initialization) + * of core::pair entries. + * \param alloc_keys a memory allocation function with prototype + * `void* alloc_keys(size_t count, size_t size)` that allocates space for + * `count` items, each with size `size`, and initializes them such that + * core::is_empty() returns `true` for each key. + */ + hash_map(const std::initializer_list>& list, + alloc_keys_func alloc_keys = calloc) : + table(list.size() * RESIZE_THRESHOLD_INVERSE + 1, alloc_keys) + { + if (!initialize_values()) { + fprintf(stderr, "hash_map ERROR: Unable to allocate memory.\n"); + exit(EXIT_FAILURE); + } + typename std::initializer_list>::iterator i; + for (i = list.begin(); i != list.end(); i++) + insert(i->key, i->value); + } + + ~hash_map() { + ::free(values); + } + + /** + * Forces the underlying hash_map::table.keys and hash_map::values to be + * resized to the requested `capacity`. + * + * **WARNING:** If `new_capacity <= hash_map::table.size`, the hashtable + * could become full during the resize process, leading to an infinite loop + * due to the linear probing collision resolution mechanism. + * + * \param alloc_keys a memory allocation function with prototype + * `void* alloc_keys(size_t count, size_t size)` that allocates space for + * `count` items, each with size `size`, and initializes them such that + * core::is_empty() returns `true` for each key. + */ + bool resize(unsigned int new_capacity, + alloc_keys_func alloc_keys = calloc) + { + K* old_keys = table.keys; + V* old_values = values; + + table.keys = (K*) alloc_keys(new_capacity, sizeof(K)); + if (table.keys == NULL) { + /* revert changes and return error */ + table.keys = old_keys; + return false; + } + + values = (V*) malloc(sizeof(V) * new_capacity); + if (values == NULL) { + /* revert changes and return error */ + ::free(table.keys); + table.keys = old_keys; + values = old_values; + return false; + } + + /* iterate through keys and re-hash the elements */ + unsigned int old_capacity = table.capacity; + table.capacity = new_capacity; + for (unsigned int i = 0; i < old_capacity; i++) { + if (!hasher::is_empty(old_keys[i])) { + unsigned int new_bucket = table.next_empty(old_keys[i]); + core::move(old_keys[i], table.keys[new_bucket]); + core::move(old_values[i], values[new_bucket]); + } + } + ::free(old_keys); + ::free(old_values); + return true; + } + + /** + * This function first determines whether `hash_map::table.size < hash_map::table.capacity * RESIZE_THRESHOLD`. + * If not, the capacities of the underlying hash_map::table and + * hash_map::values is increased by RESIZE_FACTOR until the condition is + * satisfied. This is useful to ensure the hashtable is sufficiently large + * when preparing to add new elements. + * \param alloc_keys a memory allocation function with prototype + * `void* alloc_keys(size_t count, size_t size)` that allocates space for + * `count` items, each with size `size`, and initializes them such that + * core::is_empty() returns `true` for each key. + * \returns `true` if the resize was successful, and `false` if there is insufficient memory. + */ + inline bool check_size(alloc_keys_func alloc_keys = calloc) { + return check_size(table.size, alloc_keys); + } + + /** + * For a requested number of elements `new_size`, this function first + * determines whether `new_size < hash_map::table.capacity * RESIZE_THRESHOLD`. + * If not, the capacities of the underlying hash_map::table and + * hash_map::values is increased by RESIZE_FACTOR until the condition is + * satisfied. This is useful to ensure the hashtable is sufficiently large + * when preparing to add new elements. + * \param alloc_keys a memory allocation function with prototype + * `void* alloc_keys(size_t count, size_t size)` that allocates space for + * `count` items, each with size `size`, and initializes them such that + * core::is_empty() returns `true` for each key. + * \returns `true` if the resize was successful, and `false` if there is insufficient memory. + */ + inline bool check_size(unsigned int new_size, alloc_keys_func alloc_keys = calloc) + { + while (new_size >= table.capacity * RESIZE_THRESHOLD) { + if (!resize(RESIZE_FACTOR * table.capacity, alloc_keys)) { + fprintf(stderr, "hash_map.put ERROR: Unable to resize hashtable.\n"); + return false; + } + } + return true; + } + + /** + * Adds the given key-value pair to this map. The assignment operator is + * used insert the entry, and so this function should not be used if `K` + * and `V` are not [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). + * In such a case, insertion should be performed manually by using + * hash_map::get or hash_set::index_to_insert to find the appropriate index + * and directly modifying hash_map::table.keys, hash_map::table.size, and + * hash_map::values. See the example in the hash_map description. + * \param alloc_keys a memory allocation function with prototype + * `void* alloc_keys(size_t count, size_t size)` that allocates space for + * `count` items, each with size `size`, and initializes them such that + * core::is_empty() returns `true` for each key. + * \tparam K is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). + * \tparam V is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). + */ + bool put(const K& key, const V& value, + alloc_keys_func alloc_keys = calloc) + { + if (!check_size(table.size, alloc_keys)) return false; + insert(key, value); + return true; + } + + /** + * Adds the given key-value pairs in `elements` to this map. The assignment + * operator is used insert each entry, and so this function should not be + * used if `K` and `V` are not [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). + * In such a case, insertion should be performed manually by using + * hash_map::get or hash_map::index_to_insert to find the appropriate index + * and directly modifying hash_map::table.keys, hash_map::table.size, and + * hash_map::values. See the example in the hash_map description. + * \param alloc_keys a memory allocation function with prototype + * `void* alloc_keys(size_t count, size_t size)` that allocates space for + * `count` items, each with size `size`, and initializes them such that + * core::is_empty() returns `true` for each key. + * \tparam K is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). + * \tparam V is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). + */ + bool put_all(const hash_map& elements, + alloc_keys_func alloc_keys = calloc) + { + if (!check_size(table.size + elements.table.size, alloc_keys)) + return false; + for (unsigned int i = 0; i < elements.table.capacity; i++) + if (!hasher::is_empty(elements.table.keys[i])) + insert(elements.table.keys[i], elements.values[i]); + return true; + } + + /** + * This function removes `key` and associated value from the map. This + * function does not free the removed element. + * \returns `true` if the element is removed, and `false` if the set does not contain `element`. + */ + inline bool remove(const K& key) + { + return table.remove(key, values); + } + + /** + * This function removes the entry at the bucket given by `index`. This + * function assumes that an entry is located at the given bucket with the + * correct provided hash value. This function does not free the removed key or + * value. + */ + inline void remove_at(unsigned int index) + { + table.remove_at(values, index); + } + + /** + * Retrieves the value associated with the given `key`. This function + * assumes the given key exists in the map. + */ + V& get(const K& key) const + { + return values[table.index_of(key)]; + } + + /** + * If the given `key` exists in this map, this function returns the value + * associated with the key, and sets `contains` to `true`. If `key` is not + * in the map, `contains` is set to `false`. + */ + inline V& get(const K& key, bool& contains) const + { + unsigned int index; + return get(key, contains, index); + } + + /** + * If the given `key` exists in this map, this function returns the value + * associated with the key, sets `contains` to `true`, and sets `index` to + * the bucket containing the key. If `key` is not in the map, `contains` is + * set to `false`, and `index` is set to the index where the key would be + * located, if it had existed in the map. + */ + V& get(const K& key, bool& contains, unsigned int& index) const + { + index = table.index_of(key, contains); + return values[index]; + } + + /** + * If the given `key` exists in this map, this function returns the value + * associated with the key, sets `contains` to `true`, and sets `index` to + * the bucket containing the key. If `key` is not in the map, `contains` is + * set to `false`, and `index` is set to the index where the key would be + * located, if it had existed in the map. In any case, the evaluated hash + * function of `key` is stored in `hash_value`. + */ + V& get(const K& key, bool& contains, + unsigned int& index, + unsigned int& hash_value) const + { + index = table.index_of(key, contains, hash_value); + return values[index]; + } + + /** + * Removes all entries from this hash_map. Note that this function does + * not free each entry beforehand. + */ + inline void clear() { + table.clear(); + } + + /** + * Returns a hash_map_iterator pointing to the first entry in this container. + * + * **NOTE:** Unlike the libstdc++ [unordered_set](http://en.cppreference.com/w/cpp/container/unordered_set) + * and [unordered_map](http://en.cppreference.com/w/cpp/container/unordered_map) + * iterators, we do not keep a linked list among the elements, so if rapid + * iteration over elements is more critical than rapid queries, consider + * using an array_map. + */ + inline hash_map_iterator begin() { + return hash_map_iterator(*this, table.first_empty()); + } + + /** + * Returns a hash_map_iterator pointing to the end of this container. + * + * **NOTE:** Unlike the libstdc++ [unordered_set](http://en.cppreference.com/w/cpp/container/unordered_set) + * and [unordered_map](http://en.cppreference.com/w/cpp/container/unordered_map) + * iterators, we do not keep a linked list among the elements, so if rapid + * iteration over elements is more critical than rapid queries, consider + * using an array_map. + */ + inline hash_map_iterator end() { + return hash_map_iterator(*this, table.capacity); + } + + /** + * Returns a const hash_map_iterator pointing to the first entry in this container. + * + * **NOTE:** Unlike the libstdc++ [unordered_set](http://en.cppreference.com/w/cpp/container/unordered_set) + * and [unordered_map](http://en.cppreference.com/w/cpp/container/unordered_map) + * iterators, we do not keep a linked list among the elements, so if rapid + * iteration over elements is more critical than rapid queries, consider + * using an array_map. + */ + inline hash_map_iterator begin() const { + return hash_map_iterator(*this, table.first_empty()); + } + + /** + * Returns a const hash_map_iterator pointing to the end of this container. + * + * **NOTE:** Unlike the libstdc++ [unordered_set](http://en.cppreference.com/w/cpp/container/unordered_set) + * and [unordered_map](http://en.cppreference.com/w/cpp/container/unordered_map) + * iterators, we do not keep a linked list among the elements, so if rapid + * iteration over elements is more critical than rapid queries, consider + * using an array_map. + */ + inline hash_map_iterator end() const { + return hash_map_iterator(*this, table.capacity); + } + + /** + * Swaps the contents of the hash_map `first` with that of `second`. + */ + static inline void swap(hash_map& first, hash_map& second) { + hash_set::swap(first.table, second.table); + core::swap(first.values, second.values); + } + + /** + * Moves the contents of the hash_map `src` into `dst`. Note this function + * does not copy the contents of the underlying hash_map::table or + * hash_map::values, it merely copies the pointers. + */ + static inline void move(const hash_map& src, hash_map& dst) { + hash_set::move(src.table, dst.table); + dst.values = src.values; + } + + /** + * Copies the contents of the hash_map `src` into `dst`. + * \param alloc_keys a memory allocation function with prototype + * `void* alloc_keys(size_t count, size_t size)` that allocates space for + * `count` items, each with size `size`, and initializes them such that + * core::is_empty() returns `true` for each key. + */ + static inline bool copy(const hash_map& src, hash_map& dst, alloc_keys_func alloc_keys = calloc) { + dst.table.capacity = src.table.capacity; + dst.table.size = src.table.size; + dst.table.keys = (K*) alloc_keys(src.table.capacity, sizeof(K)); + if (dst.table.keys == NULL) return false; + dst.values = (V*) malloc(sizeof(V) * src.table.capacity); + if (dst.values == NULL) { + core::free(dst.table.keys); + return false; + } + + for (unsigned int i = 0; i < src.table.capacity; i++) { + if (is_empty(src.table.keys[i])) continue; + if (!core::copy(src.table.keys[i], dst.table.keys[i])) { + free(dst); return false; + } else if (!core::copy(src.values[i], dst.values[i])) { + core::free(dst.table.keys[i]); + set_empty(dst.table.keys[i]); + free(dst); return false; + } + } + return true; + } + + template + static inline long unsigned int size_of(const hash_map& map, + const key_value_metric& metric) + { + long unsigned int sum = core::size_of(map.table.capacity) + core::size_of(map.table.size); + for (unsigned int i = 0; i < map.table.capacity; i++) { + if (is_empty(map.table.keys[i])) + sum += sizeof(K) + sizeof(V); + else sum += core::size_of(map.table.keys[i], metric.key_metric) + core::size_of(map.values[i], metric.value_metric); + } + return sum; + } + + /** + * Frees hash_map::table and hash_map::values. This should not be used if + * `map` was constructed on the stack, as the destructor will automatically + * free hash_map::table and hash_map::values. The existing entries of `map` + * are not freed. + */ + static inline void free(hash_map& map) { + core::free(map.table); + core::free(map.values); + } + +private: + /* NOTE: this function assumes table is initialized */ + bool initialize_values() { + values = (V*) malloc(sizeof(V) * table.capacity); + if (values == NULL) { + core::free(table); + return false; + } + return true; + } + + bool initialize( + unsigned int initial_capacity, + alloc_keys_func alloc_keys) + { + if (!table.initialize(initial_capacity, alloc_keys)) + return false; + return initialize_values(); + } + + inline void place(const K& key, const V& value, unsigned int index) + { + table.place(key, index); + values[index] = value; + } + + inline void insert_unique( + const K& key, const V& value) + { + place(key, value, table.next_empty(key)); + } + + inline void insert(const K& key, const V& value) + { + bool contains; + unsigned int index = table.index_to_insert(key, contains); + if (!contains) { + place(key, value, index); + table.size++; + } + } + + template + friend bool hash_map_init(hash_map& map, + unsigned int capacity, alloc_keys_func alloc_keys); +}; + +/** + * Initializes the hash_map `map` with the given initial `capacity`. + * \param alloc_keys a memory allocation function with prototype + * `void* alloc_keys(size_t count, size_t size)` that allocates space for + * `count` items, each with size `size`, and initializes them such that + * core::is_empty() returns `true` for each key. + */ +template +bool hash_map_init( + hash_map& map, + unsigned int capacity, + alloc_keys_func alloc_keys = calloc) +{ + if (!map.initialize(capacity, alloc_keys)) { + fprintf(stderr, "hash_map_init ERROR: Unable to allocate memory.\n"); + return false; + } + return true; +} + +/** + * Swaps the underlying buffers of the given hash_maps `first` and `second`. + */ +template +void swap(hash_map& first, hash_map& second) { + K* keys_swap = first.table.keys; + first.table.keys = second.table.keys; + second.table.keys = keys_swap; + + V* values_swap = first.values; + first.values = second.values; + second.values = values_swap; + + unsigned int ui_swap = first.table.size; + first.table.size = second.table.size; + second.table.size = ui_swap; + + ui_swap = first.table.capacity; + first.table.capacity = second.table.capacity; + second.table.capacity = ui_swap; +} + +/** + * A sequentially-ordered associative container that contains a set of + * key-value pairs, where the keys are unique and have type `K` and the values + * have type `V`. The keys are stored sequentially in a native array + * array_map::keys, and the values are stored sequentially in a parallel native + * array array_map::values. + * + * When inserting a key-value pair into an array_map, a linear search is + * performed to determine if the key already exists in the map. If so, its + * corresponding value is replaced by the new value. If not, the new key and + * value are inserted at the ends of their respective arrays, and + * array_map::size is incremented. + * + * **Performance:** This data structure provides linear-time access and + * modification. However, due to locality of reference, this data structure + * may perform operations more quickly than hash_map if the size of the map + * is small. + * + * A simple example of the use of array_map is given below, where the expected + * output is `a.get(3): c, a.get(4): d, 1:a 4:d 3:c `. + * + * ```{.cpp} + * #include + * #include + * using namespace core; + * + * int main() { + * array_map a = array_map(8); + * a.put(1, 'a'); a.put(2, 'b'); + * a.put(3, 'c'); a.put(4, 'd'); + * a.remove(2); + * + * printf("a.get(3): %c, ", a.get(3)); + * printf("a.get(4): %c, ", a.get(4)); + * for (const auto& entry : a) + * printf("%d:%c ", entry.key, entry.value); + * } + * ``` + * + * + * However, if `a` is not allocated on the stack, the destructor will not be + * automatically called, and so it must be freed using `core::free` or + * `hash_map::free`. In the example below, the expected output is the same as + * that of the program above: `a.get(3): c, a.get(4): d, 1:a 4:d 3:c `. + * + * ```{.cpp} + * #include + * #include + * using namespace core; + * + * int main() { + * array_map& a = *((array_map*) alloca(sizeof(array_map))); + * array_map_init(a, 8); + * a.put(1, 'a'); a.put(2, 'b'); + * a.put(3, 'c'); a.put(4, 'd'); + * a.remove(2); + * + * printf("a.get(3): %c, ", a.get(3)); + * printf("a.get(4): %c, ", a.get(4)); + * for (const auto& entry : a) + * printf("%d:%c ", entry.key, entry.value); + * free(a); + * } + * ``` + * + * + * A number of member functions require `K` and `V` to be + * [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). + * If this is not the case, those operations can be performed directly on the + * public fields, as in the following example. The expected output is + * `first:1 second:2 `. + * + * ```{.cpp} + * #include + * #include + * #include + * using namespace core; + * + * struct custom_string { + * char* buffer; + * + * static void move(const custom_string& src, custom_string& dst) { + * dst.buffer = src.buffer; + * } + * + * static void free(custom_string& s) { + * core::free(s.buffer); + * } + * }; + * + * inline bool operator == (const custom_string& first, const custom_string& second) { + * if (first.buffer == NULL) + * return second.buffer == NULL; + * return strcmp(first.buffer, second.buffer) == 0; + * } + * + * bool init(custom_string& s, const char* src) { + * s.buffer = (char*) malloc(sizeof(char) * (strlen(src) + 1)); + * if (s.buffer == NULL) + * return false; + * memcpy(s.buffer, src, sizeof(char) * (strlen(src) + 1)); + * return true; + * } + * + * int main() { + * custom_string key_one, key_two; + * int value_one = 1, value_two = 2; + * init(key_one, "first"); + * init(key_two, "second"); + * + * array_map a = array_map(8); + * a.ensure_capacity(2); + * + * unsigned int index; + * a.get(key_one, index); + * if (index == a.size) { + * core::move(key_one, a.keys[index]); + * a.values[index] = value_one; + * a.size++; + * } + * + * a.get(key_two, index); + * if (index == a.size) { + * core::move(key_two, a.keys[index]); + * a.values[index] = value_two; + * a.size++; + * } + * + * for (const auto& entry : a) + * printf("%s:%d ", entry.key.buffer, entry.value); + * for (auto entry : a) + * free(entry.key); + * } + * ``` + */ +template +struct array_map { + /** + * The type of the keys in this array_map. + */ + typedef K key_type; + + /** + * The type of the values in this array_map. + */ + typedef V value_type; + + /** + * The native array of keys. + */ + K* keys; + + /** + * The native array of values parallel to array_map::keys. + */ + V* values; + + /** + * The capacity of array_map::keys and array_map::values. + */ + size_t capacity; + + /** + * The number of entries in the array_map. + */ + size_t size; + + /** + * Constructs the hash_map with the given `initial_capacity`. + */ + array_map(size_t initial_capacity) : size(0) { + if (!initialize(initial_capacity)) { + fprintf(stderr, "array_map ERROR: Error during initialization.\n"); + exit(EXIT_FAILURE); + } + } + + ~array_map() { free(); } + + /** + * Given the requested number of elements `new_length`, this function + * determines whether array_map::capacity is sufficient. If not, it + * attempts to increase its capacity by factors of RESIZE_FACTOR. + * \returns `true` if the resize was successful, and `false` if there is insufficient memory. + */ + bool ensure_capacity(size_t new_length) { + if (new_length <= capacity) + return true; + + size_t new_capacity = capacity; + if (!expand(keys, new_capacity, new_length)) + return false; + if (!resize(values, new_capacity)) + return false; + capacity = new_capacity; + return true; + } + + /** + * Adds the given key-value pair to this map. The assignment operator is + * used insert the entry, and so this function should not be used if `K` + * and `V` are not [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). + * In such a case, insertion should be performed manually by using + * array_map::get to compute the appropriate index and directly modifying + * the the public fields. See the example in the description of array_map. + * \tparam K is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). + * \tparam V is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). + */ + bool put(const K& key, const V& value) { + size_t index = index_of(key); + if (index < size) { + values[index] = value; + return true; + } + + if (!ensure_capacity(size + 1)) + return false; + keys[size] = key; + values[size] = value; + size++; + return true; + } + + /** + * Performs a linear search to find the index of the given `key`. If the + * `key` is not in this map, array_map::size is returned. + */ + template + inline size_t index_of(const Key& key) const { + return core::index_of(key, keys, size); + } + + /** + * Performs a linear search to find the index of the given `key`, with the + * search beginning at the index `start`. If the `key` is not in this map, + * array_map::size is returned. + */ + template + inline unsigned int index_of(const Key& key, unsigned int start) const { + for (unsigned int i = start; i < size; i++) + if (keys[i] == key) + return i; + return (unsigned int) size; + } + + /** + * Performs a reverse linear search to find the index of the given `key`. + * If the `key` is not in this map, `static_cast(-1)` is + * returned. + */ + template + inline unsigned int last_index_of(const Key& key) const { + return core::last_index_of(key, keys, size); + } + + /** + * Returns `true` if the given `key` exists in the map, and `false` otherwise. + */ + inline bool contains(const K& key) const { + return index_of(key) < size; + } + + /** + * Retrieves the value associated with the given `key`. This function + * assumes the given key exists in the map. + */ + template + inline V& get(const Key& key) { + return values[index_of(key)]; + } + + /** + * Retrieves the const value associated with the given `key`. This function + * assumes the given key exists in the map. + */ + template + inline const V& get(const Key& key) const { + return values[index_of(key)]; + } + + /** + * If the given `key` exists in the map, this function returns the value + * associated with the key, and sets `index` to the index of + * array_map::keys where the key is located. If `key` does not exist in the + * map, `index` is set to array_map::size. + */ + template + inline V& get(const Key& key, unsigned int& index) { + index = index_of(key); + return values[index]; + } + + /** + * If the given `key` exists in the map, this function returns the value + * associated with the key, and sets `contains` to `true`. If `key` does + * not exist in the map, `contains` is set to `false`. + */ + template + inline V& get(const Key& key, bool& contains) { + size_t index = index_of(key); + contains = (index != size); + return values[index]; + } + + /** + * If the given `key` exists in the map, this function returns the const + * value associated with the key, and sets `contains` to `true`. If `key` + * does not exist in the map, `contains` is set to `false`. + */ + template + inline const V& get(const Key& key, bool& contains) const { + unsigned int index = index_of(key); + contains = (index != size); + return values[index]; + } + + /** + * This function removes `key` and associated value from the map. This + * function does not free the removed element. + * \returns `true` if the element is removed, and `false` if the set does not contain `element`. + */ + bool remove(const K& key) { + size_t index = index_of(key); + if (index == size) + return false; + remove_at(index); + return true; + } + + /** + * This function removes the key-value pair located at the given `index` in + * the map. This function does not free the removed element. This function + * assumes `0 <= index < array_map::size`. + * \tparam K satisfies is_moveable. + * \tparam V satisfies is_moveable. + */ + inline void remove_at(size_t index) { + size--; + if (index == size) + return; + + /* move the last item into the position of the removed item */ + core::move(keys[size], keys[index]); + core::move(values[size], values[index]); + } + + /** + * Removes all entries from this array_map. Note that this function does + * not free each entry beforehand. + */ + inline void clear() { + size = 0; + } + + /** + * Returns an array_map_iterator pointing to the first entry in this container. + */ + inline array_map_iterator begin() { + return array_map_iterator(*this, 0); + } + + /** + * Returns an array_map_iterator pointing to the end of this container. + */ + inline array_map_iterator end() { + return array_map_iterator(*this, size); + } + + /** + * Returns a const array_map_iterator pointing to the first entry in this container. + */ + inline array_map_iterator begin() const { + return array_map_iterator(*this, 0); + } + + /** + * Returns a const array_map_iterator pointing to the end of this container. + */ + inline array_map_iterator end() const { + return array_map_iterator(*this, size); + } + + /** + * Swaps the contents of the array_map `first` with that of `second`. + */ + static inline void swap(array_map& first, array_map& second) { + core::swap(first.keys, second.keys); + core::swap(first.values, second.values); + core::swap(first.capacity, second.capacity); + core::swap(first.size, second.size); + } + + /** + * Moves the contents of the array_map `src` into `dst`. Note this function + * does not copy the contents of the underlying array_map::keys or + * array_map::values, it merely copies the pointers. + */ + static inline void move(const array_map& src, array_map& dst) { + dst.keys = src.keys; + dst.values = src.values; + dst.capacity = src.capacity; + dst.size = src.size; + } + + template + static inline long unsigned int size_of(const array_map& map, + const key_value_metric& metric) + { + long unsigned int sum = core::size_of(map.capacity) + core::size_of(map.size); + for (unsigned int i = 0; i < map.size; i++) + sum += core::size_of(map.keys[i], metric.key_metric) + core::size_of(map.values[i], metric.value_metric); + return sum + (map.capacity - map.size) * (sizeof(K) + sizeof(V)); + } + + static inline long unsigned int size_of(const array_map& map, const default_metric& metric) { + return size_of(map, make_key_value_metric(default_metric(), default_metric())); + } + + /** + * Frees array_map::keys and array_map::values. This should not be used if + * `map` was constructed on the stack, as the destructor will automatically + * free array_map::keys and array_map::values. The existing entries of + * `map` are not freed. + */ + static inline void free(array_map& map) { map.free(); } + +private: + inline bool initialize(size_t initial_capacity) { + capacity = initial_capacity; + keys = (K*) malloc(sizeof(K) * capacity); + if (keys == NULL) { + fprintf(stderr, "array_map.initialize ERROR: Out of memory.\n"); + return false; + } + values = (V*) malloc(sizeof(V) * capacity); + if (values == NULL) { + core::free(keys); + fprintf(stderr, "array_map.initialize ERROR: Out of memory.\n"); + return false; + } + return true; + } + + inline void free() { + core::free(keys); + core::free(values); + } + + template + friend bool array_map_init(array_map&, size_t); +}; + +/** + * Initializes the array_map `map` with the given `initial_capacity`. + */ +template +bool array_map_init(array_map& map, size_t initial_capacity) { + map.size = 0; + return map.initialize(initial_capacity); +} + +/** + * Returns the number of elements in the given hash_set `set`. + */ +template +inline unsigned int size(const hash_set& set) { + return set.size; +} + +/** + * Returns the number of elements in the given hash_map `map`. + */ +template +inline unsigned int size(const hash_map& map) { + return map.table.size; +} + +/** + * Returns the number of elements in the given array_map `map`. + */ +template +inline unsigned int size(const array_map& map) { + return map.size; +} + +/** + * Assuming the given `map` has value type that satisfies + * [is_integral](http://en.cppreference.com/w/cpp/types/is_integral), + * and the values in the map are unique and no larger than `size(map)`, + * this function returns an array of pointers to the keys in `map`, where the + * pointer at index `i` references the key in `map` that corresponds to the + * value `i`. The caller of this function is responsible for the memory of the + * returned native array, and must call free to release it. + * \tparam MapType a map type that allows range-based for iteration, and for + * which the function `unsigned int size(const MapType&)` is defined. + */ +template +inline const typename MapType::key_type** invert(const MapType& map) { + const typename MapType::key_type** inverse = + (const typename MapType::key_type**) calloc(size(map) + 1, sizeof(typename MapType::key_type*)); + if (inverse == NULL) { + fprintf(stderr, "invert ERROR: Unable to invert map. Out of memory.\n"); + return NULL; + } + for (const auto& entry : map) + inverse[entry.value] = &entry.key; + return inverse; +} + + +inline bool hash_map_test(void) +{ + hash_map map = hash_map(4); + bool contains_four = true, contains_minus_seven = true; + map.get(4, contains_four); + map.get(-7, contains_minus_seven); + if (contains_four || contains_minus_seven) { + fprintf(stderr, "hash_map_test ERROR: Map with no inserted elements should be empty.\n"); + return false; + } + + /* test insertion and retrieval operations */ + map.put(-7, "negative seven"); + map.put(4, "four"); + + if (map.table.size != 2) { + fprintf(stderr, "hash_map_test ERROR: Map size is %d after adding two elements.\n", map.table.size); + return false; + } + if (strcmp(map.get(-7), "negative seven") + || strcmp(map.get(4), "four")) + { + fprintf(stderr, "hash_map_test ERROR: Simple hashtable insertion failed.\n"); + return false; + } + + map.put(4, "new four"); + map.put(5, "five"); + map.put(12, "twelve"); + map.put(7, "seven"); + map.put(13, "thirteen"); + + /* test automatic capacity resizing */ + if (map.table.size != 6) { + fprintf(stderr, "hash_map_test ERROR: Map size is %d after adding six elements.\n", map.table.size); + return false; + } + if (map.table.capacity != 4 * RESIZE_FACTOR * RESIZE_FACTOR) { + fprintf(stderr, "hash_map_test ERROR: Unexpected hashtable capacity.\n"); + return false; + } + if (strcmp(map.get(4), "new four") + || strcmp(map.get(12), "twelve") + || strcmp(map.get(7), "seven") + || strcmp(map.get(-7), "negative seven")) + { + fprintf(stderr, "hash_map_test ERROR: Additional hashtable insertion failed.\n"); + return false; + } + + /* test removal operation */ + if (!map.remove(4)) { + fprintf(stderr, "hash_map_test ERROR: Removal of key '4' failed.\n"); + return false; + } + if (strcmp(map.get(12), "twelve") + || strcmp(map.get(13), "thirteen") + || strcmp(map.get(5), "five") + || strcmp(map.get(-7), "negative seven") + || strcmp(map.get(7), "seven")) { + fprintf(stderr, "hash_map_test ERROR: Hashtable lookup failed after removal.\n"); + return false; + } + + bool contains = true; + map.get(4, contains); + if (contains) { + fprintf(stderr, "hash_map_test ERROR: Retrieval of removed key failed.\n"); + return false; + } + + return true; +} + +inline bool array_map_test(void) +{ + array_map map = array_map(4); + bool contains_four = true, contains_minus_seven = true; + map.get(4, contains_four); + map.get(-7, contains_minus_seven); + if (contains_four || contains_minus_seven) { + fprintf(stderr, "array_map_test ERROR: Map with no inserted elements should be empty.\n"); + return false; + } + + /* test insertion and retrieval operations */ + map.put(-7, "negative seven"); + map.put(4, "four"); + + if (map.size != 2) { + fprintf(stderr, "array_map_test ERROR: Map size is %zu after adding two elements.\n", map.size); + return false; + } + if (strcmp(map.get(-7), "negative seven") + || strcmp(map.get(4), "four")) + { + fprintf(stderr, "array_map_test ERROR: Simple hashtable insertion failed.\n"); + return false; + } + + map.put(4, "new four"); + map.put(5, "five"); + map.put(12, "twelve"); + map.put(7, "seven"); + map.put(13, "thirteen"); + + if (map.size != 6) { + fprintf(stderr, "array_map_test ERROR: Map size is %zu after adding six elements.\n", map.size); + return false; + } + if (strcmp(map.get(4), "new four") + || strcmp(map.get(12), "twelve") + || strcmp(map.get(7), "seven") + || strcmp(map.get(-7), "negative seven")) + { + fprintf(stderr, "array_map_test ERROR: Additional hashtable insertion failed.\n"); + return false; + } + + /* test removal operation */ + if (!map.remove(4)) { + fprintf(stderr, "array_map_test ERROR: Removal of key '4' failed.\n"); + return false; + } + if (strcmp(map.get(12), "twelve") + || strcmp(map.get(13), "thirteen") + || strcmp(map.get(5), "five") + || strcmp(map.get(-7), "negative seven") + || strcmp(map.get(7), "seven")) { + fprintf(stderr, "array_map_test ERROR: Hashtable lookup failed after removal.\n"); + return false; + } + + bool contains = true; + map.get(4, contains); + if (contains) { + fprintf(stderr, "array_map_test ERROR: Retrieval of removed key failed.\n"); + return false; + } + + return true; +} + +} /* namespace core */ + +#endif /* MAP_H_ */ diff --git a/reasoners/data_preprocess/logic/graph_dataset_gen/core/random.h b/reasoners/data_preprocess/logic/graph_dataset_gen/core/random.h new file mode 100644 index 000000000..6ba64e0ef --- /dev/null +++ b/reasoners/data_preprocess/logic/graph_dataset_gen/core/random.h @@ -0,0 +1,321 @@ +/** + * \file random.h + * + * This file defines the basic pseudorandom number generation functionality + * used by the core library, as well as functions to sample from a number of + * built-in distributions such as categorical, discrete uniform, continuous + * uniform, Bernoulli, geometric, beta, gamma, and Dirichlet. + * + * + */ + +#ifndef RANDOM_H_ +#define RANDOM_H_ + +#include +#include + +#include "timer.h" +#include "array.h" +#include "io.h" + +namespace core { + + +#if defined(NDEBUG) +static unsigned int seed = (unsigned int) milliseconds(); +#else +static unsigned int seed = 0; +#endif + +static thread_local std::minstd_rand engine = std::minstd_rand(seed); + +/** + * Returns the initial random seed used by all core functions that require pseudorandom number generation. + */ +inline unsigned int get_seed() { + return seed; +} + +/** + * Sets the seed of the underlying pseudorandom number generator. + */ +inline void set_seed(unsigned int new_seed) { + engine.seed(new_seed); + seed = new_seed; +} + +/** + * Reads the state of the pseudorandom number generator from `in`. This is + * useful to persist the state of the pseudorandom number generator. + */ +template +inline bool read_random_state(Stream& in) +{ + size_t length; + if (!read(length, in)) return false; + char* state = (char*) alloca(sizeof(char) * length); + if (state == NULL || !read(state, in, (unsigned int) length)) + return false; + + std::stringstream buffer(std::string(state, length)); + buffer >> engine; + return true; +} + +/** + * Writes the state of the pseudorandom number generator to `out`. This is + * useful to persist the state of the pseudorandom number generator. + */ +template +inline bool write_random_state(Stream& out) { + std::stringstream buffer; + buffer << engine; + std::string data = buffer.str(); + return write(data.length(), out) && write(data.c_str(), out, (unsigned int) data.length()); +} + +/** + * Samples from a categorical distribution, where the unnormalized probability + * of returning the index `i` is given by `probability[i]`. This function + * normalizes and overwrites `probability` with its cumulative distribution + * function. + */ +template::value>::type* = nullptr> +unsigned int sample_categorical(V* probability, unsigned int length) +{ +#if !defined(NDEBUG) + if (length == 0) + fprintf(stderr, "sample_categorical WARNING: Specified length is zero.\n"); +#endif + + for (unsigned int j = 1; j < length; j++) + probability[j] += probability[j - 1]; + + /* select the new table assignment */ + V random = probability[length - 1] * ((V) engine() / engine.max()); + unsigned int selected_table = length - 1; + for (unsigned int j = 0; j < length; j++) { + if (random < probability[j]) { + selected_table = j; + break; + } + } + return selected_table; +} + +/** + * Returns the smallest index `i` such that `random < sum from j=0 to i-1 of probability[j]`. + * Thus, this function implements the inverse cumulative distribution function + * for the categorical distribution with unnormalized probabilities given by + * `probability`. + */ +template +inline unsigned int select_categorical( + const U* probability, V random, unsigned int length) +{ +#if !defined(NDEBUG) + if (length == 0) + fprintf(stderr, "select_categorical WARNING: Specified length is zero.\n"); +#endif + + V aggregator = (V) 0; + unsigned int selected_table = length - 1; + for (unsigned int j = 0; j < length; j++) { + aggregator += probability[j]; + if (random <= aggregator) { + selected_table = j; + break; + } + } + return selected_table; +} + +/** + * Samples from a categorical distribution, where the unnormalized probability + * of returning the index `i` is given by `probability[i]` and its sum is given + * by the floating-point `sum`. This function doesn't modify `probability`. + */ +template::value>::type* = nullptr> +unsigned int sample_categorical(const V* probability, V sum, unsigned int length) +{ +#if !defined(NDEBUG) + if (length == 0) + fprintf(stderr, "sample_categorical WARNING: Specified length is zero.\n"); +#endif + + /* select the new table assignment */ + V random = sum * ((V) engine() / engine.max()); + return select_categorical(probability, random, length); +} + +/** + * Samples from a categorical distribution, where the unnormalized probability + * of returning the index `i` is given by `probability[i]` and its sum is given + * by the unsigned integer `sum`. This function doesn't modify `probability`. + */ +unsigned int sample_categorical( + const unsigned int* probability, + unsigned int sum, unsigned int length) +{ +#if !defined(NDEBUG) + if (length == 0) + fprintf(stderr, "sample_categorical WARNING: Specified length is zero.\n"); +#endif + + /* select the new table assignment */ + unsigned int random = engine() % sum; + return select_categorical(probability, random, length); +} + + +/** + */ + +/* forward declarations */ + +template struct array; + +/** + * Returns a sample from the discrete uniform distribution over `{0, ..., n - 1}`. + */ +inline unsigned int sample_uniform(unsigned int n) { + return engine() % n; +} + +/** + * Returns a sample from the discrete uniform distribution over `elements`. + */ +template +inline const T& sample_uniform(const T* elements, unsigned int length) { + return elements[engine() % length]; +} + +/** + * Returns a sample from the discrete uniform distribution over `elements`. + */ +template +inline const T& sample_uniform(const T (&elements)[N]) { + return elements[engine() % N]; +} + +/** + * Returns a sample from the discrete uniform distribution over `elements`. + */ +template +inline const T& sample_uniform(const array& elements) { + return sample_uniform(elements.data, (unsigned int) elements.length); +} + +/** + * Returns a sample from the continuous uniform distribution over [0, 1]. + */ +template +inline V sample_uniform() { + return (V) engine() / engine.max(); +} + +/** + * Returns a sample from the Bernoulli distribution: with probability 0.5, + * `true` is returned, otherwise `false` is returned. + */ +template +inline bool sample_bernoulli(const V& p) { + return sample_uniform() < p; +} + +/** + * Returns a sample from the geometric distribution with success probability `p`. + */ +template +inline unsigned int sample_geometric(const V& p) { + auto geom = std::geometric_distribution(p); + return geom(engine); +} + +/** + * Returns a sample from the negative binomial distribution with number of + * failures `r` and success probability `p`. + */ +template +inline double log_probability_negative_binomial(unsigned int x, const V& r, const V& p) { + return lgamma(x + r) - lgamma(x + 1) - lgamma(r) + r*log(1 - p) + x*log(p); +} + +/** + * Returns a sample from the beta distribution with shape parameter `alpha` and + * scale parameter 1. This function assumes `alpha > 0`. + */ +template +inline V sample_beta(const V& alpha) { + static std::gamma_distribution first_gamma = std::gamma_distribution(1.0); + std::gamma_distribution second_gamma = std::gamma_distribution(alpha); + V first = first_gamma(engine); + V second = second_gamma(engine); + return first / (first + second); +} + +/** + * Returns a sample from the beta distribution with shape parameter `alpha` and + * scale parameter `beta`. This function assumes `alpha > 0` and `beta > 0`. + */ +template +inline V sample_beta(const V& alpha, const V& beta) { + std::gamma_distribution first_gamma = std::gamma_distribution(alpha, 1.0); + std::gamma_distribution second_gamma = std::gamma_distribution(beta, 1.0); + V first = first_gamma(engine); + V second = second_gamma(engine); + return first / (first + second); +} + +/** + * Returns a sample from the gamma distribution with shape parameter `alpha` + * and rate parameter `beta`. This function assumes `alpha > 0` and `beta > 0`. + */ +template +inline V sample_gamma(const V& alpha, const V& beta) { + std::gamma_distribution gamma = std::gamma_distribution(alpha, 1.0 / beta); + return gamma(engine); +} + +/** + * Returns the log probability of drawing the observation `x` from a gamma + * distribution with shape parameter `alpha` and rate parameter `beta`. This + * function assumes `x > 0`, `alpha > 0`, and `beta > 0`. + */ +template +inline V log_probability_gamma(const V& x, const V& alpha, const V& beta) { + return alpha * log(beta) - lgamma(alpha) + (alpha - 1) * log(x) - beta * x; +} + +/** + * Samples from the Dirichlet distribution with parameter `alpha` and dimension + * `length`. The sample is written to `dst`. This function assumes `dst` has + * size at least `length`, `alpha[i] > 0` for all `i = 0, ..., length - 1`. + */ +template +inline void sample_dirichlet(V* dst, const V* alpha, unsigned int length) { + V sum = 0.0; + for (unsigned int i = 0; i < length; i++) { + if (alpha[i] == 0.0) { + dst[i] = 0.0; + } else { + std::gamma_distribution gamma_dist(alpha[i], 1.0); + V value = gamma_dist(engine); + dst[i] = value; + sum += value; + } + } + + for (unsigned int i = 0; i < length; i++) + dst[i] /= sum; +} + + +} /* namespace core */ + +#endif /* RANDOM_H_ */ diff --git a/reasoners/data_preprocess/logic/graph_dataset_gen/core/timer.h b/reasoners/data_preprocess/logic/graph_dataset_gen/core/timer.h new file mode 100644 index 000000000..306f7fd34 --- /dev/null +++ b/reasoners/data_preprocess/logic/graph_dataset_gen/core/timer.h @@ -0,0 +1,130 @@ +/** + * \file timer.h + * + * + */ + +#ifndef TIMER_H_ +#define TIMER_H_ + +#if defined(_WIN32) /* on Windows */ +#define _WINSOCKAPI_ +#include +#undef max +#undef min +#elif defined(__APPLE__) /* on Mac */ +#include +#else /* on Linux */ +#include +#endif + + +namespace core { + + +#if defined(__APPLE__) +struct timebase_info { + mach_timebase_info_data_t info; + + timebase_info() { + info = {0, 0}; + mach_timebase_info(&info); + } +}; + +inline timebase_info& get_timebase() { + static timebase_info timebase; + return timebase; +} +#endif + + +/** + * A simple structure that measures time. + */ +struct timer +{ +#if defined(_WIN32) /* on Windows */ + ULONGLONG start_time; +#elif defined(__APPLE__) /* on Mac */ + uint64_t start_time; +#else /* on Linux */ + timespec start_time; +#endif + + /** + * Constructor that starts the timer. + */ + timer() { + start(); + } + + /** + * Starts the timer. + */ + inline void start() { +#if defined(_WIN32) /* on Windows */ + start_time = GetTickCount64(); +#elif defined(__APPLE__) /* on Mac */ + start_time = mach_absolute_time(); +#else /* on Linux */ + clock_gettime(CLOCK_MONOTONIC, &start_time); +#endif + } + + /** + * Returns the number of milliseconds elapsed since the timer was last started. + */ + inline unsigned long long milliseconds() { +#if defined(_WIN32) /* on Windows */ + ULONGLONG end_time = GetTickCount64(); + return (unsigned long long) (end_time - start_time); +#elif defined(__APPLE__) /* on Mac */ + uint64_t end_time = mach_absolute_time(); + return (end_time - start_time) * get_timebase().info.numer / get_timebase().info.denom / 1000000; +#else /* on Linux */ + timespec end_time; + clock_gettime(CLOCK_MONOTONIC, &end_time); + return (unsigned long long) (end_time.tv_sec - start_time.tv_sec) * 1000 + (end_time.tv_nsec - start_time.tv_nsec) / 1000000; +#endif + } + + /** + * Returns the number of nanoseconds elapsed since the timer was last started. + */ + inline double nanoseconds() { +#if defined(_WIN32) /* on Windows */ + ULONGLONG end_time = GetTickCount64(); + return (end_time - start_time) * 1.0e6; +#elif defined(__APPLE__) /* on Mac */ + uint64_t end_time = mach_absolute_time(); + return (double) (end_time - start_time) * get_timebase().info.numer / get_timebase().info.denom; +#else /* on Linux */ + timespec end_time; + clock_gettime(CLOCK_MONOTONIC, &end_time); + return (end_time.tv_sec - start_time.tv_sec) * 1.0e9 + (end_time.tv_nsec - start_time.tv_nsec); +#endif + } +}; + +/** + * On Windows and Mac, this returns the number of milliseconds elapsed since + * the system was started. On Linux, it returns the system's best estimate of + * the current time of day. + */ +inline unsigned long long milliseconds() { +#if defined(_WIN32) /* on Windows */ + return GetTickCount64(); +#elif defined(__APPLE__) /* on Mac */ + return mach_absolute_time() * get_timebase().info.numer / get_timebase().info.denom / 1000000; +#else /* on Linux */ + timespec time; + clock_gettime(CLOCK_REALTIME, &time); + return (unsigned long long) time.tv_sec * 1000 + time.tv_nsec / 1000000; +#endif +} + + +} /* namespace core */ + +#endif /* TIMER_H_ */ diff --git a/reasoners/data_preprocess/logic/graph_dataset_gen/core/utility.h b/reasoners/data_preprocess/logic/graph_dataset_gen/core/utility.h new file mode 100644 index 000000000..25bfd99d6 --- /dev/null +++ b/reasoners/data_preprocess/logic/graph_dataset_gen/core/utility.h @@ -0,0 +1,799 @@ +/** + * \file utility.h + * + * This file contains a number of useful miscellaneous definitions and + * implementations, such as a core::string structure, a binary logarithm + * function, and a function to access the directory structure in the + * filesystem. + * + * + */ + +#ifndef UTILITY_H_ +#define UTILITY_H_ + +#include +#include "io.h" + +#if defined(_WIN32) +#if !defined(UNICODE) +#error "Unicode support required. Please compile with UNICODE preprocessor definition." +#endif +#define _WINSOCKAPI_ +#include +#include +#undef max +#undef min +#else +#include +#include +#endif + + +namespace core { + +/* a useful type trait for detecting whether the function 'print_special_string' is defined */ +namespace detail { + template static auto test_print_special_string(int32_t) -> + decltype(bool(print_special_string(0u, std::declval(), std::declval()...)), std::true_type{}); + template static auto test_print_special_string(int64_t) -> std::false_type; +} + +template struct has_print_special_string : decltype(core::detail::test_print_special_string(0)){}; + + +/** + * A basic string structure, containing a native array of `char` elements and an `unsigned int length`. + */ +struct string { + /** + * The length of the string in characters. + */ + unsigned int length; + + /** + * The native `char` array containing the string data. + */ + char* data; + + /** + * A constructor that does not initialize any fields. **WARNING:** The + * destructor will free string::data, which this constructor does not + * initialize. The user must initialize string::data using `init` or + * manually before the string is destroyed. + */ + string() { } + + /** + * Constructs the string by copying from the given null-terminated C string `src`. + */ + string(const char* src) { + if (!initialize(src, (unsigned int) strlen(src))) + exit(EXIT_FAILURE); + } + + /** + * Constructs the string by copying from the given native char array `src` + * with given `length`. + */ + string(const char* src, unsigned int length) { + if (!initialize(src, length)) + exit(EXIT_FAILURE); + } + + /** + * Constructs the string by initializing string::data with size `length` + * but without settings its contents. + */ + explicit string(unsigned int length) { + if (!initialize(length)) + exit(EXIT_FAILURE); + } + + ~string() { + core::free(data); + } + + /** + * Accesses the character at the given `index`. + */ + inline char& operator [] (unsigned int index) { + return data[index]; + } + + /** + * Accesses the character at the given `index`. + */ + inline const char& operator [] (unsigned int index) const { + return data[index]; + } + + /** + * Initializes this string by copying from `s`. Note that if this string + * was previously initialized, it is not freed. + */ + inline void operator = (const string& s) { + initialize(s.data, s.length); + } + + /** + * Appends the given null-terminated C string to this string. + */ + inline void operator += (const char* src) { + unsigned int src_length = (unsigned int) strlen(src); + char* new_data = (char*) realloc(data, sizeof(char) * (length + src_length)); + if (new_data == NULL) { + fprintf(stderr, "string.operator += ERROR: Unable to expand string.\n"); + exit(EXIT_FAILURE); + } + + data = new_data; + memcpy(data + length, src, sizeof(char) * src_length); + length += src_length; + } + + /** + * Appends the given string to this string. + */ + inline void operator += (const string& src) { + unsigned int src_length = src.length; + char* new_data = (char*) realloc(data, sizeof(char) * (length + src_length)); + if (new_data == NULL) { + fprintf(stderr, "string.operator += ERROR: Unable to expand string.\n"); + exit(EXIT_FAILURE); + } + + data = new_data; + memcpy(data + length, src.data, sizeof(char) * src_length); + length += src_length; + } + + /** + * Returns whether the current string precedes `other` in lexicographical order. + */ + inline bool operator < (const string& other) const { + for (unsigned int i = 0; ; i++) { + if (i == length) { + if (i == other.length) return false; + else return true; + } else if (i == other.length) { + return false; + } + + if (data[i] > other.data[i]) + return false; + else if (data[i] < other.data[i]) + return true; + } + } + + /** + * Returns the smallest index `i` such that `string::data[i] == c`. + */ + inline unsigned int index_of(char c) const { + for (unsigned int i = 0; i < length; i++) { + if (data[i] == c) + return i; + } + return length; + } + + /** + * Returns whether string::data is NULL. This enables strings to be used as keys in hashtables. + */ + static inline bool is_empty(const string& key) { + return key.data == NULL; + } + + /** + * Sets string::data to NULL. This enables strings to be used as keys in hashtables. + */ + static inline void set_empty(string& key) { + key.data = NULL; + } + + /** + * Sets string::data to NULL for every element in `keys`. This enables + * strings to be used as keys in hashtables. + */ + static inline void set_empty(string* keys, unsigned int length) { + memset(static_cast(keys), 0, sizeof(string) * length); + } + + /** + * Returns the hash of the given `key`. + */ + static inline unsigned int hash(const string& key) { + return default_hash(key.data, key.length); + } + + /** + * Copies the string::length and string::data pointer from `src` to `dst`. + * Note this function does not create a new pointer and copy the character + * contents, it simply copies the `char*` pointer. + */ + static inline void move(const string& src, string& dst) { + dst.length = src.length; + dst.data = src.data; + } + + /** + * Copies the string in `src` to `dst`. This function initializes a new + * string::data array in `dst` and copies the contents from `src.data`. + */ + static inline bool copy(const string& src, string& dst) { + return dst.initialize(src.data, src.length); + } + + /** + * Swaps the contents and lengths of `first` and `second`. + */ + static inline void swap(string& first, string& second) { + core::swap(first.length, second.length); + core::swap(first.data, second.data); + } + + template + static inline long unsigned int size_of(const string& str, const Metric& metric) { + return core::size_of(str.length) + sizeof(char) * str.length; + } + + /** + * Frees the underlying `char` array in `str`. + */ + static inline void free(string& str) { + core::free(str.data); + } + +private: + bool initialize(unsigned int src_length) { + length = src_length; + data = (char*) malloc(sizeof(char) * (length == 0 ? 1 : length)); + if (data == NULL) { + fprintf(stderr, "string.initialize ERROR: Unable to initialize string.\n"); + return false; + } + return true; + } + + bool initialize(const char* src, unsigned int src_length) { + if (!initialize(src_length)) + return false; + memcpy(data, src, sizeof(char) * length); + return true; + } + + friend bool init(string&, const char*, unsigned int); + friend bool init(string&, unsigned int); +}; + +/** + * Initializes the string `dst` with the given native `char` array `src` and the given `length`. + */ +inline bool init(string& dst, const char* src, unsigned int length) { + return dst.initialize(src, length); +} + +/** + * Initializes the string `dst` with the given string `src`. + */ +inline bool init(string& dst, const string& src) { + return init(dst, src.data, src.length); +} + +/** + * Initializes the string `dst` by allocating string::data with size `length`, + * but this function does not set its contents. + */ +inline bool init(string& dst, unsigned int length) { + return dst.initialize(length); +} + +/** + * Reads a string `s` from `in`. + * \param s an uninitialized string structure. This function initializes `s`, + * and the caller is responsible for its memory and must call free to + * release its memory resources. + */ +template +inline bool read(string& s, Stream& in) { + if (!read(s.length, in)) return false; + s.data = (char*) malloc(sizeof(char) * s.length); + if (s.data == NULL) + return false; + return read(s.data, in, s.length); +} + +/** + * Writes the string `s` to `out`. + */ +template +inline bool write(const string& s, Stream& out) { + if (!write(s.length, out)) return false; + return write(s.data, out, s.length); +} + +/** + * Prints the string `s` to `stream`. + */ +template +inline bool print(const string& s, Stream&& stream) { + return fwrite(s.data, sizeof(char), s.length, stream) == s.length; +} + +/** + * Compares the string `first` to the null-terminated C string `second` and + * returns `true` if they are equivalent, and `false` otherwise. + */ +inline bool operator == (const string& first, const char* second) { + for (unsigned int i = 0; i < first.length; i++) { + if (first[i] != second[i]) + return false; + } + if (second[first.length] != '\0') + return false; + return true; +} + +/** + * Compares the null-terminated C string `first` to the string `second` and + * returns `true` if they are equivalent, and `false` otherwise. + */ +inline bool operator == (const char* first, const string& second) { + return (second == first); +} + +/** + * Compares the string `first` to the string `second` and returns `true` if + * they are equivalent, and `false` otherwise. + */ +inline bool operator == (const string& first, const string& second) { + if (first.length != second.length) return false; + /* we are guaranteed that only the first may be uninitialized */ + if (first.data == NULL) + return false; + return memcmp(first.data, second.data, first.length * sizeof(char)) == 0; +} + +/** + * Compares the string `first` to the null-terminated C string `second` and + * returns `false` if they are equivalent, and `true` otherwise. + */ +inline bool operator != (const string& first, const char* second) { + return !(first == second); +} + +/** + * Compares the null-terminated C string `first` to the string `second` and + * returns `false` if they are equivalent, and `true` otherwise. + */ +inline bool operator != (const char* first, const string& second) { + return !(second == first); +} + +/** + * Compares the string `first` to the string `second` and returns `false` if + * they are equivalent, and `true` otherwise. + */ +inline bool operator != (const string& first, const string& second) { + if (first.length != second.length) return true; + if (first.data == NULL) { + if (second.data == NULL) return true; + else return false; + } else if (second.data == NULL) + return false; + return memcmp(first.data, second.data, first.length * sizeof(char)) != 0; +} + +struct sequence { + unsigned int* tokens; + unsigned int length; + + sequence(unsigned int* src, unsigned int length) : tokens(src), length(length) { } + + inline bool operator = (const sequence& src) { + return initialize(src); + } + + inline unsigned int& operator [] (unsigned int index) { + return tokens[index]; + } + + inline unsigned int operator [] (unsigned int index) const { + return tokens[index]; + } + + static inline unsigned int hash(const sequence& key) { + return default_hash(key.tokens, key.length); + } + + static inline bool is_empty(const sequence& key) { + return key.tokens == NULL; + } + + static inline void set_empty(sequence& key) { + key.tokens = NULL; + } + + static inline void move(const sequence& src, sequence& dst) { + dst.tokens = src.tokens; + dst.length = src.length; + } + + static inline bool copy(const sequence& src, sequence& dst) { + return dst.initialize(src); + } + + static inline void swap(sequence& first, sequence& second) { + core::swap(first.tokens, second.tokens); + core::swap(first.length, second.length); + } + + static inline void free(sequence& seq) { + core::free(seq.tokens); + } + +private: + inline bool initialize(const sequence& src) { + length = src.length; + tokens = (unsigned int*) malloc(sizeof(unsigned int) * length); + if (tokens == NULL) { + fprintf(stderr, "sequence.initialize ERROR: Out of memory.\n"); + return false; + } + memcpy(tokens, src.tokens, sizeof(unsigned int) * length); + return true; + } +}; + +inline bool init(sequence& seq, unsigned int length) { + seq.length = length; + seq.tokens = (unsigned int*) malloc(sizeof(unsigned int) * length); + if (seq.tokens == NULL) { + fprintf(stderr, "init ERROR: Insufficient memory for token array in sequence.\n"); + return false; + } + return true; +} + +inline bool init(sequence& seq, const sequence& src) { + return sequence::copy(src, seq); +} + +inline bool operator == (const sequence& first, const sequence& second) { + /* only the first argument can be uninitialized */ + if (first.tokens == NULL) return false; + if (first.length != second.length) return false; + for (unsigned int i = 0; i < first.length; i++) + if (first.tokens[i] != second.tokens[i]) return false; + return true; +} + +inline bool operator != (const sequence& first, const sequence& second) { + return !(first == second); +} + +inline bool operator < (const sequence& first, const sequence& second) { + if (first.length < second.length) return true; + else if (first.length > second.length) return false; + for (unsigned int i = 0; i < first.length; i++) { + if (first[i] < second[i]) return true; + else if (first[i] > second[i]) return false; + } + return false; +} + +template +inline bool read(sequence& item, Stream& in) { + if (!read(item.length, in)) return false; + item.tokens = (unsigned int*) malloc(sizeof(unsigned int) * item.length); + if (item.tokens == NULL) return false; + if (!read(item.tokens, in, item.length)) { + free(item.tokens); + return false; + } + return true; +} + +template +inline bool write(const sequence& item, Stream& out) { + return write(item.length, out) && write(item.tokens, out, item.length); +} + +template +inline bool print(const sequence& item, + Stream&& out, Printer&&... printer) +{ + if (item.length == 0) return true; + bool success = print(item[0], out, std::forward(printer)...); + for (unsigned int i = 1; i < item.length; i++) + success &= print(' ', out) && print(item[i], out, std::forward(printer)...); + return success; +} + +/** + * A scribe that maps unsigned integer indices to core::string pointers. + * + * ```{.cpp} + * #include + * using namespace core; + * + * int main() { + * string first = "first"; + * string second = "second"; + * + * string_map_scribe scribe; + * scribe.map = (const string**) malloc(sizeof(const string*) * 2); + * scribe.map[0] = &first; scribe.map[1] = &second; + * scribe.length = 2; + * + * print(0, stdout, scribe); print(' ', stdout); + * print(1, stdout, scribe); + * + * free(scribe.map); + * } + * ``` + * The expected output of this program is `first second`. + * + * Another way to construct this structure is to convert a + * hash_map into a core::string** array using the + * core::invert() function. + * + * ```{.cpp} + * #include + * using namespace core; + * + * int main() { + * hash_map token_map(16); + * token_map.put("first", 0); + * token_map.put("second", 1); + * + * string_map_scribe scribe; + * scribe.map = invert(token_map); + * scribe.length = 2; + * + * print(0, stdout, scribe); print(' ', stdout); + * print(1, stdout, scribe); + * + * free(scribe.map); + * for (auto entry : token_map) + * free(entry.key); + * } + * ``` + * The expected output of this program is `first second`. Notice that since + * hash_map does not automatically free its elements, we do it manually using a + * range-based for loop. + */ +struct string_map_scribe { + /** + * The native array of `const string*` elements that represents the + * effective map from unsigned integers to strings. + */ + const string** map; + + /** + * The length of the native array string_map_scribe::map. + */ + unsigned int length; +}; + +template::value>::type* = nullptr> +inline bool print_special_string_helper(unsigned int item, Stream& out, Printer&&... printer) { + return print_special_string(item, out, std::forward(printer)...); +} + +template::value>::type* = nullptr> +inline bool print_special_string_helper(unsigned int item, Stream& out, Printer&&... printer) { + fprintf(stderr, "print ERROR: The unsigned int %u exceeds the bounds of the " + "string_map_scribe. Did you forget to implement print_special_string?\n", item); + return true; +} + +/** + * Prints `item` to `out` using the string_map_scribe `printer`. If + * `item < printer.length`, the string at index `item` is accessed from + * string_map_scribe::map and printed to `out`. Otherwise, there are two cases: + * 1. If the function `bool print_special_string(unsigned int, Stream&)` is + * defined, this function calls it with arguments `item` and `out`. + * 2. If such a function is not defined, an error message is printed to + * [stderr](http://en.cppreference.com/w/cpp/io/c) and the function + * returns `true`. + * + * \param string_printer a scribe for which the function + * `bool print(const string&, Stream&, Printer&&...)` is defined, which is + * used to print the string itself. Note that since this is a variadic + * argument, it may be empty. + */ +template +inline bool print(unsigned int item, Stream&& out, const string_map_scribe& printer, Printer&&... string_printer) +{ + if (item < printer.length) + return print(*printer.map[item], out, std::forward(string_printer)...); + else return print_special_string_helper(item, out, std::forward(string_printer)...); +} + +/** + * Looks for the token `identifier` in the given hash_map `map`. If such a key + * exists in the map, `id` is set to its corresponding value and `true` is + * returned. If not, a new entry is added to the map with `identifier` as the + * key and `map.table.size + 1` as its value, and `id` is set to this new value. + * \returns `true` upon success, or `false` if the hash_map `map` could not be + * resized to accommodate the new entry. + */ +bool get_token(const string& identifier, unsigned int& id, hash_map& map) { + if (!map.check_size()) { + fprintf(stderr, "get_token ERROR: Unable to expand token map.\n"); + return false; + } + + bool contains; + unsigned int bucket; + unsigned int& value = map.get(identifier, contains, bucket); + if (!contains) { + map.table.keys[bucket] = identifier; + map.table.size++; + value = map.table.size; + } + id = value; + return true; +} + +/** + * Opens the file with the given `filename` with the given `mode` and returns + * either a [FILE](http://en.cppreference.com/w/c/io) pointer on success, or + * `NULL` on failure. [fclose](http://en.cppreference.com/w/c/io/fclose) should + * be used to close the stream once reading is complete. + * + * On Windows, this function uses + * [fopen_s](https://msdn.microsoft.com/en-us/library/z5hh6ee9.aspx) to open + * the stream. + */ +inline FILE* open_file(const char* filename, const char* mode) { +#if defined(_WIN32) + FILE* file; + if (fopen_s(&file, filename, mode) != 0) + return NULL; + return file; +#else + return fopen(filename, mode); +#endif +} + +/** + * Reads the contents of the file whose path is given by `filename`. If the + * file cannot be opened for reading, or if there is insufficient memory to + * allocate a buffer, `NULL` is returned. `bytes_read` is set to the number of + * bytes read, and the file contents are returned. The caller is responsible + * for the memory of the returned buffer and must call free to release its + * memory resources. + * \tparam AppendNull if `true`, a null terminating character is appended. + */ +template +inline char* read_file(const char* filename, size_t& bytes_read) +{ + FILE* fin = open_file(filename, "rb"); + if (fin == NULL || fseek(fin, 0, SEEK_END) != 0) { + if (fin != NULL) { +#if defined(_WIN32) + errno = (int) GetLastError(); +#endif + perror("read_file ERROR"); + fclose(fin); + } + return NULL; + } + + long int filesize = ftell(fin); + if (filesize == -1L) { + fprintf(stderr, "read_file ERROR: `ftell` returned error.\n"); + fclose(fin); + return NULL; + } + + if (fseek(fin, 0, SEEK_SET) != 0) { + fprintf(stderr, "read_file ERROR: `seek` returned error.\n"); + fclose(fin); + return NULL; + } + + char* data = (char*) malloc(sizeof(char) * (AppendNull ? (filesize + 1) : filesize)); + if (data == NULL) { + fprintf(stderr, "read_file ERROR: Out of memory.\n"); + fclose(fin); + return NULL; + } + bytes_read = fread(data, sizeof(char), filesize, fin); + fclose(fin); + if (AppendNull) + data[filesize] = '\0'; + return data; +} + +/** + * This function inspects the directory given by the path `directory` and adds + * the list of filenames in that directory to the string array `out`. + * \see http://stackoverflow.com/questions/306533/how-do-i-get-a-list-of-files-in-a-directory-in-c + */ +inline bool get_files_in_directory(array& out, const char* directory) +{ +#if defined(_WIN32) + HANDLE dir; + WIN32_FIND_DATA file_data; + + size_t required; + if (mbstowcs_s(&required, NULL, 0, directory, 0) != 0) + return false; + required--; /* ignore the null terminator */ + wchar_t* dir_prefix = (wchar_t*) malloc(sizeof(wchar_t) * (required + 3)); + if (dir_prefix == NULL) { + fprintf(stderr, "get_files_in_directory ERROR: Out of memory.\n"); + return false; + } + if (mbstowcs_s(&required, dir_prefix, required + 3, directory, required) != 0) { + free(dir_prefix); + return false; + } + required--; /* ignore the null terminator */ + dir_prefix[required] = '/'; + dir_prefix[required + 1] = '*'; + dir_prefix[required + 2] = '\0'; + if ((dir = FindFirstFile(dir_prefix, &file_data)) == INVALID_HANDLE_VALUE) { + free(dir_prefix); + return false; + } + free(dir_prefix); + + do { + const bool is_directory = (file_data.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY) != 0; + if (file_data.cFileName[0] == '.' || is_directory) + continue; + + if (!out.ensure_capacity(out.length + 1)) + return false; + + if (wcstombs_s(&required, NULL, 0, file_data.cFileName, 0) != 0 + || !init(out[(unsigned int) out.length], (unsigned int) required)) + return false; + out.length++; + if (wcstombs_s(&required, out.last().data, required, file_data.cFileName, wcslen(file_data.cFileName)) != 0) + return false; + out.last().length--; /* remove the null terminator */ + } while (FindNextFile(dir, &file_data)); + + FindClose(dir); + return true; +#else + DIR* dir; + dirent* ent; + struct stat st; + + dir = opendir(directory); + if (dir == NULL) + return false; + while ((ent = readdir(dir)) != NULL) { + string full_file_name = string(directory); + full_file_name += "/"; + full_file_name += ent->d_name; + full_file_name += " "; + full_file_name[full_file_name.length - 1] = '\0'; + if (ent->d_name[0] == '.' || stat(full_file_name.data, &st) == -1) + continue; + + const bool is_directory = (st.st_mode & S_IFDIR) != 0; + if (is_directory) + continue; + + if (!out.ensure_capacity(out.length + 1)) + return false; + init(out[out.length], ent->d_name); + out.length++; + } + closedir(dir); + return true; +#endif +} + +} /* namespace core */ + +#endif /* UTILITY_H_ */ diff --git a/reasoners/data_preprocess/logic/graph_dataset_gen/core/xxhash.h b/reasoners/data_preprocess/logic/graph_dataset_gen/core/xxhash.h new file mode 100644 index 000000000..08ab79457 --- /dev/null +++ b/reasoners/data_preprocess/logic/graph_dataset_gen/core/xxhash.h @@ -0,0 +1,5580 @@ +/* + * xxHash - Extremely Fast Hash algorithm + * Header File + * Copyright (C) 2012-2020 Yann Collet + * + * BSD 2-Clause License (https://www.opensource.org/licenses/bsd-license.php) + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are + * met: + * + * * Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above + * copyright notice, this list of conditions and the following disclaimer + * in the documentation and/or other materials provided with the + * distribution. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS + * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT + * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR + * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT + * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, + * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT + * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, + * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY + * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT + * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + * You can contact the author at: + * - xxHash homepage: https://www.xxhash.com + * - xxHash source repository: https://github.com/Cyan4973/xxHash + */ +/*! + * @mainpage xxHash + * + * @file xxhash.h + * xxHash prototypes and implementation + */ +/* TODO: update */ +/* Notice extracted from xxHash homepage: + +xxHash is an extremely fast hash algorithm, running at RAM speed limits. +It also successfully passes all tests from the SMHasher suite. + +Comparison (single thread, Windows Seven 32 bits, using SMHasher on a Core 2 Duo @3GHz) + +Name Speed Q.Score Author +xxHash 5.4 GB/s 10 +CrapWow 3.2 GB/s 2 Andrew +MurmurHash 3a 2.7 GB/s 10 Austin Appleby +SpookyHash 2.0 GB/s 10 Bob Jenkins +SBox 1.4 GB/s 9 Bret Mulvey +Lookup3 1.2 GB/s 9 Bob Jenkins +SuperFastHash 1.2 GB/s 1 Paul Hsieh +CityHash64 1.05 GB/s 10 Pike & Alakuijala +FNV 0.55 GB/s 5 Fowler, Noll, Vo +CRC32 0.43 GB/s 9 +MD5-32 0.33 GB/s 10 Ronald L. Rivest +SHA1-32 0.28 GB/s 10 + +Q.Score is a measure of quality of the hash function. +It depends on successfully passing SMHasher test set. +10 is a perfect score. + +Note: SMHasher's CRC32 implementation is not the fastest one. +Other speed-oriented implementations can be faster, +especially in combination with PCLMUL instruction: +https://fastcompression.blogspot.com/2019/03/presenting-xxh3.html?showComment=1552696407071#c3490092340461170735 + +A 64-bit version, named XXH64, is available since r35. +It offers much better speed, but for 64-bit applications only. +Name Speed on 64 bits Speed on 32 bits +XXH64 13.8 GB/s 1.9 GB/s +XXH32 6.8 GB/s 6.0 GB/s +*/ + +#if defined (__cplusplus) +extern "C" { +#endif + +/* **************************** + * INLINE mode + ******************************/ +/*! + * XXH_INLINE_ALL (and XXH_PRIVATE_API) + * Use these build macros to inline xxhash into the target unit. + * Inlining improves performance on small inputs, especially when the length is + * expressed as a compile-time constant: + * + * https://fastcompression.blogspot.com/2018/03/xxhash-for-small-keys-impressive-power.html + * + * It also keeps xxHash symbols private to the unit, so they are not exported. + * + * Usage: + * #define XXH_INLINE_ALL + * #include "xxhash.h" + * + * Do not compile and link xxhash.o as a separate object, as it is not useful. + */ +#if (defined(XXH_INLINE_ALL) || defined(XXH_PRIVATE_API)) \ + && !defined(XXH_INLINE_ALL_31684351384) + /* this section should be traversed only once */ +# define XXH_INLINE_ALL_31684351384 + /* give access to the advanced API, required to compile implementations */ +# undef XXH_STATIC_LINKING_ONLY /* avoid macro redef */ +# define XXH_STATIC_LINKING_ONLY + /* make all functions private */ +# undef XXH_PUBLIC_API +# if defined(__GNUC__) +# define XXH_PUBLIC_API static __inline __attribute__((unused)) +# elif defined (__cplusplus) || (defined (__STDC_VERSION__) && (__STDC_VERSION__ >= 199901L) /* C99 */) +# define XXH_PUBLIC_API static inline +# elif defined(_MSC_VER) +# define XXH_PUBLIC_API static __inline +# else + /* note: this version may generate warnings for unused static functions */ +# define XXH_PUBLIC_API static +# endif + + /* + * This part deals with the special case where a unit wants to inline xxHash, + * but "xxhash.h" has previously been included without XXH_INLINE_ALL, + * such as part of some previously included *.h header file. + * Without further action, the new include would just be ignored, + * and functions would effectively _not_ be inlined (silent failure). + * The following macros solve this situation by prefixing all inlined names, + * avoiding naming collision with previous inclusions. + */ + /* Before that, we unconditionally #undef all symbols, + * in case they were already defined with XXH_NAMESPACE. + * They will then be redefined for XXH_INLINE_ALL + */ +# undef XXH_versionNumber + /* XXH32 */ +# undef XXH32 +# undef XXH32_createState +# undef XXH32_freeState +# undef XXH32_reset +# undef XXH32_update +# undef XXH32_digest +# undef XXH32_copyState +# undef XXH32_canonicalFromHash +# undef XXH32_hashFromCanonical + /* XXH64 */ +# undef XXH64 +# undef XXH64_createState +# undef XXH64_freeState +# undef XXH64_reset +# undef XXH64_update +# undef XXH64_digest +# undef XXH64_copyState +# undef XXH64_canonicalFromHash +# undef XXH64_hashFromCanonical + /* XXH3_64bits */ +# undef XXH3_64bits +# undef XXH3_64bits_withSecret +# undef XXH3_64bits_withSeed +# undef XXH3_64bits_withSecretandSeed +# undef XXH3_createState +# undef XXH3_freeState +# undef XXH3_copyState +# undef XXH3_64bits_reset +# undef XXH3_64bits_reset_withSeed +# undef XXH3_64bits_reset_withSecret +# undef XXH3_64bits_update +# undef XXH3_64bits_digest +# undef XXH3_generateSecret + /* XXH3_128bits */ +# undef XXH128 +# undef XXH3_128bits +# undef XXH3_128bits_withSeed +# undef XXH3_128bits_withSecret +# undef XXH3_128bits_reset +# undef XXH3_128bits_reset_withSeed +# undef XXH3_128bits_reset_withSecret +# undef XXH3_128bits_reset_withSecretandSeed +# undef XXH3_128bits_update +# undef XXH3_128bits_digest +# undef XXH128_isEqual +# undef XXH128_cmp +# undef XXH128_canonicalFromHash +# undef XXH128_hashFromCanonical + /* Finally, free the namespace itself */ +# undef XXH_NAMESPACE + + /* employ the namespace for XXH_INLINE_ALL */ +# define XXH_NAMESPACE XXH_INLINE_ + /* + * Some identifiers (enums, type names) are not symbols, + * but they must nonetheless be renamed to avoid redeclaration. + * Alternative solution: do not redeclare them. + * However, this requires some #ifdefs, and has a more dispersed impact. + * Meanwhile, renaming can be achieved in a single place. + */ +# define XXH_IPREF(Id) XXH_NAMESPACE ## Id +# define XXH_OK XXH_IPREF(XXH_OK) +# define XXH_ERROR XXH_IPREF(XXH_ERROR) +# define XXH_errorcode XXH_IPREF(XXH_errorcode) +# define XXH32_canonical_t XXH_IPREF(XXH32_canonical_t) +# define XXH64_canonical_t XXH_IPREF(XXH64_canonical_t) +# define XXH128_canonical_t XXH_IPREF(XXH128_canonical_t) +# define XXH32_state_s XXH_IPREF(XXH32_state_s) +# define XXH32_state_t XXH_IPREF(XXH32_state_t) +# define XXH64_state_s XXH_IPREF(XXH64_state_s) +# define XXH64_state_t XXH_IPREF(XXH64_state_t) +# define XXH3_state_s XXH_IPREF(XXH3_state_s) +# define XXH3_state_t XXH_IPREF(XXH3_state_t) +# define XXH128_hash_t XXH_IPREF(XXH128_hash_t) + /* Ensure the header is parsed again, even if it was previously included */ +# undef XXHASH_H_5627135585666179 +# undef XXHASH_H_STATIC_13879238742 +#endif /* XXH_INLINE_ALL || XXH_PRIVATE_API */ + + + +/* **************************************************************** + * Stable API + *****************************************************************/ +#ifndef XXHASH_H_5627135585666179 +#define XXHASH_H_5627135585666179 1 + + +/*! + * @defgroup public Public API + * Contains details on the public xxHash functions. + * @{ + */ +/* specific declaration modes for Windows */ +#if !defined(XXH_INLINE_ALL) && !defined(XXH_PRIVATE_API) +# if defined(WIN32) && defined(_MSC_VER) && (defined(XXH_IMPORT) || defined(XXH_EXPORT)) +# ifdef XXH_EXPORT +# define XXH_PUBLIC_API __declspec(dllexport) +# elif XXH_IMPORT +# define XXH_PUBLIC_API __declspec(dllimport) +# endif +# else +# define XXH_PUBLIC_API /* do nothing */ +# endif +#endif + +#ifdef XXH_DOXYGEN +/*! + * @brief Emulate a namespace by transparently prefixing all symbols. + * + * If you want to include _and expose_ xxHash functions from within your own + * library, but also want to avoid symbol collisions with other libraries which + * may also include xxHash, you can use XXH_NAMESPACE to automatically prefix + * any public symbol from xxhash library with the value of XXH_NAMESPACE + * (therefore, avoid empty or numeric values). + * + * Note that no change is required within the calling program as long as it + * includes `xxhash.h`: Regular symbol names will be automatically translated + * by this header. + */ +# define XXH_NAMESPACE /* YOUR NAME HERE */ +# undef XXH_NAMESPACE +#endif + +#ifdef XXH_NAMESPACE +# define XXH_CAT(A,B) A##B +# define XXH_NAME2(A,B) XXH_CAT(A,B) +# define XXH_versionNumber XXH_NAME2(XXH_NAMESPACE, XXH_versionNumber) +/* XXH32 */ +# define XXH32 XXH_NAME2(XXH_NAMESPACE, XXH32) +# define XXH32_createState XXH_NAME2(XXH_NAMESPACE, XXH32_createState) +# define XXH32_freeState XXH_NAME2(XXH_NAMESPACE, XXH32_freeState) +# define XXH32_reset XXH_NAME2(XXH_NAMESPACE, XXH32_reset) +# define XXH32_update XXH_NAME2(XXH_NAMESPACE, XXH32_update) +# define XXH32_digest XXH_NAME2(XXH_NAMESPACE, XXH32_digest) +# define XXH32_copyState XXH_NAME2(XXH_NAMESPACE, XXH32_copyState) +# define XXH32_canonicalFromHash XXH_NAME2(XXH_NAMESPACE, XXH32_canonicalFromHash) +# define XXH32_hashFromCanonical XXH_NAME2(XXH_NAMESPACE, XXH32_hashFromCanonical) +/* XXH64 */ +# define XXH64 XXH_NAME2(XXH_NAMESPACE, XXH64) +# define XXH64_createState XXH_NAME2(XXH_NAMESPACE, XXH64_createState) +# define XXH64_freeState XXH_NAME2(XXH_NAMESPACE, XXH64_freeState) +# define XXH64_reset XXH_NAME2(XXH_NAMESPACE, XXH64_reset) +# define XXH64_update XXH_NAME2(XXH_NAMESPACE, XXH64_update) +# define XXH64_digest XXH_NAME2(XXH_NAMESPACE, XXH64_digest) +# define XXH64_copyState XXH_NAME2(XXH_NAMESPACE, XXH64_copyState) +# define XXH64_canonicalFromHash XXH_NAME2(XXH_NAMESPACE, XXH64_canonicalFromHash) +# define XXH64_hashFromCanonical XXH_NAME2(XXH_NAMESPACE, XXH64_hashFromCanonical) +/* XXH3_64bits */ +# define XXH3_64bits XXH_NAME2(XXH_NAMESPACE, XXH3_64bits) +# define XXH3_64bits_withSecret XXH_NAME2(XXH_NAMESPACE, XXH3_64bits_withSecret) +# define XXH3_64bits_withSeed XXH_NAME2(XXH_NAMESPACE, XXH3_64bits_withSeed) +# define XXH3_64bits_withSecretandSeed XXH_NAME2(XXH_NAMESPACE, XXH3_64bits_withSecretandSeed) +# define XXH3_createState XXH_NAME2(XXH_NAMESPACE, XXH3_createState) +# define XXH3_freeState XXH_NAME2(XXH_NAMESPACE, XXH3_freeState) +# define XXH3_copyState XXH_NAME2(XXH_NAMESPACE, XXH3_copyState) +# define XXH3_64bits_reset XXH_NAME2(XXH_NAMESPACE, XXH3_64bits_reset) +# define XXH3_64bits_reset_withSeed XXH_NAME2(XXH_NAMESPACE, XXH3_64bits_reset_withSeed) +# define XXH3_64bits_reset_withSecret XXH_NAME2(XXH_NAMESPACE, XXH3_64bits_reset_withSecret) +# define XXH3_64bits_reset_withSecretandSeed XXH_NAME2(XXH_NAMESPACE, XXH3_64bits_reset_withSecretandSeed) +# define XXH3_64bits_update XXH_NAME2(XXH_NAMESPACE, XXH3_64bits_update) +# define XXH3_64bits_digest XXH_NAME2(XXH_NAMESPACE, XXH3_64bits_digest) +# define XXH3_generateSecret XXH_NAME2(XXH_NAMESPACE, XXH3_generateSecret) +# define XXH3_generateSecret_fromSeed XXH_NAME2(XXH_NAMESPACE, XXH3_generateSecret_fromSeed) +/* XXH3_128bits */ +# define XXH128 XXH_NAME2(XXH_NAMESPACE, XXH128) +# define XXH3_128bits XXH_NAME2(XXH_NAMESPACE, XXH3_128bits) +# define XXH3_128bits_withSeed XXH_NAME2(XXH_NAMESPACE, XXH3_128bits_withSeed) +# define XXH3_128bits_withSecret XXH_NAME2(XXH_NAMESPACE, XXH3_128bits_withSecret) +# define XXH3_128bits_withSecretandSeed XXH_NAME2(XXH_NAMESPACE, XXH3_128bits_withSecretandSeed) +# define XXH3_128bits_reset XXH_NAME2(XXH_NAMESPACE, XXH3_128bits_reset) +# define XXH3_128bits_reset_withSeed XXH_NAME2(XXH_NAMESPACE, XXH3_128bits_reset_withSeed) +# define XXH3_128bits_reset_withSecret XXH_NAME2(XXH_NAMESPACE, XXH3_128bits_reset_withSecret) +# define XXH3_128bits_reset_withSecretandSeed XXH_NAME2(XXH_NAMESPACE, XXH3_128bits_reset_withSecretandSeed) +# define XXH3_128bits_update XXH_NAME2(XXH_NAMESPACE, XXH3_128bits_update) +# define XXH3_128bits_digest XXH_NAME2(XXH_NAMESPACE, XXH3_128bits_digest) +# define XXH128_isEqual XXH_NAME2(XXH_NAMESPACE, XXH128_isEqual) +# define XXH128_cmp XXH_NAME2(XXH_NAMESPACE, XXH128_cmp) +# define XXH128_canonicalFromHash XXH_NAME2(XXH_NAMESPACE, XXH128_canonicalFromHash) +# define XXH128_hashFromCanonical XXH_NAME2(XXH_NAMESPACE, XXH128_hashFromCanonical) +#endif + + +/* ************************************* +* Version +***************************************/ +#define XXH_VERSION_MAJOR 0 +#define XXH_VERSION_MINOR 8 +#define XXH_VERSION_RELEASE 1 +#define XXH_VERSION_NUMBER (XXH_VERSION_MAJOR *100*100 + XXH_VERSION_MINOR *100 + XXH_VERSION_RELEASE) + +/*! + * @brief Obtains the xxHash version. + * + * This is mostly useful when xxHash is compiled as a shared library, + * since the returned value comes from the library, as opposed to header file. + * + * @return `XXH_VERSION_NUMBER` of the invoked library. + */ +XXH_PUBLIC_API unsigned XXH_versionNumber (void); + + +/* **************************** +* Common basic types +******************************/ +#include /* size_t */ +typedef enum { XXH_OK=0, XXH_ERROR } XXH_errorcode; + + +/*-********************************************************************** +* 32-bit hash +************************************************************************/ +#if defined(XXH_DOXYGEN) /* Don't show include */ +/*! + * @brief An unsigned 32-bit integer. + * + * Not necessarily defined to `uint32_t` but functionally equivalent. + */ +typedef uint32_t XXH32_hash_t; + +#elif !defined (__VMS) \ + && (defined (__cplusplus) \ + || (defined (__STDC_VERSION__) && (__STDC_VERSION__ >= 199901L) /* C99 */) ) +# include + typedef uint32_t XXH32_hash_t; + +#else +# include +# if UINT_MAX == 0xFFFFFFFFUL + typedef unsigned int XXH32_hash_t; +# else +# if ULONG_MAX == 0xFFFFFFFFUL + typedef unsigned long XXH32_hash_t; +# else +# error "unsupported platform: need a 32-bit type" +# endif +# endif +#endif + +/*! + * @} + * + * @defgroup xxh32_family XXH32 family + * @ingroup public + * Contains functions used in the classic 32-bit xxHash algorithm. + * + * @note + * XXH32 is useful for older platforms, with no or poor 64-bit performance. + * Note that @ref xxh3_family provides competitive speed + * for both 32-bit and 64-bit systems, and offers true 64/128 bit hash results. + * + * @see @ref xxh64_family, @ref xxh3_family : Other xxHash families + * @see @ref xxh32_impl for implementation details + * @{ + */ + +/*! + * @brief Calculates the 32-bit hash of @p input using xxHash32. + * + * Speed on Core 2 Duo @ 3 GHz (single thread, SMHasher benchmark): 5.4 GB/s + * + * @param input The block of data to be hashed, at least @p length bytes in size. + * @param length The length of @p input, in bytes. + * @param seed The 32-bit seed to alter the hash's output predictably. + * + * @pre + * The memory between @p input and @p input + @p length must be valid, + * readable, contiguous memory. However, if @p length is `0`, @p input may be + * `NULL`. In C++, this also must be *TriviallyCopyable*. + * + * @return The calculated 32-bit hash value. + * + * @see + * XXH64(), XXH3_64bits_withSeed(), XXH3_128bits_withSeed(), XXH128(): + * Direct equivalents for the other variants of xxHash. + * @see + * XXH32_createState(), XXH32_update(), XXH32_digest(): Streaming version. + */ +XXH_PUBLIC_API XXH32_hash_t XXH32 (const void* input, size_t length, XXH32_hash_t seed); + +/*! + * Streaming functions generate the xxHash value from an incremental input. + * This method is slower than single-call functions, due to state management. + * For small inputs, prefer `XXH32()` and `XXH64()`, which are better optimized. + * + * An XXH state must first be allocated using `XXH*_createState()`. + * + * Start a new hash by initializing the state with a seed using `XXH*_reset()`. + * + * Then, feed the hash state by calling `XXH*_update()` as many times as necessary. + * + * The function returns an error code, with 0 meaning OK, and any other value + * meaning there is an error. + * + * Finally, a hash value can be produced anytime, by using `XXH*_digest()`. + * This function returns the nn-bits hash as an int or long long. + * + * It's still possible to continue inserting input into the hash state after a + * digest, and generate new hash values later on by invoking `XXH*_digest()`. + * + * When done, release the state using `XXH*_freeState()`. + * + * Example code for incrementally hashing a file: + * @code{.c} + * #include + * #include + * #define BUFFER_SIZE 256 + * + * // Note: XXH64 and XXH3 use the same interface. + * XXH32_hash_t + * hashFile(FILE* stream) + * { + * XXH32_state_t* state; + * unsigned char buf[BUFFER_SIZE]; + * size_t amt; + * XXH32_hash_t hash; + * + * state = XXH32_createState(); // Create a state + * assert(state != NULL); // Error check here + * XXH32_reset(state, 0xbaad5eed); // Reset state with our seed + * while ((amt = fread(buf, 1, sizeof(buf), stream)) != 0) { + * XXH32_update(state, buf, amt); // Hash the file in chunks + * } + * hash = XXH32_digest(state); // Finalize the hash + * XXH32_freeState(state); // Clean up + * return hash; + * } + * @endcode + */ + +/*! + * @typedef struct XXH32_state_s XXH32_state_t + * @brief The opaque state struct for the XXH32 streaming API. + * + * @see XXH32_state_s for details. + */ +typedef struct XXH32_state_s XXH32_state_t; + +/*! + * @brief Allocates an @ref XXH32_state_t. + * + * Must be freed with XXH32_freeState(). + * @return An allocated XXH32_state_t on success, `NULL` on failure. + */ +XXH_PUBLIC_API XXH32_state_t* XXH32_createState(void); +/*! + * @brief Frees an @ref XXH32_state_t. + * + * Must be allocated with XXH32_createState(). + * @param statePtr A pointer to an @ref XXH32_state_t allocated with @ref XXH32_createState(). + * @return XXH_OK. + */ +XXH_PUBLIC_API XXH_errorcode XXH32_freeState(XXH32_state_t* statePtr); +/*! + * @brief Copies one @ref XXH32_state_t to another. + * + * @param dst_state The state to copy to. + * @param src_state The state to copy from. + * @pre + * @p dst_state and @p src_state must not be `NULL` and must not overlap. + */ +XXH_PUBLIC_API void XXH32_copyState(XXH32_state_t* dst_state, const XXH32_state_t* src_state); + +/*! + * @brief Resets an @ref XXH32_state_t to begin a new hash. + * + * This function resets and seeds a state. Call it before @ref XXH32_update(). + * + * @param statePtr The state struct to reset. + * @param seed The 32-bit seed to alter the hash result predictably. + * + * @pre + * @p statePtr must not be `NULL`. + * + * @return @ref XXH_OK on success, @ref XXH_ERROR on failure. + */ +XXH_PUBLIC_API XXH_errorcode XXH32_reset (XXH32_state_t* statePtr, XXH32_hash_t seed); + +/*! + * @brief Consumes a block of @p input to an @ref XXH32_state_t. + * + * Call this to incrementally consume blocks of data. + * + * @param statePtr The state struct to update. + * @param input The block of data to be hashed, at least @p length bytes in size. + * @param length The length of @p input, in bytes. + * + * @pre + * @p statePtr must not be `NULL`. + * @pre + * The memory between @p input and @p input + @p length must be valid, + * readable, contiguous memory. However, if @p length is `0`, @p input may be + * `NULL`. In C++, this also must be *TriviallyCopyable*. + * + * @return @ref XXH_OK on success, @ref XXH_ERROR on failure. + */ +XXH_PUBLIC_API XXH_errorcode XXH32_update (XXH32_state_t* statePtr, const void* input, size_t length); + +/*! + * @brief Returns the calculated hash value from an @ref XXH32_state_t. + * + * @note + * Calling XXH32_digest() will not affect @p statePtr, so you can update, + * digest, and update again. + * + * @param statePtr The state struct to calculate the hash from. + * + * @pre + * @p statePtr must not be `NULL`. + * + * @return The calculated xxHash32 value from that state. + */ +XXH_PUBLIC_API XXH32_hash_t XXH32_digest (const XXH32_state_t* statePtr); + +/******* Canonical representation *******/ + +/* + * The default return values from XXH functions are unsigned 32 and 64 bit + * integers. + * This the simplest and fastest format for further post-processing. + * + * However, this leaves open the question of what is the order on the byte level, + * since little and big endian conventions will store the same number differently. + * + * The canonical representation settles this issue by mandating big-endian + * convention, the same convention as human-readable numbers (large digits first). + * + * When writing hash values to storage, sending them over a network, or printing + * them, it's highly recommended to use the canonical representation to ensure + * portability across a wider range of systems, present and future. + * + * The following functions allow transformation of hash values to and from + * canonical format. + */ + +/*! + * @brief Canonical (big endian) representation of @ref XXH32_hash_t. + */ +typedef struct { + unsigned char digest[4]; /*!< Hash bytes, big endian */ +} XXH32_canonical_t; + +/*! + * @brief Converts an @ref XXH32_hash_t to a big endian @ref XXH32_canonical_t. + * + * @param dst The @ref XXH32_canonical_t pointer to be stored to. + * @param hash The @ref XXH32_hash_t to be converted. + * + * @pre + * @p dst must not be `NULL`. + */ +XXH_PUBLIC_API void XXH32_canonicalFromHash(XXH32_canonical_t* dst, XXH32_hash_t hash); + +/*! + * @brief Converts an @ref XXH32_canonical_t to a native @ref XXH32_hash_t. + * + * @param src The @ref XXH32_canonical_t to convert. + * + * @pre + * @p src must not be `NULL`. + * + * @return The converted hash. + */ +XXH_PUBLIC_API XXH32_hash_t XXH32_hashFromCanonical(const XXH32_canonical_t* src); + + +#ifdef __has_attribute +# define XXH_HAS_ATTRIBUTE(x) __has_attribute(x) +#else +# define XXH_HAS_ATTRIBUTE(x) 0 +#endif + +/* C-language Attributes are added in C23. */ +#if defined(__STDC_VERSION__) && (__STDC_VERSION__ > 201710L) && defined(__has_c_attribute) +# define XXH_HAS_C_ATTRIBUTE(x) __has_c_attribute(x) +#else +# define XXH_HAS_C_ATTRIBUTE(x) 0 +#endif + +#if defined(__cplusplus) && defined(__has_cpp_attribute) +# define XXH_HAS_CPP_ATTRIBUTE(x) __has_cpp_attribute(x) +#else +# define XXH_HAS_CPP_ATTRIBUTE(x) 0 +#endif + +/* +Define XXH_FALLTHROUGH macro for annotating switch case with the 'fallthrough' attribute +introduced in CPP17 and C23. +CPP17 : https://en.cppreference.com/w/cpp/language/attributes/fallthrough +C23 : https://en.cppreference.com/w/c/language/attributes/fallthrough +*/ +#if XXH_HAS_C_ATTRIBUTE(x) +# define XXH_FALLTHROUGH [[fallthrough]] +#elif XXH_HAS_CPP_ATTRIBUTE(x) +# define XXH_FALLTHROUGH [[fallthrough]] +#elif XXH_HAS_ATTRIBUTE(__fallthrough__) +# define XXH_FALLTHROUGH __attribute__ ((fallthrough)) +#else +# define XXH_FALLTHROUGH +#endif + +/*! + * @} + * @ingroup public + * @{ + */ + +#ifndef XXH_NO_LONG_LONG +/*-********************************************************************** +* 64-bit hash +************************************************************************/ +#if defined(XXH_DOXYGEN) /* don't include */ +/*! + * @brief An unsigned 64-bit integer. + * + * Not necessarily defined to `uint64_t` but functionally equivalent. + */ +typedef uint64_t XXH64_hash_t; +#elif !defined (__VMS) \ + && (defined (__cplusplus) \ + || (defined (__STDC_VERSION__) && (__STDC_VERSION__ >= 199901L) /* C99 */) ) +# include + typedef uint64_t XXH64_hash_t; +#else +# include +# if defined(__LP64__) && ULONG_MAX == 0xFFFFFFFFFFFFFFFFULL + /* LP64 ABI says uint64_t is unsigned long */ + typedef unsigned long XXH64_hash_t; +# else + /* the following type must have a width of 64-bit */ + typedef unsigned long long XXH64_hash_t; +# endif +#endif + +/*! + * @} + * + * @defgroup xxh64_family XXH64 family + * @ingroup public + * @{ + * Contains functions used in the classic 64-bit xxHash algorithm. + * + * @note + * XXH3 provides competitive speed for both 32-bit and 64-bit systems, + * and offers true 64/128 bit hash results. + * It provides better speed for systems with vector processing capabilities. + */ + + +/*! + * @brief Calculates the 64-bit hash of @p input using xxHash64. + * + * This function usually runs faster on 64-bit systems, but slower on 32-bit + * systems (see benchmark). + * + * @param input The block of data to be hashed, at least @p length bytes in size. + * @param length The length of @p input, in bytes. + * @param seed The 64-bit seed to alter the hash's output predictably. + * + * @pre + * The memory between @p input and @p input + @p length must be valid, + * readable, contiguous memory. However, if @p length is `0`, @p input may be + * `NULL`. In C++, this also must be *TriviallyCopyable*. + * + * @return The calculated 64-bit hash. + * + * @see + * XXH32(), XXH3_64bits_withSeed(), XXH3_128bits_withSeed(), XXH128(): + * Direct equivalents for the other variants of xxHash. + * @see + * XXH64_createState(), XXH64_update(), XXH64_digest(): Streaming version. + */ +XXH_PUBLIC_API XXH64_hash_t XXH64(const void* input, size_t length, XXH64_hash_t seed); + +/******* Streaming *******/ +/*! + * @brief The opaque state struct for the XXH64 streaming API. + * + * @see XXH64_state_s for details. + */ +typedef struct XXH64_state_s XXH64_state_t; /* incomplete type */ +XXH_PUBLIC_API XXH64_state_t* XXH64_createState(void); +XXH_PUBLIC_API XXH_errorcode XXH64_freeState(XXH64_state_t* statePtr); +XXH_PUBLIC_API void XXH64_copyState(XXH64_state_t* dst_state, const XXH64_state_t* src_state); + +XXH_PUBLIC_API XXH_errorcode XXH64_reset (XXH64_state_t* statePtr, XXH64_hash_t seed); +XXH_PUBLIC_API XXH_errorcode XXH64_update (XXH64_state_t* statePtr, const void* input, size_t length); +XXH_PUBLIC_API XXH64_hash_t XXH64_digest (const XXH64_state_t* statePtr); + +/******* Canonical representation *******/ +typedef struct { unsigned char digest[sizeof(XXH64_hash_t)]; } XXH64_canonical_t; +XXH_PUBLIC_API void XXH64_canonicalFromHash(XXH64_canonical_t* dst, XXH64_hash_t hash); +XXH_PUBLIC_API XXH64_hash_t XXH64_hashFromCanonical(const XXH64_canonical_t* src); + +/*! + * @} + * ************************************************************************ + * @defgroup xxh3_family XXH3 family + * @ingroup public + * @{ + * + * XXH3 is a more recent hash algorithm featuring: + * - Improved speed for both small and large inputs + * - True 64-bit and 128-bit outputs + * - SIMD acceleration + * - Improved 32-bit viability + * + * Speed analysis methodology is explained here: + * + * https://fastcompression.blogspot.com/2019/03/presenting-xxh3.html + * + * Compared to XXH64, expect XXH3 to run approximately + * ~2x faster on large inputs and >3x faster on small ones, + * exact differences vary depending on platform. + * + * XXH3's speed benefits greatly from SIMD and 64-bit arithmetic, + * but does not require it. + * Any 32-bit and 64-bit targets that can run XXH32 smoothly + * can run XXH3 at competitive speeds, even without vector support. + * Further details are explained in the implementation. + * + * Optimized implementations are provided for AVX512, AVX2, SSE2, NEON, POWER8, + * ZVector and scalar targets. This can be controlled via the XXH_VECTOR macro. + * + * XXH3 implementation is portable: + * it has a generic C90 formulation that can be compiled on any platform, + * all implementations generage exactly the same hash value on all platforms. + * Starting from v0.8.0, it's also labelled "stable", meaning that + * any future version will also generate the same hash value. + * + * XXH3 offers 2 variants, _64bits and _128bits. + * + * When only 64 bits are needed, prefer invoking the _64bits variant, as it + * reduces the amount of mixing, resulting in faster speed on small inputs. + * It's also generally simpler to manipulate a scalar return type than a struct. + * + * The API supports one-shot hashing, streaming mode, and custom secrets. + */ + +/*-********************************************************************** +* XXH3 64-bit variant +************************************************************************/ + +/* XXH3_64bits(): + * default 64-bit variant, using default secret and default seed of 0. + * It's the fastest variant. */ +XXH_PUBLIC_API XXH64_hash_t XXH3_64bits(const void* data, size_t len); + +/* + * XXH3_64bits_withSeed(): + * This variant generates a custom secret on the fly + * based on default secret altered using the `seed` value. + * While this operation is decently fast, note that it's not completely free. + * Note: seed==0 produces the same results as XXH3_64bits(). + */ +XXH_PUBLIC_API XXH64_hash_t XXH3_64bits_withSeed(const void* data, size_t len, XXH64_hash_t seed); + +/*! + * The bare minimum size for a custom secret. + * + * @see + * XXH3_64bits_withSecret(), XXH3_64bits_reset_withSecret(), + * XXH3_128bits_withSecret(), XXH3_128bits_reset_withSecret(). + */ +#define XXH3_SECRET_SIZE_MIN 136 + +/* + * XXH3_64bits_withSecret(): + * It's possible to provide any blob of bytes as a "secret" to generate the hash. + * This makes it more difficult for an external actor to prepare an intentional collision. + * The main condition is that secretSize *must* be large enough (>= XXH3_SECRET_SIZE_MIN). + * However, the quality of the secret impacts the dispersion of the hash algorithm. + * Therefore, the secret _must_ look like a bunch of random bytes. + * Avoid "trivial" or structured data such as repeated sequences or a text document. + * Whenever in doubt about the "randomness" of the blob of bytes, + * consider employing "XXH3_generateSecret()" instead (see below). + * It will generate a proper high entropy secret derived from the blob of bytes. + * Another advantage of using XXH3_generateSecret() is that + * it guarantees that all bits within the initial blob of bytes + * will impact every bit of the output. + * This is not necessarily the case when using the blob of bytes directly + * because, when hashing _small_ inputs, only a portion of the secret is employed. + */ +XXH_PUBLIC_API XXH64_hash_t XXH3_64bits_withSecret(const void* data, size_t len, const void* secret, size_t secretSize); + + +/******* Streaming *******/ +/* + * Streaming requires state maintenance. + * This operation costs memory and CPU. + * As a consequence, streaming is slower than one-shot hashing. + * For better performance, prefer one-shot functions whenever applicable. + */ + +/*! + * @brief The state struct for the XXH3 streaming API. + * + * @see XXH3_state_s for details. + */ +typedef struct XXH3_state_s XXH3_state_t; +XXH_PUBLIC_API XXH3_state_t* XXH3_createState(void); +XXH_PUBLIC_API XXH_errorcode XXH3_freeState(XXH3_state_t* statePtr); +XXH_PUBLIC_API void XXH3_copyState(XXH3_state_t* dst_state, const XXH3_state_t* src_state); + +/* + * XXH3_64bits_reset(): + * Initialize with default parameters. + * digest will be equivalent to `XXH3_64bits()`. + */ +XXH_PUBLIC_API XXH_errorcode XXH3_64bits_reset(XXH3_state_t* statePtr); +/* + * XXH3_64bits_reset_withSeed(): + * Generate a custom secret from `seed`, and store it into `statePtr`. + * digest will be equivalent to `XXH3_64bits_withSeed()`. + */ +XXH_PUBLIC_API XXH_errorcode XXH3_64bits_reset_withSeed(XXH3_state_t* statePtr, XXH64_hash_t seed); +/* + * XXH3_64bits_reset_withSecret(): + * `secret` is referenced, it _must outlive_ the hash streaming session. + * Similar to one-shot API, `secretSize` must be >= `XXH3_SECRET_SIZE_MIN`, + * and the quality of produced hash values depends on secret's entropy + * (secret's content should look like a bunch of random bytes). + * When in doubt about the randomness of a candidate `secret`, + * consider employing `XXH3_generateSecret()` instead (see below). + */ +XXH_PUBLIC_API XXH_errorcode XXH3_64bits_reset_withSecret(XXH3_state_t* statePtr, const void* secret, size_t secretSize); + +XXH_PUBLIC_API XXH_errorcode XXH3_64bits_update (XXH3_state_t* statePtr, const void* input, size_t length); +XXH_PUBLIC_API XXH64_hash_t XXH3_64bits_digest (const XXH3_state_t* statePtr); + +/* note : canonical representation of XXH3 is the same as XXH64 + * since they both produce XXH64_hash_t values */ + + +/*-********************************************************************** +* XXH3 128-bit variant +************************************************************************/ + +/*! + * @brief The return value from 128-bit hashes. + * + * Stored in little endian order, although the fields themselves are in native + * endianness. + */ +typedef struct { + XXH64_hash_t low64; /*!< `value & 0xFFFFFFFFFFFFFFFF` */ + XXH64_hash_t high64; /*!< `value >> 64` */ +} XXH128_hash_t; + +XXH_PUBLIC_API XXH128_hash_t XXH3_128bits(const void* data, size_t len); +XXH_PUBLIC_API XXH128_hash_t XXH3_128bits_withSeed(const void* data, size_t len, XXH64_hash_t seed); +XXH_PUBLIC_API XXH128_hash_t XXH3_128bits_withSecret(const void* data, size_t len, const void* secret, size_t secretSize); + +/******* Streaming *******/ +/* + * Streaming requires state maintenance. + * This operation costs memory and CPU. + * As a consequence, streaming is slower than one-shot hashing. + * For better performance, prefer one-shot functions whenever applicable. + * + * XXH3_128bits uses the same XXH3_state_t as XXH3_64bits(). + * Use already declared XXH3_createState() and XXH3_freeState(). + * + * All reset and streaming functions have same meaning as their 64-bit counterpart. + */ + +XXH_PUBLIC_API XXH_errorcode XXH3_128bits_reset(XXH3_state_t* statePtr); +XXH_PUBLIC_API XXH_errorcode XXH3_128bits_reset_withSeed(XXH3_state_t* statePtr, XXH64_hash_t seed); +XXH_PUBLIC_API XXH_errorcode XXH3_128bits_reset_withSecret(XXH3_state_t* statePtr, const void* secret, size_t secretSize); + +XXH_PUBLIC_API XXH_errorcode XXH3_128bits_update (XXH3_state_t* statePtr, const void* input, size_t length); +XXH_PUBLIC_API XXH128_hash_t XXH3_128bits_digest (const XXH3_state_t* statePtr); + +/* Following helper functions make it possible to compare XXH128_hast_t values. + * Since XXH128_hash_t is a structure, this capability is not offered by the language. + * Note: For better performance, these functions can be inlined using XXH_INLINE_ALL */ + +/*! + * XXH128_isEqual(): + * Return: 1 if `h1` and `h2` are equal, 0 if they are not. + */ +XXH_PUBLIC_API int XXH128_isEqual(XXH128_hash_t h1, XXH128_hash_t h2); + +/*! + * XXH128_cmp(): + * + * This comparator is compatible with stdlib's `qsort()`/`bsearch()`. + * + * return: >0 if *h128_1 > *h128_2 + * =0 if *h128_1 == *h128_2 + * <0 if *h128_1 < *h128_2 + */ +XXH_PUBLIC_API int XXH128_cmp(const void* h128_1, const void* h128_2); + + +/******* Canonical representation *******/ +typedef struct { unsigned char digest[sizeof(XXH128_hash_t)]; } XXH128_canonical_t; +XXH_PUBLIC_API void XXH128_canonicalFromHash(XXH128_canonical_t* dst, XXH128_hash_t hash); +XXH_PUBLIC_API XXH128_hash_t XXH128_hashFromCanonical(const XXH128_canonical_t* src); + + +#endif /* XXH_NO_LONG_LONG */ + +/*! + * @} + */ +#endif /* XXHASH_H_5627135585666179 */ + + + +#if defined(XXH_STATIC_LINKING_ONLY) && !defined(XXHASH_H_STATIC_13879238742) +#define XXHASH_H_STATIC_13879238742 +/* **************************************************************************** + * This section contains declarations which are not guaranteed to remain stable. + * They may change in future versions, becoming incompatible with a different + * version of the library. + * These declarations should only be used with static linking. + * Never use them in association with dynamic linking! + ***************************************************************************** */ + +/* + * These definitions are only present to allow static allocation + * of XXH states, on stack or in a struct, for example. + * Never **ever** access their members directly. + */ + +/*! + * @internal + * @brief Structure for XXH32 streaming API. + * + * @note This is only defined when @ref XXH_STATIC_LINKING_ONLY, + * @ref XXH_INLINE_ALL, or @ref XXH_IMPLEMENTATION is defined. Otherwise it is + * an opaque type. This allows fields to safely be changed. + * + * Typedef'd to @ref XXH32_state_t. + * Do not access the members of this struct directly. + * @see XXH64_state_s, XXH3_state_s + */ +struct XXH32_state_s { + XXH32_hash_t total_len_32; /*!< Total length hashed, modulo 2^32 */ + XXH32_hash_t large_len; /*!< Whether the hash is >= 16 (handles @ref total_len_32 overflow) */ + XXH32_hash_t v[4]; /*!< Accumulator lanes */ + XXH32_hash_t mem32[4]; /*!< Internal buffer for partial reads. Treated as unsigned char[16]. */ + XXH32_hash_t memsize; /*!< Amount of data in @ref mem32 */ + XXH32_hash_t reserved; /*!< Reserved field. Do not read or write to it, it may be removed. */ +}; /* typedef'd to XXH32_state_t */ + + +#ifndef XXH_NO_LONG_LONG /* defined when there is no 64-bit support */ + +/*! + * @internal + * @brief Structure for XXH64 streaming API. + * + * @note This is only defined when @ref XXH_STATIC_LINKING_ONLY, + * @ref XXH_INLINE_ALL, or @ref XXH_IMPLEMENTATION is defined. Otherwise it is + * an opaque type. This allows fields to safely be changed. + * + * Typedef'd to @ref XXH64_state_t. + * Do not access the members of this struct directly. + * @see XXH32_state_s, XXH3_state_s + */ +struct XXH64_state_s { + XXH64_hash_t total_len; /*!< Total length hashed. This is always 64-bit. */ + XXH64_hash_t v[4]; /*!< Accumulator lanes */ + XXH64_hash_t mem64[4]; /*!< Internal buffer for partial reads. Treated as unsigned char[32]. */ + XXH32_hash_t memsize; /*!< Amount of data in @ref mem64 */ + XXH32_hash_t reserved32; /*!< Reserved field, needed for padding anyways*/ + XXH64_hash_t reserved64; /*!< Reserved field. Do not read or write to it, it may be removed. */ +}; /* typedef'd to XXH64_state_t */ + +#if defined(__STDC_VERSION__) && (__STDC_VERSION__ >= 201112L) /* >= C11 */ +# include +# define XXH_ALIGN(n) alignas(n) +#elif defined(__cplusplus) && (__cplusplus >= 201103L) /* >= C++11 */ +/* In C++ alignas() is a keyword */ +# define XXH_ALIGN(n) alignas(n) +#elif defined(__GNUC__) +# define XXH_ALIGN(n) __attribute__ ((aligned(n))) +#elif defined(_MSC_VER) +# define XXH_ALIGN(n) __declspec(align(n)) +#else +# define XXH_ALIGN(n) /* disabled */ +#endif + +/* Old GCC versions only accept the attribute after the type in structures. */ +#if !(defined(__STDC_VERSION__) && (__STDC_VERSION__ >= 201112L)) /* C11+ */ \ + && ! (defined(__cplusplus) && (__cplusplus >= 201103L)) /* >= C++11 */ \ + && defined(__GNUC__) +# define XXH_ALIGN_MEMBER(align, type) type XXH_ALIGN(align) +#else +# define XXH_ALIGN_MEMBER(align, type) XXH_ALIGN(align) type +#endif + +/*! + * @brief The size of the internal XXH3 buffer. + * + * This is the optimal update size for incremental hashing. + * + * @see XXH3_64b_update(), XXH3_128b_update(). + */ +#define XXH3_INTERNALBUFFER_SIZE 256 + +/*! + * @brief Default size of the secret buffer (and @ref XXH3_kSecret). + * + * This is the size used in @ref XXH3_kSecret and the seeded functions. + * + * Not to be confused with @ref XXH3_SECRET_SIZE_MIN. + */ +#define XXH3_SECRET_DEFAULT_SIZE 192 + +/*! + * @internal + * @brief Structure for XXH3 streaming API. + * + * @note This is only defined when @ref XXH_STATIC_LINKING_ONLY, + * @ref XXH_INLINE_ALL, or @ref XXH_IMPLEMENTATION is defined. + * Otherwise it is an opaque type. + * Never use this definition in combination with dynamic library. + * This allows fields to safely be changed in the future. + * + * @note ** This structure has a strict alignment requirement of 64 bytes!! ** + * Do not allocate this with `malloc()` or `new`, + * it will not be sufficiently aligned. + * Use @ref XXH3_createState() and @ref XXH3_freeState(), or stack allocation. + * + * Typedef'd to @ref XXH3_state_t. + * Do never access the members of this struct directly. + * + * @see XXH3_INITSTATE() for stack initialization. + * @see XXH3_createState(), XXH3_freeState(). + * @see XXH32_state_s, XXH64_state_s + */ +struct XXH3_state_s { + XXH_ALIGN_MEMBER(64, XXH64_hash_t acc[8]); + /*!< The 8 accumulators. Similar to `vN` in @ref XXH32_state_s::v1 and @ref XXH64_state_s */ + XXH_ALIGN_MEMBER(64, unsigned char customSecret[XXH3_SECRET_DEFAULT_SIZE]); + /*!< Used to store a custom secret generated from a seed. */ + XXH_ALIGN_MEMBER(64, unsigned char buffer[XXH3_INTERNALBUFFER_SIZE]); + /*!< The internal buffer. @see XXH32_state_s::mem32 */ + XXH32_hash_t bufferedSize; + /*!< The amount of memory in @ref buffer, @see XXH32_state_s::memsize */ + XXH32_hash_t useSeed; + /*!< Reserved field. Needed for padding on 64-bit. */ + size_t nbStripesSoFar; + /*!< Number or stripes processed. */ + XXH64_hash_t totalLen; + /*!< Total length hashed. 64-bit even on 32-bit targets. */ + size_t nbStripesPerBlock; + /*!< Number of stripes per block. */ + size_t secretLimit; + /*!< Size of @ref customSecret or @ref extSecret */ + XXH64_hash_t seed; + /*!< Seed for _withSeed variants. Must be zero otherwise, @see XXH3_INITSTATE() */ + XXH64_hash_t reserved64; + /*!< Reserved field. */ + const unsigned char* extSecret; + /*!< Reference to an external secret for the _withSecret variants, NULL + * for other variants. */ + /* note: there may be some padding at the end due to alignment on 64 bytes */ +}; /* typedef'd to XXH3_state_t */ + +#undef XXH_ALIGN_MEMBER + +/*! + * @brief Initializes a stack-allocated `XXH3_state_s`. + * + * When the @ref XXH3_state_t structure is merely emplaced on stack, + * it should be initialized with XXH3_INITSTATE() or a memset() + * in case its first reset uses XXH3_NNbits_reset_withSeed(). + * This init can be omitted if the first reset uses default or _withSecret mode. + * This operation isn't necessary when the state is created with XXH3_createState(). + * Note that this doesn't prepare the state for a streaming operation, + * it's still necessary to use XXH3_NNbits_reset*() afterwards. + */ +#define XXH3_INITSTATE(XXH3_state_ptr) { (XXH3_state_ptr)->seed = 0; } + + +/* XXH128() : + * simple alias to pre-selected XXH3_128bits variant + */ +XXH_PUBLIC_API XXH128_hash_t XXH128(const void* data, size_t len, XXH64_hash_t seed); + + +/* === Experimental API === */ +/* Symbols defined below must be considered tied to a specific library version. */ + +/* + * XXH3_generateSecret(): + * + * Derive a high-entropy secret from any user-defined content, named customSeed. + * The generated secret can be used in combination with `*_withSecret()` functions. + * The `_withSecret()` variants are useful to provide a higher level of protection than 64-bit seed, + * as it becomes much more difficult for an external actor to guess how to impact the calculation logic. + * + * The function accepts as input a custom seed of any length and any content, + * and derives from it a high-entropy secret of length @secretSize + * into an already allocated buffer @secretBuffer. + * @secretSize must be >= XXH3_SECRET_SIZE_MIN + * + * The generated secret can then be used with any `*_withSecret()` variant. + * Functions `XXH3_128bits_withSecret()`, `XXH3_64bits_withSecret()`, + * `XXH3_128bits_reset_withSecret()` and `XXH3_64bits_reset_withSecret()` + * are part of this list. They all accept a `secret` parameter + * which must be large enough for implementation reasons (>= XXH3_SECRET_SIZE_MIN) + * _and_ feature very high entropy (consist of random-looking bytes). + * These conditions can be a high bar to meet, so + * XXH3_generateSecret() can be employed to ensure proper quality. + * + * customSeed can be anything. It can have any size, even small ones, + * and its content can be anything, even "poor entropy" sources such as a bunch of zeroes. + * The resulting `secret` will nonetheless provide all required qualities. + * + * When customSeedSize > 0, supplying NULL as customSeed is undefined behavior. + */ +XXH_PUBLIC_API XXH_errorcode XXH3_generateSecret(void* secretBuffer, size_t secretSize, const void* customSeed, size_t customSeedSize); + + +/* + * XXH3_generateSecret_fromSeed(): + * + * Generate the same secret as the _withSeed() variants. + * + * The resulting secret has a length of XXH3_SECRET_DEFAULT_SIZE (necessarily). + * @secretBuffer must be already allocated, of size at least XXH3_SECRET_DEFAULT_SIZE bytes. + * + * The generated secret can be used in combination with + *`*_withSecret()` and `_withSecretandSeed()` variants. + * This generator is notably useful in combination with `_withSecretandSeed()`, + * as a way to emulate a faster `_withSeed()` variant. + */ +XXH_PUBLIC_API void XXH3_generateSecret_fromSeed(void* secretBuffer, XXH64_hash_t seed); + +/* + * *_withSecretandSeed() : + * These variants generate hash values using either + * @seed for "short" keys (< XXH3_MIDSIZE_MAX = 240 bytes) + * or @secret for "large" keys (>= XXH3_MIDSIZE_MAX). + * + * This generally benefits speed, compared to `_withSeed()` or `_withSecret()`. + * `_withSeed()` has to generate the secret on the fly for "large" keys. + * It's fast, but can be perceptible for "not so large" keys (< 1 KB). + * `_withSecret()` has to generate the masks on the fly for "small" keys, + * which requires more instructions than _withSeed() variants. + * Therefore, _withSecretandSeed variant combines the best of both worlds. + * + * When @secret has been generated by XXH3_generateSecret_fromSeed(), + * this variant produces *exactly* the same results as `_withSeed()` variant, + * hence offering only a pure speed benefit on "large" input, + * by skipping the need to regenerate the secret for every large input. + * + * Another usage scenario is to hash the secret to a 64-bit hash value, + * for example with XXH3_64bits(), which then becomes the seed, + * and then employ both the seed and the secret in _withSecretandSeed(). + * On top of speed, an added benefit is that each bit in the secret + * has a 50% chance to swap each bit in the output, + * via its impact to the seed. + * This is not guaranteed when using the secret directly in "small data" scenarios, + * because only portions of the secret are employed for small data. + */ +XXH_PUBLIC_API XXH64_hash_t +XXH3_64bits_withSecretandSeed(const void* data, size_t len, + const void* secret, size_t secretSize, + XXH64_hash_t seed); + +XXH_PUBLIC_API XXH128_hash_t +XXH3_128bits_withSecretandSeed(const void* data, size_t len, + const void* secret, size_t secretSize, + XXH64_hash_t seed64); + +XXH_PUBLIC_API XXH_errorcode +XXH3_64bits_reset_withSecretandSeed(XXH3_state_t* statePtr, + const void* secret, size_t secretSize, + XXH64_hash_t seed64); + +XXH_PUBLIC_API XXH_errorcode +XXH3_128bits_reset_withSecretandSeed(XXH3_state_t* statePtr, + const void* secret, size_t secretSize, + XXH64_hash_t seed64); + + +#endif /* XXH_NO_LONG_LONG */ +#if defined(XXH_INLINE_ALL) || defined(XXH_PRIVATE_API) +# define XXH_IMPLEMENTATION +#endif + +#endif /* defined(XXH_STATIC_LINKING_ONLY) && !defined(XXHASH_H_STATIC_13879238742) */ + + +/* ======================================================================== */ +/* ======================================================================== */ +/* ======================================================================== */ + + +/*-********************************************************************** + * xxHash implementation + *-********************************************************************** + * xxHash's implementation used to be hosted inside xxhash.c. + * + * However, inlining requires implementation to be visible to the compiler, + * hence be included alongside the header. + * Previously, implementation was hosted inside xxhash.c, + * which was then #included when inlining was activated. + * This construction created issues with a few build and install systems, + * as it required xxhash.c to be stored in /include directory. + * + * xxHash implementation is now directly integrated within xxhash.h. + * As a consequence, xxhash.c is no longer needed in /include. + * + * xxhash.c is still available and is still useful. + * In a "normal" setup, when xxhash is not inlined, + * xxhash.h only exposes the prototypes and public symbols, + * while xxhash.c can be built into an object file xxhash.o + * which can then be linked into the final binary. + ************************************************************************/ + +#if ( defined(XXH_INLINE_ALL) || defined(XXH_PRIVATE_API) \ + || defined(XXH_IMPLEMENTATION) ) && !defined(XXH_IMPLEM_13a8737387) +# define XXH_IMPLEM_13a8737387 + +/* ************************************* +* Tuning parameters +***************************************/ + +/*! + * @defgroup tuning Tuning parameters + * @{ + * + * Various macros to control xxHash's behavior. + */ +#ifdef XXH_DOXYGEN +/*! + * @brief Define this to disable 64-bit code. + * + * Useful if only using the @ref xxh32_family and you have a strict C90 compiler. + */ +# define XXH_NO_LONG_LONG +# undef XXH_NO_LONG_LONG /* don't actually */ +/*! + * @brief Controls how unaligned memory is accessed. + * + * By default, access to unaligned memory is controlled by `memcpy()`, which is + * safe and portable. + * + * Unfortunately, on some target/compiler combinations, the generated assembly + * is sub-optimal. + * + * The below switch allow selection of a different access method + * in the search for improved performance. + * + * @par Possible options: + * + * - `XXH_FORCE_MEMORY_ACCESS=0` (default): `memcpy` + * @par + * Use `memcpy()`. Safe and portable. Note that most modern compilers will + * eliminate the function call and treat it as an unaligned access. + * + * - `XXH_FORCE_MEMORY_ACCESS=1`: `__attribute__((packed))` + * @par + * Depends on compiler extensions and is therefore not portable. + * This method is safe _if_ your compiler supports it, + * and *generally* as fast or faster than `memcpy`. + * + * - `XXH_FORCE_MEMORY_ACCESS=2`: Direct cast + * @par + * Casts directly and dereferences. This method doesn't depend on the + * compiler, but it violates the C standard as it directly dereferences an + * unaligned pointer. It can generate buggy code on targets which do not + * support unaligned memory accesses, but in some circumstances, it's the + * only known way to get the most performance. + * + * - `XXH_FORCE_MEMORY_ACCESS=3`: Byteshift + * @par + * Also portable. This can generate the best code on old compilers which don't + * inline small `memcpy()` calls, and it might also be faster on big-endian + * systems which lack a native byteswap instruction. However, some compilers + * will emit literal byteshifts even if the target supports unaligned access. + * . + * + * @warning + * Methods 1 and 2 rely on implementation-defined behavior. Use these with + * care, as what works on one compiler/platform/optimization level may cause + * another to read garbage data or even crash. + * + * See http://fastcompression.blogspot.com/2015/08/accessing-unaligned-memory.html for details. + * + * Prefer these methods in priority order (0 > 3 > 1 > 2) + */ +# define XXH_FORCE_MEMORY_ACCESS 0 + +/*! + * @def XXH_FORCE_ALIGN_CHECK + * @brief If defined to non-zero, adds a special path for aligned inputs (XXH32() + * and XXH64() only). + * + * This is an important performance trick for architectures without decent + * unaligned memory access performance. + * + * It checks for input alignment, and when conditions are met, uses a "fast + * path" employing direct 32-bit/64-bit reads, resulting in _dramatically + * faster_ read speed. + * + * The check costs one initial branch per hash, which is generally negligible, + * but not zero. + * + * Moreover, it's not useful to generate an additional code path if memory + * access uses the same instruction for both aligned and unaligned + * addresses (e.g. x86 and aarch64). + * + * In these cases, the alignment check can be removed by setting this macro to 0. + * Then the code will always use unaligned memory access. + * Align check is automatically disabled on x86, x64 & arm64, + * which are platforms known to offer good unaligned memory accesses performance. + * + * This option does not affect XXH3 (only XXH32 and XXH64). + */ +# define XXH_FORCE_ALIGN_CHECK 0 + +/*! + * @def XXH_NO_INLINE_HINTS + * @brief When non-zero, sets all functions to `static`. + * + * By default, xxHash tries to force the compiler to inline almost all internal + * functions. + * + * This can usually improve performance due to reduced jumping and improved + * constant folding, but significantly increases the size of the binary which + * might not be favorable. + * + * Additionally, sometimes the forced inlining can be detrimental to performance, + * depending on the architecture. + * + * XXH_NO_INLINE_HINTS marks all internal functions as static, giving the + * compiler full control on whether to inline or not. + * + * When not optimizing (-O0), optimizing for size (-Os, -Oz), or using + * -fno-inline with GCC or Clang, this will automatically be defined. + */ +# define XXH_NO_INLINE_HINTS 0 + +/*! + * @def XXH32_ENDJMP + * @brief Whether to use a jump for `XXH32_finalize`. + * + * For performance, `XXH32_finalize` uses multiple branches in the finalizer. + * This is generally preferable for performance, + * but depending on exact architecture, a jmp may be preferable. + * + * This setting is only possibly making a difference for very small inputs. + */ +# define XXH32_ENDJMP 0 + +/*! + * @internal + * @brief Redefines old internal names. + * + * For compatibility with code that uses xxHash's internals before the names + * were changed to improve namespacing. There is no other reason to use this. + */ +# define XXH_OLD_NAMES +# undef XXH_OLD_NAMES /* don't actually use, it is ugly. */ +#endif /* XXH_DOXYGEN */ +/*! + * @} + */ + +#ifndef XXH_FORCE_MEMORY_ACCESS /* can be defined externally, on command line for example */ + /* prefer __packed__ structures (method 1) for gcc on armv7+ and mips */ +# if !defined(__clang__) && \ +( \ + (defined(__INTEL_COMPILER) && !defined(_WIN32)) || \ + ( \ + defined(__GNUC__) && ( \ + (defined(__ARM_ARCH) && __ARM_ARCH >= 7) || \ + ( \ + defined(__mips__) && \ + (__mips <= 5 || __mips_isa_rev < 6) && \ + (!defined(__mips16) || defined(__mips_mips16e2)) \ + ) \ + ) \ + ) \ +) +# define XXH_FORCE_MEMORY_ACCESS 1 +# endif +#endif + +#ifndef XXH_FORCE_ALIGN_CHECK /* can be defined externally */ +# if defined(__i386) || defined(__x86_64__) || defined(__aarch64__) \ + || defined(_M_IX86) || defined(_M_X64) || defined(_M_ARM64) /* visual */ +# define XXH_FORCE_ALIGN_CHECK 0 +# else +# define XXH_FORCE_ALIGN_CHECK 1 +# endif +#endif + +#ifndef XXH_NO_INLINE_HINTS +# if defined(__OPTIMIZE_SIZE__) /* -Os, -Oz */ \ + || defined(__NO_INLINE__) /* -O0, -fno-inline */ +# define XXH_NO_INLINE_HINTS 1 +# else +# define XXH_NO_INLINE_HINTS 0 +# endif +#endif + +#ifndef XXH32_ENDJMP +/* generally preferable for performance */ +# define XXH32_ENDJMP 0 +#endif + +/*! + * @defgroup impl Implementation + * @{ + */ + + +/* ************************************* +* Includes & Memory related functions +***************************************/ +/* + * Modify the local functions below should you wish to use + * different memory routines for malloc() and free() + */ +#include + +/*! + * @internal + * @brief Modify this function to use a different routine than malloc(). + */ +static void* XXH_malloc(size_t s) { return malloc(s); } + +/*! + * @internal + * @brief Modify this function to use a different routine than free(). + */ +static void XXH_free(void* p) { free(p); } + +#include + +/*! + * @internal + * @brief Modify this function to use a different routine than memcpy(). + */ +static void* XXH_memcpy(void* dest, const void* src, size_t size) +{ + return memcpy(dest,src,size); +} + +#include /* ULLONG_MAX */ + + +/* ************************************* +* Compiler Specific Options +***************************************/ +#ifdef _MSC_VER /* Visual Studio warning fix */ +# pragma warning(disable : 4127) /* disable: C4127: conditional expression is constant */ +#endif + +#if XXH_NO_INLINE_HINTS /* disable inlining hints */ +# if defined(__GNUC__) || defined(__clang__) +# define XXH_FORCE_INLINE static __attribute__((unused)) +# else +# define XXH_FORCE_INLINE static +# endif +# define XXH_NO_INLINE static +/* enable inlining hints */ +#elif defined(__GNUC__) || defined(__clang__) +# define XXH_FORCE_INLINE static __inline__ __attribute__((always_inline, unused)) +# define XXH_NO_INLINE static __attribute__((noinline)) +#elif defined(_MSC_VER) /* Visual Studio */ +# define XXH_FORCE_INLINE static __forceinline +# define XXH_NO_INLINE static __declspec(noinline) +#elif defined (__cplusplus) \ + || (defined (__STDC_VERSION__) && (__STDC_VERSION__ >= 199901L)) /* C99 */ +# define XXH_FORCE_INLINE static inline +# define XXH_NO_INLINE static +#else +# define XXH_FORCE_INLINE static +# define XXH_NO_INLINE static +#endif + + + +/* ************************************* +* Debug +***************************************/ +/*! + * @ingroup tuning + * @def XXH_DEBUGLEVEL + * @brief Sets the debugging level. + * + * XXH_DEBUGLEVEL is expected to be defined externally, typically via the + * compiler's command line options. The value must be a number. + */ +#ifndef XXH_DEBUGLEVEL +# ifdef DEBUGLEVEL /* backwards compat */ +# define XXH_DEBUGLEVEL DEBUGLEVEL +# else +# define XXH_DEBUGLEVEL 0 +# endif +#endif + +#if (XXH_DEBUGLEVEL>=1) +# include /* note: can still be disabled with NDEBUG */ +# define XXH_ASSERT(c) assert(c) +#else +# define XXH_ASSERT(c) ((void)0) +#endif + +/* note: use after variable declarations */ +#ifndef XXH_STATIC_ASSERT +# if defined(__STDC_VERSION__) && (__STDC_VERSION__ >= 201112L) /* C11 */ +# include +# define XXH_STATIC_ASSERT_WITH_MESSAGE(c,m) do { static_assert((c),m); } while(0) +# elif defined(__cplusplus) && (__cplusplus >= 201103L) /* C++11 */ +# define XXH_STATIC_ASSERT_WITH_MESSAGE(c,m) do { static_assert((c),m); } while(0) +# else +# define XXH_STATIC_ASSERT_WITH_MESSAGE(c,m) do { struct xxh_sa { char x[(c) ? 1 : -1]; }; } while(0) +# endif +# define XXH_STATIC_ASSERT(c) XXH_STATIC_ASSERT_WITH_MESSAGE((c),#c) +#endif + +/*! + * @internal + * @def XXH_COMPILER_GUARD(var) + * @brief Used to prevent unwanted optimizations for @p var. + * + * It uses an empty GCC inline assembly statement with a register constraint + * which forces @p var into a general purpose register (eg eax, ebx, ecx + * on x86) and marks it as modified. + * + * This is used in a few places to avoid unwanted autovectorization (e.g. + * XXH32_round()). All vectorization we want is explicit via intrinsics, + * and _usually_ isn't wanted elsewhere. + * + * We also use it to prevent unwanted constant folding for AArch64 in + * XXH3_initCustomSecret_scalar(). + */ +#if defined(__GNUC__) || defined(__clang__) +# define XXH_COMPILER_GUARD(var) __asm__ __volatile__("" : "+r" (var)) +#else +# define XXH_COMPILER_GUARD(var) ((void)0) +#endif + +/* ************************************* +* Basic Types +***************************************/ +#if !defined (__VMS) \ + && (defined (__cplusplus) \ + || (defined (__STDC_VERSION__) && (__STDC_VERSION__ >= 199901L) /* C99 */) ) +# include + typedef uint8_t xxh_u8; +#else + typedef unsigned char xxh_u8; +#endif +typedef XXH32_hash_t xxh_u32; + +#ifdef XXH_OLD_NAMES +# define BYTE xxh_u8 +# define U8 xxh_u8 +# define U32 xxh_u32 +#endif + +/* *** Memory access *** */ + +/*! + * @internal + * @fn xxh_u32 XXH_read32(const void* ptr) + * @brief Reads an unaligned 32-bit integer from @p ptr in native endianness. + * + * Affected by @ref XXH_FORCE_MEMORY_ACCESS. + * + * @param ptr The pointer to read from. + * @return The 32-bit native endian integer from the bytes at @p ptr. + */ + +/*! + * @internal + * @fn xxh_u32 XXH_readLE32(const void* ptr) + * @brief Reads an unaligned 32-bit little endian integer from @p ptr. + * + * Affected by @ref XXH_FORCE_MEMORY_ACCESS. + * + * @param ptr The pointer to read from. + * @return The 32-bit little endian integer from the bytes at @p ptr. + */ + +/*! + * @internal + * @fn xxh_u32 XXH_readBE32(const void* ptr) + * @brief Reads an unaligned 32-bit big endian integer from @p ptr. + * + * Affected by @ref XXH_FORCE_MEMORY_ACCESS. + * + * @param ptr The pointer to read from. + * @return The 32-bit big endian integer from the bytes at @p ptr. + */ + +/*! + * @internal + * @fn xxh_u32 XXH_readLE32_align(const void* ptr, XXH_alignment align) + * @brief Like @ref XXH_readLE32(), but has an option for aligned reads. + * + * Affected by @ref XXH_FORCE_MEMORY_ACCESS. + * Note that when @ref XXH_FORCE_ALIGN_CHECK == 0, the @p align parameter is + * always @ref XXH_alignment::XXH_unaligned. + * + * @param ptr The pointer to read from. + * @param align Whether @p ptr is aligned. + * @pre + * If @p align == @ref XXH_alignment::XXH_aligned, @p ptr must be 4 byte + * aligned. + * @return The 32-bit little endian integer from the bytes at @p ptr. + */ + +#if (defined(XXH_FORCE_MEMORY_ACCESS) && (XXH_FORCE_MEMORY_ACCESS==3)) +/* + * Manual byteshift. Best for old compilers which don't inline memcpy. + * We actually directly use XXH_readLE32 and XXH_readBE32. + */ +#elif (defined(XXH_FORCE_MEMORY_ACCESS) && (XXH_FORCE_MEMORY_ACCESS==2)) + +/* + * Force direct memory access. Only works on CPU which support unaligned memory + * access in hardware. + */ +static xxh_u32 XXH_read32(const void* memPtr) { return *(const xxh_u32*) memPtr; } + +#elif (defined(XXH_FORCE_MEMORY_ACCESS) && (XXH_FORCE_MEMORY_ACCESS==1)) + +/* + * __pack instructions are safer but compiler specific, hence potentially + * problematic for some compilers. + * + * Currently only defined for GCC and ICC. + */ +#ifdef XXH_OLD_NAMES +typedef union { xxh_u32 u32; } __attribute__((packed)) unalign; +#endif +static xxh_u32 XXH_read32(const void* ptr) +{ + typedef union { xxh_u32 u32; } __attribute__((packed)) xxh_unalign; + return ((const xxh_unalign*)ptr)->u32; +} + +#else + +/* + * Portable and safe solution. Generally efficient. + * see: http://fastcompression.blogspot.com/2015/08/accessing-unaligned-memory.html + */ +static xxh_u32 XXH_read32(const void* memPtr) +{ + xxh_u32 val; + XXH_memcpy(&val, memPtr, sizeof(val)); + return val; +} + +#endif /* XXH_FORCE_DIRECT_MEMORY_ACCESS */ + + +/* *** Endianness *** */ + +/*! + * @ingroup tuning + * @def XXH_CPU_LITTLE_ENDIAN + * @brief Whether the target is little endian. + * + * Defined to 1 if the target is little endian, or 0 if it is big endian. + * It can be defined externally, for example on the compiler command line. + * + * If it is not defined, + * a runtime check (which is usually constant folded) is used instead. + * + * @note + * This is not necessarily defined to an integer constant. + * + * @see XXH_isLittleEndian() for the runtime check. + */ +#ifndef XXH_CPU_LITTLE_ENDIAN +/* + * Try to detect endianness automatically, to avoid the nonstandard behavior + * in `XXH_isLittleEndian()` + */ +# if defined(_WIN32) /* Windows is always little endian */ \ + || defined(__LITTLE_ENDIAN__) \ + || (defined(__BYTE_ORDER__) && __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__) +# define XXH_CPU_LITTLE_ENDIAN 1 +# elif defined(__BIG_ENDIAN__) \ + || (defined(__BYTE_ORDER__) && __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__) +# define XXH_CPU_LITTLE_ENDIAN 0 +# else +/*! + * @internal + * @brief Runtime check for @ref XXH_CPU_LITTLE_ENDIAN. + * + * Most compilers will constant fold this. + */ +static int XXH_isLittleEndian(void) +{ + /* + * Portable and well-defined behavior. + * Don't use static: it is detrimental to performance. + */ + const union { xxh_u32 u; xxh_u8 c[4]; } one = { 1 }; + return one.c[0]; +} +# define XXH_CPU_LITTLE_ENDIAN XXH_isLittleEndian() +# endif +#endif + + + + +/* **************************************** +* Compiler-specific Functions and Macros +******************************************/ +#define XXH_GCC_VERSION (__GNUC__ * 100 + __GNUC_MINOR__) + +#ifdef __has_builtin +# define XXH_HAS_BUILTIN(x) __has_builtin(x) +#else +# define XXH_HAS_BUILTIN(x) 0 +#endif + +/*! + * @internal + * @def XXH_rotl32(x,r) + * @brief 32-bit rotate left. + * + * @param x The 32-bit integer to be rotated. + * @param r The number of bits to rotate. + * @pre + * @p r > 0 && @p r < 32 + * @note + * @p x and @p r may be evaluated multiple times. + * @return The rotated result. + */ +#if !defined(NO_CLANG_BUILTIN) && XXH_HAS_BUILTIN(__builtin_rotateleft32) \ + && XXH_HAS_BUILTIN(__builtin_rotateleft64) +# define XXH_rotl32 __builtin_rotateleft32 +# define XXH_rotl64 __builtin_rotateleft64 +/* Note: although _rotl exists for minGW (GCC under windows), performance seems poor */ +#elif defined(_MSC_VER) +# define XXH_rotl32(x,r) _rotl(x,r) +# define XXH_rotl64(x,r) _rotl64(x,r) +#else +# define XXH_rotl32(x,r) (((x) << (r)) | ((x) >> (32 - (r)))) +# define XXH_rotl64(x,r) (((x) << (r)) | ((x) >> (64 - (r)))) +#endif + +/*! + * @internal + * @fn xxh_u32 XXH_swap32(xxh_u32 x) + * @brief A 32-bit byteswap. + * + * @param x The 32-bit integer to byteswap. + * @return @p x, byteswapped. + */ +#if defined(_MSC_VER) /* Visual Studio */ +# define XXH_swap32 _byteswap_ulong +#elif XXH_GCC_VERSION >= 403 +# define XXH_swap32 __builtin_bswap32 +#else +static xxh_u32 XXH_swap32 (xxh_u32 x) +{ + return ((x << 24) & 0xff000000 ) | + ((x << 8) & 0x00ff0000 ) | + ((x >> 8) & 0x0000ff00 ) | + ((x >> 24) & 0x000000ff ); +} +#endif + + +/* *************************** +* Memory reads +*****************************/ + +/*! + * @internal + * @brief Enum to indicate whether a pointer is aligned. + */ +typedef enum { + XXH_aligned, /*!< Aligned */ + XXH_unaligned /*!< Possibly unaligned */ +} XXH_alignment; + +/* + * XXH_FORCE_MEMORY_ACCESS==3 is an endian-independent byteshift load. + * + * This is ideal for older compilers which don't inline memcpy. + */ +#if (defined(XXH_FORCE_MEMORY_ACCESS) && (XXH_FORCE_MEMORY_ACCESS==3)) + +XXH_FORCE_INLINE xxh_u32 XXH_readLE32(const void* memPtr) +{ + const xxh_u8* bytePtr = (const xxh_u8 *)memPtr; + return bytePtr[0] + | ((xxh_u32)bytePtr[1] << 8) + | ((xxh_u32)bytePtr[2] << 16) + | ((xxh_u32)bytePtr[3] << 24); +} + +XXH_FORCE_INLINE xxh_u32 XXH_readBE32(const void* memPtr) +{ + const xxh_u8* bytePtr = (const xxh_u8 *)memPtr; + return bytePtr[3] + | ((xxh_u32)bytePtr[2] << 8) + | ((xxh_u32)bytePtr[1] << 16) + | ((xxh_u32)bytePtr[0] << 24); +} + +#else +XXH_FORCE_INLINE xxh_u32 XXH_readLE32(const void* ptr) +{ + return XXH_CPU_LITTLE_ENDIAN ? XXH_read32(ptr) : XXH_swap32(XXH_read32(ptr)); +} + +static xxh_u32 XXH_readBE32(const void* ptr) +{ + return XXH_CPU_LITTLE_ENDIAN ? XXH_swap32(XXH_read32(ptr)) : XXH_read32(ptr); +} +#endif + +XXH_FORCE_INLINE xxh_u32 +XXH_readLE32_align(const void* ptr, XXH_alignment align) +{ + if (align==XXH_unaligned) { + return XXH_readLE32(ptr); + } else { + return XXH_CPU_LITTLE_ENDIAN ? *(const xxh_u32*)ptr : XXH_swap32(*(const xxh_u32*)ptr); + } +} + + +/* ************************************* +* Misc +***************************************/ +/*! @ingroup public */ +XXH_PUBLIC_API unsigned XXH_versionNumber (void) { return XXH_VERSION_NUMBER; } + + +/* ******************************************************************* +* 32-bit hash functions +*********************************************************************/ +/*! + * @} + * @defgroup xxh32_impl XXH32 implementation + * @ingroup impl + * @{ + */ + /* #define instead of static const, to be used as initializers */ +#define XXH_PRIME32_1 0x9E3779B1U /*!< 0b10011110001101110111100110110001 */ +#define XXH_PRIME32_2 0x85EBCA77U /*!< 0b10000101111010111100101001110111 */ +#define XXH_PRIME32_3 0xC2B2AE3DU /*!< 0b11000010101100101010111000111101 */ +#define XXH_PRIME32_4 0x27D4EB2FU /*!< 0b00100111110101001110101100101111 */ +#define XXH_PRIME32_5 0x165667B1U /*!< 0b00010110010101100110011110110001 */ + +#ifdef XXH_OLD_NAMES +# define PRIME32_1 XXH_PRIME32_1 +# define PRIME32_2 XXH_PRIME32_2 +# define PRIME32_3 XXH_PRIME32_3 +# define PRIME32_4 XXH_PRIME32_4 +# define PRIME32_5 XXH_PRIME32_5 +#endif + +/*! + * @internal + * @brief Normal stripe processing routine. + * + * This shuffles the bits so that any bit from @p input impacts several bits in + * @p acc. + * + * @param acc The accumulator lane. + * @param input The stripe of input to mix. + * @return The mixed accumulator lane. + */ +static xxh_u32 XXH32_round(xxh_u32 acc, xxh_u32 input) +{ + acc += input * XXH_PRIME32_2; + acc = XXH_rotl32(acc, 13); + acc *= XXH_PRIME32_1; +#if (defined(__SSE4_1__) || defined(__aarch64__)) && !defined(XXH_ENABLE_AUTOVECTORIZE) + /* + * UGLY HACK: + * A compiler fence is the only thing that prevents GCC and Clang from + * autovectorizing the XXH32 loop (pragmas and attributes don't work for some + * reason) without globally disabling SSE4.1. + * + * The reason we want to avoid vectorization is because despite working on + * 4 integers at a time, there are multiple factors slowing XXH32 down on + * SSE4: + * - There's a ridiculous amount of lag from pmulld (10 cycles of latency on + * newer chips!) making it slightly slower to multiply four integers at + * once compared to four integers independently. Even when pmulld was + * fastest, Sandy/Ivy Bridge, it is still not worth it to go into SSE + * just to multiply unless doing a long operation. + * + * - Four instructions are required to rotate, + * movqda tmp, v // not required with VEX encoding + * pslld tmp, 13 // tmp <<= 13 + * psrld v, 19 // x >>= 19 + * por v, tmp // x |= tmp + * compared to one for scalar: + * roll v, 13 // reliably fast across the board + * shldl v, v, 13 // Sandy Bridge and later prefer this for some reason + * + * - Instruction level parallelism is actually more beneficial here because + * the SIMD actually serializes this operation: While v1 is rotating, v2 + * can load data, while v3 can multiply. SSE forces them to operate + * together. + * + * This is also enabled on AArch64, as Clang autovectorizes it incorrectly + * and it is pointless writing a NEON implementation that is basically the + * same speed as scalar for XXH32. + */ + XXH_COMPILER_GUARD(acc); +#endif + return acc; +} + +/*! + * @internal + * @brief Mixes all bits to finalize the hash. + * + * The final mix ensures that all input bits have a chance to impact any bit in + * the output digest, resulting in an unbiased distribution. + * + * @param h32 The hash to avalanche. + * @return The avalanched hash. + */ +static xxh_u32 XXH32_avalanche(xxh_u32 h32) +{ + h32 ^= h32 >> 15; + h32 *= XXH_PRIME32_2; + h32 ^= h32 >> 13; + h32 *= XXH_PRIME32_3; + h32 ^= h32 >> 16; + return(h32); +} + +#define XXH_get32bits(p) XXH_readLE32_align(p, align) + +/*! + * @internal + * @brief Processes the last 0-15 bytes of @p ptr. + * + * There may be up to 15 bytes remaining to consume from the input. + * This final stage will digest them to ensure that all input bytes are present + * in the final mix. + * + * @param h32 The hash to finalize. + * @param ptr The pointer to the remaining input. + * @param len The remaining length, modulo 16. + * @param align Whether @p ptr is aligned. + * @return The finalized hash. + */ +static xxh_u32 +XXH32_finalize(xxh_u32 h32, const xxh_u8* ptr, size_t len, XXH_alignment align) +{ +#define XXH_PROCESS1 do { \ + h32 += (*ptr++) * XXH_PRIME32_5; \ + h32 = XXH_rotl32(h32, 11) * XXH_PRIME32_1; \ +} while (0) + +#define XXH_PROCESS4 do { \ + h32 += XXH_get32bits(ptr) * XXH_PRIME32_3; \ + ptr += 4; \ + h32 = XXH_rotl32(h32, 17) * XXH_PRIME32_4; \ +} while (0) + + if (ptr==NULL) XXH_ASSERT(len == 0); + + /* Compact rerolled version; generally faster */ + if (!XXH32_ENDJMP) { + len &= 15; + while (len >= 4) { + XXH_PROCESS4; + len -= 4; + } + while (len > 0) { + XXH_PROCESS1; + --len; + } + return XXH32_avalanche(h32); + } else { + switch(len&15) /* or switch(bEnd - p) */ { + case 12: XXH_PROCESS4; + XXH_FALLTHROUGH; + case 8: XXH_PROCESS4; + XXH_FALLTHROUGH; + case 4: XXH_PROCESS4; + return XXH32_avalanche(h32); + + case 13: XXH_PROCESS4; + XXH_FALLTHROUGH; + case 9: XXH_PROCESS4; + XXH_FALLTHROUGH; + case 5: XXH_PROCESS4; + XXH_PROCESS1; + return XXH32_avalanche(h32); + + case 14: XXH_PROCESS4; + XXH_FALLTHROUGH; + case 10: XXH_PROCESS4; + XXH_FALLTHROUGH; + case 6: XXH_PROCESS4; + XXH_PROCESS1; + XXH_PROCESS1; + return XXH32_avalanche(h32); + + case 15: XXH_PROCESS4; + XXH_FALLTHROUGH; + case 11: XXH_PROCESS4; + XXH_FALLTHROUGH; + case 7: XXH_PROCESS4; + XXH_FALLTHROUGH; + case 3: XXH_PROCESS1; + XXH_FALLTHROUGH; + case 2: XXH_PROCESS1; + XXH_FALLTHROUGH; + case 1: XXH_PROCESS1; + XXH_FALLTHROUGH; + case 0: return XXH32_avalanche(h32); + } + XXH_ASSERT(0); + return h32; /* reaching this point is deemed impossible */ + } +} + +#ifdef XXH_OLD_NAMES +# define PROCESS1 XXH_PROCESS1 +# define PROCESS4 XXH_PROCESS4 +#else +# undef XXH_PROCESS1 +# undef XXH_PROCESS4 +#endif + +/*! + * @internal + * @brief The implementation for @ref XXH32(). + * + * @param input , len , seed Directly passed from @ref XXH32(). + * @param align Whether @p input is aligned. + * @return The calculated hash. + */ +XXH_FORCE_INLINE xxh_u32 +XXH32_endian_align(const xxh_u8* input, size_t len, xxh_u32 seed, XXH_alignment align) +{ + xxh_u32 h32; + + if (input==NULL) XXH_ASSERT(len == 0); + + if (len>=16) { + const xxh_u8* const bEnd = input + len; + const xxh_u8* const limit = bEnd - 15; + xxh_u32 v1 = seed + XXH_PRIME32_1 + XXH_PRIME32_2; + xxh_u32 v2 = seed + XXH_PRIME32_2; + xxh_u32 v3 = seed + 0; + xxh_u32 v4 = seed - XXH_PRIME32_1; + + do { + v1 = XXH32_round(v1, XXH_get32bits(input)); input += 4; + v2 = XXH32_round(v2, XXH_get32bits(input)); input += 4; + v3 = XXH32_round(v3, XXH_get32bits(input)); input += 4; + v4 = XXH32_round(v4, XXH_get32bits(input)); input += 4; + } while (input < limit); + + h32 = XXH_rotl32(v1, 1) + XXH_rotl32(v2, 7) + + XXH_rotl32(v3, 12) + XXH_rotl32(v4, 18); + } else { + h32 = seed + XXH_PRIME32_5; + } + + h32 += (xxh_u32)len; + + return XXH32_finalize(h32, input, len&15, align); +} + +/*! @ingroup xxh32_family */ +XXH_PUBLIC_API XXH32_hash_t XXH32 (const void* input, size_t len, XXH32_hash_t seed) +{ +#if 0 + /* Simple version, good for code maintenance, but unfortunately slow for small inputs */ + XXH32_state_t state; + XXH32_reset(&state, seed); + XXH32_update(&state, (const xxh_u8*)input, len); + return XXH32_digest(&state); +#else + if (XXH_FORCE_ALIGN_CHECK) { + if ((((size_t)input) & 3) == 0) { /* Input is 4-bytes aligned, leverage the speed benefit */ + return XXH32_endian_align((const xxh_u8*)input, len, seed, XXH_aligned); + } } + + return XXH32_endian_align((const xxh_u8*)input, len, seed, XXH_unaligned); +#endif +} + + + +/******* Hash streaming *******/ +/*! + * @ingroup xxh32_family + */ +XXH_PUBLIC_API XXH32_state_t* XXH32_createState(void) +{ + return (XXH32_state_t*)XXH_malloc(sizeof(XXH32_state_t)); +} +/*! @ingroup xxh32_family */ +XXH_PUBLIC_API XXH_errorcode XXH32_freeState(XXH32_state_t* statePtr) +{ + XXH_free(statePtr); + return XXH_OK; +} + +/*! @ingroup xxh32_family */ +XXH_PUBLIC_API void XXH32_copyState(XXH32_state_t* dstState, const XXH32_state_t* srcState) +{ + XXH_memcpy(dstState, srcState, sizeof(*dstState)); +} + +/*! @ingroup xxh32_family */ +XXH_PUBLIC_API XXH_errorcode XXH32_reset(XXH32_state_t* statePtr, XXH32_hash_t seed) +{ + XXH32_state_t state; /* using a local state to memcpy() in order to avoid strict-aliasing warnings */ + memset(&state, 0, sizeof(state)); + state.v[0] = seed + XXH_PRIME32_1 + XXH_PRIME32_2; + state.v[1] = seed + XXH_PRIME32_2; + state.v[2] = seed + 0; + state.v[3] = seed - XXH_PRIME32_1; + /* do not write into reserved, planned to be removed in a future version */ + XXH_memcpy(statePtr, &state, sizeof(state) - sizeof(state.reserved)); + return XXH_OK; +} + + +/*! @ingroup xxh32_family */ +XXH_PUBLIC_API XXH_errorcode +XXH32_update(XXH32_state_t* state, const void* input, size_t len) +{ + if (input==NULL) { + XXH_ASSERT(len == 0); + return XXH_OK; + } + + { const xxh_u8* p = (const xxh_u8*)input; + const xxh_u8* const bEnd = p + len; + + state->total_len_32 += (XXH32_hash_t)len; + state->large_len |= (XXH32_hash_t)((len>=16) | (state->total_len_32>=16)); + + if (state->memsize + len < 16) { /* fill in tmp buffer */ + XXH_memcpy((xxh_u8*)(state->mem32) + state->memsize, input, len); + state->memsize += (XXH32_hash_t)len; + return XXH_OK; + } + + if (state->memsize) { /* some data left from previous update */ + XXH_memcpy((xxh_u8*)(state->mem32) + state->memsize, input, 16-state->memsize); + { const xxh_u32* p32 = state->mem32; + state->v[0] = XXH32_round(state->v[0], XXH_readLE32(p32)); p32++; + state->v[1] = XXH32_round(state->v[1], XXH_readLE32(p32)); p32++; + state->v[2] = XXH32_round(state->v[2], XXH_readLE32(p32)); p32++; + state->v[3] = XXH32_round(state->v[3], XXH_readLE32(p32)); + } + p += 16-state->memsize; + state->memsize = 0; + } + + if (p <= bEnd-16) { + const xxh_u8* const limit = bEnd - 16; + + do { + state->v[0] = XXH32_round(state->v[0], XXH_readLE32(p)); p+=4; + state->v[1] = XXH32_round(state->v[1], XXH_readLE32(p)); p+=4; + state->v[2] = XXH32_round(state->v[2], XXH_readLE32(p)); p+=4; + state->v[3] = XXH32_round(state->v[3], XXH_readLE32(p)); p+=4; + } while (p<=limit); + + } + + if (p < bEnd) { + XXH_memcpy(state->mem32, p, (size_t)(bEnd-p)); + state->memsize = (unsigned)(bEnd-p); + } + } + + return XXH_OK; +} + + +/*! @ingroup xxh32_family */ +XXH_PUBLIC_API XXH32_hash_t XXH32_digest(const XXH32_state_t* state) +{ + xxh_u32 h32; + + if (state->large_len) { + h32 = XXH_rotl32(state->v[0], 1) + + XXH_rotl32(state->v[1], 7) + + XXH_rotl32(state->v[2], 12) + + XXH_rotl32(state->v[3], 18); + } else { + h32 = state->v[2] /* == seed */ + XXH_PRIME32_5; + } + + h32 += state->total_len_32; + + return XXH32_finalize(h32, (const xxh_u8*)state->mem32, state->memsize, XXH_aligned); +} + + +/******* Canonical representation *******/ + +/*! + * @ingroup xxh32_family + * The default return values from XXH functions are unsigned 32 and 64 bit + * integers. + * + * The canonical representation uses big endian convention, the same convention + * as human-readable numbers (large digits first). + * + * This way, hash values can be written into a file or buffer, remaining + * comparable across different systems. + * + * The following functions allow transformation of hash values to and from their + * canonical format. + */ +XXH_PUBLIC_API void XXH32_canonicalFromHash(XXH32_canonical_t* dst, XXH32_hash_t hash) +{ + XXH_STATIC_ASSERT(sizeof(XXH32_canonical_t) == sizeof(XXH32_hash_t)); + if (XXH_CPU_LITTLE_ENDIAN) hash = XXH_swap32(hash); + XXH_memcpy(dst, &hash, sizeof(*dst)); +} +/*! @ingroup xxh32_family */ +XXH_PUBLIC_API XXH32_hash_t XXH32_hashFromCanonical(const XXH32_canonical_t* src) +{ + return XXH_readBE32(src); +} + + +#ifndef XXH_NO_LONG_LONG + +/* ******************************************************************* +* 64-bit hash functions +*********************************************************************/ +/*! + * @} + * @ingroup impl + * @{ + */ +/******* Memory access *******/ + +typedef XXH64_hash_t xxh_u64; + +#ifdef XXH_OLD_NAMES +# define U64 xxh_u64 +#endif + +#if (defined(XXH_FORCE_MEMORY_ACCESS) && (XXH_FORCE_MEMORY_ACCESS==3)) +/* + * Manual byteshift. Best for old compilers which don't inline memcpy. + * We actually directly use XXH_readLE64 and XXH_readBE64. + */ +#elif (defined(XXH_FORCE_MEMORY_ACCESS) && (XXH_FORCE_MEMORY_ACCESS==2)) + +/* Force direct memory access. Only works on CPU which support unaligned memory access in hardware */ +static xxh_u64 XXH_read64(const void* memPtr) +{ + return *(const xxh_u64*) memPtr; +} + +#elif (defined(XXH_FORCE_MEMORY_ACCESS) && (XXH_FORCE_MEMORY_ACCESS==1)) + +/* + * __pack instructions are safer, but compiler specific, hence potentially + * problematic for some compilers. + * + * Currently only defined for GCC and ICC. + */ +#ifdef XXH_OLD_NAMES +typedef union { xxh_u32 u32; xxh_u64 u64; } __attribute__((packed)) unalign64; +#endif +static xxh_u64 XXH_read64(const void* ptr) +{ + typedef union { xxh_u32 u32; xxh_u64 u64; } __attribute__((packed)) xxh_unalign64; + return ((const xxh_unalign64*)ptr)->u64; +} + +#else + +/* + * Portable and safe solution. Generally efficient. + * see: http://fastcompression.blogspot.com/2015/08/accessing-unaligned-memory.html + */ +static xxh_u64 XXH_read64(const void* memPtr) +{ + xxh_u64 val; + XXH_memcpy(&val, memPtr, sizeof(val)); + return val; +} + +#endif /* XXH_FORCE_DIRECT_MEMORY_ACCESS */ + +#if defined(_MSC_VER) /* Visual Studio */ +# define XXH_swap64 _byteswap_uint64 +#elif XXH_GCC_VERSION >= 403 +# define XXH_swap64 __builtin_bswap64 +#else +static xxh_u64 XXH_swap64(xxh_u64 x) +{ + return ((x << 56) & 0xff00000000000000ULL) | + ((x << 40) & 0x00ff000000000000ULL) | + ((x << 24) & 0x0000ff0000000000ULL) | + ((x << 8) & 0x000000ff00000000ULL) | + ((x >> 8) & 0x00000000ff000000ULL) | + ((x >> 24) & 0x0000000000ff0000ULL) | + ((x >> 40) & 0x000000000000ff00ULL) | + ((x >> 56) & 0x00000000000000ffULL); +} +#endif + + +/* XXH_FORCE_MEMORY_ACCESS==3 is an endian-independent byteshift load. */ +#if (defined(XXH_FORCE_MEMORY_ACCESS) && (XXH_FORCE_MEMORY_ACCESS==3)) + +XXH_FORCE_INLINE xxh_u64 XXH_readLE64(const void* memPtr) +{ + const xxh_u8* bytePtr = (const xxh_u8 *)memPtr; + return bytePtr[0] + | ((xxh_u64)bytePtr[1] << 8) + | ((xxh_u64)bytePtr[2] << 16) + | ((xxh_u64)bytePtr[3] << 24) + | ((xxh_u64)bytePtr[4] << 32) + | ((xxh_u64)bytePtr[5] << 40) + | ((xxh_u64)bytePtr[6] << 48) + | ((xxh_u64)bytePtr[7] << 56); +} + +XXH_FORCE_INLINE xxh_u64 XXH_readBE64(const void* memPtr) +{ + const xxh_u8* bytePtr = (const xxh_u8 *)memPtr; + return bytePtr[7] + | ((xxh_u64)bytePtr[6] << 8) + | ((xxh_u64)bytePtr[5] << 16) + | ((xxh_u64)bytePtr[4] << 24) + | ((xxh_u64)bytePtr[3] << 32) + | ((xxh_u64)bytePtr[2] << 40) + | ((xxh_u64)bytePtr[1] << 48) + | ((xxh_u64)bytePtr[0] << 56); +} + +#else +XXH_FORCE_INLINE xxh_u64 XXH_readLE64(const void* ptr) +{ + return XXH_CPU_LITTLE_ENDIAN ? XXH_read64(ptr) : XXH_swap64(XXH_read64(ptr)); +} + +static xxh_u64 XXH_readBE64(const void* ptr) +{ + return XXH_CPU_LITTLE_ENDIAN ? XXH_swap64(XXH_read64(ptr)) : XXH_read64(ptr); +} +#endif + +XXH_FORCE_INLINE xxh_u64 +XXH_readLE64_align(const void* ptr, XXH_alignment align) +{ + if (align==XXH_unaligned) + return XXH_readLE64(ptr); + else + return XXH_CPU_LITTLE_ENDIAN ? *(const xxh_u64*)ptr : XXH_swap64(*(const xxh_u64*)ptr); +} + + +/******* xxh64 *******/ +/*! + * @} + * @defgroup xxh64_impl XXH64 implementation + * @ingroup impl + * @{ + */ +/* #define rather that static const, to be used as initializers */ +#define XXH_PRIME64_1 0x9E3779B185EBCA87ULL /*!< 0b1001111000110111011110011011000110000101111010111100101010000111 */ +#define XXH_PRIME64_2 0xC2B2AE3D27D4EB4FULL /*!< 0b1100001010110010101011100011110100100111110101001110101101001111 */ +#define XXH_PRIME64_3 0x165667B19E3779F9ULL /*!< 0b0001011001010110011001111011000110011110001101110111100111111001 */ +#define XXH_PRIME64_4 0x85EBCA77C2B2AE63ULL /*!< 0b1000010111101011110010100111011111000010101100101010111001100011 */ +#define XXH_PRIME64_5 0x27D4EB2F165667C5ULL /*!< 0b0010011111010100111010110010111100010110010101100110011111000101 */ + +#ifdef XXH_OLD_NAMES +# define PRIME64_1 XXH_PRIME64_1 +# define PRIME64_2 XXH_PRIME64_2 +# define PRIME64_3 XXH_PRIME64_3 +# define PRIME64_4 XXH_PRIME64_4 +# define PRIME64_5 XXH_PRIME64_5 +#endif + +static xxh_u64 XXH64_round(xxh_u64 acc, xxh_u64 input) +{ + acc += input * XXH_PRIME64_2; + acc = XXH_rotl64(acc, 31); + acc *= XXH_PRIME64_1; + return acc; +} + +static xxh_u64 XXH64_mergeRound(xxh_u64 acc, xxh_u64 val) +{ + val = XXH64_round(0, val); + acc ^= val; + acc = acc * XXH_PRIME64_1 + XXH_PRIME64_4; + return acc; +} + +static xxh_u64 XXH64_avalanche(xxh_u64 h64) +{ + h64 ^= h64 >> 33; + h64 *= XXH_PRIME64_2; + h64 ^= h64 >> 29; + h64 *= XXH_PRIME64_3; + h64 ^= h64 >> 32; + return h64; +} + + +#define XXH_get64bits(p) XXH_readLE64_align(p, align) + +static xxh_u64 +XXH64_finalize(xxh_u64 h64, const xxh_u8* ptr, size_t len, XXH_alignment align) +{ + if (ptr==NULL) XXH_ASSERT(len == 0); + len &= 31; + while (len >= 8) { + xxh_u64 const k1 = XXH64_round(0, XXH_get64bits(ptr)); + ptr += 8; + h64 ^= k1; + h64 = XXH_rotl64(h64,27) * XXH_PRIME64_1 + XXH_PRIME64_4; + len -= 8; + } + if (len >= 4) { + h64 ^= (xxh_u64)(XXH_get32bits(ptr)) * XXH_PRIME64_1; + ptr += 4; + h64 = XXH_rotl64(h64, 23) * XXH_PRIME64_2 + XXH_PRIME64_3; + len -= 4; + } + while (len > 0) { + h64 ^= (*ptr++) * XXH_PRIME64_5; + h64 = XXH_rotl64(h64, 11) * XXH_PRIME64_1; + --len; + } + return XXH64_avalanche(h64); +} + +#ifdef XXH_OLD_NAMES +# define PROCESS1_64 XXH_PROCESS1_64 +# define PROCESS4_64 XXH_PROCESS4_64 +# define PROCESS8_64 XXH_PROCESS8_64 +#else +# undef XXH_PROCESS1_64 +# undef XXH_PROCESS4_64 +# undef XXH_PROCESS8_64 +#endif + +XXH_FORCE_INLINE xxh_u64 +XXH64_endian_align(const xxh_u8* input, size_t len, xxh_u64 seed, XXH_alignment align) +{ + xxh_u64 h64; + if (input==NULL) XXH_ASSERT(len == 0); + + if (len>=32) { + const xxh_u8* const bEnd = input + len; + const xxh_u8* const limit = bEnd - 31; + xxh_u64 v1 = seed + XXH_PRIME64_1 + XXH_PRIME64_2; + xxh_u64 v2 = seed + XXH_PRIME64_2; + xxh_u64 v3 = seed + 0; + xxh_u64 v4 = seed - XXH_PRIME64_1; + + do { + v1 = XXH64_round(v1, XXH_get64bits(input)); input+=8; + v2 = XXH64_round(v2, XXH_get64bits(input)); input+=8; + v3 = XXH64_round(v3, XXH_get64bits(input)); input+=8; + v4 = XXH64_round(v4, XXH_get64bits(input)); input+=8; + } while (inputtotal_len += len; + + if (state->memsize + len < 32) { /* fill in tmp buffer */ + XXH_memcpy(((xxh_u8*)state->mem64) + state->memsize, input, len); + state->memsize += (xxh_u32)len; + return XXH_OK; + } + + if (state->memsize) { /* tmp buffer is full */ + XXH_memcpy(((xxh_u8*)state->mem64) + state->memsize, input, 32-state->memsize); + state->v[0] = XXH64_round(state->v[0], XXH_readLE64(state->mem64+0)); + state->v[1] = XXH64_round(state->v[1], XXH_readLE64(state->mem64+1)); + state->v[2] = XXH64_round(state->v[2], XXH_readLE64(state->mem64+2)); + state->v[3] = XXH64_round(state->v[3], XXH_readLE64(state->mem64+3)); + p += 32 - state->memsize; + state->memsize = 0; + } + + if (p+32 <= bEnd) { + const xxh_u8* const limit = bEnd - 32; + + do { + state->v[0] = XXH64_round(state->v[0], XXH_readLE64(p)); p+=8; + state->v[1] = XXH64_round(state->v[1], XXH_readLE64(p)); p+=8; + state->v[2] = XXH64_round(state->v[2], XXH_readLE64(p)); p+=8; + state->v[3] = XXH64_round(state->v[3], XXH_readLE64(p)); p+=8; + } while (p<=limit); + + } + + if (p < bEnd) { + XXH_memcpy(state->mem64, p, (size_t)(bEnd-p)); + state->memsize = (unsigned)(bEnd-p); + } + } + + return XXH_OK; +} + + +/*! @ingroup xxh64_family */ +XXH_PUBLIC_API XXH64_hash_t XXH64_digest(const XXH64_state_t* state) +{ + xxh_u64 h64; + + if (state->total_len >= 32) { + h64 = XXH_rotl64(state->v[0], 1) + XXH_rotl64(state->v[1], 7) + XXH_rotl64(state->v[2], 12) + XXH_rotl64(state->v[3], 18); + h64 = XXH64_mergeRound(h64, state->v[0]); + h64 = XXH64_mergeRound(h64, state->v[1]); + h64 = XXH64_mergeRound(h64, state->v[2]); + h64 = XXH64_mergeRound(h64, state->v[3]); + } else { + h64 = state->v[2] /*seed*/ + XXH_PRIME64_5; + } + + h64 += (xxh_u64) state->total_len; + + return XXH64_finalize(h64, (const xxh_u8*)state->mem64, (size_t)state->total_len, XXH_aligned); +} + + +/******* Canonical representation *******/ + +/*! @ingroup xxh64_family */ +XXH_PUBLIC_API void XXH64_canonicalFromHash(XXH64_canonical_t* dst, XXH64_hash_t hash) +{ + XXH_STATIC_ASSERT(sizeof(XXH64_canonical_t) == sizeof(XXH64_hash_t)); + if (XXH_CPU_LITTLE_ENDIAN) hash = XXH_swap64(hash); + XXH_memcpy(dst, &hash, sizeof(*dst)); +} + +/*! @ingroup xxh64_family */ +XXH_PUBLIC_API XXH64_hash_t XXH64_hashFromCanonical(const XXH64_canonical_t* src) +{ + return XXH_readBE64(src); +} + +#ifndef XXH_NO_XXH3 + +/* ********************************************************************* +* XXH3 +* New generation hash designed for speed on small keys and vectorization +************************************************************************ */ +/*! + * @} + * @defgroup xxh3_impl XXH3 implementation + * @ingroup impl + * @{ + */ + +/* === Compiler specifics === */ + +#if ((defined(sun) || defined(__sun)) && __cplusplus) /* Solaris includes __STDC_VERSION__ with C++. Tested with GCC 5.5 */ +# define XXH_RESTRICT /* disable */ +#elif defined (__STDC_VERSION__) && __STDC_VERSION__ >= 199901L /* >= C99 */ +# define XXH_RESTRICT restrict +#else +/* Note: it might be useful to define __restrict or __restrict__ for some C++ compilers */ +# define XXH_RESTRICT /* disable */ +#endif + +#if (defined(__GNUC__) && (__GNUC__ >= 3)) \ + || (defined(__INTEL_COMPILER) && (__INTEL_COMPILER >= 800)) \ + || defined(__clang__) +# define XXH_likely(x) __builtin_expect(x, 1) +# define XXH_unlikely(x) __builtin_expect(x, 0) +#else +# define XXH_likely(x) (x) +# define XXH_unlikely(x) (x) +#endif + +#if defined(__GNUC__) +# if defined(__AVX2__) +# include +# elif defined(__SSE2__) +# include +# elif defined(__ARM_NEON__) || defined(__ARM_NEON) +# define inline __inline__ /* circumvent a clang bug */ +# include +# undef inline +# endif +#elif defined(_MSC_VER) +# include +#endif + +/* + * One goal of XXH3 is to make it fast on both 32-bit and 64-bit, while + * remaining a true 64-bit/128-bit hash function. + * + * This is done by prioritizing a subset of 64-bit operations that can be + * emulated without too many steps on the average 32-bit machine. + * + * For example, these two lines seem similar, and run equally fast on 64-bit: + * + * xxh_u64 x; + * x ^= (x >> 47); // good + * x ^= (x >> 13); // bad + * + * However, to a 32-bit machine, there is a major difference. + * + * x ^= (x >> 47) looks like this: + * + * x.lo ^= (x.hi >> (47 - 32)); + * + * while x ^= (x >> 13) looks like this: + * + * // note: funnel shifts are not usually cheap. + * x.lo ^= (x.lo >> 13) | (x.hi << (32 - 13)); + * x.hi ^= (x.hi >> 13); + * + * The first one is significantly faster than the second, simply because the + * shift is larger than 32. This means: + * - All the bits we need are in the upper 32 bits, so we can ignore the lower + * 32 bits in the shift. + * - The shift result will always fit in the lower 32 bits, and therefore, + * we can ignore the upper 32 bits in the xor. + * + * Thanks to this optimization, XXH3 only requires these features to be efficient: + * + * - Usable unaligned access + * - A 32-bit or 64-bit ALU + * - If 32-bit, a decent ADC instruction + * - A 32 or 64-bit multiply with a 64-bit result + * - For the 128-bit variant, a decent byteswap helps short inputs. + * + * The first two are already required by XXH32, and almost all 32-bit and 64-bit + * platforms which can run XXH32 can run XXH3 efficiently. + * + * Thumb-1, the classic 16-bit only subset of ARM's instruction set, is one + * notable exception. + * + * First of all, Thumb-1 lacks support for the UMULL instruction which + * performs the important long multiply. This means numerous __aeabi_lmul + * calls. + * + * Second of all, the 8 functional registers are just not enough. + * Setup for __aeabi_lmul, byteshift loads, pointers, and all arithmetic need + * Lo registers, and this shuffling results in thousands more MOVs than A32. + * + * A32 and T32 don't have this limitation. They can access all 14 registers, + * do a 32->64 multiply with UMULL, and the flexible operand allowing free + * shifts is helpful, too. + * + * Therefore, we do a quick sanity check. + * + * If compiling Thumb-1 for a target which supports ARM instructions, we will + * emit a warning, as it is not a "sane" platform to compile for. + * + * Usually, if this happens, it is because of an accident and you probably need + * to specify -march, as you likely meant to compile for a newer architecture. + * + * Credit: large sections of the vectorial and asm source code paths + * have been contributed by @easyaspi314 + */ +#if defined(__thumb__) && !defined(__thumb2__) && defined(__ARM_ARCH_ISA_ARM) +# warning "XXH3 is highly inefficient without ARM or Thumb-2." +#endif + +/* ========================================== + * Vectorization detection + * ========================================== */ + +#ifdef XXH_DOXYGEN +/*! + * @ingroup tuning + * @brief Overrides the vectorization implementation chosen for XXH3. + * + * Can be defined to 0 to disable SIMD or any of the values mentioned in + * @ref XXH_VECTOR_TYPE. + * + * If this is not defined, it uses predefined macros to determine the best + * implementation. + */ +# define XXH_VECTOR XXH_SCALAR +/*! + * @ingroup tuning + * @brief Possible values for @ref XXH_VECTOR. + * + * Note that these are actually implemented as macros. + * + * If this is not defined, it is detected automatically. + * @ref XXH_X86DISPATCH overrides this. + */ +enum XXH_VECTOR_TYPE /* fake enum */ { + XXH_SCALAR = 0, /*!< Portable scalar version */ + XXH_SSE2 = 1, /*!< + * SSE2 for Pentium 4, Opteron, all x86_64. + * + * @note SSE2 is also guaranteed on Windows 10, macOS, and + * Android x86. + */ + XXH_AVX2 = 2, /*!< AVX2 for Haswell and Bulldozer */ + XXH_AVX512 = 3, /*!< AVX512 for Skylake and Icelake */ + XXH_NEON = 4, /*!< NEON for most ARMv7-A and all AArch64 */ + XXH_VSX = 5, /*!< VSX and ZVector for POWER8/z13 (64-bit) */ +}; +/*! + * @ingroup tuning + * @brief Selects the minimum alignment for XXH3's accumulators. + * + * When using SIMD, this should match the alignment reqired for said vector + * type, so, for example, 32 for AVX2. + * + * Default: Auto detected. + */ +# define XXH_ACC_ALIGN 8 +#endif + +/* Actual definition */ +#ifndef XXH_DOXYGEN +# define XXH_SCALAR 0 +# define XXH_SSE2 1 +# define XXH_AVX2 2 +# define XXH_AVX512 3 +# define XXH_NEON 4 +# define XXH_VSX 5 +#endif + +#ifndef XXH_VECTOR /* can be defined on command line */ +# if defined(__AVX512F__) +# define XXH_VECTOR XXH_AVX512 +# elif defined(__AVX2__) +# define XXH_VECTOR XXH_AVX2 +# elif defined(__SSE2__) || defined(_M_AMD64) || defined(_M_X64) || (defined(_M_IX86_FP) && (_M_IX86_FP == 2)) +# define XXH_VECTOR XXH_SSE2 +# elif ( \ + defined(__ARM_NEON__) || defined(__ARM_NEON) /* gcc */ \ + || defined(_M_ARM64) || defined(_M_ARM_ARMV7VE) /* msvc */ \ + ) && ( \ + defined(_WIN32) || defined(__LITTLE_ENDIAN__) /* little endian only */ \ + || (defined(__BYTE_ORDER__) && __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__) \ + ) +# define XXH_VECTOR XXH_NEON +# elif (defined(__PPC64__) && defined(__POWER8_VECTOR__)) \ + || (defined(__s390x__) && defined(__VEC__)) \ + && defined(__GNUC__) /* TODO: IBM XL */ +# define XXH_VECTOR XXH_VSX +# else +# define XXH_VECTOR XXH_SCALAR +# endif +#endif + +/* + * Controls the alignment of the accumulator, + * for compatibility with aligned vector loads, which are usually faster. + */ +#ifndef XXH_ACC_ALIGN +# if defined(XXH_X86DISPATCH) +# define XXH_ACC_ALIGN 64 /* for compatibility with avx512 */ +# elif XXH_VECTOR == XXH_SCALAR /* scalar */ +# define XXH_ACC_ALIGN 8 +# elif XXH_VECTOR == XXH_SSE2 /* sse2 */ +# define XXH_ACC_ALIGN 16 +# elif XXH_VECTOR == XXH_AVX2 /* avx2 */ +# define XXH_ACC_ALIGN 32 +# elif XXH_VECTOR == XXH_NEON /* neon */ +# define XXH_ACC_ALIGN 16 +# elif XXH_VECTOR == XXH_VSX /* vsx */ +# define XXH_ACC_ALIGN 16 +# elif XXH_VECTOR == XXH_AVX512 /* avx512 */ +# define XXH_ACC_ALIGN 64 +# endif +#endif + +#if defined(XXH_X86DISPATCH) || XXH_VECTOR == XXH_SSE2 \ + || XXH_VECTOR == XXH_AVX2 || XXH_VECTOR == XXH_AVX512 +# define XXH_SEC_ALIGN XXH_ACC_ALIGN +#else +# define XXH_SEC_ALIGN 8 +#endif + +/* + * UGLY HACK: + * GCC usually generates the best code with -O3 for xxHash. + * + * However, when targeting AVX2, it is overzealous in its unrolling resulting + * in code roughly 3/4 the speed of Clang. + * + * There are other issues, such as GCC splitting _mm256_loadu_si256 into + * _mm_loadu_si128 + _mm256_inserti128_si256. This is an optimization which + * only applies to Sandy and Ivy Bridge... which don't even support AVX2. + * + * That is why when compiling the AVX2 version, it is recommended to use either + * -O2 -mavx2 -march=haswell + * or + * -O2 -mavx2 -mno-avx256-split-unaligned-load + * for decent performance, or to use Clang instead. + * + * Fortunately, we can control the first one with a pragma that forces GCC into + * -O2, but the other one we can't control without "failed to inline always + * inline function due to target mismatch" warnings. + */ +#if XXH_VECTOR == XXH_AVX2 /* AVX2 */ \ + && defined(__GNUC__) && !defined(__clang__) /* GCC, not Clang */ \ + && defined(__OPTIMIZE__) && !defined(__OPTIMIZE_SIZE__) /* respect -O0 and -Os */ +# pragma GCC push_options +# pragma GCC optimize("-O2") +#endif + + +#if XXH_VECTOR == XXH_NEON +/* + * NEON's setup for vmlal_u32 is a little more complicated than it is on + * SSE2, AVX2, and VSX. + * + * While PMULUDQ and VMULEUW both perform a mask, VMLAL.U32 performs an upcast. + * + * To do the same operation, the 128-bit 'Q' register needs to be split into + * two 64-bit 'D' registers, performing this operation:: + * + * [ a | b ] + * | '---------. .--------' | + * | x | + * | .---------' '--------. | + * [ a & 0xFFFFFFFF | b & 0xFFFFFFFF ],[ a >> 32 | b >> 32 ] + * + * Due to significant changes in aarch64, the fastest method for aarch64 is + * completely different than the fastest method for ARMv7-A. + * + * ARMv7-A treats D registers as unions overlaying Q registers, so modifying + * D11 will modify the high half of Q5. This is similar to how modifying AH + * will only affect bits 8-15 of AX on x86. + * + * VZIP takes two registers, and puts even lanes in one register and odd lanes + * in the other. + * + * On ARMv7-A, this strangely modifies both parameters in place instead of + * taking the usual 3-operand form. + * + * Therefore, if we want to do this, we can simply use a D-form VZIP.32 on the + * lower and upper halves of the Q register to end up with the high and low + * halves where we want - all in one instruction. + * + * vzip.32 d10, d11 @ d10 = { d10[0], d11[0] }; d11 = { d10[1], d11[1] } + * + * Unfortunately we need inline assembly for this: Instructions modifying two + * registers at once is not possible in GCC or Clang's IR, and they have to + * create a copy. + * + * aarch64 requires a different approach. + * + * In order to make it easier to write a decent compiler for aarch64, many + * quirks were removed, such as conditional execution. + * + * NEON was also affected by this. + * + * aarch64 cannot access the high bits of a Q-form register, and writes to a + * D-form register zero the high bits, similar to how writes to W-form scalar + * registers (or DWORD registers on x86_64) work. + * + * The formerly free vget_high intrinsics now require a vext (with a few + * exceptions) + * + * Additionally, VZIP was replaced by ZIP1 and ZIP2, which are the equivalent + * of PUNPCKL* and PUNPCKH* in SSE, respectively, in order to only modify one + * operand. + * + * The equivalent of the VZIP.32 on the lower and upper halves would be this + * mess: + * + * ext v2.4s, v0.4s, v0.4s, #2 // v2 = { v0[2], v0[3], v0[0], v0[1] } + * zip1 v1.2s, v0.2s, v2.2s // v1 = { v0[0], v2[0] } + * zip2 v0.2s, v0.2s, v1.2s // v0 = { v0[1], v2[1] } + * + * Instead, we use a literal downcast, vmovn_u64 (XTN), and vshrn_n_u64 (SHRN): + * + * shrn v1.2s, v0.2d, #32 // v1 = (uint32x2_t)(v0 >> 32); + * xtn v0.2s, v0.2d // v0 = (uint32x2_t)(v0 & 0xFFFFFFFF); + * + * This is available on ARMv7-A, but is less efficient than a single VZIP.32. + */ + +/*! + * Function-like macro: + * void XXH_SPLIT_IN_PLACE(uint64x2_t &in, uint32x2_t &outLo, uint32x2_t &outHi) + * { + * outLo = (uint32x2_t)(in & 0xFFFFFFFF); + * outHi = (uint32x2_t)(in >> 32); + * in = UNDEFINED; + * } + */ +# if !defined(XXH_NO_VZIP_HACK) /* define to disable */ \ + && defined(__GNUC__) \ + && !defined(__aarch64__) && !defined(__arm64__) && !defined(_M_ARM64) +# define XXH_SPLIT_IN_PLACE(in, outLo, outHi) \ + do { \ + /* Undocumented GCC/Clang operand modifier: %e0 = lower D half, %f0 = upper D half */ \ + /* https://github.com/gcc-mirror/gcc/blob/38cf91e5/gcc/config/arm/arm.c#L22486 */ \ + /* https://github.com/llvm-mirror/llvm/blob/2c4ca683/lib/Target/ARM/ARMAsmPrinter.cpp#L399 */ \ + __asm__("vzip.32 %e0, %f0" : "+w" (in)); \ + (outLo) = vget_low_u32 (vreinterpretq_u32_u64(in)); \ + (outHi) = vget_high_u32(vreinterpretq_u32_u64(in)); \ + } while (0) +# else +# define XXH_SPLIT_IN_PLACE(in, outLo, outHi) \ + do { \ + (outLo) = vmovn_u64 (in); \ + (outHi) = vshrn_n_u64 ((in), 32); \ + } while (0) +# endif +#endif /* XXH_VECTOR == XXH_NEON */ + +/* + * VSX and Z Vector helpers. + * + * This is very messy, and any pull requests to clean this up are welcome. + * + * There are a lot of problems with supporting VSX and s390x, due to + * inconsistent intrinsics, spotty coverage, and multiple endiannesses. + */ +#if XXH_VECTOR == XXH_VSX +# if defined(__s390x__) +# include +# else +/* gcc's altivec.h can have the unwanted consequence to unconditionally + * #define bool, vector, and pixel keywords, + * with bad consequences for programs already using these keywords for other purposes. + * The paragraph defining these macros is skipped when __APPLE_ALTIVEC__ is defined. + * __APPLE_ALTIVEC__ is _generally_ defined automatically by the compiler, + * but it seems that, in some cases, it isn't. + * Force the build macro to be defined, so that keywords are not altered. + */ +# if defined(__GNUC__) && !defined(__APPLE_ALTIVEC__) +# define __APPLE_ALTIVEC__ +# endif +# include +# endif + +typedef __vector unsigned long long xxh_u64x2; +typedef __vector unsigned char xxh_u8x16; +typedef __vector unsigned xxh_u32x4; + +# ifndef XXH_VSX_BE +# if defined(__BIG_ENDIAN__) \ + || (defined(__BYTE_ORDER__) && __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__) +# define XXH_VSX_BE 1 +# elif defined(__VEC_ELEMENT_REG_ORDER__) && __VEC_ELEMENT_REG_ORDER__ == __ORDER_BIG_ENDIAN__ +# warning "-maltivec=be is not recommended. Please use native endianness." +# define XXH_VSX_BE 1 +# else +# define XXH_VSX_BE 0 +# endif +# endif /* !defined(XXH_VSX_BE) */ + +# if XXH_VSX_BE +# if defined(__POWER9_VECTOR__) || (defined(__clang__) && defined(__s390x__)) +# define XXH_vec_revb vec_revb +# else +/*! + * A polyfill for POWER9's vec_revb(). + */ +XXH_FORCE_INLINE xxh_u64x2 XXH_vec_revb(xxh_u64x2 val) +{ + xxh_u8x16 const vByteSwap = { 0x07, 0x06, 0x05, 0x04, 0x03, 0x02, 0x01, 0x00, + 0x0F, 0x0E, 0x0D, 0x0C, 0x0B, 0x0A, 0x09, 0x08 }; + return vec_perm(val, val, vByteSwap); +} +# endif +# endif /* XXH_VSX_BE */ + +/*! + * Performs an unaligned vector load and byte swaps it on big endian. + */ +XXH_FORCE_INLINE xxh_u64x2 XXH_vec_loadu(const void *ptr) +{ + xxh_u64x2 ret; + XXH_memcpy(&ret, ptr, sizeof(xxh_u64x2)); +# if XXH_VSX_BE + ret = XXH_vec_revb(ret); +# endif + return ret; +} + +/* + * vec_mulo and vec_mule are very problematic intrinsics on PowerPC + * + * These intrinsics weren't added until GCC 8, despite existing for a while, + * and they are endian dependent. Also, their meaning swap depending on version. + * */ +# if defined(__s390x__) + /* s390x is always big endian, no issue on this platform */ +# define XXH_vec_mulo vec_mulo +# define XXH_vec_mule vec_mule +# elif defined(__clang__) && XXH_HAS_BUILTIN(__builtin_altivec_vmuleuw) +/* Clang has a better way to control this, we can just use the builtin which doesn't swap. */ +# define XXH_vec_mulo __builtin_altivec_vmulouw +# define XXH_vec_mule __builtin_altivec_vmuleuw +# else +/* gcc needs inline assembly */ +/* Adapted from https://github.com/google/highwayhash/blob/master/highwayhash/hh_vsx.h. */ +XXH_FORCE_INLINE xxh_u64x2 XXH_vec_mulo(xxh_u32x4 a, xxh_u32x4 b) +{ + xxh_u64x2 result; + __asm__("vmulouw %0, %1, %2" : "=v" (result) : "v" (a), "v" (b)); + return result; +} +XXH_FORCE_INLINE xxh_u64x2 XXH_vec_mule(xxh_u32x4 a, xxh_u32x4 b) +{ + xxh_u64x2 result; + __asm__("vmuleuw %0, %1, %2" : "=v" (result) : "v" (a), "v" (b)); + return result; +} +# endif /* XXH_vec_mulo, XXH_vec_mule */ +#endif /* XXH_VECTOR == XXH_VSX */ + + +/* prefetch + * can be disabled, by declaring XXH_NO_PREFETCH build macro */ +#if defined(XXH_NO_PREFETCH) +# define XXH_PREFETCH(ptr) (void)(ptr) /* disabled */ +#else +# if defined(_MSC_VER) && (defined(_M_X64) || defined(_M_IX86)) /* _mm_prefetch() not defined outside of x86/x64 */ +# include /* https://msdn.microsoft.com/fr-fr/library/84szxsww(v=vs.90).aspx */ +# define XXH_PREFETCH(ptr) _mm_prefetch((const char*)(ptr), _MM_HINT_T0) +# elif defined(__GNUC__) && ( (__GNUC__ >= 4) || ( (__GNUC__ == 3) && (__GNUC_MINOR__ >= 1) ) ) +# define XXH_PREFETCH(ptr) __builtin_prefetch((ptr), 0 /* rw==read */, 3 /* locality */) +# else +# define XXH_PREFETCH(ptr) (void)(ptr) /* disabled */ +# endif +#endif /* XXH_NO_PREFETCH */ + + +/* ========================================== + * XXH3 default settings + * ========================================== */ + +#define XXH_SECRET_DEFAULT_SIZE 192 /* minimum XXH3_SECRET_SIZE_MIN */ + +#if (XXH_SECRET_DEFAULT_SIZE < XXH3_SECRET_SIZE_MIN) +# error "default keyset is not large enough" +#endif + +/*! Pseudorandom secret taken directly from FARSH. */ +XXH_ALIGN(64) static const xxh_u8 XXH3_kSecret[XXH_SECRET_DEFAULT_SIZE] = { + 0xb8, 0xfe, 0x6c, 0x39, 0x23, 0xa4, 0x4b, 0xbe, 0x7c, 0x01, 0x81, 0x2c, 0xf7, 0x21, 0xad, 0x1c, + 0xde, 0xd4, 0x6d, 0xe9, 0x83, 0x90, 0x97, 0xdb, 0x72, 0x40, 0xa4, 0xa4, 0xb7, 0xb3, 0x67, 0x1f, + 0xcb, 0x79, 0xe6, 0x4e, 0xcc, 0xc0, 0xe5, 0x78, 0x82, 0x5a, 0xd0, 0x7d, 0xcc, 0xff, 0x72, 0x21, + 0xb8, 0x08, 0x46, 0x74, 0xf7, 0x43, 0x24, 0x8e, 0xe0, 0x35, 0x90, 0xe6, 0x81, 0x3a, 0x26, 0x4c, + 0x3c, 0x28, 0x52, 0xbb, 0x91, 0xc3, 0x00, 0xcb, 0x88, 0xd0, 0x65, 0x8b, 0x1b, 0x53, 0x2e, 0xa3, + 0x71, 0x64, 0x48, 0x97, 0xa2, 0x0d, 0xf9, 0x4e, 0x38, 0x19, 0xef, 0x46, 0xa9, 0xde, 0xac, 0xd8, + 0xa8, 0xfa, 0x76, 0x3f, 0xe3, 0x9c, 0x34, 0x3f, 0xf9, 0xdc, 0xbb, 0xc7, 0xc7, 0x0b, 0x4f, 0x1d, + 0x8a, 0x51, 0xe0, 0x4b, 0xcd, 0xb4, 0x59, 0x31, 0xc8, 0x9f, 0x7e, 0xc9, 0xd9, 0x78, 0x73, 0x64, + 0xea, 0xc5, 0xac, 0x83, 0x34, 0xd3, 0xeb, 0xc3, 0xc5, 0x81, 0xa0, 0xff, 0xfa, 0x13, 0x63, 0xeb, + 0x17, 0x0d, 0xdd, 0x51, 0xb7, 0xf0, 0xda, 0x49, 0xd3, 0x16, 0x55, 0x26, 0x29, 0xd4, 0x68, 0x9e, + 0x2b, 0x16, 0xbe, 0x58, 0x7d, 0x47, 0xa1, 0xfc, 0x8f, 0xf8, 0xb8, 0xd1, 0x7a, 0xd0, 0x31, 0xce, + 0x45, 0xcb, 0x3a, 0x8f, 0x95, 0x16, 0x04, 0x28, 0xaf, 0xd7, 0xfb, 0xca, 0xbb, 0x4b, 0x40, 0x7e, +}; + + +#ifdef XXH_OLD_NAMES +# define kSecret XXH3_kSecret +#endif + +#ifdef XXH_DOXYGEN +/*! + * @brief Calculates a 32-bit to 64-bit long multiply. + * + * Implemented as a macro. + * + * Wraps `__emulu` on MSVC x86 because it tends to call `__allmul` when it doesn't + * need to (but it shouldn't need to anyways, it is about 7 instructions to do + * a 64x64 multiply...). Since we know that this will _always_ emit `MULL`, we + * use that instead of the normal method. + * + * If you are compiling for platforms like Thumb-1 and don't have a better option, + * you may also want to write your own long multiply routine here. + * + * @param x, y Numbers to be multiplied + * @return 64-bit product of the low 32 bits of @p x and @p y. + */ +XXH_FORCE_INLINE xxh_u64 +XXH_mult32to64(xxh_u64 x, xxh_u64 y) +{ + return (x & 0xFFFFFFFF) * (y & 0xFFFFFFFF); +} +#elif defined(_MSC_VER) && defined(_M_IX86) +# include +# define XXH_mult32to64(x, y) __emulu((unsigned)(x), (unsigned)(y)) +#else +/* + * Downcast + upcast is usually better than masking on older compilers like + * GCC 4.2 (especially 32-bit ones), all without affecting newer compilers. + * + * The other method, (x & 0xFFFFFFFF) * (y & 0xFFFFFFFF), will AND both operands + * and perform a full 64x64 multiply -- entirely redundant on 32-bit. + */ +# define XXH_mult32to64(x, y) ((xxh_u64)(xxh_u32)(x) * (xxh_u64)(xxh_u32)(y)) +#endif + +/*! + * @brief Calculates a 64->128-bit long multiply. + * + * Uses `__uint128_t` and `_umul128` if available, otherwise uses a scalar + * version. + * + * @param lhs , rhs The 64-bit integers to be multiplied + * @return The 128-bit result represented in an @ref XXH128_hash_t. + */ +static XXH128_hash_t +XXH_mult64to128(xxh_u64 lhs, xxh_u64 rhs) +{ + /* + * GCC/Clang __uint128_t method. + * + * On most 64-bit targets, GCC and Clang define a __uint128_t type. + * This is usually the best way as it usually uses a native long 64-bit + * multiply, such as MULQ on x86_64 or MUL + UMULH on aarch64. + * + * Usually. + * + * Despite being a 32-bit platform, Clang (and emscripten) define this type + * despite not having the arithmetic for it. This results in a laggy + * compiler builtin call which calculates a full 128-bit multiply. + * In that case it is best to use the portable one. + * https://github.com/Cyan4973/xxHash/issues/211#issuecomment-515575677 + */ +#if defined(__GNUC__) && !defined(__wasm__) \ + && defined(__SIZEOF_INT128__) \ + || (defined(_INTEGRAL_MAX_BITS) && _INTEGRAL_MAX_BITS >= 128) + + __uint128_t const product = (__uint128_t)lhs * (__uint128_t)rhs; + XXH128_hash_t r128; + r128.low64 = (xxh_u64)(product); + r128.high64 = (xxh_u64)(product >> 64); + return r128; + + /* + * MSVC for x64's _umul128 method. + * + * xxh_u64 _umul128(xxh_u64 Multiplier, xxh_u64 Multiplicand, xxh_u64 *HighProduct); + * + * This compiles to single operand MUL on x64. + */ +#elif defined(_M_X64) || defined(_M_IA64) + +#ifndef _MSC_VER +# pragma intrinsic(_umul128) +#endif + xxh_u64 product_high; + xxh_u64 const product_low = _umul128(lhs, rhs, &product_high); + XXH128_hash_t r128; + r128.low64 = product_low; + r128.high64 = product_high; + return r128; + + /* + * MSVC for ARM64's __umulh method. + * + * This compiles to the same MUL + UMULH as GCC/Clang's __uint128_t method. + */ +#elif defined(_M_ARM64) + +#ifndef _MSC_VER +# pragma intrinsic(__umulh) +#endif + XXH128_hash_t r128; + r128.low64 = lhs * rhs; + r128.high64 = __umulh(lhs, rhs); + return r128; + +#else + /* + * Portable scalar method. Optimized for 32-bit and 64-bit ALUs. + * + * This is a fast and simple grade school multiply, which is shown below + * with base 10 arithmetic instead of base 0x100000000. + * + * 9 3 // D2 lhs = 93 + * x 7 5 // D2 rhs = 75 + * ---------- + * 1 5 // D2 lo_lo = (93 % 10) * (75 % 10) = 15 + * 4 5 | // D2 hi_lo = (93 / 10) * (75 % 10) = 45 + * 2 1 | // D2 lo_hi = (93 % 10) * (75 / 10) = 21 + * + 6 3 | | // D2 hi_hi = (93 / 10) * (75 / 10) = 63 + * --------- + * 2 7 | // D2 cross = (15 / 10) + (45 % 10) + 21 = 27 + * + 6 7 | | // D2 upper = (27 / 10) + (45 / 10) + 63 = 67 + * --------- + * 6 9 7 5 // D4 res = (27 * 10) + (15 % 10) + (67 * 100) = 6975 + * + * The reasons for adding the products like this are: + * 1. It avoids manual carry tracking. Just like how + * (9 * 9) + 9 + 9 = 99, the same applies with this for UINT64_MAX. + * This avoids a lot of complexity. + * + * 2. It hints for, and on Clang, compiles to, the powerful UMAAL + * instruction available in ARM's Digital Signal Processing extension + * in 32-bit ARMv6 and later, which is shown below: + * + * void UMAAL(xxh_u32 *RdLo, xxh_u32 *RdHi, xxh_u32 Rn, xxh_u32 Rm) + * { + * xxh_u64 product = (xxh_u64)*RdLo * (xxh_u64)*RdHi + Rn + Rm; + * *RdLo = (xxh_u32)(product & 0xFFFFFFFF); + * *RdHi = (xxh_u32)(product >> 32); + * } + * + * This instruction was designed for efficient long multiplication, and + * allows this to be calculated in only 4 instructions at speeds + * comparable to some 64-bit ALUs. + * + * 3. It isn't terrible on other platforms. Usually this will be a couple + * of 32-bit ADD/ADCs. + */ + + /* First calculate all of the cross products. */ + xxh_u64 const lo_lo = XXH_mult32to64(lhs & 0xFFFFFFFF, rhs & 0xFFFFFFFF); + xxh_u64 const hi_lo = XXH_mult32to64(lhs >> 32, rhs & 0xFFFFFFFF); + xxh_u64 const lo_hi = XXH_mult32to64(lhs & 0xFFFFFFFF, rhs >> 32); + xxh_u64 const hi_hi = XXH_mult32to64(lhs >> 32, rhs >> 32); + + /* Now add the products together. These will never overflow. */ + xxh_u64 const cross = (lo_lo >> 32) + (hi_lo & 0xFFFFFFFF) + lo_hi; + xxh_u64 const upper = (hi_lo >> 32) + (cross >> 32) + hi_hi; + xxh_u64 const lower = (cross << 32) | (lo_lo & 0xFFFFFFFF); + + XXH128_hash_t r128; + r128.low64 = lower; + r128.high64 = upper; + return r128; +#endif +} + +/*! + * @brief Calculates a 64-bit to 128-bit multiply, then XOR folds it. + * + * The reason for the separate function is to prevent passing too many structs + * around by value. This will hopefully inline the multiply, but we don't force it. + * + * @param lhs , rhs The 64-bit integers to multiply + * @return The low 64 bits of the product XOR'd by the high 64 bits. + * @see XXH_mult64to128() + */ +static xxh_u64 +XXH3_mul128_fold64(xxh_u64 lhs, xxh_u64 rhs) +{ + XXH128_hash_t product = XXH_mult64to128(lhs, rhs); + return product.low64 ^ product.high64; +} + +/*! Seems to produce slightly better code on GCC for some reason. */ +XXH_FORCE_INLINE xxh_u64 XXH_xorshift64(xxh_u64 v64, int shift) +{ + XXH_ASSERT(0 <= shift && shift < 64); + return v64 ^ (v64 >> shift); +} + +/* + * This is a fast avalanche stage, + * suitable when input bits are already partially mixed + */ +static XXH64_hash_t XXH3_avalanche(xxh_u64 h64) +{ + h64 = XXH_xorshift64(h64, 37); + h64 *= 0x165667919E3779F9ULL; + h64 = XXH_xorshift64(h64, 32); + return h64; +} + +/* + * This is a stronger avalanche, + * inspired by Pelle Evensen's rrmxmx + * preferable when input has not been previously mixed + */ +static XXH64_hash_t XXH3_rrmxmx(xxh_u64 h64, xxh_u64 len) +{ + /* this mix is inspired by Pelle Evensen's rrmxmx */ + h64 ^= XXH_rotl64(h64, 49) ^ XXH_rotl64(h64, 24); + h64 *= 0x9FB21C651E98DF25ULL; + h64 ^= (h64 >> 35) + len ; + h64 *= 0x9FB21C651E98DF25ULL; + return XXH_xorshift64(h64, 28); +} + + +/* ========================================== + * Short keys + * ========================================== + * One of the shortcomings of XXH32 and XXH64 was that their performance was + * sub-optimal on short lengths. It used an iterative algorithm which strongly + * favored lengths that were a multiple of 4 or 8. + * + * Instead of iterating over individual inputs, we use a set of single shot + * functions which piece together a range of lengths and operate in constant time. + * + * Additionally, the number of multiplies has been significantly reduced. This + * reduces latency, especially when emulating 64-bit multiplies on 32-bit. + * + * Depending on the platform, this may or may not be faster than XXH32, but it + * is almost guaranteed to be faster than XXH64. + */ + +/* + * At very short lengths, there isn't enough input to fully hide secrets, or use + * the entire secret. + * + * There is also only a limited amount of mixing we can do before significantly + * impacting performance. + * + * Therefore, we use different sections of the secret and always mix two secret + * samples with an XOR. This should have no effect on performance on the + * seedless or withSeed variants because everything _should_ be constant folded + * by modern compilers. + * + * The XOR mixing hides individual parts of the secret and increases entropy. + * + * This adds an extra layer of strength for custom secrets. + */ +XXH_FORCE_INLINE XXH64_hash_t +XXH3_len_1to3_64b(const xxh_u8* input, size_t len, const xxh_u8* secret, XXH64_hash_t seed) +{ + XXH_ASSERT(input != NULL); + XXH_ASSERT(1 <= len && len <= 3); + XXH_ASSERT(secret != NULL); + /* + * len = 1: combined = { input[0], 0x01, input[0], input[0] } + * len = 2: combined = { input[1], 0x02, input[0], input[1] } + * len = 3: combined = { input[2], 0x03, input[0], input[1] } + */ + { xxh_u8 const c1 = input[0]; + xxh_u8 const c2 = input[len >> 1]; + xxh_u8 const c3 = input[len - 1]; + xxh_u32 const combined = ((xxh_u32)c1 << 16) | ((xxh_u32)c2 << 24) + | ((xxh_u32)c3 << 0) | ((xxh_u32)len << 8); + xxh_u64 const bitflip = (XXH_readLE32(secret) ^ XXH_readLE32(secret+4)) + seed; + xxh_u64 const keyed = (xxh_u64)combined ^ bitflip; + return XXH64_avalanche(keyed); + } +} + +XXH_FORCE_INLINE XXH64_hash_t +XXH3_len_4to8_64b(const xxh_u8* input, size_t len, const xxh_u8* secret, XXH64_hash_t seed) +{ + XXH_ASSERT(input != NULL); + XXH_ASSERT(secret != NULL); + XXH_ASSERT(4 <= len && len <= 8); + seed ^= (xxh_u64)XXH_swap32((xxh_u32)seed) << 32; + { xxh_u32 const input1 = XXH_readLE32(input); + xxh_u32 const input2 = XXH_readLE32(input + len - 4); + xxh_u64 const bitflip = (XXH_readLE64(secret+8) ^ XXH_readLE64(secret+16)) - seed; + xxh_u64 const input64 = input2 + (((xxh_u64)input1) << 32); + xxh_u64 const keyed = input64 ^ bitflip; + return XXH3_rrmxmx(keyed, len); + } +} + +XXH_FORCE_INLINE XXH64_hash_t +XXH3_len_9to16_64b(const xxh_u8* input, size_t len, const xxh_u8* secret, XXH64_hash_t seed) +{ + XXH_ASSERT(input != NULL); + XXH_ASSERT(secret != NULL); + XXH_ASSERT(9 <= len && len <= 16); + { xxh_u64 const bitflip1 = (XXH_readLE64(secret+24) ^ XXH_readLE64(secret+32)) + seed; + xxh_u64 const bitflip2 = (XXH_readLE64(secret+40) ^ XXH_readLE64(secret+48)) - seed; + xxh_u64 const input_lo = XXH_readLE64(input) ^ bitflip1; + xxh_u64 const input_hi = XXH_readLE64(input + len - 8) ^ bitflip2; + xxh_u64 const acc = len + + XXH_swap64(input_lo) + input_hi + + XXH3_mul128_fold64(input_lo, input_hi); + return XXH3_avalanche(acc); + } +} + +XXH_FORCE_INLINE XXH64_hash_t +XXH3_len_0to16_64b(const xxh_u8* input, size_t len, const xxh_u8* secret, XXH64_hash_t seed) +{ + XXH_ASSERT(len <= 16); + { if (XXH_likely(len > 8)) return XXH3_len_9to16_64b(input, len, secret, seed); + if (XXH_likely(len >= 4)) return XXH3_len_4to8_64b(input, len, secret, seed); + if (len) return XXH3_len_1to3_64b(input, len, secret, seed); + return XXH64_avalanche(seed ^ (XXH_readLE64(secret+56) ^ XXH_readLE64(secret+64))); + } +} + +/* + * DISCLAIMER: There are known *seed-dependent* multicollisions here due to + * multiplication by zero, affecting hashes of lengths 17 to 240. + * + * However, they are very unlikely. + * + * Keep this in mind when using the unseeded XXH3_64bits() variant: As with all + * unseeded non-cryptographic hashes, it does not attempt to defend itself + * against specially crafted inputs, only random inputs. + * + * Compared to classic UMAC where a 1 in 2^31 chance of 4 consecutive bytes + * cancelling out the secret is taken an arbitrary number of times (addressed + * in XXH3_accumulate_512), this collision is very unlikely with random inputs + * and/or proper seeding: + * + * This only has a 1 in 2^63 chance of 8 consecutive bytes cancelling out, in a + * function that is only called up to 16 times per hash with up to 240 bytes of + * input. + * + * This is not too bad for a non-cryptographic hash function, especially with + * only 64 bit outputs. + * + * The 128-bit variant (which trades some speed for strength) is NOT affected + * by this, although it is always a good idea to use a proper seed if you care + * about strength. + */ +XXH_FORCE_INLINE xxh_u64 XXH3_mix16B(const xxh_u8* XXH_RESTRICT input, + const xxh_u8* XXH_RESTRICT secret, xxh_u64 seed64) +{ +#if defined(__GNUC__) && !defined(__clang__) /* GCC, not Clang */ \ + && defined(__i386__) && defined(__SSE2__) /* x86 + SSE2 */ \ + && !defined(XXH_ENABLE_AUTOVECTORIZE) /* Define to disable like XXH32 hack */ + /* + * UGLY HACK: + * GCC for x86 tends to autovectorize the 128-bit multiply, resulting in + * slower code. + * + * By forcing seed64 into a register, we disrupt the cost model and + * cause it to scalarize. See `XXH32_round()` + * + * FIXME: Clang's output is still _much_ faster -- On an AMD Ryzen 3600, + * XXH3_64bits @ len=240 runs at 4.6 GB/s with Clang 9, but 3.3 GB/s on + * GCC 9.2, despite both emitting scalar code. + * + * GCC generates much better scalar code than Clang for the rest of XXH3, + * which is why finding a more optimal codepath is an interest. + */ + XXH_COMPILER_GUARD(seed64); +#endif + { xxh_u64 const input_lo = XXH_readLE64(input); + xxh_u64 const input_hi = XXH_readLE64(input+8); + return XXH3_mul128_fold64( + input_lo ^ (XXH_readLE64(secret) + seed64), + input_hi ^ (XXH_readLE64(secret+8) - seed64) + ); + } +} + +/* For mid range keys, XXH3 uses a Mum-hash variant. */ +XXH_FORCE_INLINE XXH64_hash_t +XXH3_len_17to128_64b(const xxh_u8* XXH_RESTRICT input, size_t len, + const xxh_u8* XXH_RESTRICT secret, size_t secretSize, + XXH64_hash_t seed) +{ + XXH_ASSERT(secretSize >= XXH3_SECRET_SIZE_MIN); (void)secretSize; + XXH_ASSERT(16 < len && len <= 128); + + { xxh_u64 acc = len * XXH_PRIME64_1; + if (len > 32) { + if (len > 64) { + if (len > 96) { + acc += XXH3_mix16B(input+48, secret+96, seed); + acc += XXH3_mix16B(input+len-64, secret+112, seed); + } + acc += XXH3_mix16B(input+32, secret+64, seed); + acc += XXH3_mix16B(input+len-48, secret+80, seed); + } + acc += XXH3_mix16B(input+16, secret+32, seed); + acc += XXH3_mix16B(input+len-32, secret+48, seed); + } + acc += XXH3_mix16B(input+0, secret+0, seed); + acc += XXH3_mix16B(input+len-16, secret+16, seed); + + return XXH3_avalanche(acc); + } +} + +#define XXH3_MIDSIZE_MAX 240 + +XXH_NO_INLINE XXH64_hash_t +XXH3_len_129to240_64b(const xxh_u8* XXH_RESTRICT input, size_t len, + const xxh_u8* XXH_RESTRICT secret, size_t secretSize, + XXH64_hash_t seed) +{ + XXH_ASSERT(secretSize >= XXH3_SECRET_SIZE_MIN); (void)secretSize; + XXH_ASSERT(128 < len && len <= XXH3_MIDSIZE_MAX); + + #define XXH3_MIDSIZE_STARTOFFSET 3 + #define XXH3_MIDSIZE_LASTOFFSET 17 + + { xxh_u64 acc = len * XXH_PRIME64_1; + int const nbRounds = (int)len / 16; + int i; + for (i=0; i<8; i++) { + acc += XXH3_mix16B(input+(16*i), secret+(16*i), seed); + } + acc = XXH3_avalanche(acc); + XXH_ASSERT(nbRounds >= 8); +#if defined(__clang__) /* Clang */ \ + && (defined(__ARM_NEON) || defined(__ARM_NEON__)) /* NEON */ \ + && !defined(XXH_ENABLE_AUTOVECTORIZE) /* Define to disable */ + /* + * UGLY HACK: + * Clang for ARMv7-A tries to vectorize this loop, similar to GCC x86. + * In everywhere else, it uses scalar code. + * + * For 64->128-bit multiplies, even if the NEON was 100% optimal, it + * would still be slower than UMAAL (see XXH_mult64to128). + * + * Unfortunately, Clang doesn't handle the long multiplies properly and + * converts them to the nonexistent "vmulq_u64" intrinsic, which is then + * scalarized into an ugly mess of VMOV.32 instructions. + * + * This mess is difficult to avoid without turning autovectorization + * off completely, but they are usually relatively minor and/or not + * worth it to fix. + * + * This loop is the easiest to fix, as unlike XXH32, this pragma + * _actually works_ because it is a loop vectorization instead of an + * SLP vectorization. + */ + #pragma clang loop vectorize(disable) +#endif + for (i=8 ; i < nbRounds; i++) { + acc += XXH3_mix16B(input+(16*i), secret+(16*(i-8)) + XXH3_MIDSIZE_STARTOFFSET, seed); + } + /* last bytes */ + acc += XXH3_mix16B(input + len - 16, secret + XXH3_SECRET_SIZE_MIN - XXH3_MIDSIZE_LASTOFFSET, seed); + return XXH3_avalanche(acc); + } +} + + +/* ======= Long Keys ======= */ + +#define XXH_STRIPE_LEN 64 +#define XXH_SECRET_CONSUME_RATE 8 /* nb of secret bytes consumed at each accumulation */ +#define XXH_ACC_NB (XXH_STRIPE_LEN / sizeof(xxh_u64)) + +#ifdef XXH_OLD_NAMES +# define STRIPE_LEN XXH_STRIPE_LEN +# define ACC_NB XXH_ACC_NB +#endif + +XXH_FORCE_INLINE void XXH_writeLE64(void* dst, xxh_u64 v64) +{ + if (!XXH_CPU_LITTLE_ENDIAN) v64 = XXH_swap64(v64); + XXH_memcpy(dst, &v64, sizeof(v64)); +} + +/* Several intrinsic functions below are supposed to accept __int64 as argument, + * as documented in https://software.intel.com/sites/landingpage/IntrinsicsGuide/ . + * However, several environments do not define __int64 type, + * requiring a workaround. + */ +#if !defined (__VMS) \ + && (defined (__cplusplus) \ + || (defined (__STDC_VERSION__) && (__STDC_VERSION__ >= 199901L) /* C99 */) ) + typedef int64_t xxh_i64; +#else + /* the following type must have a width of 64-bit */ + typedef long long xxh_i64; +#endif + +/* + * XXH3_accumulate_512 is the tightest loop for long inputs, and it is the most optimized. + * + * It is a hardened version of UMAC, based off of FARSH's implementation. + * + * This was chosen because it adapts quite well to 32-bit, 64-bit, and SIMD + * implementations, and it is ridiculously fast. + * + * We harden it by mixing the original input to the accumulators as well as the product. + * + * This means that in the (relatively likely) case of a multiply by zero, the + * original input is preserved. + * + * On 128-bit inputs, we swap 64-bit pairs when we add the input to improve + * cross-pollination, as otherwise the upper and lower halves would be + * essentially independent. + * + * This doesn't matter on 64-bit hashes since they all get merged together in + * the end, so we skip the extra step. + * + * Both XXH3_64bits and XXH3_128bits use this subroutine. + */ + +#if (XXH_VECTOR == XXH_AVX512) \ + || (defined(XXH_DISPATCH_AVX512) && XXH_DISPATCH_AVX512 != 0) + +#ifndef XXH_TARGET_AVX512 +# define XXH_TARGET_AVX512 /* disable attribute target */ +#endif + +XXH_FORCE_INLINE XXH_TARGET_AVX512 void +XXH3_accumulate_512_avx512(void* XXH_RESTRICT acc, + const void* XXH_RESTRICT input, + const void* XXH_RESTRICT secret) +{ + __m512i* const xacc = (__m512i *) acc; + XXH_ASSERT((((size_t)acc) & 63) == 0); + XXH_STATIC_ASSERT(XXH_STRIPE_LEN == sizeof(__m512i)); + + { + /* data_vec = input[0]; */ + __m512i const data_vec = _mm512_loadu_si512 (input); + /* key_vec = secret[0]; */ + __m512i const key_vec = _mm512_loadu_si512 (secret); + /* data_key = data_vec ^ key_vec; */ + __m512i const data_key = _mm512_xor_si512 (data_vec, key_vec); + /* data_key_lo = data_key >> 32; */ + __m512i const data_key_lo = _mm512_shuffle_epi32 (data_key, (_MM_PERM_ENUM)_MM_SHUFFLE(0, 3, 0, 1)); + /* product = (data_key & 0xffffffff) * (data_key_lo & 0xffffffff); */ + __m512i const product = _mm512_mul_epu32 (data_key, data_key_lo); + /* xacc[0] += swap(data_vec); */ + __m512i const data_swap = _mm512_shuffle_epi32(data_vec, (_MM_PERM_ENUM)_MM_SHUFFLE(1, 0, 3, 2)); + __m512i const sum = _mm512_add_epi64(*xacc, data_swap); + /* xacc[0] += product; */ + *xacc = _mm512_add_epi64(product, sum); + } +} + +/* + * XXH3_scrambleAcc: Scrambles the accumulators to improve mixing. + * + * Multiplication isn't perfect, as explained by Google in HighwayHash: + * + * // Multiplication mixes/scrambles bytes 0-7 of the 64-bit result to + * // varying degrees. In descending order of goodness, bytes + * // 3 4 2 5 1 6 0 7 have quality 228 224 164 160 100 96 36 32. + * // As expected, the upper and lower bytes are much worse. + * + * Source: https://github.com/google/highwayhash/blob/0aaf66b/highwayhash/hh_avx2.h#L291 + * + * Since our algorithm uses a pseudorandom secret to add some variance into the + * mix, we don't need to (or want to) mix as often or as much as HighwayHash does. + * + * This isn't as tight as XXH3_accumulate, but still written in SIMD to avoid + * extraction. + * + * Both XXH3_64bits and XXH3_128bits use this subroutine. + */ + +XXH_FORCE_INLINE XXH_TARGET_AVX512 void +XXH3_scrambleAcc_avx512(void* XXH_RESTRICT acc, const void* XXH_RESTRICT secret) +{ + XXH_ASSERT((((size_t)acc) & 63) == 0); + XXH_STATIC_ASSERT(XXH_STRIPE_LEN == sizeof(__m512i)); + { __m512i* const xacc = (__m512i*) acc; + const __m512i prime32 = _mm512_set1_epi32((int)XXH_PRIME32_1); + + /* xacc[0] ^= (xacc[0] >> 47) */ + __m512i const acc_vec = *xacc; + __m512i const shifted = _mm512_srli_epi64 (acc_vec, 47); + __m512i const data_vec = _mm512_xor_si512 (acc_vec, shifted); + /* xacc[0] ^= secret; */ + __m512i const key_vec = _mm512_loadu_si512 (secret); + __m512i const data_key = _mm512_xor_si512 (data_vec, key_vec); + + /* xacc[0] *= XXH_PRIME32_1; */ + __m512i const data_key_hi = _mm512_shuffle_epi32 (data_key, (_MM_PERM_ENUM)_MM_SHUFFLE(0, 3, 0, 1)); + __m512i const prod_lo = _mm512_mul_epu32 (data_key, prime32); + __m512i const prod_hi = _mm512_mul_epu32 (data_key_hi, prime32); + *xacc = _mm512_add_epi64(prod_lo, _mm512_slli_epi64(prod_hi, 32)); + } +} + +XXH_FORCE_INLINE XXH_TARGET_AVX512 void +XXH3_initCustomSecret_avx512(void* XXH_RESTRICT customSecret, xxh_u64 seed64) +{ + XXH_STATIC_ASSERT((XXH_SECRET_DEFAULT_SIZE & 63) == 0); + XXH_STATIC_ASSERT(XXH_SEC_ALIGN == 64); + XXH_ASSERT(((size_t)customSecret & 63) == 0); + (void)(&XXH_writeLE64); + { int const nbRounds = XXH_SECRET_DEFAULT_SIZE / sizeof(__m512i); + __m512i const seed = _mm512_mask_set1_epi64(_mm512_set1_epi64((xxh_i64)seed64), 0xAA, (xxh_i64)(0U - seed64)); + + const __m512i* const src = (const __m512i*) ((const void*) XXH3_kSecret); + __m512i* const dest = ( __m512i*) customSecret; + int i; + XXH_ASSERT(((size_t)src & 63) == 0); /* control alignment */ + XXH_ASSERT(((size_t)dest & 63) == 0); + for (i=0; i < nbRounds; ++i) { + /* GCC has a bug, _mm512_stream_load_si512 accepts 'void*', not 'void const*', + * this will warn "discards 'const' qualifier". */ + union { + const __m512i* cp; + void* p; + } remote_const_void; + remote_const_void.cp = src + i; + dest[i] = _mm512_add_epi64(_mm512_stream_load_si512(remote_const_void.p), seed); + } } +} + +#endif + +#if (XXH_VECTOR == XXH_AVX2) \ + || (defined(XXH_DISPATCH_AVX2) && XXH_DISPATCH_AVX2 != 0) + +#ifndef XXH_TARGET_AVX2 +# define XXH_TARGET_AVX2 /* disable attribute target */ +#endif + +XXH_FORCE_INLINE XXH_TARGET_AVX2 void +XXH3_accumulate_512_avx2( void* XXH_RESTRICT acc, + const void* XXH_RESTRICT input, + const void* XXH_RESTRICT secret) +{ + XXH_ASSERT((((size_t)acc) & 31) == 0); + { __m256i* const xacc = (__m256i *) acc; + /* Unaligned. This is mainly for pointer arithmetic, and because + * _mm256_loadu_si256 requires a const __m256i * pointer for some reason. */ + const __m256i* const xinput = (const __m256i *) input; + /* Unaligned. This is mainly for pointer arithmetic, and because + * _mm256_loadu_si256 requires a const __m256i * pointer for some reason. */ + const __m256i* const xsecret = (const __m256i *) secret; + + size_t i; + for (i=0; i < XXH_STRIPE_LEN/sizeof(__m256i); i++) { + /* data_vec = xinput[i]; */ + __m256i const data_vec = _mm256_loadu_si256 (xinput+i); + /* key_vec = xsecret[i]; */ + __m256i const key_vec = _mm256_loadu_si256 (xsecret+i); + /* data_key = data_vec ^ key_vec; */ + __m256i const data_key = _mm256_xor_si256 (data_vec, key_vec); + /* data_key_lo = data_key >> 32; */ + __m256i const data_key_lo = _mm256_shuffle_epi32 (data_key, _MM_SHUFFLE(0, 3, 0, 1)); + /* product = (data_key & 0xffffffff) * (data_key_lo & 0xffffffff); */ + __m256i const product = _mm256_mul_epu32 (data_key, data_key_lo); + /* xacc[i] += swap(data_vec); */ + __m256i const data_swap = _mm256_shuffle_epi32(data_vec, _MM_SHUFFLE(1, 0, 3, 2)); + __m256i const sum = _mm256_add_epi64(xacc[i], data_swap); + /* xacc[i] += product; */ + xacc[i] = _mm256_add_epi64(product, sum); + } } +} + +XXH_FORCE_INLINE XXH_TARGET_AVX2 void +XXH3_scrambleAcc_avx2(void* XXH_RESTRICT acc, const void* XXH_RESTRICT secret) +{ + XXH_ASSERT((((size_t)acc) & 31) == 0); + { __m256i* const xacc = (__m256i*) acc; + /* Unaligned. This is mainly for pointer arithmetic, and because + * _mm256_loadu_si256 requires a const __m256i * pointer for some reason. */ + const __m256i* const xsecret = (const __m256i *) secret; + const __m256i prime32 = _mm256_set1_epi32((int)XXH_PRIME32_1); + + size_t i; + for (i=0; i < XXH_STRIPE_LEN/sizeof(__m256i); i++) { + /* xacc[i] ^= (xacc[i] >> 47) */ + __m256i const acc_vec = xacc[i]; + __m256i const shifted = _mm256_srli_epi64 (acc_vec, 47); + __m256i const data_vec = _mm256_xor_si256 (acc_vec, shifted); + /* xacc[i] ^= xsecret; */ + __m256i const key_vec = _mm256_loadu_si256 (xsecret+i); + __m256i const data_key = _mm256_xor_si256 (data_vec, key_vec); + + /* xacc[i] *= XXH_PRIME32_1; */ + __m256i const data_key_hi = _mm256_shuffle_epi32 (data_key, _MM_SHUFFLE(0, 3, 0, 1)); + __m256i const prod_lo = _mm256_mul_epu32 (data_key, prime32); + __m256i const prod_hi = _mm256_mul_epu32 (data_key_hi, prime32); + xacc[i] = _mm256_add_epi64(prod_lo, _mm256_slli_epi64(prod_hi, 32)); + } + } +} + +XXH_FORCE_INLINE XXH_TARGET_AVX2 void XXH3_initCustomSecret_avx2(void* XXH_RESTRICT customSecret, xxh_u64 seed64) +{ + XXH_STATIC_ASSERT((XXH_SECRET_DEFAULT_SIZE & 31) == 0); + XXH_STATIC_ASSERT((XXH_SECRET_DEFAULT_SIZE / sizeof(__m256i)) == 6); + XXH_STATIC_ASSERT(XXH_SEC_ALIGN <= 64); + (void)(&XXH_writeLE64); + XXH_PREFETCH(customSecret); + { __m256i const seed = _mm256_set_epi64x((xxh_i64)(0U - seed64), (xxh_i64)seed64, (xxh_i64)(0U - seed64), (xxh_i64)seed64); + + const __m256i* const src = (const __m256i*) ((const void*) XXH3_kSecret); + __m256i* dest = ( __m256i*) customSecret; + +# if defined(__GNUC__) || defined(__clang__) + /* + * On GCC & Clang, marking 'dest' as modified will cause the compiler: + * - do not extract the secret from sse registers in the internal loop + * - use less common registers, and avoid pushing these reg into stack + */ + XXH_COMPILER_GUARD(dest); +# endif + XXH_ASSERT(((size_t)src & 31) == 0); /* control alignment */ + XXH_ASSERT(((size_t)dest & 31) == 0); + + /* GCC -O2 need unroll loop manually */ + dest[0] = _mm256_add_epi64(_mm256_stream_load_si256(src+0), seed); + dest[1] = _mm256_add_epi64(_mm256_stream_load_si256(src+1), seed); + dest[2] = _mm256_add_epi64(_mm256_stream_load_si256(src+2), seed); + dest[3] = _mm256_add_epi64(_mm256_stream_load_si256(src+3), seed); + dest[4] = _mm256_add_epi64(_mm256_stream_load_si256(src+4), seed); + dest[5] = _mm256_add_epi64(_mm256_stream_load_si256(src+5), seed); + } +} + +#endif + +/* x86dispatch always generates SSE2 */ +#if (XXH_VECTOR == XXH_SSE2) || defined(XXH_X86DISPATCH) + +#ifndef XXH_TARGET_SSE2 +# define XXH_TARGET_SSE2 /* disable attribute target */ +#endif + +XXH_FORCE_INLINE XXH_TARGET_SSE2 void +XXH3_accumulate_512_sse2( void* XXH_RESTRICT acc, + const void* XXH_RESTRICT input, + const void* XXH_RESTRICT secret) +{ + /* SSE2 is just a half-scale version of the AVX2 version. */ + XXH_ASSERT((((size_t)acc) & 15) == 0); + { __m128i* const xacc = (__m128i *) acc; + /* Unaligned. This is mainly for pointer arithmetic, and because + * _mm_loadu_si128 requires a const __m128i * pointer for some reason. */ + const __m128i* const xinput = (const __m128i *) input; + /* Unaligned. This is mainly for pointer arithmetic, and because + * _mm_loadu_si128 requires a const __m128i * pointer for some reason. */ + const __m128i* const xsecret = (const __m128i *) secret; + + size_t i; + for (i=0; i < XXH_STRIPE_LEN/sizeof(__m128i); i++) { + /* data_vec = xinput[i]; */ + __m128i const data_vec = _mm_loadu_si128 (xinput+i); + /* key_vec = xsecret[i]; */ + __m128i const key_vec = _mm_loadu_si128 (xsecret+i); + /* data_key = data_vec ^ key_vec; */ + __m128i const data_key = _mm_xor_si128 (data_vec, key_vec); + /* data_key_lo = data_key >> 32; */ + __m128i const data_key_lo = _mm_shuffle_epi32 (data_key, _MM_SHUFFLE(0, 3, 0, 1)); + /* product = (data_key & 0xffffffff) * (data_key_lo & 0xffffffff); */ + __m128i const product = _mm_mul_epu32 (data_key, data_key_lo); + /* xacc[i] += swap(data_vec); */ + __m128i const data_swap = _mm_shuffle_epi32(data_vec, _MM_SHUFFLE(1,0,3,2)); + __m128i const sum = _mm_add_epi64(xacc[i], data_swap); + /* xacc[i] += product; */ + xacc[i] = _mm_add_epi64(product, sum); + } } +} + +XXH_FORCE_INLINE XXH_TARGET_SSE2 void +XXH3_scrambleAcc_sse2(void* XXH_RESTRICT acc, const void* XXH_RESTRICT secret) +{ + XXH_ASSERT((((size_t)acc) & 15) == 0); + { __m128i* const xacc = (__m128i*) acc; + /* Unaligned. This is mainly for pointer arithmetic, and because + * _mm_loadu_si128 requires a const __m128i * pointer for some reason. */ + const __m128i* const xsecret = (const __m128i *) secret; + const __m128i prime32 = _mm_set1_epi32((int)XXH_PRIME32_1); + + size_t i; + for (i=0; i < XXH_STRIPE_LEN/sizeof(__m128i); i++) { + /* xacc[i] ^= (xacc[i] >> 47) */ + __m128i const acc_vec = xacc[i]; + __m128i const shifted = _mm_srli_epi64 (acc_vec, 47); + __m128i const data_vec = _mm_xor_si128 (acc_vec, shifted); + /* xacc[i] ^= xsecret[i]; */ + __m128i const key_vec = _mm_loadu_si128 (xsecret+i); + __m128i const data_key = _mm_xor_si128 (data_vec, key_vec); + + /* xacc[i] *= XXH_PRIME32_1; */ + __m128i const data_key_hi = _mm_shuffle_epi32 (data_key, _MM_SHUFFLE(0, 3, 0, 1)); + __m128i const prod_lo = _mm_mul_epu32 (data_key, prime32); + __m128i const prod_hi = _mm_mul_epu32 (data_key_hi, prime32); + xacc[i] = _mm_add_epi64(prod_lo, _mm_slli_epi64(prod_hi, 32)); + } + } +} + +XXH_FORCE_INLINE XXH_TARGET_SSE2 void XXH3_initCustomSecret_sse2(void* XXH_RESTRICT customSecret, xxh_u64 seed64) +{ + XXH_STATIC_ASSERT((XXH_SECRET_DEFAULT_SIZE & 15) == 0); + (void)(&XXH_writeLE64); + { int const nbRounds = XXH_SECRET_DEFAULT_SIZE / sizeof(__m128i); + +# if defined(_MSC_VER) && defined(_M_IX86) && _MSC_VER < 1900 + /* MSVC 32bit mode does not support _mm_set_epi64x before 2015 */ + XXH_ALIGN(16) const xxh_i64 seed64x2[2] = { (xxh_i64)seed64, (xxh_i64)(0U - seed64) }; + __m128i const seed = _mm_load_si128((__m128i const*)seed64x2); +# else + __m128i const seed = _mm_set_epi64x((xxh_i64)(0U - seed64), (xxh_i64)seed64); +# endif + int i; + + const void* const src16 = XXH3_kSecret; + __m128i* dst16 = (__m128i*) customSecret; +# if defined(__GNUC__) || defined(__clang__) + /* + * On GCC & Clang, marking 'dest' as modified will cause the compiler: + * - do not extract the secret from sse registers in the internal loop + * - use less common registers, and avoid pushing these reg into stack + */ + XXH_COMPILER_GUARD(dst16); +# endif + XXH_ASSERT(((size_t)src16 & 15) == 0); /* control alignment */ + XXH_ASSERT(((size_t)dst16 & 15) == 0); + + for (i=0; i < nbRounds; ++i) { + dst16[i] = _mm_add_epi64(_mm_load_si128((const __m128i *)src16+i), seed); + } } +} + +#endif + +#if (XXH_VECTOR == XXH_NEON) + +XXH_FORCE_INLINE void +XXH3_accumulate_512_neon( void* XXH_RESTRICT acc, + const void* XXH_RESTRICT input, + const void* XXH_RESTRICT secret) +{ + XXH_ASSERT((((size_t)acc) & 15) == 0); + { + uint64x2_t* const xacc = (uint64x2_t *) acc; + /* We don't use a uint32x4_t pointer because it causes bus errors on ARMv7. */ + uint8_t const* const xinput = (const uint8_t *) input; + uint8_t const* const xsecret = (const uint8_t *) secret; + + size_t i; + for (i=0; i < XXH_STRIPE_LEN / sizeof(uint64x2_t); i++) { + /* data_vec = xinput[i]; */ + uint8x16_t data_vec = vld1q_u8(xinput + (i * 16)); + /* key_vec = xsecret[i]; */ + uint8x16_t key_vec = vld1q_u8(xsecret + (i * 16)); + uint64x2_t data_key; + uint32x2_t data_key_lo, data_key_hi; + /* xacc[i] += swap(data_vec); */ + uint64x2_t const data64 = vreinterpretq_u64_u8(data_vec); + uint64x2_t const swapped = vextq_u64(data64, data64, 1); + xacc[i] = vaddq_u64 (xacc[i], swapped); + /* data_key = data_vec ^ key_vec; */ + data_key = vreinterpretq_u64_u8(veorq_u8(data_vec, key_vec)); + /* data_key_lo = (uint32x2_t) (data_key & 0xFFFFFFFF); + * data_key_hi = (uint32x2_t) (data_key >> 32); + * data_key = UNDEFINED; */ + XXH_SPLIT_IN_PLACE(data_key, data_key_lo, data_key_hi); + /* xacc[i] += (uint64x2_t) data_key_lo * (uint64x2_t) data_key_hi; */ + xacc[i] = vmlal_u32 (xacc[i], data_key_lo, data_key_hi); + + } + } +} + +XXH_FORCE_INLINE void +XXH3_scrambleAcc_neon(void* XXH_RESTRICT acc, const void* XXH_RESTRICT secret) +{ + XXH_ASSERT((((size_t)acc) & 15) == 0); + + { uint64x2_t* xacc = (uint64x2_t*) acc; + uint8_t const* xsecret = (uint8_t const*) secret; + uint32x2_t prime = vdup_n_u32 (XXH_PRIME32_1); + + size_t i; + for (i=0; i < XXH_STRIPE_LEN/sizeof(uint64x2_t); i++) { + /* xacc[i] ^= (xacc[i] >> 47); */ + uint64x2_t acc_vec = xacc[i]; + uint64x2_t shifted = vshrq_n_u64 (acc_vec, 47); + uint64x2_t data_vec = veorq_u64 (acc_vec, shifted); + + /* xacc[i] ^= xsecret[i]; */ + uint8x16_t key_vec = vld1q_u8 (xsecret + (i * 16)); + uint64x2_t data_key = veorq_u64 (data_vec, vreinterpretq_u64_u8(key_vec)); + + /* xacc[i] *= XXH_PRIME32_1 */ + uint32x2_t data_key_lo, data_key_hi; + /* data_key_lo = (uint32x2_t) (xacc[i] & 0xFFFFFFFF); + * data_key_hi = (uint32x2_t) (xacc[i] >> 32); + * xacc[i] = UNDEFINED; */ + XXH_SPLIT_IN_PLACE(data_key, data_key_lo, data_key_hi); + { /* + * prod_hi = (data_key >> 32) * XXH_PRIME32_1; + * + * Avoid vmul_u32 + vshll_n_u32 since Clang 6 and 7 will + * incorrectly "optimize" this: + * tmp = vmul_u32(vmovn_u64(a), vmovn_u64(b)); + * shifted = vshll_n_u32(tmp, 32); + * to this: + * tmp = "vmulq_u64"(a, b); // no such thing! + * shifted = vshlq_n_u64(tmp, 32); + * + * However, unlike SSE, Clang lacks a 64-bit multiply routine + * for NEON, and it scalarizes two 64-bit multiplies instead. + * + * vmull_u32 has the same timing as vmul_u32, and it avoids + * this bug completely. + * See https://bugs.llvm.org/show_bug.cgi?id=39967 + */ + uint64x2_t prod_hi = vmull_u32 (data_key_hi, prime); + /* xacc[i] = prod_hi << 32; */ + xacc[i] = vshlq_n_u64(prod_hi, 32); + /* xacc[i] += (prod_hi & 0xFFFFFFFF) * XXH_PRIME32_1; */ + xacc[i] = vmlal_u32(xacc[i], data_key_lo, prime); + } + } } +} + +#endif + +#if (XXH_VECTOR == XXH_VSX) + +XXH_FORCE_INLINE void +XXH3_accumulate_512_vsx( void* XXH_RESTRICT acc, + const void* XXH_RESTRICT input, + const void* XXH_RESTRICT secret) +{ + /* presumed aligned */ + unsigned long long* const xacc = (unsigned long long*) acc; + xxh_u64x2 const* const xinput = (xxh_u64x2 const*) input; /* no alignment restriction */ + xxh_u64x2 const* const xsecret = (xxh_u64x2 const*) secret; /* no alignment restriction */ + xxh_u64x2 const v32 = { 32, 32 }; + size_t i; + for (i = 0; i < XXH_STRIPE_LEN / sizeof(xxh_u64x2); i++) { + /* data_vec = xinput[i]; */ + xxh_u64x2 const data_vec = XXH_vec_loadu(xinput + i); + /* key_vec = xsecret[i]; */ + xxh_u64x2 const key_vec = XXH_vec_loadu(xsecret + i); + xxh_u64x2 const data_key = data_vec ^ key_vec; + /* shuffled = (data_key << 32) | (data_key >> 32); */ + xxh_u32x4 const shuffled = (xxh_u32x4)vec_rl(data_key, v32); + /* product = ((xxh_u64x2)data_key & 0xFFFFFFFF) * ((xxh_u64x2)shuffled & 0xFFFFFFFF); */ + xxh_u64x2 const product = XXH_vec_mulo((xxh_u32x4)data_key, shuffled); + /* acc_vec = xacc[i]; */ + xxh_u64x2 acc_vec = vec_xl(0, xacc + 2 * i); + acc_vec += product; + + /* swap high and low halves */ +#ifdef __s390x__ + acc_vec += vec_permi(data_vec, data_vec, 2); +#else + acc_vec += vec_xxpermdi(data_vec, data_vec, 2); +#endif + /* xacc[i] = acc_vec; */ + vec_xst(acc_vec, 0, xacc + 2 * i); + } +} + +XXH_FORCE_INLINE void +XXH3_scrambleAcc_vsx(void* XXH_RESTRICT acc, const void* XXH_RESTRICT secret) +{ + XXH_ASSERT((((size_t)acc) & 15) == 0); + + { xxh_u64x2* const xacc = (xxh_u64x2*) acc; + const xxh_u64x2* const xsecret = (const xxh_u64x2*) secret; + /* constants */ + xxh_u64x2 const v32 = { 32, 32 }; + xxh_u64x2 const v47 = { 47, 47 }; + xxh_u32x4 const prime = { XXH_PRIME32_1, XXH_PRIME32_1, XXH_PRIME32_1, XXH_PRIME32_1 }; + size_t i; + for (i = 0; i < XXH_STRIPE_LEN / sizeof(xxh_u64x2); i++) { + /* xacc[i] ^= (xacc[i] >> 47); */ + xxh_u64x2 const acc_vec = xacc[i]; + xxh_u64x2 const data_vec = acc_vec ^ (acc_vec >> v47); + + /* xacc[i] ^= xsecret[i]; */ + xxh_u64x2 const key_vec = XXH_vec_loadu(xsecret + i); + xxh_u64x2 const data_key = data_vec ^ key_vec; + + /* xacc[i] *= XXH_PRIME32_1 */ + /* prod_lo = ((xxh_u64x2)data_key & 0xFFFFFFFF) * ((xxh_u64x2)prime & 0xFFFFFFFF); */ + xxh_u64x2 const prod_even = XXH_vec_mule((xxh_u32x4)data_key, prime); + /* prod_hi = ((xxh_u64x2)data_key >> 32) * ((xxh_u64x2)prime >> 32); */ + xxh_u64x2 const prod_odd = XXH_vec_mulo((xxh_u32x4)data_key, prime); + xacc[i] = prod_odd + (prod_even << v32); + } } +} + +#endif + +/* scalar variants - universal */ + +XXH_FORCE_INLINE void +XXH3_accumulate_512_scalar(void* XXH_RESTRICT acc, + const void* XXH_RESTRICT input, + const void* XXH_RESTRICT secret) +{ + xxh_u64* const xacc = (xxh_u64*) acc; /* presumed aligned */ + const xxh_u8* const xinput = (const xxh_u8*) input; /* no alignment restriction */ + const xxh_u8* const xsecret = (const xxh_u8*) secret; /* no alignment restriction */ + size_t i; + XXH_ASSERT(((size_t)acc & (XXH_ACC_ALIGN-1)) == 0); + for (i=0; i < XXH_ACC_NB; i++) { + xxh_u64 const data_val = XXH_readLE64(xinput + 8*i); + xxh_u64 const data_key = data_val ^ XXH_readLE64(xsecret + i*8); + xacc[i ^ 1] += data_val; /* swap adjacent lanes */ + xacc[i] += XXH_mult32to64(data_key & 0xFFFFFFFF, data_key >> 32); + } +} + +XXH_FORCE_INLINE void +XXH3_scrambleAcc_scalar(void* XXH_RESTRICT acc, const void* XXH_RESTRICT secret) +{ + xxh_u64* const xacc = (xxh_u64*) acc; /* presumed aligned */ + const xxh_u8* const xsecret = (const xxh_u8*) secret; /* no alignment restriction */ + size_t i; + XXH_ASSERT((((size_t)acc) & (XXH_ACC_ALIGN-1)) == 0); + for (i=0; i < XXH_ACC_NB; i++) { + xxh_u64 const key64 = XXH_readLE64(xsecret + 8*i); + xxh_u64 acc64 = xacc[i]; + acc64 = XXH_xorshift64(acc64, 47); + acc64 ^= key64; + acc64 *= XXH_PRIME32_1; + xacc[i] = acc64; + } +} + +XXH_FORCE_INLINE void +XXH3_initCustomSecret_scalar(void* XXH_RESTRICT customSecret, xxh_u64 seed64) +{ + /* + * We need a separate pointer for the hack below, + * which requires a non-const pointer. + * Any decent compiler will optimize this out otherwise. + */ + const xxh_u8* kSecretPtr = XXH3_kSecret; + XXH_STATIC_ASSERT((XXH_SECRET_DEFAULT_SIZE & 15) == 0); + +#if defined(__clang__) && defined(__aarch64__) + /* + * UGLY HACK: + * Clang generates a bunch of MOV/MOVK pairs for aarch64, and they are + * placed sequentially, in order, at the top of the unrolled loop. + * + * While MOVK is great for generating constants (2 cycles for a 64-bit + * constant compared to 4 cycles for LDR), long MOVK chains stall the + * integer pipelines: + * I L S + * MOVK + * MOVK + * MOVK + * MOVK + * ADD + * SUB STR + * STR + * By forcing loads from memory (as the asm line causes Clang to assume + * that XXH3_kSecretPtr has been changed), the pipelines are used more + * efficiently: + * I L S + * LDR + * ADD LDR + * SUB STR + * STR + * XXH3_64bits_withSeed, len == 256, Snapdragon 835 + * without hack: 2654.4 MB/s + * with hack: 3202.9 MB/s + */ + XXH_COMPILER_GUARD(kSecretPtr); +#endif + /* + * Note: in debug mode, this overrides the asm optimization + * and Clang will emit MOVK chains again. + */ + XXH_ASSERT(kSecretPtr == XXH3_kSecret); + + { int const nbRounds = XXH_SECRET_DEFAULT_SIZE / 16; + int i; + for (i=0; i < nbRounds; i++) { + /* + * The asm hack causes Clang to assume that kSecretPtr aliases with + * customSecret, and on aarch64, this prevented LDP from merging two + * loads together for free. Putting the loads together before the stores + * properly generates LDP. + */ + xxh_u64 lo = XXH_readLE64(kSecretPtr + 16*i) + seed64; + xxh_u64 hi = XXH_readLE64(kSecretPtr + 16*i + 8) - seed64; + XXH_writeLE64((xxh_u8*)customSecret + 16*i, lo); + XXH_writeLE64((xxh_u8*)customSecret + 16*i + 8, hi); + } } +} + + +typedef void (*XXH3_f_accumulate_512)(void* XXH_RESTRICT, const void*, const void*); +typedef void (*XXH3_f_scrambleAcc)(void* XXH_RESTRICT, const void*); +typedef void (*XXH3_f_initCustomSecret)(void* XXH_RESTRICT, xxh_u64); + + +#if (XXH_VECTOR == XXH_AVX512) + +#define XXH3_accumulate_512 XXH3_accumulate_512_avx512 +#define XXH3_scrambleAcc XXH3_scrambleAcc_avx512 +#define XXH3_initCustomSecret XXH3_initCustomSecret_avx512 + +#elif (XXH_VECTOR == XXH_AVX2) + +#define XXH3_accumulate_512 XXH3_accumulate_512_avx2 +#define XXH3_scrambleAcc XXH3_scrambleAcc_avx2 +#define XXH3_initCustomSecret XXH3_initCustomSecret_avx2 + +#elif (XXH_VECTOR == XXH_SSE2) + +#define XXH3_accumulate_512 XXH3_accumulate_512_sse2 +#define XXH3_scrambleAcc XXH3_scrambleAcc_sse2 +#define XXH3_initCustomSecret XXH3_initCustomSecret_sse2 + +#elif (XXH_VECTOR == XXH_NEON) + +#define XXH3_accumulate_512 XXH3_accumulate_512_neon +#define XXH3_scrambleAcc XXH3_scrambleAcc_neon +#define XXH3_initCustomSecret XXH3_initCustomSecret_scalar + +#elif (XXH_VECTOR == XXH_VSX) + +#define XXH3_accumulate_512 XXH3_accumulate_512_vsx +#define XXH3_scrambleAcc XXH3_scrambleAcc_vsx +#define XXH3_initCustomSecret XXH3_initCustomSecret_scalar + +#else /* scalar */ + +#define XXH3_accumulate_512 XXH3_accumulate_512_scalar +#define XXH3_scrambleAcc XXH3_scrambleAcc_scalar +#define XXH3_initCustomSecret XXH3_initCustomSecret_scalar + +#endif + + + +#ifndef XXH_PREFETCH_DIST +# ifdef __clang__ +# define XXH_PREFETCH_DIST 320 +# else +# if (XXH_VECTOR == XXH_AVX512) +# define XXH_PREFETCH_DIST 512 +# else +# define XXH_PREFETCH_DIST 384 +# endif +# endif /* __clang__ */ +#endif /* XXH_PREFETCH_DIST */ + +/* + * XXH3_accumulate() + * Loops over XXH3_accumulate_512(). + * Assumption: nbStripes will not overflow the secret size + */ +XXH_FORCE_INLINE void +XXH3_accumulate( xxh_u64* XXH_RESTRICT acc, + const xxh_u8* XXH_RESTRICT input, + const xxh_u8* XXH_RESTRICT secret, + size_t nbStripes, + XXH3_f_accumulate_512 f_acc512) +{ + size_t n; + for (n = 0; n < nbStripes; n++ ) { + const xxh_u8* const in = input + n*XXH_STRIPE_LEN; + XXH_PREFETCH(in + XXH_PREFETCH_DIST); + f_acc512(acc, + in, + secret + n*XXH_SECRET_CONSUME_RATE); + } +} + +XXH_FORCE_INLINE void +XXH3_hashLong_internal_loop(xxh_u64* XXH_RESTRICT acc, + const xxh_u8* XXH_RESTRICT input, size_t len, + const xxh_u8* XXH_RESTRICT secret, size_t secretSize, + XXH3_f_accumulate_512 f_acc512, + XXH3_f_scrambleAcc f_scramble) +{ + size_t const nbStripesPerBlock = (secretSize - XXH_STRIPE_LEN) / XXH_SECRET_CONSUME_RATE; + size_t const block_len = XXH_STRIPE_LEN * nbStripesPerBlock; + size_t const nb_blocks = (len - 1) / block_len; + + size_t n; + + XXH_ASSERT(secretSize >= XXH3_SECRET_SIZE_MIN); + + for (n = 0; n < nb_blocks; n++) { + XXH3_accumulate(acc, input + n*block_len, secret, nbStripesPerBlock, f_acc512); + f_scramble(acc, secret + secretSize - XXH_STRIPE_LEN); + } + + /* last partial block */ + XXH_ASSERT(len > XXH_STRIPE_LEN); + { size_t const nbStripes = ((len - 1) - (block_len * nb_blocks)) / XXH_STRIPE_LEN; + XXH_ASSERT(nbStripes <= (secretSize / XXH_SECRET_CONSUME_RATE)); + XXH3_accumulate(acc, input + nb_blocks*block_len, secret, nbStripes, f_acc512); + + /* last stripe */ + { const xxh_u8* const p = input + len - XXH_STRIPE_LEN; +#define XXH_SECRET_LASTACC_START 7 /* not aligned on 8, last secret is different from acc & scrambler */ + f_acc512(acc, p, secret + secretSize - XXH_STRIPE_LEN - XXH_SECRET_LASTACC_START); + } } +} + +XXH_FORCE_INLINE xxh_u64 +XXH3_mix2Accs(const xxh_u64* XXH_RESTRICT acc, const xxh_u8* XXH_RESTRICT secret) +{ + return XXH3_mul128_fold64( + acc[0] ^ XXH_readLE64(secret), + acc[1] ^ XXH_readLE64(secret+8) ); +} + +static XXH64_hash_t +XXH3_mergeAccs(const xxh_u64* XXH_RESTRICT acc, const xxh_u8* XXH_RESTRICT secret, xxh_u64 start) +{ + xxh_u64 result64 = start; + size_t i = 0; + + for (i = 0; i < 4; i++) { + result64 += XXH3_mix2Accs(acc+2*i, secret + 16*i); +#if defined(__clang__) /* Clang */ \ + && (defined(__arm__) || defined(__thumb__)) /* ARMv7 */ \ + && (defined(__ARM_NEON) || defined(__ARM_NEON__)) /* NEON */ \ + && !defined(XXH_ENABLE_AUTOVECTORIZE) /* Define to disable */ + /* + * UGLY HACK: + * Prevent autovectorization on Clang ARMv7-a. Exact same problem as + * the one in XXH3_len_129to240_64b. Speeds up shorter keys > 240b. + * XXH3_64bits, len == 256, Snapdragon 835: + * without hack: 2063.7 MB/s + * with hack: 2560.7 MB/s + */ + XXH_COMPILER_GUARD(result64); +#endif + } + + return XXH3_avalanche(result64); +} + +#define XXH3_INIT_ACC { XXH_PRIME32_3, XXH_PRIME64_1, XXH_PRIME64_2, XXH_PRIME64_3, \ + XXH_PRIME64_4, XXH_PRIME32_2, XXH_PRIME64_5, XXH_PRIME32_1 } + +XXH_FORCE_INLINE XXH64_hash_t +XXH3_hashLong_64b_internal(const void* XXH_RESTRICT input, size_t len, + const void* XXH_RESTRICT secret, size_t secretSize, + XXH3_f_accumulate_512 f_acc512, + XXH3_f_scrambleAcc f_scramble) +{ + XXH_ALIGN(XXH_ACC_ALIGN) xxh_u64 acc[XXH_ACC_NB] = XXH3_INIT_ACC; + + XXH3_hashLong_internal_loop(acc, (const xxh_u8*)input, len, (const xxh_u8*)secret, secretSize, f_acc512, f_scramble); + + /* converge into final hash */ + XXH_STATIC_ASSERT(sizeof(acc) == 64); + /* do not align on 8, so that the secret is different from the accumulator */ +#define XXH_SECRET_MERGEACCS_START 11 + XXH_ASSERT(secretSize >= sizeof(acc) + XXH_SECRET_MERGEACCS_START); + return XXH3_mergeAccs(acc, (const xxh_u8*)secret + XXH_SECRET_MERGEACCS_START, (xxh_u64)len * XXH_PRIME64_1); +} + +/* + * It's important for performance to transmit secret's size (when it's static) + * so that the compiler can properly optimize the vectorized loop. + * This makes a big performance difference for "medium" keys (<1 KB) when using AVX instruction set. + */ +XXH_FORCE_INLINE XXH64_hash_t +XXH3_hashLong_64b_withSecret(const void* XXH_RESTRICT input, size_t len, + XXH64_hash_t seed64, const xxh_u8* XXH_RESTRICT secret, size_t secretLen) +{ + (void)seed64; + return XXH3_hashLong_64b_internal(input, len, secret, secretLen, XXH3_accumulate_512, XXH3_scrambleAcc); +} + +/* + * It's preferable for performance that XXH3_hashLong is not inlined, + * as it results in a smaller function for small data, easier to the instruction cache. + * Note that inside this no_inline function, we do inline the internal loop, + * and provide a statically defined secret size to allow optimization of vector loop. + */ +XXH_NO_INLINE XXH64_hash_t +XXH3_hashLong_64b_default(const void* XXH_RESTRICT input, size_t len, + XXH64_hash_t seed64, const xxh_u8* XXH_RESTRICT secret, size_t secretLen) +{ + (void)seed64; (void)secret; (void)secretLen; + return XXH3_hashLong_64b_internal(input, len, XXH3_kSecret, sizeof(XXH3_kSecret), XXH3_accumulate_512, XXH3_scrambleAcc); +} + +/* + * XXH3_hashLong_64b_withSeed(): + * Generate a custom key based on alteration of default XXH3_kSecret with the seed, + * and then use this key for long mode hashing. + * + * This operation is decently fast but nonetheless costs a little bit of time. + * Try to avoid it whenever possible (typically when seed==0). + * + * It's important for performance that XXH3_hashLong is not inlined. Not sure + * why (uop cache maybe?), but the difference is large and easily measurable. + */ +XXH_FORCE_INLINE XXH64_hash_t +XXH3_hashLong_64b_withSeed_internal(const void* input, size_t len, + XXH64_hash_t seed, + XXH3_f_accumulate_512 f_acc512, + XXH3_f_scrambleAcc f_scramble, + XXH3_f_initCustomSecret f_initSec) +{ + if (seed == 0) + return XXH3_hashLong_64b_internal(input, len, + XXH3_kSecret, sizeof(XXH3_kSecret), + f_acc512, f_scramble); + { XXH_ALIGN(XXH_SEC_ALIGN) xxh_u8 secret[XXH_SECRET_DEFAULT_SIZE]; + f_initSec(secret, seed); + return XXH3_hashLong_64b_internal(input, len, secret, sizeof(secret), + f_acc512, f_scramble); + } +} + +/* + * It's important for performance that XXH3_hashLong is not inlined. + */ +XXH_NO_INLINE XXH64_hash_t +XXH3_hashLong_64b_withSeed(const void* input, size_t len, + XXH64_hash_t seed, const xxh_u8* secret, size_t secretLen) +{ + (void)secret; (void)secretLen; + return XXH3_hashLong_64b_withSeed_internal(input, len, seed, + XXH3_accumulate_512, XXH3_scrambleAcc, XXH3_initCustomSecret); +} + + +typedef XXH64_hash_t (*XXH3_hashLong64_f)(const void* XXH_RESTRICT, size_t, + XXH64_hash_t, const xxh_u8* XXH_RESTRICT, size_t); + +XXH_FORCE_INLINE XXH64_hash_t +XXH3_64bits_internal(const void* XXH_RESTRICT input, size_t len, + XXH64_hash_t seed64, const void* XXH_RESTRICT secret, size_t secretLen, + XXH3_hashLong64_f f_hashLong) +{ + XXH_ASSERT(secretLen >= XXH3_SECRET_SIZE_MIN); + /* + * If an action is to be taken if `secretLen` condition is not respected, + * it should be done here. + * For now, it's a contract pre-condition. + * Adding a check and a branch here would cost performance at every hash. + * Also, note that function signature doesn't offer room to return an error. + */ + if (len <= 16) + return XXH3_len_0to16_64b((const xxh_u8*)input, len, (const xxh_u8*)secret, seed64); + if (len <= 128) + return XXH3_len_17to128_64b((const xxh_u8*)input, len, (const xxh_u8*)secret, secretLen, seed64); + if (len <= XXH3_MIDSIZE_MAX) + return XXH3_len_129to240_64b((const xxh_u8*)input, len, (const xxh_u8*)secret, secretLen, seed64); + return f_hashLong(input, len, seed64, (const xxh_u8*)secret, secretLen); +} + + +/* === Public entry point === */ + +/*! @ingroup xxh3_family */ +XXH_PUBLIC_API XXH64_hash_t XXH3_64bits(const void* input, size_t len) +{ + return XXH3_64bits_internal(input, len, 0, XXH3_kSecret, sizeof(XXH3_kSecret), XXH3_hashLong_64b_default); +} + +/*! @ingroup xxh3_family */ +XXH_PUBLIC_API XXH64_hash_t +XXH3_64bits_withSecret(const void* input, size_t len, const void* secret, size_t secretSize) +{ + return XXH3_64bits_internal(input, len, 0, secret, secretSize, XXH3_hashLong_64b_withSecret); +} + +/*! @ingroup xxh3_family */ +XXH_PUBLIC_API XXH64_hash_t +XXH3_64bits_withSeed(const void* input, size_t len, XXH64_hash_t seed) +{ + return XXH3_64bits_internal(input, len, seed, XXH3_kSecret, sizeof(XXH3_kSecret), XXH3_hashLong_64b_withSeed); +} + +XXH_PUBLIC_API XXH64_hash_t +XXH3_64bits_withSecretandSeed(const void* input, size_t len, const void* secret, size_t secretSize, XXH64_hash_t seed) +{ + if (len <= XXH3_MIDSIZE_MAX) + return XXH3_64bits_internal(input, len, seed, XXH3_kSecret, sizeof(XXH3_kSecret), NULL); + return XXH3_hashLong_64b_withSecret(input, len, seed, (const xxh_u8*)secret, secretSize); +} + + +/* === XXH3 streaming === */ + +/* + * Malloc's a pointer that is always aligned to align. + * + * This must be freed with `XXH_alignedFree()`. + * + * malloc typically guarantees 16 byte alignment on 64-bit systems and 8 byte + * alignment on 32-bit. This isn't enough for the 32 byte aligned loads in AVX2 + * or on 32-bit, the 16 byte aligned loads in SSE2 and NEON. + * + * This underalignment previously caused a rather obvious crash which went + * completely unnoticed due to XXH3_createState() not actually being tested. + * Credit to RedSpah for noticing this bug. + * + * The alignment is done manually: Functions like posix_memalign or _mm_malloc + * are avoided: To maintain portability, we would have to write a fallback + * like this anyways, and besides, testing for the existence of library + * functions without relying on external build tools is impossible. + * + * The method is simple: Overallocate, manually align, and store the offset + * to the original behind the returned pointer. + * + * Align must be a power of 2 and 8 <= align <= 128. + */ +static void* XXH_alignedMalloc(size_t s, size_t align) +{ + XXH_ASSERT(align <= 128 && align >= 8); /* range check */ + XXH_ASSERT((align & (align-1)) == 0); /* power of 2 */ + XXH_ASSERT(s != 0 && s < (s + align)); /* empty/overflow */ + { /* Overallocate to make room for manual realignment and an offset byte */ + xxh_u8* base = (xxh_u8*)XXH_malloc(s + align); + if (base != NULL) { + /* + * Get the offset needed to align this pointer. + * + * Even if the returned pointer is aligned, there will always be + * at least one byte to store the offset to the original pointer. + */ + size_t offset = align - ((size_t)base & (align - 1)); /* base % align */ + /* Add the offset for the now-aligned pointer */ + xxh_u8* ptr = base + offset; + + XXH_ASSERT((size_t)ptr % align == 0); + + /* Store the offset immediately before the returned pointer. */ + ptr[-1] = (xxh_u8)offset; + return ptr; + } + return NULL; + } +} +/* + * Frees an aligned pointer allocated by XXH_alignedMalloc(). Don't pass + * normal malloc'd pointers, XXH_alignedMalloc has a specific data layout. + */ +static void XXH_alignedFree(void* p) +{ + if (p != NULL) { + xxh_u8* ptr = (xxh_u8*)p; + /* Get the offset byte we added in XXH_malloc. */ + xxh_u8 offset = ptr[-1]; + /* Free the original malloc'd pointer */ + xxh_u8* base = ptr - offset; + XXH_free(base); + } +} +/*! @ingroup xxh3_family */ +XXH_PUBLIC_API XXH3_state_t* XXH3_createState(void) +{ + XXH3_state_t* const state = (XXH3_state_t*)XXH_alignedMalloc(sizeof(XXH3_state_t), 64); + if (state==NULL) return NULL; + XXH3_INITSTATE(state); + return state; +} + +/*! @ingroup xxh3_family */ +XXH_PUBLIC_API XXH_errorcode XXH3_freeState(XXH3_state_t* statePtr) +{ + XXH_alignedFree(statePtr); + return XXH_OK; +} + +/*! @ingroup xxh3_family */ +XXH_PUBLIC_API void +XXH3_copyState(XXH3_state_t* dst_state, const XXH3_state_t* src_state) +{ + XXH_memcpy(dst_state, src_state, sizeof(*dst_state)); +} + +static void +XXH3_reset_internal(XXH3_state_t* statePtr, + XXH64_hash_t seed, + const void* secret, size_t secretSize) +{ + size_t const initStart = offsetof(XXH3_state_t, bufferedSize); + size_t const initLength = offsetof(XXH3_state_t, nbStripesPerBlock) - initStart; + XXH_ASSERT(offsetof(XXH3_state_t, nbStripesPerBlock) > initStart); + XXH_ASSERT(statePtr != NULL); + /* set members from bufferedSize to nbStripesPerBlock (excluded) to 0 */ + memset((char*)statePtr + initStart, 0, initLength); + statePtr->acc[0] = XXH_PRIME32_3; + statePtr->acc[1] = XXH_PRIME64_1; + statePtr->acc[2] = XXH_PRIME64_2; + statePtr->acc[3] = XXH_PRIME64_3; + statePtr->acc[4] = XXH_PRIME64_4; + statePtr->acc[5] = XXH_PRIME32_2; + statePtr->acc[6] = XXH_PRIME64_5; + statePtr->acc[7] = XXH_PRIME32_1; + statePtr->seed = seed; + statePtr->useSeed = (seed != 0); + statePtr->extSecret = (const unsigned char*)secret; + XXH_ASSERT(secretSize >= XXH3_SECRET_SIZE_MIN); + statePtr->secretLimit = secretSize - XXH_STRIPE_LEN; + statePtr->nbStripesPerBlock = statePtr->secretLimit / XXH_SECRET_CONSUME_RATE; +} + +/*! @ingroup xxh3_family */ +XXH_PUBLIC_API XXH_errorcode +XXH3_64bits_reset(XXH3_state_t* statePtr) +{ + if (statePtr == NULL) return XXH_ERROR; + XXH3_reset_internal(statePtr, 0, XXH3_kSecret, XXH_SECRET_DEFAULT_SIZE); + return XXH_OK; +} + +/*! @ingroup xxh3_family */ +XXH_PUBLIC_API XXH_errorcode +XXH3_64bits_reset_withSecret(XXH3_state_t* statePtr, const void* secret, size_t secretSize) +{ + if (statePtr == NULL) return XXH_ERROR; + XXH3_reset_internal(statePtr, 0, secret, secretSize); + if (secret == NULL) return XXH_ERROR; + if (secretSize < XXH3_SECRET_SIZE_MIN) return XXH_ERROR; + return XXH_OK; +} + +/*! @ingroup xxh3_family */ +XXH_PUBLIC_API XXH_errorcode +XXH3_64bits_reset_withSeed(XXH3_state_t* statePtr, XXH64_hash_t seed) +{ + if (statePtr == NULL) return XXH_ERROR; + if (seed==0) return XXH3_64bits_reset(statePtr); + if ((seed != statePtr->seed) || (statePtr->extSecret != NULL)) + XXH3_initCustomSecret(statePtr->customSecret, seed); + XXH3_reset_internal(statePtr, seed, NULL, XXH_SECRET_DEFAULT_SIZE); + return XXH_OK; +} + +/*! @ingroup xxh3_family */ +XXH_PUBLIC_API XXH_errorcode +XXH3_64bits_reset_withSecretandSeed(XXH3_state_t* statePtr, const void* secret, size_t secretSize, XXH64_hash_t seed64) +{ + if (statePtr == NULL) return XXH_ERROR; + if (secret == NULL) return XXH_ERROR; + if (secretSize < XXH3_SECRET_SIZE_MIN) return XXH_ERROR; + XXH3_reset_internal(statePtr, seed64, secret, secretSize); + statePtr->useSeed = 1; /* always, even if seed64==0 */ + return XXH_OK; +} + +/* Note : when XXH3_consumeStripes() is invoked, + * there must be a guarantee that at least one more byte must be consumed from input + * so that the function can blindly consume all stripes using the "normal" secret segment */ +XXH_FORCE_INLINE void +XXH3_consumeStripes(xxh_u64* XXH_RESTRICT acc, + size_t* XXH_RESTRICT nbStripesSoFarPtr, size_t nbStripesPerBlock, + const xxh_u8* XXH_RESTRICT input, size_t nbStripes, + const xxh_u8* XXH_RESTRICT secret, size_t secretLimit, + XXH3_f_accumulate_512 f_acc512, + XXH3_f_scrambleAcc f_scramble) +{ + XXH_ASSERT(nbStripes <= nbStripesPerBlock); /* can handle max 1 scramble per invocation */ + XXH_ASSERT(*nbStripesSoFarPtr < nbStripesPerBlock); + if (nbStripesPerBlock - *nbStripesSoFarPtr <= nbStripes) { + /* need a scrambling operation */ + size_t const nbStripesToEndofBlock = nbStripesPerBlock - *nbStripesSoFarPtr; + size_t const nbStripesAfterBlock = nbStripes - nbStripesToEndofBlock; + XXH3_accumulate(acc, input, secret + nbStripesSoFarPtr[0] * XXH_SECRET_CONSUME_RATE, nbStripesToEndofBlock, f_acc512); + f_scramble(acc, secret + secretLimit); + XXH3_accumulate(acc, input + nbStripesToEndofBlock * XXH_STRIPE_LEN, secret, nbStripesAfterBlock, f_acc512); + *nbStripesSoFarPtr = nbStripesAfterBlock; + } else { + XXH3_accumulate(acc, input, secret + nbStripesSoFarPtr[0] * XXH_SECRET_CONSUME_RATE, nbStripes, f_acc512); + *nbStripesSoFarPtr += nbStripes; + } +} + +#ifndef XXH3_STREAM_USE_STACK +# ifndef __clang__ /* clang doesn't need additional stack space */ +# define XXH3_STREAM_USE_STACK 1 +# endif +#endif +/* + * Both XXH3_64bits_update and XXH3_128bits_update use this routine. + */ +XXH_FORCE_INLINE XXH_errorcode +XXH3_update(XXH3_state_t* XXH_RESTRICT const state, + const xxh_u8* XXH_RESTRICT input, size_t len, + XXH3_f_accumulate_512 f_acc512, + XXH3_f_scrambleAcc f_scramble) +{ + if (input==NULL) { + XXH_ASSERT(len == 0); + return XXH_OK; + } + + XXH_ASSERT(state != NULL); + { const xxh_u8* const bEnd = input + len; + const unsigned char* const secret = (state->extSecret == NULL) ? state->customSecret : state->extSecret; +#if defined(XXH3_STREAM_USE_STACK) && XXH3_STREAM_USE_STACK >= 1 + /* For some reason, gcc and MSVC seem to suffer greatly + * when operating accumulators directly into state. + * Operating into stack space seems to enable proper optimization. + * clang, on the other hand, doesn't seem to need this trick */ + XXH_ALIGN(XXH_ACC_ALIGN) xxh_u64 acc[8]; memcpy(acc, state->acc, sizeof(acc)); +#else + xxh_u64* XXH_RESTRICT const acc = state->acc; +#endif + state->totalLen += len; + XXH_ASSERT(state->bufferedSize <= XXH3_INTERNALBUFFER_SIZE); + + /* small input : just fill in tmp buffer */ + if (state->bufferedSize + len <= XXH3_INTERNALBUFFER_SIZE) { + XXH_memcpy(state->buffer + state->bufferedSize, input, len); + state->bufferedSize += (XXH32_hash_t)len; + return XXH_OK; + } + + /* total input is now > XXH3_INTERNALBUFFER_SIZE */ + #define XXH3_INTERNALBUFFER_STRIPES (XXH3_INTERNALBUFFER_SIZE / XXH_STRIPE_LEN) + XXH_STATIC_ASSERT(XXH3_INTERNALBUFFER_SIZE % XXH_STRIPE_LEN == 0); /* clean multiple */ + + /* + * Internal buffer is partially filled (always, except at beginning) + * Complete it, then consume it. + */ + if (state->bufferedSize) { + size_t const loadSize = XXH3_INTERNALBUFFER_SIZE - state->bufferedSize; + XXH_memcpy(state->buffer + state->bufferedSize, input, loadSize); + input += loadSize; + XXH3_consumeStripes(acc, + &state->nbStripesSoFar, state->nbStripesPerBlock, + state->buffer, XXH3_INTERNALBUFFER_STRIPES, + secret, state->secretLimit, + f_acc512, f_scramble); + state->bufferedSize = 0; + } + XXH_ASSERT(input < bEnd); + + /* large input to consume : ingest per full block */ + if ((size_t)(bEnd - input) > state->nbStripesPerBlock * XXH_STRIPE_LEN) { + size_t nbStripes = (size_t)(bEnd - 1 - input) / XXH_STRIPE_LEN; + XXH_ASSERT(state->nbStripesPerBlock >= state->nbStripesSoFar); + /* join to current block's end */ + { size_t const nbStripesToEnd = state->nbStripesPerBlock - state->nbStripesSoFar; + XXH_ASSERT(nbStripes <= nbStripes); + XXH3_accumulate(acc, input, secret + state->nbStripesSoFar * XXH_SECRET_CONSUME_RATE, nbStripesToEnd, f_acc512); + f_scramble(acc, secret + state->secretLimit); + state->nbStripesSoFar = 0; + input += nbStripesToEnd * XXH_STRIPE_LEN; + nbStripes -= nbStripesToEnd; + } + /* consume per entire blocks */ + while(nbStripes >= state->nbStripesPerBlock) { + XXH3_accumulate(acc, input, secret, state->nbStripesPerBlock, f_acc512); + f_scramble(acc, secret + state->secretLimit); + input += state->nbStripesPerBlock * XXH_STRIPE_LEN; + nbStripes -= state->nbStripesPerBlock; + } + /* consume last partial block */ + XXH3_accumulate(acc, input, secret, nbStripes, f_acc512); + input += nbStripes * XXH_STRIPE_LEN; + XXH_ASSERT(input < bEnd); /* at least some bytes left */ + state->nbStripesSoFar = nbStripes; + /* buffer predecessor of last partial stripe */ + XXH_memcpy(state->buffer + sizeof(state->buffer) - XXH_STRIPE_LEN, input - XXH_STRIPE_LEN, XXH_STRIPE_LEN); + XXH_ASSERT(bEnd - input <= XXH_STRIPE_LEN); + } else { + /* content to consume <= block size */ + /* Consume input by a multiple of internal buffer size */ + if (bEnd - input > XXH3_INTERNALBUFFER_SIZE) { + const xxh_u8* const limit = bEnd - XXH3_INTERNALBUFFER_SIZE; + do { + XXH3_consumeStripes(acc, + &state->nbStripesSoFar, state->nbStripesPerBlock, + input, XXH3_INTERNALBUFFER_STRIPES, + secret, state->secretLimit, + f_acc512, f_scramble); + input += XXH3_INTERNALBUFFER_SIZE; + } while (inputbuffer + sizeof(state->buffer) - XXH_STRIPE_LEN, input - XXH_STRIPE_LEN, XXH_STRIPE_LEN); + } + } + + /* Some remaining input (always) : buffer it */ + XXH_ASSERT(input < bEnd); + XXH_ASSERT(bEnd - input <= XXH3_INTERNALBUFFER_SIZE); + XXH_ASSERT(state->bufferedSize == 0); + XXH_memcpy(state->buffer, input, (size_t)(bEnd-input)); + state->bufferedSize = (XXH32_hash_t)(bEnd-input); +#if defined(XXH3_STREAM_USE_STACK) && XXH3_STREAM_USE_STACK >= 1 + /* save stack accumulators into state */ + memcpy(state->acc, acc, sizeof(acc)); +#endif + } + + return XXH_OK; +} + +/*! @ingroup xxh3_family */ +XXH_PUBLIC_API XXH_errorcode +XXH3_64bits_update(XXH3_state_t* state, const void* input, size_t len) +{ + return XXH3_update(state, (const xxh_u8*)input, len, + XXH3_accumulate_512, XXH3_scrambleAcc); +} + + +XXH_FORCE_INLINE void +XXH3_digest_long (XXH64_hash_t* acc, + const XXH3_state_t* state, + const unsigned char* secret) +{ + /* + * Digest on a local copy. This way, the state remains unaltered, and it can + * continue ingesting more input afterwards. + */ + XXH_memcpy(acc, state->acc, sizeof(state->acc)); + if (state->bufferedSize >= XXH_STRIPE_LEN) { + size_t const nbStripes = (state->bufferedSize - 1) / XXH_STRIPE_LEN; + size_t nbStripesSoFar = state->nbStripesSoFar; + XXH3_consumeStripes(acc, + &nbStripesSoFar, state->nbStripesPerBlock, + state->buffer, nbStripes, + secret, state->secretLimit, + XXH3_accumulate_512, XXH3_scrambleAcc); + /* last stripe */ + XXH3_accumulate_512(acc, + state->buffer + state->bufferedSize - XXH_STRIPE_LEN, + secret + state->secretLimit - XXH_SECRET_LASTACC_START); + } else { /* bufferedSize < XXH_STRIPE_LEN */ + xxh_u8 lastStripe[XXH_STRIPE_LEN]; + size_t const catchupSize = XXH_STRIPE_LEN - state->bufferedSize; + XXH_ASSERT(state->bufferedSize > 0); /* there is always some input buffered */ + XXH_memcpy(lastStripe, state->buffer + sizeof(state->buffer) - catchupSize, catchupSize); + XXH_memcpy(lastStripe + catchupSize, state->buffer, state->bufferedSize); + XXH3_accumulate_512(acc, + lastStripe, + secret + state->secretLimit - XXH_SECRET_LASTACC_START); + } +} + +/*! @ingroup xxh3_family */ +XXH_PUBLIC_API XXH64_hash_t XXH3_64bits_digest (const XXH3_state_t* state) +{ + const unsigned char* const secret = (state->extSecret == NULL) ? state->customSecret : state->extSecret; + if (state->totalLen > XXH3_MIDSIZE_MAX) { + XXH_ALIGN(XXH_ACC_ALIGN) XXH64_hash_t acc[XXH_ACC_NB]; + XXH3_digest_long(acc, state, secret); + return XXH3_mergeAccs(acc, + secret + XXH_SECRET_MERGEACCS_START, + (xxh_u64)state->totalLen * XXH_PRIME64_1); + } + /* totalLen <= XXH3_MIDSIZE_MAX: digesting a short input */ + if (state->useSeed) + return XXH3_64bits_withSeed(state->buffer, (size_t)state->totalLen, state->seed); + return XXH3_64bits_withSecret(state->buffer, (size_t)(state->totalLen), + secret, state->secretLimit + XXH_STRIPE_LEN); +} + + + +/* ========================================== + * XXH3 128 bits (a.k.a XXH128) + * ========================================== + * XXH3's 128-bit variant has better mixing and strength than the 64-bit variant, + * even without counting the significantly larger output size. + * + * For example, extra steps are taken to avoid the seed-dependent collisions + * in 17-240 byte inputs (See XXH3_mix16B and XXH128_mix32B). + * + * This strength naturally comes at the cost of some speed, especially on short + * lengths. Note that longer hashes are about as fast as the 64-bit version + * due to it using only a slight modification of the 64-bit loop. + * + * XXH128 is also more oriented towards 64-bit machines. It is still extremely + * fast for a _128-bit_ hash on 32-bit (it usually clears XXH64). + */ + +XXH_FORCE_INLINE XXH128_hash_t +XXH3_len_1to3_128b(const xxh_u8* input, size_t len, const xxh_u8* secret, XXH64_hash_t seed) +{ + /* A doubled version of 1to3_64b with different constants. */ + XXH_ASSERT(input != NULL); + XXH_ASSERT(1 <= len && len <= 3); + XXH_ASSERT(secret != NULL); + /* + * len = 1: combinedl = { input[0], 0x01, input[0], input[0] } + * len = 2: combinedl = { input[1], 0x02, input[0], input[1] } + * len = 3: combinedl = { input[2], 0x03, input[0], input[1] } + */ + { xxh_u8 const c1 = input[0]; + xxh_u8 const c2 = input[len >> 1]; + xxh_u8 const c3 = input[len - 1]; + xxh_u32 const combinedl = ((xxh_u32)c1 <<16) | ((xxh_u32)c2 << 24) + | ((xxh_u32)c3 << 0) | ((xxh_u32)len << 8); + xxh_u32 const combinedh = XXH_rotl32(XXH_swap32(combinedl), 13); + xxh_u64 const bitflipl = (XXH_readLE32(secret) ^ XXH_readLE32(secret+4)) + seed; + xxh_u64 const bitfliph = (XXH_readLE32(secret+8) ^ XXH_readLE32(secret+12)) - seed; + xxh_u64 const keyed_lo = (xxh_u64)combinedl ^ bitflipl; + xxh_u64 const keyed_hi = (xxh_u64)combinedh ^ bitfliph; + XXH128_hash_t h128; + h128.low64 = XXH64_avalanche(keyed_lo); + h128.high64 = XXH64_avalanche(keyed_hi); + return h128; + } +} + +XXH_FORCE_INLINE XXH128_hash_t +XXH3_len_4to8_128b(const xxh_u8* input, size_t len, const xxh_u8* secret, XXH64_hash_t seed) +{ + XXH_ASSERT(input != NULL); + XXH_ASSERT(secret != NULL); + XXH_ASSERT(4 <= len && len <= 8); + seed ^= (xxh_u64)XXH_swap32((xxh_u32)seed) << 32; + { xxh_u32 const input_lo = XXH_readLE32(input); + xxh_u32 const input_hi = XXH_readLE32(input + len - 4); + xxh_u64 const input_64 = input_lo + ((xxh_u64)input_hi << 32); + xxh_u64 const bitflip = (XXH_readLE64(secret+16) ^ XXH_readLE64(secret+24)) + seed; + xxh_u64 const keyed = input_64 ^ bitflip; + + /* Shift len to the left to ensure it is even, this avoids even multiplies. */ + XXH128_hash_t m128 = XXH_mult64to128(keyed, XXH_PRIME64_1 + (len << 2)); + + m128.high64 += (m128.low64 << 1); + m128.low64 ^= (m128.high64 >> 3); + + m128.low64 = XXH_xorshift64(m128.low64, 35); + m128.low64 *= 0x9FB21C651E98DF25ULL; + m128.low64 = XXH_xorshift64(m128.low64, 28); + m128.high64 = XXH3_avalanche(m128.high64); + return m128; + } +} + +XXH_FORCE_INLINE XXH128_hash_t +XXH3_len_9to16_128b(const xxh_u8* input, size_t len, const xxh_u8* secret, XXH64_hash_t seed) +{ + XXH_ASSERT(input != NULL); + XXH_ASSERT(secret != NULL); + XXH_ASSERT(9 <= len && len <= 16); + { xxh_u64 const bitflipl = (XXH_readLE64(secret+32) ^ XXH_readLE64(secret+40)) - seed; + xxh_u64 const bitfliph = (XXH_readLE64(secret+48) ^ XXH_readLE64(secret+56)) + seed; + xxh_u64 const input_lo = XXH_readLE64(input); + xxh_u64 input_hi = XXH_readLE64(input + len - 8); + XXH128_hash_t m128 = XXH_mult64to128(input_lo ^ input_hi ^ bitflipl, XXH_PRIME64_1); + /* + * Put len in the middle of m128 to ensure that the length gets mixed to + * both the low and high bits in the 128x64 multiply below. + */ + m128.low64 += (xxh_u64)(len - 1) << 54; + input_hi ^= bitfliph; + /* + * Add the high 32 bits of input_hi to the high 32 bits of m128, then + * add the long product of the low 32 bits of input_hi and XXH_PRIME32_2 to + * the high 64 bits of m128. + * + * The best approach to this operation is different on 32-bit and 64-bit. + */ + if (sizeof(void *) < sizeof(xxh_u64)) { /* 32-bit */ + /* + * 32-bit optimized version, which is more readable. + * + * On 32-bit, it removes an ADC and delays a dependency between the two + * halves of m128.high64, but it generates an extra mask on 64-bit. + */ + m128.high64 += (input_hi & 0xFFFFFFFF00000000ULL) + XXH_mult32to64((xxh_u32)input_hi, XXH_PRIME32_2); + } else { + /* + * 64-bit optimized (albeit more confusing) version. + * + * Uses some properties of addition and multiplication to remove the mask: + * + * Let: + * a = input_hi.lo = (input_hi & 0x00000000FFFFFFFF) + * b = input_hi.hi = (input_hi & 0xFFFFFFFF00000000) + * c = XXH_PRIME32_2 + * + * a + (b * c) + * Inverse Property: x + y - x == y + * a + (b * (1 + c - 1)) + * Distributive Property: x * (y + z) == (x * y) + (x * z) + * a + (b * 1) + (b * (c - 1)) + * Identity Property: x * 1 == x + * a + b + (b * (c - 1)) + * + * Substitute a, b, and c: + * input_hi.hi + input_hi.lo + ((xxh_u64)input_hi.lo * (XXH_PRIME32_2 - 1)) + * + * Since input_hi.hi + input_hi.lo == input_hi, we get this: + * input_hi + ((xxh_u64)input_hi.lo * (XXH_PRIME32_2 - 1)) + */ + m128.high64 += input_hi + XXH_mult32to64((xxh_u32)input_hi, XXH_PRIME32_2 - 1); + } + /* m128 ^= XXH_swap64(m128 >> 64); */ + m128.low64 ^= XXH_swap64(m128.high64); + + { /* 128x64 multiply: h128 = m128 * XXH_PRIME64_2; */ + XXH128_hash_t h128 = XXH_mult64to128(m128.low64, XXH_PRIME64_2); + h128.high64 += m128.high64 * XXH_PRIME64_2; + + h128.low64 = XXH3_avalanche(h128.low64); + h128.high64 = XXH3_avalanche(h128.high64); + return h128; + } } +} + +/* + * Assumption: `secret` size is >= XXH3_SECRET_SIZE_MIN + */ +XXH_FORCE_INLINE XXH128_hash_t +XXH3_len_0to16_128b(const xxh_u8* input, size_t len, const xxh_u8* secret, XXH64_hash_t seed) +{ + XXH_ASSERT(len <= 16); + { if (len > 8) return XXH3_len_9to16_128b(input, len, secret, seed); + if (len >= 4) return XXH3_len_4to8_128b(input, len, secret, seed); + if (len) return XXH3_len_1to3_128b(input, len, secret, seed); + { XXH128_hash_t h128; + xxh_u64 const bitflipl = XXH_readLE64(secret+64) ^ XXH_readLE64(secret+72); + xxh_u64 const bitfliph = XXH_readLE64(secret+80) ^ XXH_readLE64(secret+88); + h128.low64 = XXH64_avalanche(seed ^ bitflipl); + h128.high64 = XXH64_avalanche( seed ^ bitfliph); + return h128; + } } +} + +/* + * A bit slower than XXH3_mix16B, but handles multiply by zero better. + */ +XXH_FORCE_INLINE XXH128_hash_t +XXH128_mix32B(XXH128_hash_t acc, const xxh_u8* input_1, const xxh_u8* input_2, + const xxh_u8* secret, XXH64_hash_t seed) +{ + acc.low64 += XXH3_mix16B (input_1, secret+0, seed); + acc.low64 ^= XXH_readLE64(input_2) + XXH_readLE64(input_2 + 8); + acc.high64 += XXH3_mix16B (input_2, secret+16, seed); + acc.high64 ^= XXH_readLE64(input_1) + XXH_readLE64(input_1 + 8); + return acc; +} + + +XXH_FORCE_INLINE XXH128_hash_t +XXH3_len_17to128_128b(const xxh_u8* XXH_RESTRICT input, size_t len, + const xxh_u8* XXH_RESTRICT secret, size_t secretSize, + XXH64_hash_t seed) +{ + XXH_ASSERT(secretSize >= XXH3_SECRET_SIZE_MIN); (void)secretSize; + XXH_ASSERT(16 < len && len <= 128); + + { XXH128_hash_t acc; + acc.low64 = len * XXH_PRIME64_1; + acc.high64 = 0; + if (len > 32) { + if (len > 64) { + if (len > 96) { + acc = XXH128_mix32B(acc, input+48, input+len-64, secret+96, seed); + } + acc = XXH128_mix32B(acc, input+32, input+len-48, secret+64, seed); + } + acc = XXH128_mix32B(acc, input+16, input+len-32, secret+32, seed); + } + acc = XXH128_mix32B(acc, input, input+len-16, secret, seed); + { XXH128_hash_t h128; + h128.low64 = acc.low64 + acc.high64; + h128.high64 = (acc.low64 * XXH_PRIME64_1) + + (acc.high64 * XXH_PRIME64_4) + + ((len - seed) * XXH_PRIME64_2); + h128.low64 = XXH3_avalanche(h128.low64); + h128.high64 = (XXH64_hash_t)0 - XXH3_avalanche(h128.high64); + return h128; + } + } +} + +XXH_NO_INLINE XXH128_hash_t +XXH3_len_129to240_128b(const xxh_u8* XXH_RESTRICT input, size_t len, + const xxh_u8* XXH_RESTRICT secret, size_t secretSize, + XXH64_hash_t seed) +{ + XXH_ASSERT(secretSize >= XXH3_SECRET_SIZE_MIN); (void)secretSize; + XXH_ASSERT(128 < len && len <= XXH3_MIDSIZE_MAX); + + { XXH128_hash_t acc; + int const nbRounds = (int)len / 32; + int i; + acc.low64 = len * XXH_PRIME64_1; + acc.high64 = 0; + for (i=0; i<4; i++) { + acc = XXH128_mix32B(acc, + input + (32 * i), + input + (32 * i) + 16, + secret + (32 * i), + seed); + } + acc.low64 = XXH3_avalanche(acc.low64); + acc.high64 = XXH3_avalanche(acc.high64); + XXH_ASSERT(nbRounds >= 4); + for (i=4 ; i < nbRounds; i++) { + acc = XXH128_mix32B(acc, + input + (32 * i), + input + (32 * i) + 16, + secret + XXH3_MIDSIZE_STARTOFFSET + (32 * (i - 4)), + seed); + } + /* last bytes */ + acc = XXH128_mix32B(acc, + input + len - 16, + input + len - 32, + secret + XXH3_SECRET_SIZE_MIN - XXH3_MIDSIZE_LASTOFFSET - 16, + 0ULL - seed); + + { XXH128_hash_t h128; + h128.low64 = acc.low64 + acc.high64; + h128.high64 = (acc.low64 * XXH_PRIME64_1) + + (acc.high64 * XXH_PRIME64_4) + + ((len - seed) * XXH_PRIME64_2); + h128.low64 = XXH3_avalanche(h128.low64); + h128.high64 = (XXH64_hash_t)0 - XXH3_avalanche(h128.high64); + return h128; + } + } +} + +XXH_FORCE_INLINE XXH128_hash_t +XXH3_hashLong_128b_internal(const void* XXH_RESTRICT input, size_t len, + const xxh_u8* XXH_RESTRICT secret, size_t secretSize, + XXH3_f_accumulate_512 f_acc512, + XXH3_f_scrambleAcc f_scramble) +{ + XXH_ALIGN(XXH_ACC_ALIGN) xxh_u64 acc[XXH_ACC_NB] = XXH3_INIT_ACC; + + XXH3_hashLong_internal_loop(acc, (const xxh_u8*)input, len, secret, secretSize, f_acc512, f_scramble); + + /* converge into final hash */ + XXH_STATIC_ASSERT(sizeof(acc) == 64); + XXH_ASSERT(secretSize >= sizeof(acc) + XXH_SECRET_MERGEACCS_START); + { XXH128_hash_t h128; + h128.low64 = XXH3_mergeAccs(acc, + secret + XXH_SECRET_MERGEACCS_START, + (xxh_u64)len * XXH_PRIME64_1); + h128.high64 = XXH3_mergeAccs(acc, + secret + secretSize + - sizeof(acc) - XXH_SECRET_MERGEACCS_START, + ~((xxh_u64)len * XXH_PRIME64_2)); + return h128; + } +} + +/* + * It's important for performance that XXH3_hashLong is not inlined. + */ +XXH_NO_INLINE XXH128_hash_t +XXH3_hashLong_128b_default(const void* XXH_RESTRICT input, size_t len, + XXH64_hash_t seed64, + const void* XXH_RESTRICT secret, size_t secretLen) +{ + (void)seed64; (void)secret; (void)secretLen; + return XXH3_hashLong_128b_internal(input, len, XXH3_kSecret, sizeof(XXH3_kSecret), + XXH3_accumulate_512, XXH3_scrambleAcc); +} + +/* + * It's important for performance to pass @secretLen (when it's static) + * to the compiler, so that it can properly optimize the vectorized loop. + */ +XXH_FORCE_INLINE XXH128_hash_t +XXH3_hashLong_128b_withSecret(const void* XXH_RESTRICT input, size_t len, + XXH64_hash_t seed64, + const void* XXH_RESTRICT secret, size_t secretLen) +{ + (void)seed64; + return XXH3_hashLong_128b_internal(input, len, (const xxh_u8*)secret, secretLen, + XXH3_accumulate_512, XXH3_scrambleAcc); +} + +XXH_FORCE_INLINE XXH128_hash_t +XXH3_hashLong_128b_withSeed_internal(const void* XXH_RESTRICT input, size_t len, + XXH64_hash_t seed64, + XXH3_f_accumulate_512 f_acc512, + XXH3_f_scrambleAcc f_scramble, + XXH3_f_initCustomSecret f_initSec) +{ + if (seed64 == 0) + return XXH3_hashLong_128b_internal(input, len, + XXH3_kSecret, sizeof(XXH3_kSecret), + f_acc512, f_scramble); + { XXH_ALIGN(XXH_SEC_ALIGN) xxh_u8 secret[XXH_SECRET_DEFAULT_SIZE]; + f_initSec(secret, seed64); + return XXH3_hashLong_128b_internal(input, len, (const xxh_u8*)secret, sizeof(secret), + f_acc512, f_scramble); + } +} + +/* + * It's important for performance that XXH3_hashLong is not inlined. + */ +XXH_NO_INLINE XXH128_hash_t +XXH3_hashLong_128b_withSeed(const void* input, size_t len, + XXH64_hash_t seed64, const void* XXH_RESTRICT secret, size_t secretLen) +{ + (void)secret; (void)secretLen; + return XXH3_hashLong_128b_withSeed_internal(input, len, seed64, + XXH3_accumulate_512, XXH3_scrambleAcc, XXH3_initCustomSecret); +} + +typedef XXH128_hash_t (*XXH3_hashLong128_f)(const void* XXH_RESTRICT, size_t, + XXH64_hash_t, const void* XXH_RESTRICT, size_t); + +XXH_FORCE_INLINE XXH128_hash_t +XXH3_128bits_internal(const void* input, size_t len, + XXH64_hash_t seed64, const void* XXH_RESTRICT secret, size_t secretLen, + XXH3_hashLong128_f f_hl128) +{ + XXH_ASSERT(secretLen >= XXH3_SECRET_SIZE_MIN); + /* + * If an action is to be taken if `secret` conditions are not respected, + * it should be done here. + * For now, it's a contract pre-condition. + * Adding a check and a branch here would cost performance at every hash. + */ + if (len <= 16) + return XXH3_len_0to16_128b((const xxh_u8*)input, len, (const xxh_u8*)secret, seed64); + if (len <= 128) + return XXH3_len_17to128_128b((const xxh_u8*)input, len, (const xxh_u8*)secret, secretLen, seed64); + if (len <= XXH3_MIDSIZE_MAX) + return XXH3_len_129to240_128b((const xxh_u8*)input, len, (const xxh_u8*)secret, secretLen, seed64); + return f_hl128(input, len, seed64, secret, secretLen); +} + + +/* === Public XXH128 API === */ + +/*! @ingroup xxh3_family */ +XXH_PUBLIC_API XXH128_hash_t XXH3_128bits(const void* input, size_t len) +{ + return XXH3_128bits_internal(input, len, 0, + XXH3_kSecret, sizeof(XXH3_kSecret), + XXH3_hashLong_128b_default); +} + +/*! @ingroup xxh3_family */ +XXH_PUBLIC_API XXH128_hash_t +XXH3_128bits_withSecret(const void* input, size_t len, const void* secret, size_t secretSize) +{ + return XXH3_128bits_internal(input, len, 0, + (const xxh_u8*)secret, secretSize, + XXH3_hashLong_128b_withSecret); +} + +/*! @ingroup xxh3_family */ +XXH_PUBLIC_API XXH128_hash_t +XXH3_128bits_withSeed(const void* input, size_t len, XXH64_hash_t seed) +{ + return XXH3_128bits_internal(input, len, seed, + XXH3_kSecret, sizeof(XXH3_kSecret), + XXH3_hashLong_128b_withSeed); +} + +/*! @ingroup xxh3_family */ +XXH_PUBLIC_API XXH128_hash_t +XXH3_128bits_withSecretandSeed(const void* input, size_t len, const void* secret, size_t secretSize, XXH64_hash_t seed) +{ + if (len <= XXH3_MIDSIZE_MAX) + return XXH3_128bits_internal(input, len, seed, XXH3_kSecret, sizeof(XXH3_kSecret), NULL); + return XXH3_hashLong_128b_withSecret(input, len, seed, secret, secretSize); +} + +/*! @ingroup xxh3_family */ +XXH_PUBLIC_API XXH128_hash_t +XXH128(const void* input, size_t len, XXH64_hash_t seed) +{ + return XXH3_128bits_withSeed(input, len, seed); +} + + +/* === XXH3 128-bit streaming === */ + +/* + * All initialization and update functions are identical to 64-bit streaming variant. + * The only difference is the finalization routine. + */ + +/*! @ingroup xxh3_family */ +XXH_PUBLIC_API XXH_errorcode +XXH3_128bits_reset(XXH3_state_t* statePtr) +{ + return XXH3_64bits_reset(statePtr); +} + +/*! @ingroup xxh3_family */ +XXH_PUBLIC_API XXH_errorcode +XXH3_128bits_reset_withSecret(XXH3_state_t* statePtr, const void* secret, size_t secretSize) +{ + return XXH3_64bits_reset_withSecret(statePtr, secret, secretSize); +} + +/*! @ingroup xxh3_family */ +XXH_PUBLIC_API XXH_errorcode +XXH3_128bits_reset_withSeed(XXH3_state_t* statePtr, XXH64_hash_t seed) +{ + return XXH3_64bits_reset_withSeed(statePtr, seed); +} + +/*! @ingroup xxh3_family */ +XXH_PUBLIC_API XXH_errorcode +XXH3_128bits_reset_withSecretandSeed(XXH3_state_t* statePtr, const void* secret, size_t secretSize, XXH64_hash_t seed) +{ + return XXH3_64bits_reset_withSecretandSeed(statePtr, secret, secretSize, seed); +} + +/*! @ingroup xxh3_family */ +XXH_PUBLIC_API XXH_errorcode +XXH3_128bits_update(XXH3_state_t* state, const void* input, size_t len) +{ + return XXH3_update(state, (const xxh_u8*)input, len, + XXH3_accumulate_512, XXH3_scrambleAcc); +} + +/*! @ingroup xxh3_family */ +XXH_PUBLIC_API XXH128_hash_t XXH3_128bits_digest (const XXH3_state_t* state) +{ + const unsigned char* const secret = (state->extSecret == NULL) ? state->customSecret : state->extSecret; + if (state->totalLen > XXH3_MIDSIZE_MAX) { + XXH_ALIGN(XXH_ACC_ALIGN) XXH64_hash_t acc[XXH_ACC_NB]; + XXH3_digest_long(acc, state, secret); + XXH_ASSERT(state->secretLimit + XXH_STRIPE_LEN >= sizeof(acc) + XXH_SECRET_MERGEACCS_START); + { XXH128_hash_t h128; + h128.low64 = XXH3_mergeAccs(acc, + secret + XXH_SECRET_MERGEACCS_START, + (xxh_u64)state->totalLen * XXH_PRIME64_1); + h128.high64 = XXH3_mergeAccs(acc, + secret + state->secretLimit + XXH_STRIPE_LEN + - sizeof(acc) - XXH_SECRET_MERGEACCS_START, + ~((xxh_u64)state->totalLen * XXH_PRIME64_2)); + return h128; + } + } + /* len <= XXH3_MIDSIZE_MAX : short code */ + if (state->seed) + return XXH3_128bits_withSeed(state->buffer, (size_t)state->totalLen, state->seed); + return XXH3_128bits_withSecret(state->buffer, (size_t)(state->totalLen), + secret, state->secretLimit + XXH_STRIPE_LEN); +} + +/* 128-bit utility functions */ + +#include /* memcmp, memcpy */ + +/* return : 1 is equal, 0 if different */ +/*! @ingroup xxh3_family */ +XXH_PUBLIC_API int XXH128_isEqual(XXH128_hash_t h1, XXH128_hash_t h2) +{ + /* note : XXH128_hash_t is compact, it has no padding byte */ + return !(memcmp(&h1, &h2, sizeof(h1))); +} + +/* This prototype is compatible with stdlib's qsort(). + * return : >0 if *h128_1 > *h128_2 + * <0 if *h128_1 < *h128_2 + * =0 if *h128_1 == *h128_2 */ +/*! @ingroup xxh3_family */ +XXH_PUBLIC_API int XXH128_cmp(const void* h128_1, const void* h128_2) +{ + XXH128_hash_t const h1 = *(const XXH128_hash_t*)h128_1; + XXH128_hash_t const h2 = *(const XXH128_hash_t*)h128_2; + int const hcmp = (h1.high64 > h2.high64) - (h2.high64 > h1.high64); + /* note : bets that, in most cases, hash values are different */ + if (hcmp) return hcmp; + return (h1.low64 > h2.low64) - (h2.low64 > h1.low64); +} + + +/*====== Canonical representation ======*/ +/*! @ingroup xxh3_family */ +XXH_PUBLIC_API void +XXH128_canonicalFromHash(XXH128_canonical_t* dst, XXH128_hash_t hash) +{ + XXH_STATIC_ASSERT(sizeof(XXH128_canonical_t) == sizeof(XXH128_hash_t)); + if (XXH_CPU_LITTLE_ENDIAN) { + hash.high64 = XXH_swap64(hash.high64); + hash.low64 = XXH_swap64(hash.low64); + } + XXH_memcpy(dst, &hash.high64, sizeof(hash.high64)); + XXH_memcpy((char*)dst + sizeof(hash.high64), &hash.low64, sizeof(hash.low64)); +} + +/*! @ingroup xxh3_family */ +XXH_PUBLIC_API XXH128_hash_t +XXH128_hashFromCanonical(const XXH128_canonical_t* src) +{ + XXH128_hash_t h; + h.high64 = XXH_readBE64(src); + h.low64 = XXH_readBE64(src->digest + 8); + return h; +} + + + +/* ========================================== + * Secret generators + * ========================================== + */ +#define XXH_MIN(x, y) (((x) > (y)) ? (y) : (x)) + +static void XXH3_combine16(void* dst, XXH128_hash_t h128) +{ + XXH_writeLE64( dst, XXH_readLE64(dst) ^ h128.low64 ); + XXH_writeLE64( (char*)dst+8, XXH_readLE64((char*)dst+8) ^ h128.high64 ); +} + +/*! @ingroup xxh3_family */ +XXH_PUBLIC_API XXH_errorcode +XXH3_generateSecret(void* secretBuffer, size_t secretSize, const void* customSeed, size_t customSeedSize) +{ + XXH_ASSERT(secretBuffer != NULL); + if (secretBuffer == NULL) return XXH_ERROR; + XXH_ASSERT(secretSize >= XXH3_SECRET_SIZE_MIN); + if (secretSize < XXH3_SECRET_SIZE_MIN) return XXH_ERROR; + if (customSeedSize == 0) { + customSeed = XXH3_kSecret; + customSeedSize = XXH_SECRET_DEFAULT_SIZE; + } + XXH_ASSERT(customSeed != NULL); + if (customSeed == NULL) return XXH_ERROR; + + /* Fill secretBuffer with a copy of customSeed - repeat as needed */ + { size_t pos = 0; + while (pos < secretSize) { + size_t const toCopy = XXH_MIN((secretSize - pos), customSeedSize); + memcpy((char*)secretBuffer + pos, customSeed, toCopy); + pos += toCopy; + } } + + { size_t const nbSeg16 = secretSize / 16; + size_t n; + XXH128_canonical_t scrambler; + XXH128_canonicalFromHash(&scrambler, XXH128(customSeed, customSeedSize, 0)); + for (n=0; n +#include +#include +#include + +using namespace core; + +namespace py = pybind11; +using namespace py::literals; + +static constexpr unsigned int RESERVED_INDICES[] = { 0 }; +static const py::bool_ py_true(true); + +template +inline const T& choice(const T* items, unsigned int length) { + return sample_uniform(items, length); +} + +inline unsigned int randrange(unsigned int start, unsigned int end) { + return start + sample_uniform(end - start); +} + +inline unsigned int randrange(unsigned int end) { + return sample_uniform(end); +} + +struct node { + unsigned int id; + core::array children; + core::array parents; + + node(unsigned int id) : id(id), children(8), parents(8) { } + + static inline void free(node& n) { + core::free(n.children); + core::free(n.parents); + } +}; + +inline bool init(node& n, unsigned int id) { + if (!array_init(n.children, 8)) { + return false; + } else if (!array_init(n.parents, 8)) { + core::free(n.children); + return false; + } + n.id = id; + return true; +} + +inline bool operator == (const node& first, const node& second) { + return first.id == second.id; +} + +bool generate_graph(array& vertices, node*& start, node*& end, unsigned int num_vertices, unsigned int max_num_parents, unsigned int max_vertex_id) +{ + if (!vertices.ensure_capacity(num_vertices)) + return false; + for (unsigned int i = 0; i < num_vertices; i++) + if (!init(vertices[i], i)) return false; + vertices.length = num_vertices; + + /* sample a random DAG */ + for (unsigned int i = 1; i < num_vertices; i++) { + /* sample the number of parent vertices */ + unsigned int num_parents; + if (randrange(2) == 0) + num_parents = 1; + else + num_parents = randrange(1, max_num_parents); + num_parents = min(num_parents, i); + + array available_parent_ids(i); + for (unsigned int j = 0; j < i; j++) + available_parent_ids.add(j); + for (unsigned int j = 0; j < num_parents; j++) { + unsigned int u = randrange(available_parent_ids.length); + unsigned int parent_id = available_parent_ids[u]; + vertices[parent_id].children.add(&vertices[i]); + vertices[i].parents.add(&vertices[parent_id]); + available_parent_ids.remove(u); + } + } + + /* remove any correlation between graph topology and vertex IDs by shuffling the vertices */ + unsigned int* new_indices = (unsigned int*) alloca(sizeof(unsigned int) * (max_vertex_id + 1)); + for (unsigned int i = 0; i < max_vertex_id + 1; i++) new_indices[i] = i; + shuffle(new_indices, max_vertex_id + 1); + unsigned int src_index = 0; + for (unsigned int i = 0; i < vertices.length; i++) { + bool is_reserved = false; + for (unsigned int j = 0; j < array_length(RESERVED_INDICES); j++) { + if (new_indices[src_index] == RESERVED_INDICES[j]) { + is_reserved = true; + break; + } + } + if (is_reserved) + src_index++; + vertices[i].id = new_indices[src_index]; + src_index++; + } + + /* randomly select two vertices */ + start = &vertices[randrange(vertices.length - 1)]; + do { + end = &vertices[randrange(vertices.length - 1)]; + } while (end == start); + return true; +} + +/* computes the number of lookahead steps to find the answer */ +unsigned int lookahead_depth(const node* vertex, const node* next_vertex, const node* goal) +{ + array> frontier(8); + for (const node* v : vertex->children) + frontier.add(make_pair(v, v)); + array visited(16); + visited.append(vertex->children.data, vertex->children.length); + unsigned int lookahead = 0; + while (frontier.length != 0) { + bool frontier_is_next_vertex = true; + for (pair& entry : frontier) { + if (entry.value != next_vertex) { + frontier_is_next_vertex = false; + break; + } + } + if (frontier_is_next_vertex) + return lookahead; + + lookahead++; + array> new_frontier(8); + for (const pair& entry : frontier) { + const node* v = entry.key; + const node* branch = entry.value; + + if (v == goal) + return lookahead; + for (const node* child : v->children) { + if (!visited.contains(child)) { + new_frontier.add(make_pair(child, branch)); + visited.add(child); + } else if (branch == next_vertex) { + for (unsigned int i = 0; i < new_frontier.length; i++) + if (new_frontier[i].key == child) + new_frontier[i].value = branch; + } + } + } + core::swap(frontier, new_frontier); + } + return lookahead; +} + +void get_descendants(const node& vertex, array& descendants) { + array queue(8); + array visited(16); + queue[0] = &vertex; + queue.length = 1; + while (queue.length != 0) { + const node* current = queue.pop(); + visited.add(current); + for (const node* child : current->children) { + if (!descendants.contains(child)) + descendants.add(child); + if (visited.contains(child)) + continue; + queue.add(child); + } + } +} + +bool has_cycles(array& vertices) { + for (const node& vertex : vertices) { + array descendants(8); + get_descendants(vertex, descendants); + if (descendants.contains(&vertex)) + return true; + } + return false; +} + +bool generate_graph_with_lookahead(array& vertices, node*& start, node*& end, unsigned int num_vertices, unsigned int max_num_parents, unsigned int max_vertex_id, unsigned int lookahead, unsigned int num_paths, unsigned int max_prefix_vertices) +{ + num_vertices = std::max(std::max(2u, num_vertices), 1 + num_paths * lookahead); + + if (!vertices.ensure_capacity(num_vertices)) + return false; + for (unsigned int i = 0; i < num_vertices; i++) + if (!init(vertices[i], i)) return false; + vertices.length = num_vertices; + + vertices[1].parents.add(&vertices[0]); + vertices[0].children.add(&vertices[1]); + for (unsigned int i = 1; i < lookahead; i++) { + vertices[1 + i].parents.add(&vertices[i]); + vertices[i].children.add(&vertices[1 + i]); + } + unsigned int index; + if (lookahead == 0) { + index = 2; + } else { + index = 1 + lookahead; + for (unsigned int j = 0; j < num_paths - 1; j++) { + vertices[index].parents.add(&vertices[0]); + vertices[0].children.add(&vertices[index]); + index++; + unsigned int other_branch_length = lookahead + randrange(std::min(2u, num_vertices - index - (num_paths - j - 1) * lookahead + 2)); + for (unsigned int i = 1; i < other_branch_length; i++) { + vertices[index].parents.add(&vertices[index - 1]); + vertices[index - 1].children.add(&vertices[index]); + index++; + } + } + } + + unsigned int num_prefix_vertices = randrange(min(max_prefix_vertices + 1, num_vertices - index + 1)); + node* prev_vertex = &vertices[0]; + for (unsigned int i = 0; i < num_prefix_vertices; i++) { + vertices[index].children.add(prev_vertex); + prev_vertex->parents.add(&vertices[index]); + prev_vertex = &vertices[index]; + index++; + } + + start = &vertices[0]; + end = &vertices[std::max(1u, lookahead)]; + + /* sample some parent/ancestor vertices */ + constexpr float ALPHA = 0.5f; + unsigned int* in_degrees = (unsigned int*) calloc(num_vertices, sizeof(unsigned int)); + unsigned int* out_degrees = (unsigned int*) calloc(num_vertices, sizeof(unsigned int)); + if (in_degrees == nullptr || out_degrees == nullptr) { + if (in_degrees != nullptr) free(in_degrees); + return false; + } + for (unsigned int i = 0; i < num_vertices; i++) { + in_degrees[i] = vertices[i].parents.length; + out_degrees[i] = vertices[i].children.length; + } + for (unsigned int i = index; i < num_vertices; i++) { + /* sample the number of child and parent vertices */ + unsigned int num_children = randrange(0, max_num_parents); + unsigned int num_parents = randrange(num_children == 0 ? 1 : 0, max_num_parents); + num_children = std::min(num_children, i); + num_parents = std::min(num_parents, i); + + /* sample the children of this new node */ + array probabilities(index); + float total_probability = 0.0f; + for (unsigned int j = 0; j < index; j++) { + probabilities[j] = ALPHA + in_degrees[j]; + total_probability += probabilities[j]; + } + probabilities.length = index; + + array sampled_children(std::max(1u, num_children)); + for (unsigned int j = 0; j < num_children; j++) { + unsigned int u = sample_categorical(probabilities.data, total_probability, probabilities.length); + sampled_children.add(u); + total_probability -= probabilities[u]; + probabilities[u] = 0.0f; + } + + for (unsigned int child_id : sampled_children) { + vertices[index].children.add(&vertices[child_id]); + vertices[child_id].parents.add(&vertices[index]); + in_degrees[child_id] += 1; + } + + /* sample the parents of this new node */ + total_probability = 0.0f; + for (unsigned int j = 0; j < index; j++) { + probabilities[j] = ALPHA + out_degrees[j]; + total_probability += probabilities[j]; + } + + /* to avoid creating a cycle, we have to remove any descendants from the possible parents */ + array descendants(8); + get_descendants(vertices[index], descendants); + for (const node* descendant : descendants) { + total_probability -= probabilities[descendant->id]; + probabilities[descendant->id] = 0.0f; + } + if (total_probability != 0.0f) { + num_parents = std::min(num_parents, index - (unsigned int) descendants.length); + + array sampled_parents(std::max(1u, num_parents)); + for (unsigned int j = 0; j < num_parents; j++) { + unsigned int u = sample_categorical(probabilities.data, total_probability, probabilities.length); + sampled_parents.add(u); + total_probability -= probabilities[u]; + probabilities[u] = 0.0f; + } + + for (unsigned int parent_id : sampled_parents) { + vertices[parent_id].children.add(&vertices[index]); + vertices[index].parents.add(&vertices[parent_id]); + out_degrees[parent_id] += 1; + } + } + index += 1; + } + free(in_degrees); + free(out_degrees); + + /* remove any correlation between graph topology and vertex IDs by shuffling the vertices */ + unsigned int* new_indices = (unsigned int*) alloca(sizeof(unsigned int) * (max_vertex_id + 1)); + for (unsigned int i = 0; i < max_vertex_id + 1; i++) new_indices[i] = i; + shuffle(new_indices, max_vertex_id + 1); + unsigned int src_index = 0; + for (unsigned int i = 0; i < vertices.length; i++) { + bool is_reserved = false; + for (unsigned int j = 0; j < array_length(RESERVED_INDICES); j++) { + if (new_indices[src_index] == RESERVED_INDICES[j]) { + is_reserved = true; + break; + } + } + if (is_reserved) + src_index++; + vertices[i].id = new_indices[src_index]; + src_index++; + } + return true; +} + +bool generate_example(array& vertices, node*& start, node*& end, array>& paths, unsigned int num_vertices, unsigned int max_num_parents, unsigned int max_vertex_id, bool get_shortest_paths, int lookahead, unsigned int num_paths, unsigned int max_prefix_vertices) +{ + if (lookahead == -1) { + if (!generate_graph(vertices, start, end, num_vertices, max_num_parents, max_vertex_id)) + return false; + } else { + if (!generate_graph_with_lookahead(vertices, start, end, num_vertices, max_num_parents, max_vertex_id, lookahead, num_paths, max_prefix_vertices)) + return false; + } + + /* find the shortest paths from `start` to `end` */ + array> queue(16); + queue[0].key = start; + queue[0].value = 0; + queue.length = 1; + array_map> reverse_pointers(16); + while (queue.length != 0) { + pair item = queue.pop(); + node* current = item.key; + unsigned int distance = item.value; + + for (node* child : current->children) { + if (!reverse_pointers.ensure_capacity(reverse_pointers.size + 1)) { + for (const auto& entry : reverse_pointers) core::free(entry.value); + return false; + } + bool contains; + array_map& value = reverse_pointers.get(child, contains); + if (!contains) { + if (!array_map_init(value, 4)) { + for (const auto& entry : reverse_pointers) core::free(entry.value); + return false; + } + reverse_pointers.keys[reverse_pointers.size++] = child; + } + + if (!value.ensure_capacity(value.size + 1)) { + for (const auto& entry : reverse_pointers) core::free(entry.value); + return false; + } + unsigned int& distance_value = value.get(current, contains); + if (!contains) { + distance_value = distance + 1; + value.keys[value.size++] = current; + } else if (distance_value > distance + 1) { + distance_value = distance + 1; + } else { + continue; + } + + bool found_child = false; + for (unsigned int j = 0; j < queue.length; j++) { + if (queue[j].key == child) { + queue[j].value = std::min(queue[j].value, distance + 1); + found_child = true; + break; + } + } + if (!found_child) + queue.add(make_pair(child, distance + 1)); + } + } + + if (!reverse_pointers.contains(end)) { + for (const auto& entry : reverse_pointers) core::free(entry.value); + return false; + } + + array_map> forward_pointers(16); + array dist_queue(16); + dist_queue[0] = end; + dist_queue.length = 1; + while (dist_queue.length != 0) { + node* current = dist_queue.pop(); + if (current == start) + continue; + + array prev_nodes(8); + const array_map& value = reverse_pointers.get(current); + if (get_shortest_paths) { + unsigned int min_distance = value.values[0]; + for (unsigned int i = 1; i < value.size; i++) + if (value.values[i] < min_distance) min_distance = value.values[i]; + + for (unsigned int i = 0; i < value.size; i++) + if (value.values[i] == min_distance) prev_nodes.add(value.keys[i]); + } else { + prev_nodes.append(value.keys, value.size); + } + for (const node* prev : prev_nodes) { + if (!forward_pointers.ensure_capacity(forward_pointers.size + 1)) { + for (const auto& entry : reverse_pointers) core::free(entry.value); + for (const auto& entry : forward_pointers) core::free(entry.value); + return false; + } + bool contains; + array& fptrs = forward_pointers.get(prev, contains); + if (!contains) { + if (!array_init(fptrs, 4)) { + for (const auto& entry : reverse_pointers) core::free(entry.value); + for (const auto& entry : forward_pointers) core::free(entry.value); + return false; + } + forward_pointers.keys[forward_pointers.size++] = prev; + } + fptrs.add(current); + } + dist_queue.append(prev_nodes.data, prev_nodes.length); + } + + /* free `reverse_pointers` */ + for (const auto& entry : reverse_pointers) core::free(entry.value); + + /* construct the shortest paths from the forward pointers */ + array> path_queue(8); + if (!array_init(path_queue[0], 1)) { + for (const auto& entry : forward_pointers) core::free(entry.value); + return false; + } + path_queue[0].add(start); + path_queue.length = 1; + while (path_queue.length != 0) { + array& partial_path = *((array*) alloca(sizeof(array))); + core::move(path_queue.last(), partial_path); + path_queue.length--; + + if (partial_path.last() == end) { + if (!paths.ensure_capacity(paths.length + 1)) { + for (const auto& entry : forward_pointers) core::free(entry.value); + for (auto& p : path_queue) core::free(p); + core::free(partial_path); + return false; + } + core::move(partial_path, paths[paths.length++]); + if (paths.length > 64) { + for (const auto& entry : forward_pointers) core::free(entry.value); + for (auto& p : path_queue) core::free(p); + return false; + } + continue; + } + for (node* next : forward_pointers.get(partial_path.last())) { + if (!path_queue.ensure_capacity(path_queue.length + 1) + || !array_init(path_queue[path_queue.length], partial_path.length + 1)) + { + for (const auto& entry : forward_pointers) core::free(entry.value); + for (auto& p : path_queue) core::free(p); + core::free(partial_path); + return false; + } + path_queue.length++; + path_queue.last().append(partial_path.data, partial_path.length); + path_queue.last().add(next); + } + core::free(partial_path); + } + + for (const auto& entry : forward_pointers) core::free(entry.value); + return true; +} + +bool has_path(const node* start, const node* end) +{ + array stack(8); + stack[0] = start; + stack.length = 1; + array visited(16); + while (stack.length != 0) { + const node* v = stack.pop(); + if (v == end) + return true; + for (node* child : v->children) { + if (!visited.contains(child)) { + visited.add(child); + stack.add(child); + } + } + } + return false; +} + +py::tuple generate_training_set(const unsigned int max_input_size, const uint64_t dataset_size, const int max_lookahead, const unsigned int max_edges, const py::object& reserved_inputs, const int distance_from_start, const int max_prefix_vertices, const bool quiet=false, const int num_paths=2) +{ + const unsigned int QUERY_PREFIX_TOKEN = (max_input_size-5) / 3 + 4; + const unsigned int PADDING_TOKEN = (max_input_size-5) / 3 + 3; + const unsigned int EDGE_PREFIX_TOKEN = (max_input_size-5) / 3 + 2; + const unsigned int PATH_PREFIX_TOKEN = (max_input_size-5) / 3 + 1; + + unsigned int num_generated = 0; + unsigned int num_collisions = 0; + unsigned int ntokens = (max_input_size - 5) / 3 + 5; + size_t input_shape[2]{dataset_size, max_input_size}; + size_t output_shape[2]{dataset_size, ntokens}; + size_t label_shape[1]{dataset_size}; + py::array_t inputs(input_shape); + py::array_t outputs(output_shape); + py::array_t labels(label_shape); + auto inputs_mem = inputs.mutable_unchecked<2>(); + auto outputs_mem = outputs.mutable_unchecked<2>(); + auto labels_mem = labels.mutable_unchecked<1>(); + unsigned int* lookahead_step_histogram = (unsigned int*) alloca(sizeof(unsigned int) * max_input_size); + unsigned int* path_length_histogram = (unsigned int*) alloca(sizeof(unsigned int) * max_input_size); + for (unsigned int i = 0; i < max_input_size; i++) { + lookahead_step_histogram[i] = 0; + path_length_histogram[i] = 0; + } + float* MAX_FREQS_PER_BUCKET = (float*) alloca(sizeof(float) * max_input_size); + if (max_lookahead == -1) { + for (unsigned int i = 0; i < max_input_size; i++) + MAX_FREQS_PER_BUCKET[i] = 1.0; + } else { + for (unsigned int i = 0; i < (unsigned) max_lookahead + 1; i++) + MAX_FREQS_PER_BUCKET[i] = 1.0 / (max_lookahead+1); + for (unsigned int i = max_lookahead + 1; i < max_input_size; i++) + MAX_FREQS_PER_BUCKET[i] = 0.0; + MAX_FREQS_PER_BUCKET[max_lookahead] += 0.05; + } + + unsigned int* potential_lookaheads = (unsigned int*) alloca(max((size_t) 1, sizeof(unsigned int) * (max_lookahead + 1))); + unsigned int potential_lookahead_count = 0; + while (num_generated < dataset_size) { + array g(32); + node* start; node* end; + array> paths(8); + while (true) { + if (max_lookahead == -1) { + unsigned int num_vertices = randrange(3, (max_input_size - 5) / 3); + if (!generate_example(g, start, end, paths, num_vertices, 4, (max_input_size - 5) / 3, true, -1, 0, max_prefix_vertices == -1 ? max_input_size : max_prefix_vertices)) { + for (node& n : g) core::free(n); + for (array& a : paths) core::free(a); + g.length = 0; paths.length = 0; + continue; + } + } else { + potential_lookahead_count = 0; + for (unsigned int i = 0; i < (unsigned) max_lookahead + 1; i++) + if (num_generated == 0 || lookahead_step_histogram[i] / num_generated < MAX_FREQS_PER_BUCKET[i]) + potential_lookaheads[potential_lookahead_count++] = i; +// unsigned int lookahead = choice(potential_lookaheads, potential_lookahead_count); + unsigned int lookahead = max_lookahead; + +// unsigned int num_paths; +// if (lookahead == 0) { +// num_paths = randrange(1, 3); +// } else { +// unsigned int max_num_paths = (max_edges - 1) / lookahead; +// num_paths = randrange(2, max_num_paths + 1); +// } + + unsigned int num_vertices = std::min(std::min(lookahead * num_paths + 1 + randrange(0, 6), (max_input_size-5) / 3), max_edges + 1); + if (!generate_example(g, start, end, paths, num_vertices, 4, (max_input_size - 5) / 3, true, lookahead, num_paths, max_prefix_vertices == -1 ? max_input_size : max_prefix_vertices)) { + for (node& n : g) core::free(n); + for (array& a : paths) core::free(a); + g.length = 0; paths.length = 0; + continue; + } + + unsigned int shortest_path_length = paths[0].length; + for (unsigned int i = 1; i < paths.length; i++) + if (paths[i].length < shortest_path_length) + shortest_path_length = paths[i].length; + + // Count how many paths have the shortest length. + unsigned int shortest_path_count = 0; + for (unsigned int i = 0; i < paths.length; i++) { + if (paths[i].length == shortest_path_length) + shortest_path_count++; + } + + // Only accept graphs with a unique (one) shortest path. + if (shortest_path_length <= 1 || shortest_path_count != 1) { + for (node& n : g) core::free(n); + for (array& a : paths) core::free(a); + g.length = 0; paths.length = 0; + continue; + } + } + unsigned int shortest_path_length = paths[0].length; + for (unsigned int i = 1; i < paths.length; i++) + if (paths[i].length < shortest_path_length) + shortest_path_length = paths[i].length; + if (shortest_path_length > 1) + break; + for (node& n : g) core::free(n); + for (array& a : paths) core::free(a); + g.length = 0; paths.length = 0; + } + + array> edges(8); + for (node& vertex : g) + for (node* child : vertex.children) + edges.add(make_pair(vertex.id, child->id)); + if (edges.length > max_edges || edges.length * 3 + 4 > max_input_size) { + for (node& n : g) core::free(n); + for (array& a : paths) core::free(a); + g.length = 0; paths.length = 0; + continue; + } + shuffle(edges); + + array prefix(max_input_size); + for (auto& entry : edges) { + prefix[prefix.length++] = EDGE_PREFIX_TOKEN; + prefix[prefix.length++] = entry.key; + prefix[prefix.length++] = entry.value; + } + prefix[prefix.length++] = QUERY_PREFIX_TOKEN; + prefix[prefix.length++] = start->id; + prefix[prefix.length++] = end->id; + prefix[prefix.length++] = PATH_PREFIX_TOKEN; + + for (const array& path : paths) { + if (path.length == 1) + continue; + for (unsigned int j = 1; j < path.length; j++) { + if (distance_from_start != -1 && j != (unsigned int) distance_from_start) + continue; + array example(prefix.length + j); + for (unsigned int i = 0; i < prefix.length; i++) + example[i] = prefix[i]; + for (unsigned int i = 0; i < j; i++) + example[prefix.length + i] = path[i]->id; + example.length = prefix.length + j; + if (example.length > max_input_size) + continue; + + unsigned int lookahead_steps = lookahead_depth(path[j-1], path[j], end); + array useful_steps(8); + for (node* v : path[j-1]->children) + if (has_path(v, end)) useful_steps.add(v); + + /* check if this input is reserved */ + py::object contains = reserved_inputs.attr("__contains__"); + py::tuple example_tuple(example.length); + for (unsigned int i = 0; i < example.length; i++) + example_tuple[i] = example[i]; + if (contains(example_tuple).is(py_true)) { + num_collisions += 1; + continue; + } + + if (num_generated != 0 && lookahead_step_histogram[lookahead_steps] / num_generated >= MAX_FREQS_PER_BUCKET[lookahead_steps]) + continue; + lookahead_step_histogram[lookahead_steps] += 1; + path_length_histogram[j] += 1; + + for (unsigned int i = 0; i < max_input_size - example.length; i++) + inputs_mem(num_generated, i) = PADDING_TOKEN; + for (unsigned int i = 0; i < example.length; i++) + inputs_mem(num_generated, max_input_size - example.length + i) = example[i]; + for (unsigned int i = 0; i < ntokens; i++) + outputs_mem(num_generated, i) = 0.0f; + for (unsigned int i = 0; i < useful_steps.length; i++) + outputs_mem(num_generated, useful_steps[i]->id) = 1.0f; + labels_mem(num_generated) = choice(useful_steps.data, useful_steps.length)->id; + num_generated++; + if (num_generated == dataset_size) + break; + } + if (num_generated == dataset_size) + break; + } + + if (!quiet && num_generated > 0 && (num_generated % 1000 == 0 || num_generated >= dataset_size)) { + printf("%d examples generated.\n", num_generated); + + printf("Lookahead steps histogram:\n"); + printf("["); + bool first = true; + for (unsigned int i = 0; i < max_input_size; i++) { + if (lookahead_step_histogram[i] == 0) + continue; + if (!first) printf(", "); + printf("%d:%.2f", i, (float) lookahead_step_histogram[i] / num_generated + 1e-9); + first = false; + } + printf("]\n"); + + printf("Path length histogram:\n"); + printf("["); + first = true; + for (unsigned int i = 0; i < max_input_size; i++) { + if (path_length_histogram[i] == 0) + continue; + if (!first) printf(", "); + printf("%d:%.2f", i, (float) path_length_histogram[i] / num_generated + 1e-9); + first = false; + } + printf("]\n"); + fflush(stdout); + } + + for (node& n : g) core::free(n); + for (array& a : paths) core::free(a); + g.length = 0; paths.length = 0; + continue; + } + + return py::make_tuple(inputs, outputs, labels, num_collisions); +} + +py::array_t lookahead_histogram(const unsigned int max_input_size, const uint64_t num_samples, const unsigned int max_edges, const int distance_from_start, const bool quiet=false) +{ + const unsigned int QUERY_PREFIX_TOKEN = (max_input_size-5) / 3 + 4; + const unsigned int EDGE_PREFIX_TOKEN = (max_input_size-5) / 3 + 2; + const unsigned int PATH_PREFIX_TOKEN = (max_input_size-5) / 3 + 1; + + unsigned int num_generated = 0; + unsigned int max_lookahead = ((max_input_size - 5) / 3 - 1) / 2; + size_t histogram_shape[1]{max_lookahead}; + py::array_t histogram(histogram_shape); + auto histogram_mem = histogram.mutable_unchecked<1>(); + for (unsigned int i = 0; i < max_lookahead; i++) + histogram_mem(i) = 0; + + while (num_generated < num_samples) { + array g(32); + node* start; node* end; + array> paths(8); + while (true) { + unsigned int num_vertices = randrange(3, (max_input_size - 5) / 3); + if (!generate_example(g, start, end, paths, num_vertices, 4, (max_input_size - 5) / 3, true, -1, 0, -1)) { + for (node& n : g) core::free(n); + for (array& a : paths) core::free(a); + g.length = 0; paths.length = 0; + continue; + } + unsigned int shortest_path_length = paths[0].length; + for (unsigned int i = 1; i < paths.length; i++) + if (paths[i].length < shortest_path_length) + shortest_path_length = paths[i].length; + if (shortest_path_length > 1) + break; + for (node& n : g) core::free(n); + for (array& a : paths) core::free(a); + g.length = 0; paths.length = 0; + } + + array> edges(8); + for (node& vertex : g) + for (node* child : vertex.children) + edges.add(make_pair(vertex.id, child->id)); + if (edges.length > max_edges || edges.length * 3 + 4 > max_input_size) { + for (node& n : g) core::free(n); + for (array& a : paths) core::free(a); + g.length = 0; paths.length = 0; + continue; + } + shuffle(edges); + + array prefix(max_input_size); + for (auto& entry : edges) { + prefix[prefix.length++] = EDGE_PREFIX_TOKEN; + prefix[prefix.length++] = entry.key; + prefix[prefix.length++] = entry.value; + } + prefix[prefix.length++] = QUERY_PREFIX_TOKEN; + prefix[prefix.length++] = start->id; + prefix[prefix.length++] = end->id; + prefix[prefix.length++] = PATH_PREFIX_TOKEN; + + for (const array& path : paths) { + if (path.length == 1) + continue; + for (unsigned int j = 1; j < path.length; j++) { + if (distance_from_start != -1 && j != (unsigned int) distance_from_start) + continue; + array example(prefix.length + j); + for (unsigned int i = 0; i < prefix.length; i++) + example[i] = prefix[i]; + for (unsigned int i = 0; i < j; i++) + example[prefix.length + i] = path[i]->id; + example.length = prefix.length + j; + if (example.length > max_input_size) + continue; + + unsigned int lookahead_steps = lookahead_depth(path[j-1], path[j], end); + histogram_mem(lookahead_steps) += 1; + num_generated++; + if (num_generated == num_samples) + break; + } + if (num_generated == num_samples) + break; + } + + if (!quiet && num_generated > 0 && (num_generated % 1000 == 0 || num_generated >= num_samples)) { + printf("%d examples generated.\n", num_generated); + + printf("Lookahead steps histogram:\n"); + printf("["); + bool first = true; + for (unsigned int i = 0; i < max_lookahead; i++) { + if (histogram_mem(i) == 0) + continue; + if (!first) printf(", "); + printf("%d:%.2f", i, (float) histogram_mem(i) / num_generated + 1e-9); + first = false; + } + printf("]\n"); + } + + for (node& n : g) core::free(n); + for (array& a : paths) core::free(a); + g.length = 0; paths.length = 0; + continue; + } + + return histogram; +} + +py::tuple generate_reachable_training_set(const unsigned int max_input_size, const uint64_t dataset_size, const unsigned int lookahead, const unsigned int max_edges, const py::object& reserved_inputs, const int distance_from_start, const int reachable_distance, const unsigned int start_vertex_index, const bool exclude_start_vertex) +{ + const unsigned int QUERY_PREFIX_TOKEN = (max_input_size-5) / 3 + 4; + const unsigned int PADDING_TOKEN = (max_input_size-5) / 3 + 3; + const unsigned int EDGE_PREFIX_TOKEN = (max_input_size-5) / 3 + 2; + const unsigned int PATH_PREFIX_TOKEN = (max_input_size-5) / 3 + 1; + + unsigned int num_generated = 0; + unsigned int num_collisions = 0; + size_t input_shape[2]{dataset_size, max_input_size}; + size_t output_shape[2]{dataset_size, max_input_size}; + py::array_t inputs(input_shape); + py::array_t outputs(output_shape); + auto inputs_mem = inputs.mutable_unchecked<2>(); + auto outputs_mem = outputs.mutable_unchecked<2>(); + py::list valid_outputs; + + unsigned int max_vertex_id = (max_input_size - 5) / 3; + while (num_generated < dataset_size) { + array g(32); + node* start; node* end; + array> paths(8); + while (true) { + unsigned int num_paths; + if (lookahead == 0) { + num_paths = randrange(1, 3); + } else { + unsigned int max_num_paths = (max_edges - 1) / lookahead; + num_paths = randrange(2, max_num_paths + 1); + } + + unsigned int num_vertices = std::min(std::min(lookahead * num_paths + 1 + randrange(0, 6), (max_input_size-5) / 3), max_edges + 1); + if (!generate_example(g, start, end, paths, num_vertices, 4, max_vertex_id, true, lookahead, num_paths, -1)) { + for (node& n : g) core::free(n); + for (array& a : paths) core::free(a); + g.length = 0; paths.length = 0; + continue; + } + unsigned int shortest_path_length = paths[0].length; + for (unsigned int i = 1; i < paths.length; i++) + if (paths[i].length < shortest_path_length) + shortest_path_length = paths[i].length; + if (shortest_path_length > 1) + break; + for (node& n : g) core::free(n); + for (array& a : paths) core::free(a); + g.length = 0; paths.length = 0; + } + + array> edges(8); + for (node& vertex : g) + for (node* child : vertex.children) + edges.add(make_pair(vertex.id, child->id)); + if (edges.length > max_edges || edges.length * 3 + 4 > max_input_size) { + for (node& n : g) core::free(n); + for (array& a : paths) core::free(a); + g.length = 0; paths.length = 0; + continue; + } + shuffle(edges); + + array prefix(max_input_size); + for (auto& entry : edges) { + prefix[prefix.length++] = EDGE_PREFIX_TOKEN; + prefix[prefix.length++] = entry.key; + prefix[prefix.length++] = entry.value; + } + prefix[prefix.length++] = QUERY_PREFIX_TOKEN; + prefix[prefix.length++] = start->id; + prefix[prefix.length++] = end->id; + prefix[prefix.length++] = PATH_PREFIX_TOKEN; + + for (const array& path : paths) { + if (path.length == 1) + continue; + for (unsigned int j = 1; j < path.length; j++) { + if (distance_from_start != -1 && j != (unsigned int) distance_from_start) + continue; + array example(prefix.length + j); + for (unsigned int i = 0; i < prefix.length; i++) + example[i] = prefix[i]; + for (unsigned int i = 0; i < j; i++) + example[prefix.length + i] = path[i]->id; + example.length = prefix.length + j; + if (example.length > max_input_size) + continue; + + /* compute the set of reachable vertices */ + node** vertex_id_map = (node**) calloc(max_vertex_id + 1, sizeof(node*)); + for (unsigned int i = 0; i < g.length; i++) + vertex_id_map[g[i].id] = &g[i]; + array reachable(16); + array> stack(16); + unsigned int start_vertex; + if (example.length < start_vertex_index) + start_vertex = start->id; + else start_vertex = example[example.length - start_vertex_index]; + stack.add(make_pair(start_vertex, 0u)); + while (stack.length != 0) { + pair entry = stack.pop(); + unsigned int current_vertex = entry.key; + unsigned int current_distance = entry.value; + if (!reachable.contains(current_vertex)) + reachable.add(current_vertex); + if (reachable_distance > 0 && current_distance + 1 <= (unsigned int) reachable_distance) { + for (node* child : vertex_id_map[current_vertex]->children) + stack.add(make_pair(child->id, current_distance + 1)); + } else if (reachable_distance < 0 && current_distance + 1 <= (unsigned int) -reachable_distance) { + for (node* parent : vertex_id_map[current_vertex]->parents) + stack.add(make_pair(parent->id, current_distance + 1)); + } + } + if (exclude_start_vertex) + reachable.remove(reachable.index_of(start_vertex)); + + array useful_steps(8); + for (node* v : path[j-1]->children) + if (has_path(v, end)) useful_steps.add(v); + + /* check if this input is reserved */ + py::object contains = reserved_inputs.attr("__contains__"); + py::tuple example_tuple(example.length); + for (unsigned int i = 0; i < example.length; i++) + example_tuple[i] = example[i]; + if (contains(example_tuple).is(py_true)) { + num_collisions += 1; + continue; + } + + for (unsigned int i = 0; i < max_input_size - example.length; i++) + inputs_mem(num_generated, i) = PADDING_TOKEN; + for (unsigned int i = 0; i < example.length; i++) + inputs_mem(num_generated, max_input_size - example.length + i) = example[i]; + for (unsigned int i = 0; i < max_input_size - example.length; i++) + outputs_mem(num_generated, i) = 0; + for (unsigned int i = 0; i < example.length; i++) + outputs_mem(num_generated, max_input_size - example.length + i) = reachable.contains(example[i]) ? 1 : 0; + num_generated++; + if (num_generated == dataset_size) + break; + } + if (num_generated == dataset_size) + break; + } + + for (node& n : g) core::free(n); + for (array& a : paths) core::free(a); + g.length = 0; paths.length = 0; + continue; + } + + return py::make_tuple(inputs, outputs, num_collisions); +} + +bool generate_dfs_example(array& vertices, const node*& start, const node*& end, unsigned int& current_node_index, array& path, unsigned int num_vertices, unsigned int max_num_parents, unsigned int max_vertex_id, const int requested_backtrack, unsigned int max_edges) +{ + if (!vertices.ensure_capacity(num_vertices)) + return false; + for (unsigned int i = 0; i < num_vertices; i++) + if (!init(vertices[i], i)) return false; + vertices.length = num_vertices; + + /* sample some parent/ancestor vertices */ + constexpr float ALPHA = 1.0f; + unsigned int* out_degrees = (unsigned int*) calloc(num_vertices, sizeof(unsigned int)); + if (out_degrees == nullptr) + return false; + for (unsigned int i = 1; i < num_vertices; i++) { + /* sample the number of child and parent vertices */ + unsigned int num_parents = randrange(1, max_num_parents); + num_parents = std::min(num_parents, i); + + /* sample the parents of this new node */ + float total_probability = 0.0f; + array probabilities(i); + for (unsigned int j = 0; j < i; j++) { + probabilities[j] = ALPHA + out_degrees[j]; + total_probability += probabilities[j]; + } + probabilities.length = i; + + array sampled_parents(std::max(1u, num_parents)); + for (unsigned int j = 0; j < num_parents; j++) { + unsigned int u = sample_categorical(probabilities.data, total_probability, probabilities.length); + sampled_parents.add(u); + total_probability -= probabilities[u]; + probabilities[u] = 0.0f; + } + + for (unsigned int parent_id : sampled_parents) { + vertices[parent_id].children.add(&vertices[i]); + vertices[i].parents.add(&vertices[parent_id]); + out_degrees[parent_id] += 1; + } + } + free(out_degrees); + + /* create a list of edges that we can remove */ + unsigned int new_edge_count = 0; + array> all_edges(16); + for (const node& vertex : vertices) + for (const node* child : vertex.children) + all_edges.add(make_pair(vertex.id, child->id)); + unsigned int total_edge_count = all_edges.length; + + array> removable_edges(16); + removable_edges.append(all_edges.data, all_edges.length); + while (true) { + /* select a start and goal vertex uniformly at random */ + unsigned int start_index; + unsigned int end_index; + if (requested_backtrack == 0 || requested_backtrack == -1) { + start_index = randrange(vertices.length - 1); + end_index = start_index + 1 + randrange(vertices.length - start_index - 1); + } else { + /* if `requested_backtrack` is specified, divide the vertices after the + start vertex into two subgraphs: the first subgraph will contain the goal, + and the second subgraph must have size at least `requested_backtrack` */ + const node* second_subgraph_root = &vertices[vertices.length - requested_backtrack]; + if (second_subgraph_root->parents.length == 1 && second_subgraph_root->parents[0]->id == second_subgraph_root->id - 1) { + start_index = randrange(vertices.length - requested_backtrack - 1); + } else { + do { + start_index = choice(second_subgraph_root->parents.data, second_subgraph_root->parents.length)->id; + } while (start_index == second_subgraph_root->id - 1); + } + end_index = start_index + 1 + randrange(vertices.length - start_index - requested_backtrack - 1); + } + start = &vertices[start_index]; + end = &vertices[end_index]; + + if (requested_backtrack != 0 && requested_backtrack != -1) { + /* if `requested_backtrack` is specified, make sure there is a path from the + start vertex to all vertices in the first and second subgraphs */ + for (unsigned int i = start_index + 1; i < vertices.length; i++) { + bool has_path_to_start = false; + for (const node* parent : vertices[i].parents) { + if (i >= vertices.length - requested_backtrack && (parent->id == start_index || parent->id >= vertices.length - requested_backtrack)) { + has_path_to_start = true; + break; + } else if (i < vertices.length - requested_backtrack && parent->id >= start_index) { + has_path_to_start = true; + break; + } + } + if (!has_path_to_start) { + vertices[start_index].children.add(&vertices[i]); + vertices[i].parents.add(&vertices[start_index]); + new_edge_count++; + total_edge_count++; + } + } + } + + /* perform DFS from the start vertex */ + array> queue(8); + if (requested_backtrack == 0 || requested_backtrack == -1) { + queue[0] = make_pair((const node*) nullptr, start); + queue.length = 1; + } else { + /* if `requested_backtrack` is specified, start the DFS in the second subgraph */ + path.add(start); + const node** first_subgraph_children = (const node**) malloc(sizeof(const node*) * start->children.length); + const node** second_subgraph_children = (const node**) malloc(sizeof(const node*) * start->children.length); + unsigned int num_first_subgraph_children = 0; + unsigned int num_second_subgraph_children = 0; + for (const node* child : start->children) { + if (child->id >= vertices.length - requested_backtrack) { + second_subgraph_children[num_second_subgraph_children++] = child; + } else { + first_subgraph_children[num_first_subgraph_children++] = child; + } + } + shuffle(first_subgraph_children, num_first_subgraph_children); + shuffle(second_subgraph_children, num_second_subgraph_children); + for (unsigned int i = 0; i < num_first_subgraph_children; i++) + queue.add(make_pair(start, first_subgraph_children[i])); + for (unsigned int i = 0; i < num_second_subgraph_children; i++) + queue.add(make_pair(start, second_subgraph_children[i])); + free(first_subgraph_children); + free(second_subgraph_children); + } + bool found_goal = false; + while (queue.length != 0) { + pair state = queue.pop(); + const node* parent = state.key; + const node* current = state.value; + if (path.contains(current)) + continue; + path.add(current); + if (parent != nullptr) { + unsigned int index = removable_edges.index_of(make_pair(parent->id, current->id)); + if (index < removable_edges.length) + removable_edges.remove(index); + } + + if (current == end) { + found_goal = true; + break; + } else if (current->children.contains(end)) { + found_goal = true; + path.add(end); + unsigned int index = removable_edges.index_of(make_pair(current->id, end->id)); + if (index < removable_edges.length) + removable_edges.remove(index); + break; + } + + if (current->children.length == 0) + continue; + const node** children = (const node**) malloc(sizeof(const node*) * current->children.length); + for (unsigned int i = 0; i < current->children.length; i++) + children[i] = current->children[i]; + shuffle(children, current->children.length); + for (unsigned int i = 0; i < current->children.length; i++) { + if (path.contains(children[i])) continue; + queue.add(make_pair(current, children[i])); + } + free(children); + } + + /* check if the goal vertex is reachable from the start vertex */ + if (found_goal) + break; + removable_edges.clear(); + removable_edges.append(all_edges.data, all_edges.length); + path.clear(); + } + + if (requested_backtrack == -1) + current_node_index = randrange(path.length - 1); + else + current_node_index = requested_backtrack; + + /* remove edges to avoid generating a graph with too many edges */ + unsigned int edges_to_remove; + if (total_edge_count > max_edges) + edges_to_remove = total_edge_count - max_edges; + else + edges_to_remove = new_edge_count; + edges_to_remove = min((unsigned int) removable_edges.length, edges_to_remove); + for (unsigned int i = 0; i < edges_to_remove; i++) { + unsigned int u = randrange(removable_edges.length); + unsigned int parent_id = removable_edges[u].key; + unsigned int child_id = removable_edges[u].value; + vertices[parent_id].children.remove(vertices[parent_id].children.index_of(&vertices[child_id])); + vertices[child_id].parents.remove(vertices[child_id].parents.index_of(&vertices[parent_id])); + removable_edges.remove(u); + } + + /* remove any correlation between graph topology and vertex IDs by shuffling the vertices */ + unsigned int* new_indices = (unsigned int*) alloca(sizeof(unsigned int) * (max_vertex_id + 1)); + for (unsigned int i = 0; i < max_vertex_id + 1; i++) new_indices[i] = i; + shuffle(new_indices, max_vertex_id + 1); + unsigned int src_index = 0; + for (unsigned int i = 0; i < vertices.length; i++) { + bool is_reserved = false; + for (unsigned int j = 0; j < array_length(RESERVED_INDICES); j++) { + if (new_indices[src_index] == RESERVED_INDICES[j]) { + is_reserved = true; + break; + } + } + if (is_reserved) + src_index++; + vertices[i].id = new_indices[src_index]; + src_index++; + } + + return true; +} + +py::tuple generate_dfs_training_set(const unsigned int max_input_size, const uint64_t dataset_size, const py::object& reserved_inputs, const int requested_backtrack, const bool random_padding, const bool uniform, const bool quiet=false) +{ + const unsigned int QUERY_PREFIX_TOKEN = (max_input_size-5) / 3 + 4; + const unsigned int PADDING_TOKEN = (max_input_size-5) / 3 + 3; + const unsigned int EDGE_PREFIX_TOKEN = (max_input_size-5) / 3 + 2; + const unsigned int PATH_PREFIX_TOKEN = (max_input_size-5) / 3 + 1; + + unsigned int longest_path_length = (max_input_size - 4) / 4; + + unsigned int num_generated = 0; + unsigned int num_collisions = 0; + unsigned int ntokens = (max_input_size - 5) / 3 + 5; + size_t input_shape[2]{dataset_size, max_input_size}; + size_t output_shape[2]{dataset_size, ntokens}; + size_t label_shape[1]{dataset_size}; + py::array_t inputs(input_shape); + py::array_t outputs(output_shape); + py::array_t labels(label_shape); + auto inputs_mem = inputs.mutable_unchecked<2>(); + auto outputs_mem = outputs.mutable_unchecked<2>(); + auto labels_mem = labels.mutable_unchecked<1>(); + + unsigned int* backtrack_distance_histogram = (unsigned int*) alloca(sizeof(unsigned int) * max_input_size); + for (unsigned int i = 0; i < max_input_size; i++) + backtrack_distance_histogram[i] = 0; + + float* MAX_FREQS_PER_BUCKET = (float*) alloca(sizeof(float) * max_input_size); + if (requested_backtrack == -1) { + for (unsigned int i = 0; i < max_input_size; i++) + MAX_FREQS_PER_BUCKET[i] = 1.0; + } else { + for (unsigned int i = 0; i < (unsigned) requested_backtrack + 1; i++) + MAX_FREQS_PER_BUCKET[i] = 1.0 / (requested_backtrack+1); + for (unsigned int i = requested_backtrack + 1; i < max_input_size; i++) + MAX_FREQS_PER_BUCKET[i] = 0.0; + MAX_FREQS_PER_BUCKET[requested_backtrack] += 0.05; + } + + array path(32); + unsigned int* potential_backtracks = (unsigned int*) alloca(max((size_t) 1, sizeof(unsigned int) * (requested_backtrack + 1))); + unsigned int potential_backtrack_count = 0; + unsigned int num_attempts = 0; + while (num_generated < dataset_size) { + if (num_attempts >= 10000000) + break; + num_attempts++; + array g(32); + const node* start; const node* end; + unsigned int current_node_index; + while (true) { + unsigned int num_vertices = std::max(2u, randrange(longest_path_length + 1)); + int backtrack = requested_backtrack; + if (requested_backtrack != -1) { + if (uniform) { + potential_backtrack_count = 0; + for (unsigned int i = 0; i < (unsigned) requested_backtrack + 1; i++) + if (num_generated == 0 || backtrack_distance_histogram[i] / num_generated < MAX_FREQS_PER_BUCKET[i]) + potential_backtracks[potential_backtrack_count++] = i; + backtrack = choice(potential_backtracks, potential_backtrack_count); + } + num_vertices = std::max((unsigned int) backtrack + 2, num_vertices); + } + if (!generate_dfs_example(g, start, end, current_node_index, path, num_vertices, max_input_size / 24 + 1, (max_input_size - 5) / 3, backtrack, longest_path_length)) { + for (node& n : g) core::free(n); + g.length = 0; path.length = 0; + continue; + } + break; + } + + array> edges(8); + for (node& vertex : g) + for (node* child : vertex.children) + edges.add(make_pair(vertex.id, child->id)); + if (edges.length > longest_path_length) { + for (node& n : g) core::free(n); + g.length = 0; path.length = 0; + continue; + } + shuffle(edges); + + array prefix(max_input_size); + for (auto& entry : edges) { + prefix[prefix.length++] = EDGE_PREFIX_TOKEN; + prefix[prefix.length++] = entry.key; + prefix[prefix.length++] = entry.value; + } + prefix[prefix.length++] = QUERY_PREFIX_TOKEN; + prefix[prefix.length++] = start->id; + prefix[prefix.length++] = end->id; + prefix[prefix.length++] = PATH_PREFIX_TOKEN; + + /* randomly select a vertex in the DFS trace */ + path.length = current_node_index + 1; + + array unvisited(4); + unsigned int backtrack_distance = max_input_size - 1; + for (unsigned int j = current_node_index + 1; j > 0; j--) { + for (const node* child : path[j-1]->children) { + if (!path.contains(child)) + unvisited.add(child); + } + if (unvisited.length != 0) { + backtrack_distance = current_node_index + 1 - j; + break; + } + } + + if (random_padding) { + /* sample padding randomly until the path is the correct length */ + unsigned int* padding_lengths = (unsigned int*) calloc(path.length + 1, sizeof(unsigned int)); + unsigned int available_padding = longest_path_length - path.length; + for (unsigned int j = path.length; j > 0; j--) { + padding_lengths[j] = randrange(available_padding + 1); + available_padding -= padding_lengths[j]; + } + padding_lengths[0] = available_padding; + + for (unsigned int j = 0; j < path.length; j++) { + for (unsigned int k = 0; k < padding_lengths[j]; k++) + prefix[prefix.length++] = PATH_PREFIX_TOKEN; + prefix[prefix.length++] = path[j]->id; + } + for (unsigned int k = 0; k < padding_lengths[path.length]; k++) + prefix[prefix.length++] = PATH_PREFIX_TOKEN; + unsigned int effective_backtrack = backtrack_distance; + for (unsigned int j = path.length; j > 0; j--) { + effective_backtrack += padding_lengths[j]; + if (path.length - j == backtrack_distance) + break; + } + backtrack_distance = effective_backtrack; + free(padding_lengths); + } else { + for (unsigned int j = path.length; j < longest_path_length; j++) + prefix[prefix.length++] = PATH_PREFIX_TOKEN; + for (unsigned int j = 0; j < path.length; j++) + prefix[prefix.length++] = path[j]->id; + } + + if ((requested_backtrack != -1 && !uniform && (unsigned int) requested_backtrack != backtrack_distance) + || (uniform && num_generated != 0 && backtrack_distance_histogram[backtrack_distance] / num_generated >= MAX_FREQS_PER_BUCKET[backtrack_distance])) + { + for (node& n : g) core::free(n); + g.length = 0; path.length = 0; + continue; + } + + /* check if this input is reserved */ + py::object contains = reserved_inputs.attr("__contains__"); + py::tuple example_tuple(prefix.length); + for (unsigned int i = 0; i < prefix.length; i++) + example_tuple[i] = prefix[i]; + if (contains(example_tuple).is(py_true)) { + for (node& n : g) core::free(n); + g.length = 0; path.length = 0; + num_collisions += 1; + continue; + } + + backtrack_distance_histogram[backtrack_distance]++; + + for (unsigned int i = 0; i < max_input_size - prefix.length; i++) + inputs_mem(num_generated, i) = PADDING_TOKEN; + for (unsigned int i = 0; i < prefix.length; i++) + inputs_mem(num_generated, max_input_size - prefix.length + i) = prefix[i]; + for (unsigned int i = 0; i < ntokens; i++) + outputs_mem(num_generated, i) = 0.0f; + for (unsigned int i = 0; i < unvisited.length; i++) + outputs_mem(num_generated, unvisited[i]->id) = 1.0f; + labels_mem(num_generated) = path[current_node_index+1]->id; + num_generated++; + + if (!quiet && num_generated > 0 && (num_generated % 1000 == 0 || num_generated >= dataset_size)) { + printf("%d examples generated.\n", num_generated); + fflush(stdout); + + printf("Backtrack distance histogram: (log frequencies)\n"); + printf("["); + bool first = true; + for (unsigned int i = 0; i < max_input_size; i++) { + if (backtrack_distance_histogram[i] == 0) + continue; + if (!first) printf(", "); + printf("%d:%.2f", i, log(backtrack_distance_histogram[i]) - log(num_generated)); + first = false; + } + printf("]\n"); + } + + for (node& n : g) core::free(n); + g.length = 0; path.length = 0; + continue; + } + + return py::make_tuple(inputs, outputs, labels, num_collisions); +} + +bool generate_si_example(array& vertices, const node*& start, const node*& end, unsigned int& current_node_index, array& path, unsigned int num_vertices, unsigned int max_num_parents, unsigned int max_vertex_id, unsigned int max_edges, unsigned int frontier_size, unsigned int branch_size) +{ + if (!vertices.ensure_capacity(num_vertices)) + return false; + for (unsigned int i = 0; i < num_vertices; i++) + if (!init(vertices[i], i)) return false; + vertices.length = num_vertices; + + /* the current vertex must have at least `branch_size` vertices following it */ + current_node_index = randrange(num_vertices - branch_size); + node* current_node = &vertices[current_node_index]; + + /* randomly sample the start and goal vertices */ + unsigned int start_index = randrange(min(current_node_index + 1, num_vertices - frontier_size)); + unsigned int end_index = randrange(max(current_node_index + 1, start_index + frontier_size), num_vertices); + start = &vertices[start_index]; end = &vertices[end_index]; + + /* sample some parent/ancestor vertices */ + constexpr float ALPHA = 1.0f; + unsigned int* out_degrees = (unsigned int*) calloc(num_vertices, sizeof(unsigned int)); + if (out_degrees == nullptr) + return false; + for (unsigned int i = 1; i < num_vertices; i++) { + /* sample the number of child and parent vertices */ + unsigned int num_parents = randrange(1, max_num_parents); + num_parents = std::min(num_parents, i); + + /* sample the parents of this new node */ + float total_probability = 0.0f; + array probabilities(i); + for (unsigned int j = 0; j < i; j++) { + probabilities[j] = ALPHA + out_degrees[j]; + total_probability += probabilities[j]; + } + probabilities.length = i; + if (current_node->children.length == branch_size) { + total_probability -= probabilities[current_node_index]; + probabilities[current_node_index] = 0.0f; + } + + array sampled_parents(std::max(1u, num_parents)); + for (unsigned int j = 0; j < num_parents; j++) { + unsigned int u = sample_categorical(probabilities.data, total_probability, probabilities.length); + sampled_parents.add(u); + total_probability -= probabilities[u]; + probabilities[u] = 0.0f; + if (total_probability < 1.0e-6f) + break; + } + + for (unsigned int parent_id : sampled_parents) { + vertices[parent_id].children.add(&vertices[i]); + vertices[i].parents.add(&vertices[parent_id]); + out_degrees[parent_id] += 1; + } + } + free(out_degrees); + + /* randomly sample the frontier (visited vertices with unvisited child nodes) */ + array frontier(8); + array available_nodes(8); + frontier.add(&vertices[start_index]); + if (start_index != current_node_index) + frontier.add(&vertices[current_node_index]); + for (unsigned int i = start_index + 1; i < vertices.length - 1; i++) + if (i != end_index && !frontier.contains(&vertices[i])) + available_nodes.add(&vertices[i]); + while (frontier.length < frontier_size) { + unsigned int r = randrange(available_nodes.length); + frontier.add(available_nodes[r]); + available_nodes.remove(r); + } + + unsigned int new_edge_count = 0; + array> removable_edges(16); + for (const node& vertex : vertices) + for (const node* child : vertex.children) + removable_edges.add(make_pair(vertex.id, child->id)); + unsigned int total_edge_count = removable_edges.length; + + /* construct a path such that all vertices in `frontier` are visited */ + array> queue(8); + for (node* child : start->children) + if (frontier.contains(child)) queue.add(make_pair(start, child)); + while (queue.length != 0) { + unsigned int r = randrange(queue.length); + const node* prev = queue[r].key; + const node* next = queue[r].value; + queue.remove(r); + if (path.contains(next)) + continue; + path.add(prev); + path.add(next); + removable_edges.remove(removable_edges.index_of(make_pair(prev->id, next->id))); + + for (node* child : next->children) + if (frontier.contains(child)) queue.add(make_pair(next, child)); + } + + /* get set of vertices reachable from `start` */ + array stack(8); + array reachable(8); + stack[stack.length++] = &vertices[start_index]; + while (stack.length != 0) { + node* current = stack.pop(); + if (reachable.contains(current)) + continue; + reachable.add(current); + + for (node* child : current->children) { + if (!reachable.contains(child) && !stack.contains(child)) + stack.add(child); + } + } + + /* for the unvisited vertices in `frontier`, move edges until they have an incoming edge from a visited vertex */ + for (node* frontier_vertex : frontier) { + if (frontier_vertex == start || path.index_of(frontier_vertex) % 2 == 1) + continue; + + /* get the most recent visited ancestors */ + array stack(8); + array visited_ancestors(8); + stack.append(frontier_vertex->parents.data, frontier_vertex->parents.length); + while (stack.length != 0) { + node* current = stack.pop(); + + bool has_visited_parent = false; + for (node* parent : current->parents) { + if (parent == start || path.contains(parent)) { + if (!visited_ancestors.contains(parent)) + visited_ancestors.add(parent); + has_visited_parent = true; + } + } + if (!has_visited_parent) + stack.append(current->parents.data, current->parents.length); + } + + if (current_node->children.length == branch_size && visited_ancestors.contains(current_node)) + visited_ancestors.remove(visited_ancestors.index_of(current_node)); + if (visited_ancestors.length != 0) { + node* new_parent = choice(visited_ancestors.data, visited_ancestors.length); + + /* replace an existing parent with `new_parent` */ + unsigned int r = randrange(frontier_vertex->parents.length); + node* old_parent = frontier_vertex->parents[r]; + frontier_vertex->parents[r] = new_parent; + old_parent->children.remove(old_parent->children.index_of(frontier_vertex)); + new_parent->children.add(frontier_vertex); + path.add(new_parent); + path.add(frontier_vertex); + removable_edges.remove(removable_edges.index_of(make_pair(old_parent->id, frontier_vertex->id))); + + /* move the new edge in `path` to a random position */ + unsigned int insert_path_index = path.index_of(new_parent); + if (path.length == 0) insert_path_index = 0; + else if (insert_path_index % 2 == 1) insert_path_index++; + unsigned int insert_index = randrange(insert_path_index, path.length); + insert_index -= insert_index % 2; + for (unsigned int j = path.length - 1; j >= insert_index + 2; j--) + path[j] = path[j - 2]; + path[insert_index] = new_parent; + path[insert_index + 1] = frontier_vertex; + } else { + /* there is no path from `start` to `frontier_vertex`, so add a new parent that is reachable from `start` */ + array candidate_parents(8); + for (node* n : reachable) { + if (n->id < frontier_vertex->id) + candidate_parents.add(n); + } + + if (current_node->children.length == branch_size && candidate_parents.contains(current_node) && candidate_parents.length != 1) + candidate_parents.remove(candidate_parents.index_of(current_node)); + + node* new_parent = choice(candidate_parents.data, candidate_parents.length); + frontier_vertex->parents.add(new_parent); + new_parent->children.add(frontier_vertex); + path.add(new_parent); + path.add(frontier_vertex); + + /* move the new edge in `path` to a random position */ + unsigned int insert_path_index = path.index_of(new_parent); + if (path.length == 0) insert_path_index = 0; + else if (insert_path_index % 2 == 1) insert_path_index++; + unsigned int insert_index = randrange(insert_path_index, path.length); + insert_index -= insert_index % 2; + for (unsigned int j = path.length - 1; j >= insert_index + 2; j--) + path[j] = path[j - 2]; + path[insert_index] = new_parent; + path[insert_index + 1] = frontier_vertex; + } + } + + /* next we have to make sure every vertex in `frontier` has an unvisited child */ + for (node* frontier_vertex : frontier) { + bool has_unvisited_child = false; + for (const node* child : frontier_vertex->children) { + if (!path.contains(child)) { + has_unvisited_child = true; + break; + } + } + if (has_unvisited_child) continue; + + /* randomly select an unvisited vertex after `frontier_vertex` */ + array unvisited(8); + for (unsigned int i = frontier_vertex->id + 1; i < vertices.length; i++) + if (!path.contains(&vertices[i])) unvisited.add(i); + node* new_child = &vertices[choice(unvisited.data, unvisited.length)]; + + /* see if we can replace an unvisited parent node of `new_child` */ + unvisited.clear(); + for (unsigned int i = 0; i < new_child->parents.length; i++) + if (new_child->parents[i] != start && !path.contains(new_child->parents[i])) + unvisited.add(i); + if (unvisited.length != 0) { + unsigned int r = choice(unvisited.data, unvisited.length); + node* old_parent = new_child->parents[r]; + old_parent->children.remove(old_parent->children.index_of(new_child)); + frontier_vertex->children.add(new_child); + new_child->parents[r] = frontier_vertex; + removable_edges.remove(removable_edges.index_of(make_pair(old_parent->id, new_child->id))); + } else { + /* we can't remove any parents from `new_child` so simply add an edge */ + frontier_vertex->children.add(new_child); + new_child->parents.add(frontier_vertex); + total_edge_count++; new_edge_count++; + } + } + + /* next, we want to make sure `vertices[current_node_index]` has an edge to `branch_size` vertices that follows it */ + unsigned int current_branch_size = 0; + available_nodes.clear(); + array available_frontier_branches(8); + for (unsigned int i = current_node_index + 1; i < vertices.length; i++) { + available_nodes.add(&vertices[i]); + + if (frontier.contains(&vertices[i])) + available_frontier_branches.add(&vertices[i]); + } + unsigned int num_frontier_branches = 0; + while (available_frontier_branches.length > 0 && num_frontier_branches < max(2*frontier_size + branch_size, max_edges) - max_edges) { + unsigned int r = randrange(available_frontier_branches.length); + node* new_child = available_frontier_branches[r]; + available_frontier_branches.remove(r); + available_nodes.remove(available_nodes.index_of(new_child)); + + /* find the edge in the path that leads to `new_child` */ + unsigned int path_index; + for (path_index = 0; path_index < path.length; path_index += 2) + if (path[path_index + 1] == new_child) break; + + if (path[path_index] == current_node) { + /* this vertex is already visited from current */ + unsigned int index = removable_edges.index_of(make_pair(current_node_index, new_child->id)); + if (index < removable_edges.length) + removable_edges.remove(index); + current_branch_size++; + num_frontier_branches++; + continue; + } + + /* connect this vertex directly to `current_node` */ + array visited(8); + node* old_parent = (node*) path[path_index]; + old_parent->children.remove(old_parent->children.index_of(new_child)); + if (current_node->children.contains(new_child)) { + new_child->parents.remove(new_child->parents.index_of(old_parent)); + } else { + current_node->children.add(new_child); + new_child->parents[new_child->parents.index_of(old_parent)] = current_node; + } + unsigned int index = removable_edges.index_of(make_pair(old_parent->id, new_child->id)); + if (index < removable_edges.length) + removable_edges.remove(index); + index = removable_edges.index_of(make_pair(current_node_index, new_child->id)); + if (index < removable_edges.length) + removable_edges.remove(index); + + /* replace the corresponding edge from `path` */ + path[path_index] = current_node; + + current_branch_size++; + num_frontier_branches++; + } + while (current_branch_size < branch_size) { + unsigned int r = randrange(available_nodes.length); + node* new_child = available_nodes[r]; + available_nodes.remove(r); + + if (current_node->children.contains(new_child)){ + unsigned int index = removable_edges.index_of(make_pair(current_node_index, new_child->id)); + if (index < removable_edges.length) + removable_edges.remove(index); + current_branch_size++; + continue; + } + + /* see if we can replace an unvisited parent node of `new_child` */ + array unvisited(8); + for (unsigned int i = 0; i < new_child->parents.length; i++) + if (new_child->parents[i] != start && !path.contains(new_child->parents[i])) + unvisited.add(i); + if (unvisited.length != 0) { + unsigned int r = choice(unvisited.data, unvisited.length); + node* old_parent = new_child->parents[r]; + old_parent->children.remove(old_parent->children.index_of(new_child)); + current_node->children.add(new_child); + new_child->parents[r] = current_node; + removable_edges.remove(removable_edges.index_of(make_pair(old_parent->id, new_child->id))); + } else { + /* we can't remove any parents from `new_child` so simply add an edge */ + current_node->children.add(new_child); + new_child->parents.add(current_node); + total_edge_count++; new_edge_count++; + } + current_branch_size++; + } + + /* visit some of the children of `current_node` */ + //unsigned int visited_branches = randrange(branch_size); + //array available_nodes(8); + //available_nodes.append(current_node->children.data, current_node->children.length); + //unsigned int index = available_nodes.index_of(end); + //if (index != available_nodes.length) + // available_nodes.remove(index); /* make sure we don't visit the goal vertex prematurely */ + //unsigned int insert_path_index = path.index_of(current_node) + 1; + //if (path.length == 0) insert_path_index = 0; + //else if (insert_path_index % 2 == 1) insert_path_index++; + //for (unsigned int i = 0; i < visited_branches && available_nodes.length != 0; i++) { + /* count the number of branches with unvisited children */ + // unsigned int r = randrange(available_nodes.length); + // node* child = available_nodes[r]; + // available_nodes.remove(r); + + // path.add(current_node); + // path.add(child); + // unsigned int insert_index = randrange(insert_path_index, path.length); + // insert_index -= insert_index % 2; + // for (unsigned int j = path.length - 1; j >= insert_index + 2; j--) + // path[j] = path[j - 2]; + // path[insert_index] = current_node; + // path[insert_index + 1] = child; + //} + + /* make sure `end` is reachable from `start` */ + array> reachability_stack(8); + for (node* child : start->children) + reachability_stack.add(make_pair(child, start)); + array_map reverse_ptrs(8); + while (reachability_stack.length != 0) { + pair entry = reachability_stack.pop(); + node* next = entry.key; + const node* parent = entry.value; + if (next->id > end_index || reverse_ptrs.contains(next)) + continue; + reverse_ptrs.put(next, parent); + for (node* child : next->children) + if (!reachability_stack.contains(make_pair(child, next))) + reachability_stack.add(make_pair(child, next)); + } + const node* current; + const node* parent; + if (reverse_ptrs.contains(&vertices[end_index])) { + current = &vertices[end_index]; + parent = reverse_ptrs.get(current); + } else { + reverse_ptrs.put(&vertices[start_index], nullptr); + node* new_parent = choice(reverse_ptrs.keys, reverse_ptrs.size); + new_parent->children.add(&vertices[end_index]); + vertices[end_index].parents.add(new_parent); + total_edge_count++; new_edge_count++; + + current = &vertices[end_index]; + parent = new_parent; + } + /* make sure none of the edges on the path from `start` to `end` are removable */ + while (true) { + pair entry = make_pair(parent->id, current->id); + unsigned int index = removable_edges.index_of(entry); + if (index != removable_edges.length) + removable_edges.remove(index); + if (parent == start) + break; + current = parent; + parent = reverse_ptrs.get(current); + } + + /* remove edges to avoid generating a graph with too many edges */ + unsigned int edges_to_remove; + if (total_edge_count > max_edges) + edges_to_remove = total_edge_count - max_edges; + else + edges_to_remove = new_edge_count; + for (unsigned int i = 0; i < edges_to_remove && removable_edges.length != 0; i++) { + unsigned int u = randrange(removable_edges.length); + unsigned int parent_id = removable_edges[u].key; + unsigned int child_id = removable_edges[u].value; + removable_edges.remove(u); + + /* make sure that removing this edge will not remove an unvisited child node from any frontier vertex */ + if (frontier.contains(&vertices[parent_id])) { + bool can_remove_edge = false; + for (const node* child : vertices[parent_id].children) { + if (child->id != child_id && !path.contains(child)) { + can_remove_edge = true; + break; + } + } + if (!can_remove_edge) { + i--; + continue; + } + } + + vertices[parent_id].children.remove(vertices[parent_id].children.index_of(&vertices[child_id])); + vertices[child_id].parents.remove(vertices[child_id].parents.index_of(&vertices[parent_id])); + } + + /* visit some additional vertices */ + array path_length_probabilities(8); + for (unsigned int i = path.length / 2; i < max_edges + 1; i++) + path_length_probabilities.add((float) i); + unsigned int target_path_length = path.length / 2 + sample_categorical(path_length_probabilities.data, path_length_probabilities.length); + available_nodes.clear(); + for (node* frontier_vertex : frontier) { + unsigned int unvisited_child_count = 0; + for (const node* child : frontier_vertex->children) + if (!path.contains(child)) unvisited_child_count++; + + if (unvisited_child_count > 1) + available_nodes.add(frontier_vertex); + } + while (available_nodes.length != 0 && path.length < target_path_length * 2) { + unsigned int u = randrange(available_nodes.length); + node* selected_vertex = available_nodes[u]; + available_nodes.remove(u); + + array unvisited_children(4); + for (node* child : selected_vertex->children) { + if (path.contains(child)) continue; + + /* make sure this is not the unvisited child of another frontier vertex */ + bool can_visit = true; + for (const node* parent : child->parents) { + if (!frontier.contains(parent)) continue; + unsigned int unvisited_child_count = 0; + for (const node* other_child : parent->children) + if (!path.contains(other_child)) unvisited_child_count++; + if (unvisited_child_count == 1) { + can_visit = false; + break; + } + } + if (can_visit) + unvisited_children.add(child); + } + if (unvisited_children.length == 0) + continue; + u = randrange(unvisited_children.length); + node* selected_child = unvisited_children[u]; + path.add(selected_vertex); + path.add(selected_child); + if (frontier.contains(selected_vertex) && unvisited_children.length - 1 > 1) + available_nodes.add(selected_vertex); + else if (!frontier.contains(selected_vertex) && unvisited_children.length - 1 > 0) + available_nodes.add(selected_vertex); + + /* check if `selected_child` has any unvisited child nodes */ + unsigned int unvisited_child_count = 0; + for (const node* child : selected_child->children) + if (!path.contains(child)) unvisited_child_count++; + if (unvisited_child_count > 0) + available_nodes.add(selected_child); + } + + /* remove any correlation between graph topology and vertex IDs by shuffling the vertices */ + unsigned int* new_indices = (unsigned int*) alloca(sizeof(unsigned int) * (max_vertex_id + 1)); + for (unsigned int i = 0; i < max_vertex_id + 1; i++) new_indices[i] = i; + shuffle(new_indices, max_vertex_id + 1); + unsigned int src_index = 0; + for (unsigned int i = 0; i < vertices.length; i++) { + bool is_reserved = false; + for (unsigned int j = 0; j < array_length(RESERVED_INDICES); j++) { + if (new_indices[src_index] == RESERVED_INDICES[j]) { + is_reserved = true; + break; + } + } + if (is_reserved) + src_index++; + vertices[i].id = new_indices[src_index]; + src_index++; + } + + return true; +} + +py::tuple generate_si_training_set(const unsigned int max_input_size, const uint64_t dataset_size, const py::object& reserved_inputs, const int requested_frontier_size, const int requested_branch_size, bool uniform, const bool quiet=false) +{ + const unsigned int QUERY_PREFIX_TOKEN = (max_input_size-5) / 3 + 4; + const unsigned int PADDING_TOKEN = (max_input_size-5) / 3 + 3; + const unsigned int EDGE_PREFIX_TOKEN = (max_input_size-5) / 3 + 2; + const unsigned int PATH_PREFIX_TOKEN = (max_input_size-5) / 3 + 1; + + unsigned int max_edges = (max_input_size - 2) / 6; + + unsigned int num_generated = 0; + unsigned int num_collisions = 0; + unsigned int ntokens = (max_input_size - 5) / 3 + 5; + size_t input_shape[2]{dataset_size, max_input_size}; + size_t output_shape[2]{dataset_size, ntokens}; + size_t label_shape[1]{dataset_size}; + py::array_t inputs(input_shape); + py::array_t outputs(output_shape); + py::array_t labels(label_shape); + auto inputs_mem = inputs.mutable_unchecked<2>(); + auto outputs_mem = outputs.mutable_unchecked<2>(); + auto labels_mem = labels.mutable_unchecked<1>(); + + unsigned int* frontier_branch_histogram = (unsigned int*) alloca(sizeof(unsigned int) * (max_edges + 1) * (max_edges + 1)); + for (unsigned int i = 0; i < (max_edges + 1) * (max_edges + 1); i++) + frontier_branch_histogram[i] = 0; + unsigned int* visited_edges_histogram = (unsigned int*) alloca(sizeof(unsigned int) * (max_edges + 1)); + for (unsigned int i = 0; i < max_edges + 1; i++) + visited_edges_histogram[i] = 0; + + float* MAX_FREQS_PER_BUCKET = (float*) alloca(sizeof(float) * (max_edges + 1) * (max_edges + 1)); + unsigned int nonzero_buckets = 0; + for (unsigned int i = 0; i < max_edges + 1; i++) { + /* frontier_size is i */ + for (unsigned int j = 0; j < max_edges + 1; j++) { + /* branch_size is j */ + if (i == 0 || j == 0 || i > (unsigned) requested_frontier_size || j > (unsigned) requested_branch_size || i + j > max_edges + 1) { + MAX_FREQS_PER_BUCKET[i*(max_edges+1) + j] = 0.0; + } else { + MAX_FREQS_PER_BUCKET[i*(max_edges+1) + j] = 1.0; + nonzero_buckets++; + } + } + } + for (unsigned int i = 0; i < (max_edges + 1) * (max_edges + 1); i++) + if (MAX_FREQS_PER_BUCKET[i] != 0.0) MAX_FREQS_PER_BUCKET[i] /= nonzero_buckets; + if ((unsigned) requested_branch_size < max_edges + 1 - requested_frontier_size) + MAX_FREQS_PER_BUCKET[requested_frontier_size*(max_edges+1) + requested_branch_size] += 0.05 / nonzero_buckets; + else + MAX_FREQS_PER_BUCKET[requested_frontier_size*(max_edges+1) + max_edges + 1 - requested_frontier_size] += 0.05 / nonzero_buckets; + + array path(32); + pair* potential_frontier_branches = (pair*) alloca(max((size_t) 1, sizeof(pair) * ((max_edges + 1) * (max_edges + 1)))); + unsigned int potential_frontier_branch_count = 0; + unsigned int num_attempts = 0; + while (num_generated < dataset_size) { + if (num_attempts >= 10000000) + break; + num_attempts++; + array g(32); + const node* start; const node* end; + unsigned int current_node_index; + while (true) { + unsigned int frontier_size; unsigned int branch_size; + if (uniform) { + potential_frontier_branch_count = 0; + for (unsigned int i = 1; i < max_edges + 1; i++) { + for (unsigned int j = 1; j < max_edges + 1; j++) { + if ((num_generated == 0 && MAX_FREQS_PER_BUCKET[i*(max_edges+1) + j] != 0.0) || (float) frontier_branch_histogram[i*(max_edges+1) + j] / num_generated < MAX_FREQS_PER_BUCKET[i*(max_edges+1) + j]) + potential_frontier_branches[potential_frontier_branch_count++] = make_pair(i, j); + } + } + pair frontier_branch = choice(potential_frontier_branches, potential_frontier_branch_count); + frontier_size = frontier_branch.key; + branch_size = frontier_branch.value; + } else { + frontier_size = requested_frontier_size; + branch_size = requested_branch_size; + } + unsigned int num_vertices = std::max(2u, randrange(max_edges - frontier_size + 1)); + num_vertices = std::max(frontier_size + branch_size + 1 - std::min(frontier_size, branch_size), num_vertices); + if (!generate_si_example(g, start, end, current_node_index, path, num_vertices, max_input_size / 24 + 1, (max_input_size - 2) / 6 + 2, max_edges, frontier_size, branch_size)) { + for (node& n : g) core::free(n); + g.length = 0; path.length = 0; + continue; + } + break; + } + + array> edges(8); + for (node& vertex : g) + for (node* child : vertex.children) + edges.add(make_pair(vertex.id, child->id)); + if (edges.length > max_edges) { + for (node& n : g) core::free(n); + g.length = 0; path.length = 0; + continue; + } + shuffle(edges); + + array prefix(max_input_size); + for (auto& entry : edges) { + prefix[prefix.length++] = EDGE_PREFIX_TOKEN; + prefix[prefix.length++] = entry.key; + prefix[prefix.length++] = entry.value; + } + prefix[prefix.length++] = QUERY_PREFIX_TOKEN; + prefix[prefix.length++] = start->id; + prefix[prefix.length++] = end->id; + prefix[prefix.length++] = PATH_PREFIX_TOKEN; + + /* compute the frontier size and branch size of this example */ + array frontier(8); + for (unsigned int i = 0; i < path.length; i++) { + bool has_unvisited_children = false; + for (const node* child : path[i]->children) { + if (!path.contains(child)) { + has_unvisited_children = true; + break; + } + } + if (has_unvisited_children && !frontier.contains(path[i])) + frontier.add(path[i]); + } + if (frontier.length == 0) + frontier.add(start); + const node* current_node = &g[current_node_index]; + unsigned int branch_size = current_node->children.length; + + bool is_selection_step; + if (path.length == 0) + is_selection_step = false; + else + is_selection_step = (randrange(2) == 0); + if (3*(path.length/2) + (is_selection_step ? 1 : 2) > 3*(max_edges - 1) + 1) { + /* we have just barely too many edges */ + for (node& n : g) core::free(n); + g.length = 0; path.length = 0; + continue; + } + for (unsigned int j = 3*(path.length/2) + (is_selection_step ? 1 : 2); j < 3*(max_edges - 1) + 1; j++) + prefix[prefix.length++] = PATH_PREFIX_TOKEN; + for (unsigned int j = 0; j < path.length; j += 2) { + prefix[prefix.length++] = PATH_PREFIX_TOKEN; + prefix[prefix.length++] = path[j]->id; + prefix[prefix.length++] = path[j+1]->id; + } + prefix[prefix.length++] = PATH_PREFIX_TOKEN; + if (!is_selection_step) + prefix[prefix.length++] = current_node->id; + + if ((requested_frontier_size != -1 && !uniform && (unsigned int) requested_frontier_size != (unsigned int) frontier.length) + || (requested_branch_size != -1 && !uniform && (unsigned int) requested_branch_size != branch_size) + || (uniform && num_generated != 0 && (float) frontier_branch_histogram[frontier.length*(max_edges+1) + branch_size] / num_generated >= MAX_FREQS_PER_BUCKET[frontier.length*(max_edges+1) + branch_size])) + { + for (node& n : g) core::free(n); + g.length = 0; path.length = 0; + continue; + } + + /* check if this input is reserved */ + py::object contains = reserved_inputs.attr("__contains__"); + py::tuple example_tuple(prefix.length); + for (unsigned int i = 0; i < prefix.length; i++) + example_tuple[i] = prefix[i]; + if (contains(example_tuple).is(py_true)) { + for (node& n : g) core::free(n); + g.length = 0; path.length = 0; + num_collisions += 1; + continue; + } + + frontier_branch_histogram[frontier.length*(max_edges+1) + branch_size]++; + visited_edges_histogram[path.length/2]++; + + for (unsigned int i = 0; i < max_input_size - prefix.length; i++) + inputs_mem(num_generated, i) = PADDING_TOKEN; + for (unsigned int i = 0; i < prefix.length; i++) + inputs_mem(num_generated, max_input_size - prefix.length + i) = prefix[i]; + for (unsigned int i = 0; i < ntokens; i++) + outputs_mem(num_generated, i) = 0.0f; + if (is_selection_step) { + for (const node* frontier_vertex : frontier) + outputs_mem(num_generated, frontier_vertex->id) = 1.0f; + labels_mem(num_generated) = choice(frontier.data, frontier.length)->id; + } else { + array correct_answers(8); + for (const node* child : current_node->children) { + if (!path.contains(child)) { + outputs_mem(num_generated, child->id) = 1.0f; + correct_answers.add(child->id); + } + } + labels_mem(num_generated) = choice(correct_answers.data, correct_answers.length); + } + num_generated++; + + if (!quiet && num_generated > 0 && (num_generated % 1000 == 0 || num_generated >= dataset_size)) { + printf("%d examples generated.\n", num_generated); + fflush(stdout); + + printf("Frontier-branch size histogram: (log frequencies)\n"); + printf("["); + bool first = true; + for (unsigned int i = 0; i < max_edges + 1; i++) { + for (unsigned int j = 0; j < max_edges + 1; j++) { + if (frontier_branch_histogram[i*(max_edges+1) + j] == 0) + continue; + if (!first) printf(", "); + printf("(%d,%d):%.2f", i, j, log(frontier_branch_histogram[i*(max_edges+1) + j]) - log(num_generated)); + first = false; + } + } + printf("]\n"); + + printf("Visited edge count histogram: (log frequencies)\n"); + printf("["); + first = true; + for (unsigned int i = 0; i < max_edges + 1; i++) { + if (visited_edges_histogram[i] == 0) + continue; + if (!first) printf(", "); + printf("%d:%.2f", i, log(visited_edges_histogram[i]) - log(num_generated)); + first = false; + } + printf("]\n"); + } + + for (node& n : g) core::free(n); + g.length = 0; path.length = 0; + continue; + } + + return py::make_tuple(inputs, outputs, labels, num_collisions); +} + +PYBIND11_MODULE(generator, m) { + m.def("generate_training_set", &generate_training_set); + m.def("generate_reachable_training_set", &generate_reachable_training_set); + m.def("generate_dfs_training_set", &generate_dfs_training_set); + m.def("generate_si_training_set", &generate_si_training_set); + m.def("lookahead_histogram", &lookahead_histogram); + m.def("set_seed", &core::set_seed); +} diff --git a/reasoners/data_preprocess/logic/graph_dataset_gen/jobfile.sub b/reasoners/data_preprocess/logic/graph_dataset_gen/jobfile.sub new file mode 100644 index 000000000..8dfaf35f8 --- /dev/null +++ b/reasoners/data_preprocess/logic/graph_dataset_gen/jobfile.sub @@ -0,0 +1,11 @@ +#!/bin/bash +# FILENAME: jobfile.sub +#SBATCH --nodes=1 +#SBATCH --time=4:00:00 +#SBATCH --job-name gen_graph_dataset +#SBATCH --gpus-per-node=1 +module load conda +source ~/.bashrc +conda activate llm_env +module list +python logic.py --num_samples 10000 --output_file graph_dataset.json \ No newline at end of file diff --git a/reasoners/data_preprocess/logic/graph_dataset_gen/logic.py b/reasoners/data_preprocess/logic/graph_dataset_gen/logic.py new file mode 100644 index 000000000..04aa92cf8 --- /dev/null +++ b/reasoners/data_preprocess/logic/graph_dataset_gen/logic.py @@ -0,0 +1,359 @@ +import os +import random +import sysconfig +from io import StringIO +from faker import Faker +from pybind11.__main__ import print_includes +import math +import argparse +import json + +def build_module(name): + import sys + old_stdout = sys.stdout + try: + sys.stdout = StringIO() + print_includes() + includes = sys.stdout.getvalue().strip() + sys.stdout.close() + sys.stdout = old_stdout + except Exception as e: + raise e + finally: + sys.stdout = old_stdout + + python_extension_suffix = sysconfig.get_config_var("EXT_SUFFIX") + + if sys.platform == "darwin": + # macOS command + command = ( + f"g++ -std=c++11 -Ofast -DNDEBUG -fno-stack-protector " + f"-Wall -Wpedantic -undefined dynamic_lookup -shared -fPIC " + f"{includes} -I. {name}.cpp -o {name}{python_extension_suffix}" + ) + else: + # Non-macOS command + command = ( + f"g++ -Ofast -std=c++11 -DNDEBUG -fno-stack-protector " + f"-Wall -Wpedantic -shared -fPIC " + f"{includes} -I. {name}.cpp -o {name}{python_extension_suffix}" + ) + + print(command) + if os.system(command) != 0: + print(f"ERROR: Unable to compile `{name}.cpp`.") + sys.exit(1) + +try: + from os.path import getmtime + from importlib.util import find_spec + + generator_spec = find_spec('generator') + if generator_spec == None: + raise ModuleNotFoundError + if getmtime(generator_spec.origin) < getmtime('generator.cpp'): + print("C++ module `generator` is out-of-date. Compiling from source...") + build_module("generator") + import generator +except ModuleNotFoundError: + print("C++ module `generator` not found. Compiling from source...") + build_module("generator") + import generator +except ImportError: + print("Error loading C++ module `generator`. Compiling from source...") + build_module("generator") + import generator +print("C++ module `generator` loaded.") + + +def generate_graph_text( + max_input_size, + num_vertices, + max_num_parents, + lookahead, + num_paths, +): + + # Get graph using the C++ generator + inputs, outputs, labels, num_collisions = generator.generate_training_set( + max_input_size, + 1, # batch size + lookahead, + max_num_parents * num_vertices, # max edges + set(), # reserved vertices + -1, # distance from start + 0, # max_prefix_vertices + True, # quiet + num_paths # number of paths + ) + + # If the random DAG fails to generate a valid example, return None + if inputs is None or outputs is None: + return None + + # Token IDs + PADDING_TOKEN = (max_input_size - 5) // 3 + 3 + EDGE_PREFIX_TOKEN = (max_input_size - 5) // 3 + 2 + PATH_PREFIX_TOKEN = (max_input_size - 5) // 3 + 1 + QUERY_PREFIX_TOKEN = (max_input_size - 5) // 3 + 4 + + # Function for mapping tokens to letters + def token_to_str(token): + if token == EDGE_PREFIX_TOKEN: + return "E" + elif token == QUERY_PREFIX_TOKEN: + return "Q" + elif token == PATH_PREFIX_TOKEN: + return "P" + # return "" + elif token == PADDING_TOKEN: + return "" + else: + return str(token) + + # Convert all tokens to letters + text_tokens = [token_to_str(t) for t in inputs[0] if token_to_str(t) != ""] + + final_str = " ".join(text_tokens) + return final_str, labels[0] + + +def generate_name(n): + """ + Generates n random names using the Faker library. + """ + fake = Faker() + # return [fake.unique.first_name() for i in range(n)] + return [fake.unique.name() for i in range(n)] + + +def random_syllable(): + """ + Generates a random syllable based on a randomly chosen pattern. + The pattern is built from 'C' (consonant) and 'V' (vowel). + """ + vowels = "aeiou" + consonants = "bcdfghjklmnpqrstvwxyz" + # Syllable patterns (e.g. "CV" means a consonant followed by a vowel) + patterns = ["CV", "VC", "CVC", "CVV", "CCV", "VCV", "VCC"] + pattern = random.choice(patterns) + + syllable = "" + for char in pattern: + if char == "C": + syllable += random.choice(consonants) + elif char == "V": + syllable += random.choice(vowels) + return syllable + + +def generate_fake_noun(): + """ + Generates a fake noun by concatenating two randomly generated syllables. + """ + return random_syllable() + random_syllable() + + +def generate_fake_nouns(n): + """ + Returns a list of n fake nouns. + """ + nouns = set() + while len(nouns) < n: + nouns.add(generate_fake_noun()) + return list(nouns) + + +# def generate_words(n): +# fake = Faker() +# return fake.words(nb=n) + + +def logic_paragraph_from_tokens(tokens_str: str, next_step: int, use_diff_names=True): + # Split string into tokens + tokens = tokens_str.strip().split() + + edges = [] + queries = [] + + # Keep track of nodes in a set + node_ids = set() + + i = 0 + while i < len(tokens): + t = tokens[i] + + if t == "E": + # Expect "E A B" + if i + 2 >= len(tokens): + break # malformed + A = int(tokens[i + 1]) + B = int(tokens[i + 2]) + edges.append((A, B)) + node_ids.update([A, B]) + i += 3 + + elif t == "Q": + # Expect "Q X Y" + if i + 2 >= len(tokens): + break + X = int(tokens[i + 1]) + Y = int(tokens[i + 2]) + queries.append((X, Y)) + node_ids.update([X, Y]) + i += 3 + + elif t == "P": + # "P" is just a path prefix token; ignore + i += 1 + + else: + # Possibly a stray token or something else (like a vertex ID alone). + i += 1 + + # Generate fake names and fake adjectives + num_nodes = len(node_ids) + all_names = generate_name(num_nodes) + all_adjs = generate_fake_nouns(num_nodes) + + sorted_nodes = sorted(node_ids) + id_to_pair = {} # node_id -> (name, adjective) + # Assign all nodes the same name if use_diff_names == False + for idx, node_id in enumerate(sorted_nodes): + id_to_pair[node_id] = (all_names[idx if use_diff_names else 0], all_adjs[idx]) + + # Lines of logic + lines = ["Given the following list of predicates:"] + + def get_logic_line(name_a: str, name_b: str, adj_a: str, adj_b: str) -> str: + choices = [ + f"If {name_a} is {adj_a}, then {name_b} is {adj_b}.", + f"{name_a} is {adj_a} implies {name_b} is {adj_b}.", + f"{name_b} is {adj_b} is true if {name_a} is {adj_a}.", + # f"{adj_b} is true if {adj_a} is true.", + # f"If {adj_a} then {adj_b} is true.", + # f"If {adj_a} is true then {adj_b}.", + f"Given {name_a} is {adj_a} then {name_b} is {adj_b}.", + ] + + sentence = random.choice(choices) + return sentence[0].upper() + sentence[1:] + + # For each edge: E A B => "If name(A) is adj(A), then name(A) is adj(B)." + for (A, B) in edges: + name_A, adj_A = id_to_pair[A] + name_B, adj_B = id_to_pair[B] + lines.append(get_logic_line(name_A, name_B, adj_A, adj_B)) + + # For each query: Q X Y => "If name(X) is adj(X), prove that name(X) is adj(Y)." + for (X, Y) in queries: + name_x, adj_x = id_to_pair[X] + name_y, adj_y = id_to_pair[Y] + lines.append(f"\n\nIf {name_x} is {adj_x}, what is the next step to prove that {name_y} is {adj_y}?") + + # Join all lines into one paragraph + paragraph = " ".join(lines) + + # Get the correct adjective for the next step + next_step_adj = id_to_pair[next_step] + next_step_adj = id_to_pair[next_step][1] + + return paragraph, next_step_adj + + + +def main(samples_per_test: int = 3, lookahead_range: list = range(1, 5), num_paths: int = 2, max_num_parents: int = 3, logic: bool = False, seed: int = None, verbose: bool = True, print_prompts: bool = False, output_file:str ="graph_search_data.json", model: str = "gpt-4o"): + if seed is not None: + random.seed(seed) + generator.set_seed(seed) + + graph_search_puzzles = [] + correct_responses = [] + look_ahead_values = [] + counter = -1 + for look_ahead in lookahead_range: + for _ in range(samples_per_test): + txt, next_step = generate_graph_text( + max_input_size=max_num_parents * look_ahead * num_paths * 4, + num_vertices=max_num_parents * look_ahead, + max_num_parents=max_num_parents, + lookahead=look_ahead, + num_paths=num_paths, + ) + counter += 1 + # Create the prompt for the graph search + prompt = (f"{txt}\nAbove is a representation of a directed graph search problem, " + f"where E A B represents an edge from A to B, and Q X Y represents starting from X and ending at Y, " + f"find the shortest path. The vertex after P indicates our current position. Respond with only the " + f"next vertex on the shortest path from X to Y and nothing else.") + + # Change the prompt to a logic puzzle if the logic option is enabled + if logic: + logic, next_step_adj = logic_paragraph_from_tokens(txt, next_step) + prompt = f"{logic} Respond with only the trait of the next step." + + + correct_response = next_step_adj if logic else next_step + puzzle = { + "id": counter, + "prompt": prompt, + "correct_response": correct_response, + "look_ahead": look_ahead + } + graph_search_puzzles.append(puzzle) + if print_prompts: + print(f"Prompt: {prompt}\n") + print(f"Correct: {next_step_adj if logic else next_step}\n") + + with open(output_file, "w") as json_file: + json.dump(graph_search_puzzles, json_file, indent=2) + + +def multiplicative_range(start, stop, step): + result = [] + current = start + while (step > 1 and current < stop) or (step < 1 and current > stop): + result.append(current) + current *= step + return result + +def exponential_range(base, start, end): + res = [round(base ** n) for n in range(start, end + 1)] + res = list(set(res)) + return sorted(res) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Generate graph search problems.") + parser.add_argument("--output_dir", type=str, default="../../../data/raw", help="output directory for generated problems") + parser.add_argument("--output_file", type=str, default="graph_search.json", help="output file for generated problems") + parser.add_argument("--num_samples", type=int, default=10000, help="number of samples to generate") + args = parser.parse_args() + # Set the output directory and file + # base_dir = os.path.dirname(os.path.abspath(__file__)) + output_dir = args.output_dir + # Create the output directory if it doesn't exist + if not os.path.exists(output_dir): + os.makedirs(output_dir) + output_file = os.path.join(output_dir, args.output_file) + # Set the random seed for reproducibility + seed = 42 + random.seed(seed) + look_ahead = exponential_range(1.3, 1, 24) + print(f"look_ahead range={look_ahead}") + samples_per_test = round(args.num_samples/len(look_ahead)) + print(f"samples_per_test={samples_per_test}") + main( + samples_per_test=samples_per_test, + lookahead_range=look_ahead, + num_paths=9, + logic=True, + verbose=False, + print_prompts=False, + output_file = output_file, + seed=seed, + ) + +# example usage: +# python graph_search.py --output_dir ../data/graph_dataset --output_file graph_search.json --num_samples 10000 + diff --git a/reasoners/data_preprocess/logic/process_graph_dataset.py b/reasoners/data_preprocess/logic/process_graph_dataset.py new file mode 100644 index 000000000..01279a61c --- /dev/null +++ b/reasoners/data_preprocess/logic/process_graph_dataset.py @@ -0,0 +1,159 @@ +"""Downloads, processes, and saves graph logical reasoning dataset.""" + +import os +import datasets +import argparse +import re +from sklearn.model_selection import train_test_split +from transformers import AutoTokenizer + +from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset + +# Constants +INSTRUCTION_FOLLOW = "Please put your answer within and tags, for example fdebme ." + +def get_dataset(json_path): + """ + Load the graph dataset from a JSON file. + + Args: + json_path (str): Path to the JSON file + + Returns: + Dataset: Loaded dataset + """ + print(f"Loading the dataset from {json_path}...") + return datasets.load_dataset('json', data_files=json_path)['train'] + +def make_map_fn(split: str, data_source: str) -> callable: + """ + Create a mapping function for processing dataset examples. + + Args: + split (str): Data split ('train' or 'test') + data_source (str): Name of the data source + + Returns: + callable: Function to map over the dataset + """ + def process_fn(example, idx): + question = example['prompt'] + lookahead = example['look_ahead'] + answer = example['correct_response'] + + # Create user message with instructions + formatted_question = question + " " + INSTRUCTION_FOLLOW + + data = { + "data_source": data_source, + "prompt": [{ + "role": "user", + "content": formatted_question, + }], + "ability": "logical_reasoning", + "apply_chat_template": True, + "reward_model": { + "style": "rule", + "ground_truth": answer + }, + "extra_info": { + 'id': example['id'] if 'id' in example else str(idx), + 'lookahead': lookahead, + 'split': split + } + } + + if idx == 0: + print(f"data_source: {data_source}, split: {split}, idx: {idx}") + print("\n" + "=" * 100 + f"{data_source} {split} {idx}" + "=" * 10) + print(data) + + return data + + return process_fn + +if __name__ == '__main__': + """Main script execution: parse args, load, process, and save datasets.""" + parser = argparse.ArgumentParser() + parser.add_argument('--json_path', default='data/raw/graph_search.json', help='Path to json file') + parser.add_argument('--output_dir', default='data', help='Directory to save processed data') + parser.add_argument('--hdfs_dir', default=None, help='HDFS directory (optional)') + parser.add_argument('--train_size', type=float, default=0.8, help='Proportion of data for train set') + parser.add_argument('--test_size', type=float, default=0.2, help='Proportion of data for test set') + parser.add_argument('--data_source', default='graph_logical_dataset', help='Name of data source') + parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility when splitting data') + parser.add_argument('--train_sample_size', type=int, default=None, help='Number of samples to use from train. If None, use all.') + args = parser.parse_args() + + # Set random seed for reproducibility + set_seed(args.seed) + + # Validate train and test sizes + if args.train_size + args.test_size > 1.0: + raise ValueError(f"The sum of train_size ({args.train_size}) and test_size ({args.test_size}) cannot exceed 1.0") + + # Load the dataset + dataset = get_dataset(args.json_path) + + # Split dataset into train and test + train_indices, test_indices = train_test_split( + range(len(dataset)), + train_size=args.train_size, + test_size=args.test_size, + random_state=args.seed + ) + + print(f"Train set size: {len(train_indices)}, Test set size: {len(test_indices)}") + + # Process the datasets + process_train_fn = make_map_fn('train', args.data_source) + process_test_fn = make_map_fn('test', args.data_source) + + train_dataset = dataset.select(train_indices).map(function=process_train_fn, with_indices=True) + test_dataset = dataset.select(test_indices).map(function=process_test_fn, with_indices=True) + + # Store the original training dataset size + original_train_size = len(train_dataset) + + # Sample the training dataset if needed + train_dataset = sample_dataset(train_dataset, args.train_sample_size) + + # Create output directories + train_output_dir = os.path.join(args.output_dir, "train") + test_output_dir = os.path.join(args.output_dir, "test") + os.makedirs(train_output_dir, exist_ok=True) + os.makedirs(test_output_dir, exist_ok=True) + + # Save train dataset + train_output_path = save_dataset( + dataset=train_dataset, + output_dir=train_output_dir, + filename_prefix=f"logic__{args.data_source}", + sample_size=args.train_sample_size if args.train_sample_size else len(train_dataset) + ) + + # Save test dataset + test_output_path = save_dataset( + dataset=test_dataset, + output_dir=test_output_dir, + filename_prefix=f"logic__{args.data_source}", + sample_size=len(test_dataset) + ) + + # Copy to HDFS if specified + if args.hdfs_dir is not None: + try: + from verl.utils.hdfs_io import copy, makedirs + makedirs(args.hdfs_dir) + copy(src=args.output_dir, dst=args.hdfs_dir) + print(f"Data copied to HDFS: {args.hdfs_dir}") + except ImportError: + print("HDFS utilities not available. Install verl package for HDFS support.") + + print(f"Done! \n" + f"Train data saved to {train_output_path}\n" + f"Test data saved to {test_output_path}") + print(f"Original train set size: {original_train_size} examples") + print(f"Final train set size: {len(train_dataset)} examples") + print(f"Test set: {len(test_dataset)} examples") + diff --git a/reasoners/data_preprocess/logic/process_puzzles_dataset.py b/reasoners/data_preprocess/logic/process_puzzles_dataset.py new file mode 100644 index 000000000..21f8f8083 --- /dev/null +++ b/reasoners/data_preprocess/logic/process_puzzles_dataset.py @@ -0,0 +1,198 @@ +# Copyright 2024 +# +# 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. +""" +Preprocess a puzzle dataset to parquet format. +""" +import os +import datasets +import argparse +import re +from sklearn.model_selection import train_test_split + +from transformers import ( + LlamaForCausalLM, + LlamaTokenizer, + AutoTokenizer, +) + +from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset + +# Constants +INSTRUCTION_FOLLOW = "Please put your answer within and tags, for example ['pigeon', 'sparrow', 'quail'] ." + +def make_prefix(dp): + """ + Create a suitable prompt prefix. + + Args: + dp (dict): Data point containing input, instruction, etc. + + Returns: + str: Formatted prompt + """ + constraints = dp['input'] + result = dp['ground_truth'] + instruction = dp['instruction'] + prefix = f"{instruction}. The constraints are: {constraints}." + return prefix + + +def extract_from_ground_truth(text): + """ + Extract answer from the ground truth text. + + Args: + text: Ground truth answer, could be string or list + + Returns: + list: Extracted answer + """ + if isinstance(text, list): + return text + else: + return eval(text) + + +def make_map_fn(split: str, data_source: str) -> callable: + """ + Create a mapping function for processing dataset examples. + + Args: + split (str): Data split ('train' or 'test') + data_source (str): Name of the data source + + Returns: + callable: Function to map over the dataset + """ + def process_fn(example, idx): + # Generate the appropriate question format based on model type + question = make_prefix(example) + num_objects = example['num_objects'] + final_arrangement = extract_from_ground_truth(example['ground_truth']) + + data = { + "data_source": data_source, + "prompt": [{ + "role": "user", + "content": question + " " + INSTRUCTION_FOLLOW + }], + "ability": "logical_reasoning", + "reward_model": { + "style": "rule", + "ground_truth": final_arrangement + }, + "apply_chat_template": True, + "extra_info": { + 'id': example['id'] if 'id' in example else str(idx), + 'raw_instruction': example['instruction'], + 'raw_input': example['input'], + 'num_objects': num_objects, + 'split': split, + } + } + + if idx == 0: + print(f"data_source: {data_source}, split: {split}, idx: {idx}") + print("\n" + "=" * 100 + f"{data_source} {split} {idx}" + "=" * 10) + print(data) + return data + + return process_fn + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--json_path', default='data/raw/puzzles_dataset.json', help='Path to json file') + parser.add_argument('--output_dir', default='data', help='Directory to save processed data') + parser.add_argument('--hdfs_dir', default=None, help='HDFS directory (optional)') + parser.add_argument('--train_size', type=float, default=0.8, help='Proportion of data for train set') + parser.add_argument('--test_size', type=float, default=0.2, help='Proportion of data for test set') + parser.add_argument('--data_source', default='ordering_puzzle_dataset', help='Name of data source') + parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility') + parser.add_argument('--train_sample_size', type=int, default=None, help='Number of samples to use from train. If None, use all.') + args = parser.parse_args() + + # Set seed for reproducibility + set_seed(args.seed) + + # Load dataset from JSON + dataset = datasets.load_dataset('json', data_files=args.json_path)['train'] + print(f"Loaded dataset with {len(dataset)} examples") + + # Validate train and test sizes + if args.train_size + args.test_size > 1.0: + raise ValueError(f"The sum of train_size ({args.train_size}) and test_size ({args.test_size}) cannot exceed 1.0") + + # Transform dataset + process_train_fn = make_map_fn('train', args.data_source) + processed_dataset = dataset.map(function=process_train_fn, with_indices=True) + + # Split dataset into train and test + train_indices, test_indices = train_test_split( + range(len(dataset)), + train_size=args.train_size, + test_size=args.test_size, + random_state=args.seed + ) + + print(f"Train set size: {len(train_indices)}, Test set size: {len(test_indices)}") + + # Create train and test datasets + train_dataset = processed_dataset.select(train_indices) + test_dataset = processed_dataset.select(test_indices) + + # Store the original training dataset size + original_train_size = len(train_dataset) + + # Sample the training dataset if needed + train_dataset = sample_dataset(train_dataset, args.train_sample_size) + + # Create output directories + train_output_dir = os.path.join(args.output_dir, "train") + test_output_dir = os.path.join(args.output_dir, "test") + os.makedirs(train_output_dir, exist_ok=True) + os.makedirs(test_output_dir, exist_ok=True) + + # Save train dataset + train_output_path = save_dataset( + dataset=train_dataset, + output_dir=train_output_dir, + filename_prefix=f"logic__{args.data_source}", + sample_size=args.train_sample_size if args.train_sample_size else len(train_dataset) + ) + + # Save test dataset + test_output_path = save_dataset( + dataset=test_dataset, + output_dir=test_output_dir, + filename_prefix=f"logic__{args.data_source}", + sample_size=len(test_dataset) + ) + + # Copy to HDFS if specified + if args.hdfs_dir is not None: + try: + from verl.utils.hdfs_io import copy, makedirs + makedirs(args.hdfs_dir) + copy(src=args.output_dir, dst=args.hdfs_dir) + print(f"Data copied to HDFS: {args.hdfs_dir}") + except ImportError: + print("HDFS utilities not available. Install verl package for HDFS support.") + + print(f"Done! \n" + f"Train data saved to {train_output_path}\n" + f"Test data saved to {test_output_path}") + print(f"Original train set size: {original_train_size} examples") + print(f"Final train set size: {len(train_dataset)} examples") + print(f"Test set: {len(test_dataset)} examples") \ No newline at end of file diff --git a/reasoners/data_preprocess/logic/process_zebrapuzzle_dataset.py b/reasoners/data_preprocess/logic/process_zebrapuzzle_dataset.py new file mode 100644 index 000000000..04c11d29d --- /dev/null +++ b/reasoners/data_preprocess/logic/process_zebrapuzzle_dataset.py @@ -0,0 +1,143 @@ +import os +import datasets +import argparse +from sklearn.model_selection import train_test_split + +from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset + +# Instruction for zebra puzzle answers +InstructionFollow = "Output the grid in the form of a dictionary with keys as header containing a list of the attributes and rows denoting each row of the final grid. Please return the final answer in tags, for example {\"header\": [\"Position\", \"Nationality\", \"Job\"], \"rows\": [[\"1\", \"british\", \"plumber\"], [\"2\", \"polish\", \"carpenter\"]]} ." + +def make_prefix(dp): + clues = dp['clues'] + if isinstance(clues, list): + clues = "\n".join(clues) + result = dp['ground_truth'] + instruction = dp['instruction'] + prefix = f"{instruction} The clues are: {clues}." + return prefix + + +def extract_from_ground_truth(text): + if isinstance(text, dict): + return text + else: + return eval(text) + +def make_map_fn(split, data_source): + def process_fn(example, idx): + question = make_prefix(example) + grid_size = str(example['config']["cols"]) + "x" + str(example['config']["rows"]) + final_grid = extract_from_ground_truth(example['ground_truth']) + + data = { + "data_source": data_source, + "prompt": [{ + "role": "user", + "content": question + " " + InstructionFollow + }], + "ability": "logical_reasoning", + "reward_model": { + "style": "rule", + "ground_truth": final_grid, + }, + "apply_chat_template": True, + "extra_info": { + 'id': example['id'] if 'id' in example else str(idx), + "grid_size": grid_size, + 'raw_instruction': example['instruction'], + 'raw_input': example['clues'], + 'split': split, + } + } + + if idx == 0: + print(f"data_source: {data_source}, split: {split}, idx: {idx}") + print("\n" + "=" * 100 + f"{data_source} {split} {idx}" + "=" * 10) + print(data) + return data + + return process_fn + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--json_path', default='data/raw/zebra_puzzles.json', help='Path to json file') + parser.add_argument('--output_dir', default='data', help='Directory to save processed data') + parser.add_argument('--hdfs_dir', default=None, help='HDFS directory (optional)') + parser.add_argument('--train_size', type=float, default=0.8, help='Proportion of data for train set') + parser.add_argument('--test_size', type=float, default=0.2, help='Proportion of data for test set') + parser.add_argument('--data_source', default='zebra_puzzle_dataset', help='Name of data source') + parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility') + parser.add_argument('--train_sample_size', type=int, default=None, help='Number of samples to use from train. If None, use all.') + args = parser.parse_args() + + # Set seed for reproducibility + set_seed(args.seed) + + # Load dataset from JSON + dataset = datasets.load_dataset('json', data_files=args.json_path)['train'] + + # Transform dataset + process_train_fn = make_map_fn('train', args.data_source) + processed_dataset = dataset.map(function=process_train_fn, with_indices=True) + + if args.train_size + args.test_size > 1.0: + raise ValueError(f"The sum of train_size ({args.train_size}) and test_size ({args.test_size}) cannot exceed 1.0") + + # Split dataset into train and test + train_indices, test_indices = train_test_split( + range(len(dataset)), + train_size=args.train_size, + test_size=args.test_size, + random_state=args.seed + ) + + # Create train and test datasets + train_dataset = processed_dataset.select(train_indices) + test_dataset = processed_dataset.select(test_indices) + + # Store the original training dataset size + original_train_size = len(train_dataset) + + # Sample the training dataset if needed + train_dataset = sample_dataset(train_dataset, args.train_sample_size) + + # Create output directories + train_output_dir = os.path.join(args.output_dir, "train") + test_output_dir = os.path.join(args.output_dir, "test") + os.makedirs(train_output_dir, exist_ok=True) + os.makedirs(test_output_dir, exist_ok=True) + + # Save train dataset + train_output_path = save_dataset( + dataset=train_dataset, + output_dir=train_output_dir, + filename_prefix=f"logic__{args.data_source}", + sample_size=args.train_sample_size if args.train_sample_size else len(train_dataset) + ) + + # Save test dataset + test_output_path = save_dataset( + dataset=test_dataset, + output_dir=test_output_dir, + filename_prefix=f"logic__{args.data_source}", + sample_size=len(test_dataset) + ) + + # Copy to HDFS if specified + if args.hdfs_dir is not None: + try: + from verl.utils.hdfs_io import copy, makedirs + makedirs(args.hdfs_dir) + copy(src=args.output_dir, dst=args.hdfs_dir) + print(f"Data copied to HDFS: {args.hdfs_dir}") + except ImportError: + print("HDFS utilities not available. Install verl package for HDFS support.") + + print(f"Done! \n" + f"Train data saved to {train_output_path}\n" + f"Test data saved to {test_output_path}") + print(f"Original train set size: {original_train_size} examples") + print(f"Final train set size: {len(train_dataset)} examples") + print(f"Test set: {len(test_dataset)} examples") \ No newline at end of file diff --git a/reasoners/data_preprocess/logic/puzzle_gen.py b/reasoners/data_preprocess/logic/puzzle_gen.py new file mode 100644 index 000000000..655ce7651 --- /dev/null +++ b/reasoners/data_preprocess/logic/puzzle_gen.py @@ -0,0 +1,318 @@ +import random +from itertools import permutations +from typing import List, Set, Dict, Tuple, Optional +from dataclasses import dataclass +import os +from datetime import datetime +import random +import gen_proof as gen_proof +import csv +import argparse +import json + +@dataclass +class Constraint: + text: str # Human-readable constraint + object1: Optional[int] # First object's encoded value (if applicable) + object2: Optional[int] # Second object's encoded value (if applicable) + constraint_type: str # Type of mathematical constraint + value: Optional[int] # Additional value needed for constraint (e.g., number of birds between) + + def evaluate(self, positions: List[int]) -> bool: + """Evaluate if the constraint is satisfied for given positions.""" + if self.constraint_type == "left_of": + return positions.index(self.object1) < positions.index(self.object2) + elif self.constraint_type == "right_of": + return positions.index(self.object1) > positions.index(self.object2) + elif self.constraint_type == "adjacent_left": + return positions.index(self.object1) + 1 == positions.index(self.object2) + elif self.constraint_type == "adjacent_right": + return positions.index(self.object1) - 1 == positions.index(self.object2) + elif self.constraint_type == "birds_between": + return abs(positions.index(self.object1) - positions.index(self.object2)) - 1 == self.value + elif self.constraint_type == "absolute_position_left": + return positions.index(self.object1) == self.value - 1 + elif self.constraint_type == "absolute_position_right": + return positions.index(self.object1) == len(positions) - self.value + return False + +class BirdPuzzleGenerator: + def __init__(self, objects: List[str]): + self.birds = objects + self.num_positions = len(objects) + # Encode birds as integers (0 to n-1) + self.bird_encoding = {bird: i for i, bird in enumerate(objects)} + self.bird_decoding = {i: bird for i, bird in enumerate(objects)} + + # Generate ground truth + self.ground_truth_birds = list(objects) + random.shuffle(self.ground_truth_birds) + # Store encoded ground truth positions + self.ground_truth = [self.bird_encoding[bird] for bird in self.ground_truth_birds] + + def generate_all_constraints(self) -> List[Constraint]: + """Generate all possible true constraints based on ground truth.""" + absolute_constraints = [] + relative_constraints = [] + # Generate all types of constraints + for i in range(self.num_positions): + bird1_code = self.ground_truth[i] + bird1_name = self.bird_decoding[bird1_code] + + # Absolute position constraints + pos_suffix = lambda n: 'th' if 11 <= n % 100 <= 13 else {1: 'st', 2: 'nd', 3: 'rd'}.get(n % 10, 'th') + left_pos = i + 1 + right_pos = self.num_positions - i + + absolute_constraints.append(Constraint( + text=f"The {bird1_name} is the {left_pos}{pos_suffix(left_pos)} from the left", + object1=bird1_code, + object2=None, + constraint_type="absolute_position_left", + value=left_pos + )) + + absolute_constraints.append(Constraint( + text=f"The {bird1_name} is the {right_pos}{pos_suffix(right_pos)} from the right", + object1=bird1_code, + object2=None, + constraint_type="absolute_position_right", + value=right_pos + )) + + # Relative position constraints + for j in range(i + 1, self.num_positions): + bird2_code = self.ground_truth[j] + bird2_name = self.bird_decoding[bird2_code] + + # Left/Right relationships + relative_constraints.append(Constraint( + text=f"The {bird1_name} is to the left of the {bird2_name}", + object1=bird1_code, + object2=bird2_code, + constraint_type="left_of", + value=None + )) + + relative_constraints.append(Constraint( + text=f"The {bird2_name} is to the right of the {bird1_name}", + object1=bird2_code, + object2=bird1_code, + constraint_type="right_of", + value=None + )) + + # Birds between + birds_between = j - i - 1 + if birds_between > 0: + plural = "s" if birds_between > 1 else "" + relative_constraints.append(Constraint( + text=f"There {['is', 'are'][birds_between > 1]} {birds_between} bird{plural} between the {bird1_name} and the {bird2_name}", + object1=bird1_code, + object2=bird2_code, + constraint_type="birds_between", + value=birds_between + )) + + # Adjacent relationships + if birds_between == 0: + relative_constraints.append(Constraint( + text=f"The {bird1_name} is immediately to the left of the {bird2_name}", + object1=bird1_code, + object2=bird2_code, + constraint_type="adjacent_left", + value=None + )) + + relative_constraints.append(Constraint( + text=f"The {bird2_name} is immediately to the right of the {bird1_name}", + object1=bird2_code, + object2=bird1_code, + constraint_type="adjacent_right", + value=None + )) + + return absolute_constraints, relative_constraints + + def verify_arrangement(self, constraints: List[Constraint], test_arrangement: List[int]) -> bool: + """Verify if a given arrangement satisfies all constraints using mathematical evaluation.""" + return all(constraint.evaluate(test_arrangement) for constraint in constraints) + + def find_minimal_constraints(self) -> List[Constraint]: + """Find a minimal set of constraints that uniquely determines the ground truth. + Uses incremental constraint addition until solution space reduces to one arrangement.""" + absolute_constraints, relative_constraints = self.generate_all_constraints() + # Shuffle constraints to randomize selection order + abs_constraints = list(absolute_constraints) + random.shuffle(abs_constraints) + rel_constraints = list(relative_constraints) + random.shuffle(rel_constraints) + + selected_constraints = [] + + # num_absolute = random.randint(0, int(self.num_positions*0.4)) + num_absolute = 0 + while len(selected_constraints) < num_absolute: + selected_constraints.append(abs_constraints.pop(0)) + + while rel_constraints and len(selected_constraints) < len(rel_constraints): + # Take the next constraint + current_constraint = rel_constraints.pop(0) + selected_constraints.append(current_constraint) + + # Find all valid arrangements given current constraints + valid_arrangements = [] + for perm in permutations(range(self.num_positions)): + if self.verify_arrangement(selected_constraints, list(perm)): + valid_arrangements.append(list(perm)) + # if len(valid_arrangements) > 1: + # break + + # If we've found a unique solution matching ground truth + if len(valid_arrangements) == 1: + if valid_arrangements[0] == self.ground_truth: + print("Found minimal constraint set!") + return selected_constraints + else: + # Our constraints led to a wrong unique solution + print("Warning: Constraints led to incorrect unique solution.") + # Remove the last constraint and try a different one + selected_constraints.pop() + continue + + # If we have no valid arrangements, the last constraint was incompatible + if len(valid_arrangements) == 0: + print("Warning: No valid arrangements with these constraints.") + # Remove the incompatible constraint and continue + selected_constraints.pop() + continue + + print("Warning: Could not find a minimal constraint set.") + return selected_constraints + + def generate_puzzle(self) -> Tuple[List[Constraint], List[List[int]]]: + """Generate a puzzle with constraints and wrong arrangements.""" + constraints = self.find_minimal_constraints() + return constraints + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Generate ordering puzzles for evaluating reasoning in LLMs") + parser.add_argument("--num_puzzles", type=int, default=10000, help="Number of puzzles to generate") + parser.add_argument("--output_dir", type=str, default="data/raw", help="Output directory for generated puzzles") + parser.add_argument("--output_file", type=str, default="puzzles_dataset.json", help="Output filename for generated puzzles") + parser.add_argument("--test", type=bool, default=False, help="Generate set without solution traces") + args = parser.parse_args() + + birds = ["hawk", "hummingbird", "quail", "owl", "crow", "robin", "cardinal", "sparrow", "bluejay", "pigeon", + "eagle", "falcon", "finch", "woodpecker", "duck", "goose", "swan", "seagull", "pelican", "flamingo", "parrot", + "macaw", "peacock", "penguin", "ostrich", "emu", "vulture", "raven", "magpie", "starling", "warbler", "chickadee", + "swallow", "martin", "oriole", "dove", "pheasant", "turkey", "toucan", "kingfisher", "heron", "egret", "stork", + "kiwi", "albatross", "condor"] + + mammals = ["lion", "tiger", "bear", "wolf", "fox", "deer", "elephant", "giraffe", "zebra", "rhino", + "hippo", "monkey", "gorilla", "chimpanzee", "kangaroo", "koala", "panda", "raccoon", "squirrel", + "rabbit", "mouse", "rat", "bat", "whale", "dolphin", "seal", "walrus", "otter", "beaver", + "hedgehog", "armadillo", "sloth", "opossum", "weasel", "badger", "skunk", "moose", "bison", + "camel", "llama", "alpaca", "horse", "donkey", "cow", "goat", "sheep", "pig"] + + fruits = ["apple", "orange", "banana", "grape", "strawberry", "blueberry", "raspberry", "blackberry", + "watermelon", "cantaloupe", "honeydew", "pineapple", "mango", "papaya", "kiwi", "pear", + "peach", "plum", "apricot", "cherry", "lemon", "lime", "grapefruit", "coconut", "fig", + "date", "guava", "passion fruit", "pomegranate", "dragonfruit", "lychee", "persimmon", + "avocado", "starfruit", "durian", "jackfruit", "nectarine", "tangerine", "clementine"] + + vehicles = ["car", "truck", "bus", "motorcycle", "scooter", "bicycle", "train", "tram", "subway", + "helicopter", "airplane", "jet", "rocket", "spaceship", "submarine", "ship", "yacht", + "sailboat", "speedboat", "canoe", "kayak", "van", "ambulance", "firetruck", "tractor", + "bulldozer", "crane", "forklift", "excavator", "tank", "snowmobile", "golf cart", + "limousine", "taxi", "trolley", "rickshaw", "hovercraft", "blimp", "hot air balloon"] + + instruments = ["guitar", "piano", "violin", "cello", "bass", "harp", "flute", "clarinet", "saxophone", + "trumpet", "trombone", "tuba", "oboe", "bassoon", "piccolo", "recorder", + "harmonica", "accordion", "banjo", "mandolin", "ukulele", "drums", "bongos", "congas", + "tambourine", "xylophone", "marimba", "vibraphone", "glockenspiel", "triangle", "cymbals", + "bagpipes", "sitar", "shamisen", "erhu", "theremin", "synthesizer", "organ", "harpsichord"] + + gemstones = ["diamond", "ruby", "emerald", "sapphire", "amethyst", "topaz", "opal", "pearl", "jade", + "garnet", "aquamarine", "turquoise", "moonstone", "amber", "onyx", "quartz", "citrine", + "peridot", "tanzanite", "agate", "alexandrite", "beryl", "malachite", "obsidian", "jasper", + "zircon", "spinel", "sunstone", "tourmaline", "lapis lazuli", "hematite", "pyrite", + "rhodochrosite", "fluorite", "calcite", "azurite", "chrysoberyl", "sodalite", "carnelian"] + + people = ["Emma", "Liam", "Wei", "Priya", "Mohammed", "Sofia", "Jamal", "Yuki", + "Isabella", "Raj", "Noah", "Ling", "Santiago", "Amara", "Daniel", + "Fatima", "Chen", "Miguel", "Aisha", "Sanjay", "Harper", "Takashi", + "Kwame", "Elena", "Omar", "Zoe", "Hiroshi", "Nia", "Gabriel", "Jin", + "Maya", "Mateo", "Layla", "Arjun", "Clara", "Ahmed", "Mei", "Luis", + "Imani", "Alejandro", "Divya", "Kenji", "Chioma", "Paulo", "Ananya", + "Hassan", "Valentina", "Kofi", "Sakura", "Amir"] + + random.seed(42) + num_puzzles = args.num_puzzles + base_dir = os.path.dirname(os.path.abspath(__file__)) + + # Create the "generated_QAs" folder at the same level as the Python file if it doesn't exist + # output_folder = os.path.join(base_dir, args.output_dir) + output_folder = args.output_dir + if not os.path.exists(output_folder): + os.makedirs(output_folder) + + output_file = os.path.join(output_folder, args.output_file) + + if args.test: + output_file_test = os.path.join(output_folder, args.output_file.replace(".json", "_test.json")) + else: + output_file_test = os.path.join(output_folder, args.output_file.replace(".json", "_unsolved.json")) + + puzzles_data = [] + puzzles_test_data = [] + object_types = { + 'birds': birds, + 'mammals': mammals, + 'fruits': fruits, + 'vehicles': vehicles, + 'instruments': instruments, + 'gemstones': gemstones, + 'people': people + } + for id in range(num_puzzles): + num_objects = random.randint(3, 7) + + type_name = random.choice(list(object_types.keys())) + type_objects = object_types[type_name] + curr_birds = random.sample(type_objects, num_objects) + birds_str = ", ".join(curr_birds) + instruction = f"Solve the following puzzle to determine the order of the {type_name} from left to right. The {type_name} are {birds_str}." + generator = BirdPuzzleGenerator(curr_birds) + + constraints = generator.generate_all_constraints() + constraints = generator.generate_puzzle() + input = "" + input_constraints = [] + for i, constraint in enumerate(constraints, 1): + input_constraints.append(constraint.text) + input += f"{i}. {constraint.text}\n" + if not args.test: + output = gen_proof.solve_puzzle(ground_truth = generator.ground_truth_birds, raw_constraints = input_constraints) + if output: + output = "\n".join(output) + puzzle_data = {'id': id, 'instruction': instruction, 'input': input, 'output': output, 'ground_truth': generator.ground_truth_birds, 'num_objects': num_objects} + puzzles_data.append(puzzle_data) + else: + puzzle_data = {'id': id, 'instruction':instruction, 'input': input, 'ground_truth': generator.ground_truth_birds, 'num_objects': num_objects} + puzzles_test_data.append(puzzle_data) + else: + puzzle_data = {'id': id, 'instruction':instruction, 'input': input, 'ground_truth': generator.ground_truth_birds, 'num_objects': num_objects} + puzzles_test_data.append(puzzle_data) + if args.test: + with open(output_file, 'w') as json_file: + json.dump(puzzles_test_data, json_file, indent=2) + else: + with open(output_file, 'w') as json_file: + json.dump(puzzles_data, json_file, indent=2) + with open(output_file_test, 'w') as json_file_test: + json.dump(puzzles_test_data, json_file_test, indent=2) + +# example usage +# python puzzle_gen.py --num_puzzles 10000 --output_dir data/puzzles_dataset --output_file puzzles_dataset.json --test True + diff --git a/reasoners/data_preprocess/logic/zebrapuzzle_gen/generator.py b/reasoners/data_preprocess/logic/zebrapuzzle_gen/generator.py new file mode 100644 index 000000000..6c7615701 --- /dev/null +++ b/reasoners/data_preprocess/logic/zebrapuzzle_gen/generator.py @@ -0,0 +1,488 @@ +import random +import collections +import time +from typing import Literal, List, Set, Tuple, Callable + + +def format_table(header: List[str], table: List[List[str]], + top_format='{:^{}}', left_format=' {:<{}}', cell_format='{:<{}}', + col_delim=' | ', row_delim='\n', prefix_format='|', postfix_format='|'): + table = [[''] + header + ['']] + [row for row in table] + table_format = [[prefix_format + left_format] + len(header) * [top_format]] \ + + len(table) * [[prefix_format + left_format] + len(header) * [cell_format]] + col_widths = [max(len(format.format(cell, 0)) + for format, cell in zip(col_format, col)) + for col_format, col in zip(zip(*table_format), zip(*table))] + return row_delim.join( + col_delim.join( + format.format(cell, width) + for format, cell, width in zip(row_format, row, col_widths)) + f" {postfix_format}" + for row_format, row in zip(table_format, table)) + + +def update_range(wns: List[str], rns: List[List[Set[str]]], cmp: Callable): + changed = False + for rn in rns: + classified_words = set() + for n_col, set_of_words in enumerate(rn): + if len(set_of_words) == 1: + classified_words.add(next(iter(set_of_words))) + word_to_cols = dict() + for n_col, set_of_words in enumerate(rn): + if len(set_of_words) != 1: + prev_length = len(set_of_words) + set_of_words.difference_update(classified_words) + changed |= prev_length != len(set_of_words) + for word in set_of_words: + word_to_cols.setdefault(word, set()).add(n_col) + for word, cols in word_to_cols.items(): + if len(cols) == 1: + x = rn[next(iter(cols))] + if len(x) != 1: + x.clear() + x.add(word) + changed = True + + new_rns = [[{x for x in xs if x != wn} for xs in rn] for wn, rn in zip(wns, rns)] + pairs = [] + for wn, rn in zip(wns, rns): + new_pairs = [] + break_condition = True + for cn, setn in enumerate(rn): + if wn in setn: + break_condition = False + if not pairs: + pairs = [[]] + for v in pairs: + new_pairs.append([*v, cn]) + pairs = new_pairs + if break_condition: + break + for pair in pairs: + if cmp(*pair): + for nrn, cn, wn in zip(new_rns, pair, wns): + nrn[cn].add(wn) + changed |= any(rn != new_rn for rn, new_rn in zip(rns, new_rns)) + if changed: + for rn, new_rn in zip(rns, new_rns): + for old, new in zip(rn, new_rn): + old.intersection_update(new) + return changed + + +def update_ranges(relations: List[Tuple[List[int], List[str], Callable, ...]], + ranges: List[List[Set[str]]]): + changed = False + for ins, wns, callable_object, *_ in relations: + changed |= update_range(wns, [ranges[i] for i in ins], callable_object) + return changed + + +def generate_puzzle(table: List[List[str]], *, + level: Literal[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20], + minimal_conditions: bool = False, max_seconds_for_minimizing: float = None, + tries: int = 1): + if level not in range(1, 20 + 1): + raise ValueError('level must be >= 1 and <= 20') + + table_wo_left = [row[1:] for row in table] + n_attributes = len(table_wo_left) + m_objects = len(table_wo_left[0]) + + if level >= 19 and m_objects == 2: + raise ValueError('too few objects for level >= 19') + elif m_objects <= 1: + raise ValueError('m_objects must be >= 2') + elif n_attributes <= 0: + raise ValueError('n_attributes must be >= 1') + + center = m_objects // 2 + except_flag = True + rules_for_relations = [ + (2, lambda j1, j2: j1 == j2, ['{0}:{1} == {2}:{3}', '{2}:{3} == {0}:{1}']), + (2, lambda j1, j2: j1 == j2 - 1, ['{0}:{1} is on the left of {2}:{3}']), + (2, lambda j1, j2: j1 == j2 + 1, ['{0}:{1} is on the right of {2}:{3}']), + (1, lambda j1: j1 == 0, ['{0}:{1} is on the far left']), + (1, lambda j1, last_index=m_objects - 1: j1 == last_index, ['{0}:{1} is on the far right']), + ] + (m_objects % 2 != 0) * [(1, lambda j1, mid=center: j1 == mid, ['{0}:{1} is in the middle'])] + if level >= 2: + rules_for_relations += [ + (3, lambda j1, j2, j3: j2 + 1 == j1 == j3 - 1 or j3 + 1 == j1 == j2 - 1, + ['{0}:{1} is between {2}:{3} and {4}:{5}', '{0}:{1} is between {4}:{5} and {2}:{3}']), + ] + if level >= 3: + rules_for_relations += [ + (2, lambda j1, j2: j1 == j2 - 1 or j1 == j2 + 1, + ['{0}:{1} is on the left or right of {2}:{3}']), + (1, lambda j1, last_index=m_objects - 1: j1 == 0 or j1 == last_index, + ['{0}:{1} is on the far left or far right']), + ] + if level >= 4: + rules_for_relations += [ + (1, lambda j1: (j1 + 1) % 2 != 0, ['{0}:{1} is in an odd position']), + (1, lambda j1: (j1 + 1) % 2 == 0, ['{0}:{1} is in an even position']), + ] + if level >= 5: + rules_for_relations += [ + (2, lambda j1, j2: j1 < j2, ['{0}:{1} is somewhere to the left of {2}:{3}']), + (2, lambda j1, j2: j1 > j2, ['{0}:{1} is somewhere to the right of {2}:{3}']), + ] + if level >= 6: + rules_for_relations += [ + (2, lambda j1, j2: j1 != j2, ['{0}:{1} != {2}:{3}', '{2}:{3} != {0}:{1}'], except_flag), + ] + if level >= 7: + rules_for_relations += [ + (3, lambda j1, j2, j3: j2 < j1 < j3 or j3 < j1 < j2, + ['{0}:{1} is somewhere between {2}:{3} and {4}:{5}', + '{0}:{1} is somewhere between {4}:{5} and {2}:{3}']), + ] + if level >= 8: + rules_for_relations += [ + (2, lambda j1, j2: j1 >= j2, ['{0}:{1} is not to the left of {2}:{3}']), + (2, lambda j1, j2: j1 <= j2, ['{0}:{1} is not to the right of {2}:{3}']), + ] + if level >= 9: + rules_for_relations += [ + (2, lambda j1, j2: j1 % 2 != j2 % 2, + ['{0}:{1} and {2}:{3} have different parity positions', + '{2}:{3} and {0}:{1} have different parity positions'], except_flag), + (2, lambda j1, j2: j1 % 2 == j2 % 2, + ['{0}:{1} and {2}:{3} have the same parity positions', + '{2}:{3} and {0}:{1} have the same parity positions'], except_flag), + ] + if level >= 10: + rules_for_relations += [ + (3, lambda j1, j2, j3: (j1 == j2 and j1 != j3) or (j1 != j2 and j1 == j3), + ['{0}:{1} == {2}:{3} or {0}:{1} == {4}:{5}, but not both', + '{0}:{1} == {4}:{5} or {0}:{1} == {2}:{3}, but not both'], except_flag), + (3, lambda j1, j2, j3: (j1 == j2 and j2 != j3) or (j1 != j2 and j2 == j3), + ['{0}:{1} == {2}:{3} or {2}:{3} == {4}:{5}, but not both', + '{2}:{3} == {4}:{5} or {0}:{1} == {2}:{3}, but not both'], except_flag), + ] + if level >= 11: + rules_for_relations += [ + (3, lambda j1, j2, j3: j1 == j2 or j1 == j3, + ['{0}:{1} == {2}:{3} or {0}:{1} == {4}:{5} or both', + '{0}:{1} == {4}:{5} or {0}:{1} == {2}:{3} or both'], except_flag), + (3, lambda j1, j2, j3: j1 == j2 or j2 == j3, + ['{0}:{1} == {2}:{3} or {2}:{3} == {4}:{5} or both', + '{2}:{3} == {4}:{5} or {0}:{1} == {2}:{3} or both'], except_flag), + ] + if level >= 12: + rules_for_relations += [ + (3, lambda j1, j2, j3: j1 != j2 or j1 != j3, + ['{0}:{1} != {2}:{3} or {0}:{1} != {4}:{5} or both', + '{0}:{1} != {4}:{5} or {0}:{1} != {2}:{3} or both'], except_flag), + (3, lambda j1, j2, j3: j1 != j2 or j2 != j3, + ['{0}:{1} != {2}:{3} or {2}:{3} != {4}:{5} or both', + '{2}:{3} != {4}:{5} or {0}:{1} != {2}:{3} or both'], except_flag), + ] + if level >= 13: + rules_for_relations.pop(0) # pop '==' + if level >= 14: + rules_for_relations.pop(0) # pop 'is on the left of' + rules_for_relations.pop(0) # pop 'is on the right of' + if level >= 15: + rules_for_relations.pop(0) # pop 'is on the far left' + rules_for_relations.pop(0) # pop 'is on the far right' + if m_objects % 2 != 0: + rules_for_relations.pop(0) # pop 'is in the middle' + if level >= 16: + rules_for_relations.pop(0) # pop 'is between' + if level >= 17: + rules_for_relations.pop(0) # pop 'is on the left or right of' + rules_for_relations.pop(0) # pop 'is on the far left or far right' + if level >= 18: + rules_for_relations.pop(0) # pop 'is in an odd position' + rules_for_relations.pop(0) # pop 'is in an even position' + if level >= 19: + rules_for_relations.pop(0) # pop 'is somewhere to the left of' + rules_for_relations.pop(0) # pop 'is somewhere to the right of' + if level >= 20: + rules_for_relations.pop(0) # pop '!=' + is_minimized = False + time_elapsed = False + min_relations = None + while True: + ranges = [[set(table_wo_left[i]) for _ in range(len(table_wo_left[i]))] for i in range(len(table_wo_left))] + relations = list() + fail = False + while not fail: + needs_clarification = list() + no_solutions = False + solved = True + for i, rng in enumerate(ranges): + for j, rs in enumerate(rng): + if len(rs) == 0: + no_solutions = True + solved = False + break + elif len(rs) > 1: + solved = False + needs_clarification.append((i, j)) + if no_solutions: + break + if solved or min_relations is not None and len(relations) >= len(min_relations): + tries -= 1 + if min_relations is None or len(relations) < len(min_relations): + min_relations = relations + if tries > 0: + fail = True + continue + if tries <= 0: + relations = min_relations + if not minimal_conditions: + break + number_of_relations_min = len(relations) + number_of_relations_before = len(relations) + start_time = time.monotonic() + main_q = collections.deque([relations]) + while main_q: + current_relations = main_q.popleft() + for k in range(len(current_relations)): + new_ranges = [[set(table_wo_left[i]) for _ in range(len(table_wo_left[i]))] + for i in range(len(table_wo_left))] + new_relations = current_relations.copy() + new_relations.pop(k) + changed = True + while changed: + changed = update_ranges(new_relations, new_ranges) + + q = collections.deque([new_ranges]) + possible_solutions = [] + while q: + current_ranges = q.popleft() + + no_solutions = False + solved = True + for rng in current_ranges: + for rs in rng: + if len(rs) == 0: + no_solutions = True + solved = False + break + elif len(rs) > 1: + solved = False + if no_solutions or not solved: + break + if no_solutions: + continue + if solved: + if current_ranges not in possible_solutions: + possible_solutions.append(current_ranges) + if len(possible_solutions) >= 2: + break + continue + + for n_group, rng in enumerate(current_ranges): + founded = False + for n_x, rs in enumerate(rng): + if len(rs) > 1: + founded = True + for r in rs: + new_ranges = [[x.copy() for x in row] for row in current_ranges] + new_ranges[n_group][n_x] = {r} + changed = True + while changed: + changed = update_ranges(new_relations, new_ranges) + q.appendleft(new_ranges) + break + if founded: + break + if len(possible_solutions) == 1: + number_of_relations_after = len(new_relations) + if number_of_relations_min > number_of_relations_after: + number_of_relations_min = number_of_relations_after + relations = new_relations + main_q.append(new_relations) + if max_seconds_for_minimizing is not None and \ + time.monotonic() >= start_time + max_seconds_for_minimizing: + time_elapsed = True + break + if time_elapsed: + break + is_minimized = number_of_relations_min < number_of_relations_before or not time_elapsed + break + if no_solutions or not needs_clarification: + fail = True + continue + + i, j = item = random.choice(needs_clarification) + next2_i, next2_j = None, None + if level >= 2 and len(needs_clarification) > 1: + needs_clarification.remove(item) + next2_i, next2_j = random.choice(needs_clarification) + + neighbours = [] + right_neighbours = [] + for dj in range(-1, 1 + 1): + if not (0 <= j + dj < m_objects): + continue + for new_i in range(0, n_attributes): + if new_i == i and dj == 0: + continue + new_item = (new_i, j + dj) + neighbours.append(new_item) + if level >= 2 and dj == 1: + right_neighbours.append(new_item) + if not neighbours: + continue + next_i, next_j = random.choice(neighbours) + if level >= 2 and next2_i is None and right_neighbours: + next2_i, next2_j = random.choice(right_neighbours) + + permutations3 = [ + ((i, j), (next_i, next_j), (next2_i, next2_j)), ((i, j), (next2_i, next2_j), (next_i, next_j)), + ((next_i, next_j), (i, j), (next2_i, next2_j)), ((next_i, next_j), (next2_i, next2_j), (i, j)), + ((next2_i, next2_j), (i, j), (next_i, next_j)), ((next2_i, next2_j), (next_i, next_j), (i, j)) + ] if next2_i is not None else [] + permutations2 = [ + ((i, j), (next_i, next_j)), ((next_i, next_j), (next2_i, next2_j)), ((i, j), (next2_i, next2_j)), + ((next_i, next_j), (i, j)), ((next2_i, next2_j), (next_i, next_j)), ((next2_i, next2_j), (i, j)), + ] if next2_i is not None else [ + ((i, j), (next_i, next_j)), ((next_i, next_j), (i, j)) + ] + possible_variants = [] + for (n_args, cmp_function, str_variants, *flags) in rules_for_relations: + if n_args == 3: + for items in permutations3: + (ti, tj), (t_next_i, t_next_j), (t_next2_i, t_next2_j) = items + if flags and flags[0] and (ti == t_next_i or ti == t_next2_i or t_next_i == t_next2_i): + continue + if cmp_function(tj, t_next_j, t_next2_j): + possible_variants.append((n_args, items, cmp_function, random.choice(str_variants))) + elif n_args == 2: + for items in permutations2: + (ti, tj), (t_next_i, t_next_j) = items + if flags and flags[0] and ti == t_next_i: + continue + if cmp_function(tj, t_next_j): + possible_variants.append((n_args, items, cmp_function, random.choice(str_variants))) + elif n_args == 1 and cmp_function(j): + possible_variants.append((n_args, [(i, j)], cmp_function, random.choice(str_variants))) + if not possible_variants: + continue + + n_args, list_of_ij, cmp_function, string_format = random.choice(possible_variants) + list_for_format = [] + ins, wns = [], [] + for i, j in list_of_ij: + list_for_format.extend([table[i][0], table_wo_left[i][j]]) + ins.append(i) + wns.append(table_wo_left[i][j]) + relations.append((ins, wns, cmp_function, string_format.format(*list_for_format))) + + changed = True + while changed: + changed = update_ranges(relations, ranges) + + if not fail: + if minimal_conditions and not is_minimized and not time_elapsed: + continue + break + + premises = [t[-1] for t in relations] + random.shuffle(premises) + return premises + +def table_to_dict(table: List[List[str]]) -> dict: + res = {} + number_of_rows = len(table[0])-1 + res["header"] = ["Position"] + res["rows"] = [[str(i+1)] for i in range(number_of_rows)] + for i, row in enumerate(table): + res['header'].append(row[0]) + for j in range(1, len(row)): + res['rows'][j-1].append(row[j]) + return res + + + +def generate_table(n_attributes: int, m_objects: int) -> List[List[str]]: + kinds_dict = { + "Nationality": { + "american", "argentine", "australian", "brazilian", "british", + "canadian", "chinese", "colombian", "dutch", "egyptian", + "french", "german", "indian", "indonesian", "italian", + "japanese", "malaysian", "mexican", "nigerian", "pakistani", + "polish", "russian", "spanish", "thai", "turkish", + }, + "Food": { + "apple", "apricot", "artichoke", "asparagus", "avocado", + "banana", "blueberry", "broccoli", "cabbage", "carrot", + "cauliflower", "cherry", "corn", "cranberry", "cucumber", + "eggplant", "garlic", "grapefruit", "grapes", "kale", + "kiwi", "lemon", "lettuce", "lime", "mango", + "nectarine", "onion", "orange", "papaya", "peach", + "pear", "peas", "pepper", "pineapple", "plum", + "pomegranate", "potato", "pumpkin", "radish", "raspberry", + "spinach", "strawberry", "tomato", "watermelon", "zucchini", + }, + "Pet": { + "bird", "cat", "chinchilla", "dog", "ferret", + "fish", "frog", "goat", "goldfish", "guinea-pig", + "hamster", "hedgehog", "horse", "lizard", "mouse", + "pony", "rabbit", "rat", "snake", "turtle", + }, + "Job": { + "accountant", "analyst", "architect", "bartender", "chef", + "coach", "dancer", "designer", "doctor", "dressmaker", + "electrician", "engineer", "entrepreneur", "firefighter", "fisherman", + "freelancer", "journalist", "lawyer", "librarian", "manager", + "mechanic", "musician", "nurse", "paramedic", "photographer", + "pilot", "police-officer", "project-manager", "scientist", "security-guard", + "social-worker", "software-developer", "teacher", "videographer", "writer", + }, + "Beverage": { + "7up", "almond-milk", "coffee", "cola", "fanta", + "hot-chocolate", "iced-tea", "juice", "lemonade", "milk", + "mirinda", "soy-milk", "sprite", "tea", "water", + }, + "Transport": { + "airplane", "bike", "boat", "bus", "car", + "helicopter", "jet-ski", "motorbike", "quad-bike", "roller", + "scooter", "ship", "skateboard", "snowmobile", + "subway", "taxi", "train", "tram", "trike", "van", + }, + "Music-Genre": { + "ambient", "blues", "classical", "country", "d&b", + "disco", "dubstep", "electronic", "folk", "funk", + "gospel", "hip-hop", "house", "indie", "jazz", + "metal", "pop", "punk", "r&b", "reggae", + "rock", "salsa", "soul", "techno", "trance", + }, + "Movie-Genre": { + "action", "adventure", "animation", "comedy", "crime", + "disaster", "documentary", "drama", "epic", "family", + "fantasy", "horror", "martial-arts", "musical", "mystery", + "romance", "satire", "scientific", "sports", "spy", + "superhero", "thriller", "time-travel", "western", "zombie", + }, + "Sport": { + "badminton", "baseball", "basketball", "biathlon", "climbing", + "cricket", "cycling", "golf", "handball", "ice-hockey", + "lacrosse", "parkour", "rowing", "rugby", "sailing", + "skateboarding", "skiing", "snowboarding", "soccer", "surfing", + "swimming", "tennis", "volleyball", "water-polo", "weightlifting", + }, + "Hobby": { + "baking", "board-games", "camping", "card-games", "chess", + "collecting", "cooking", "dancing", "drawing", "filmmaking", + "fishing", "gardening", "hiking", "magic-tricks", "photography", + "puzzles", "reading", "rock-climbing", "singing", "skydiving", + "sudoku", "traveling", "video-games", "woodworking", "writing", + } + } + kinds = sorted(kinds_dict) + + # Check + assert n_attributes <= len(kinds_dict),\ + f'Not enough attributes: actual {len(kinds_dict)}, expected {n_attributes}' + assert all(m_objects <= len(v) for k, v in kinds_dict.items()), 'Not enough objects: ' +\ + f'actual {next(f"{k}={len(v)}" for k, v in kinds_dict.items() if m_objects > len(v))}, expected {m_objects}' + + chosen_kinds = sorted(random.sample(kinds, k=n_attributes)) + table = [[kind] + random.sample(sorted(kinds_dict[kind]), k=m_objects) for kind in chosen_kinds] + return table + diff --git a/reasoners/data_preprocess/logic/zebrapuzzle_gen/puzzle_generator.py b/reasoners/data_preprocess/logic/zebrapuzzle_gen/puzzle_generator.py new file mode 100644 index 000000000..1c47964d3 --- /dev/null +++ b/reasoners/data_preprocess/logic/zebrapuzzle_gen/puzzle_generator.py @@ -0,0 +1,202 @@ +import threading +import queue +import time +import random +from typing import List, Literal +import concurrent.futures +from generator import generate_puzzle, generate_table, table_to_dict +import multiprocessing +import argparse + + +# diverse set of configs, rows = num_attributes, cols = num_objects +# The level is a number from 1 to 20, where 1 is the easiest and 20 is the hardest +# minimal_conditions is a boolean indicating whether to use minimal conditions +PUZZLE_CONFIGS = [ + {"rows": 2, "cols": 2, "level": 5, "minimal_conditions": True}, + {"rows": 2, "cols": 4, "level": 12, "minimal_conditions": True}, + {"rows": 2, "cols": 5, "level": 3, "minimal_conditions": False}, + {"rows": 2, "cols": 6, "level": 15, "minimal_conditions": True}, + {"rows": 2, "cols": 8, "level": 7, "minimal_conditions": False}, + {"rows": 2, "cols": 10, "level": 18, "minimal_conditions": True}, + {"rows": 2, "cols": 12, "level": 10, "minimal_conditions": False}, + {"rows": 2, "cols": 15, "level": 20, "minimal_conditions": True}, + + {"rows": 3, "cols": 2, "level": 14, "minimal_conditions": False}, + {"rows": 3, "cols": 3, "level": 6, "minimal_conditions": True}, + {"rows": 3, "cols": 5, "level": 19, "minimal_conditions": False}, + {"rows": 3, "cols": 7, "level": 11, "minimal_conditions": True}, + {"rows": 3, "cols": 9, "level": 4, "minimal_conditions": False}, + {"rows": 3, "cols": 12, "level": 16, "minimal_conditions": True}, + {"rows": 3, "cols": 15, "level": 9, "minimal_conditions": False}, + + {"rows": 4, "cols": 2, "level": 17, "minimal_conditions": True}, + {"rows": 4, "cols": 4, "level": 1, "minimal_conditions": False}, + {"rows": 4, "cols": 6, "level": 13, "minimal_conditions": True}, + {"rows": 4, "cols": 8, "level": 20, "minimal_conditions": False}, + {"rows": 4, "cols": 10, "level": 5, "minimal_conditions": True}, + {"rows": 4, "cols": 13, "level": 15, "minimal_conditions": False}, + + {"rows": 5, "cols": 3, "level": 8, "minimal_conditions": True}, + {"rows": 5, "cols": 5, "level": 18, "minimal_conditions": False}, + {"rows": 5, "cols": 7, "level": 3, "minimal_conditions": True}, + {"rows": 5, "cols": 9, "level": 12, "minimal_conditions": False}, + {"rows": 5, "cols": 11, "level": 7, "minimal_conditions": True}, + {"rows": 5, "cols": 14, "level": 19, "minimal_conditions": False}, + + {"rows": 6, "cols": 2, "level": 10, "minimal_conditions": False}, + {"rows": 6, "cols": 4, "level": 6, "minimal_conditions": True}, + {"rows": 6, "cols": 6, "level": 16, "minimal_conditions": False}, + {"rows": 6, "cols": 8, "level": 2, "minimal_conditions": True}, + {"rows": 6, "cols": 10, "level": 14, "minimal_conditions": False}, + {"rows": 6, "cols": 12, "level": 9, "minimal_conditions": True}, + {"rows": 6, "cols": 15, "level": 4, "minimal_conditions": False}, + + {"rows": 7, "cols": 3, "level": 11, "minimal_conditions": True}, + {"rows": 7, "cols": 5, "level": 5, "minimal_conditions": False}, + {"rows": 7, "cols": 7, "level": 17, "minimal_conditions": True}, + {"rows": 7, "cols": 9, "level": 8, "minimal_conditions": False}, + {"rows": 7, "cols": 11, "level": 20, "minimal_conditions": True}, + {"rows": 7, "cols": 13, "level": 3, "minimal_conditions": False}, + + {"rows": 8, "cols": 2, "level": 15, "minimal_conditions": True}, + {"rows": 8, "cols": 4, "level": 7, "minimal_conditions": False}, + {"rows": 8, "cols": 6, "level": 19, "minimal_conditions": True}, + {"rows": 8, "cols": 8, "level": 10, "minimal_conditions": False}, + {"rows": 8, "cols": 10, "level": 1, "minimal_conditions": True}, + {"rows": 8, "cols": 12, "level": 13, "minimal_conditions": False}, + {"rows": 8, "cols": 14, "level": 6, "minimal_conditions": True}, + + {"rows": 9, "cols": 3, "level": 18, "minimal_conditions": False}, + {"rows": 9, "cols": 5, "level": 4, "minimal_conditions": True}, + {"rows": 9, "cols": 7, "level": 16, "minimal_conditions": False}, + {"rows": 9, "cols": 9, "level": 9, "minimal_conditions": True}, + {"rows": 9, "cols": 11, "level": 2, "minimal_conditions": False}, + {"rows": 9, "cols": 13, "level": 14, "minimal_conditions": True}, + + {"rows": 10, "cols": 2, "level": 8, "minimal_conditions": False}, + {"rows": 10, "cols": 4, "level": 20, "minimal_conditions": True}, + {"rows": 10, "cols": 6, "level": 5, "minimal_conditions": False}, + {"rows": 10, "cols": 8, "level": 17, "minimal_conditions": True}, + {"rows": 10, "cols": 10, "level": 10, "minimal_conditions": False}, + {"rows": 10, "cols": 12, "level": 3, "minimal_conditions": True}, + {"rows": 10, "cols": 14, "level": 18, "minimal_conditions": False}, + {"rows": 10, "cols": 15, "level": 20, "minimal_conditions": True}, +] + +INSTRUCTION = """Solve the following puzzle where you are given {{n_objects}} objects and their attributes like job, nationality etc. +The list of attributes and their unique values are: {{table}}. The goal is to determine the correct attribute for each object based on the provided clues.""" + + +def worker_function(config_and_id): + config, worker_id = config_and_id + rows, cols = config["rows"], config["cols"] + level = config["level"] + minimal_conditions = config["minimal_conditions"] + instance_id = config.get("instance_id", 0) + + print(f"Worker {worker_id}: Starting puzzle generation with {rows}x{cols} table, level {level}, instance {instance_id}") + + # Create a table for the puzzle + table = generate_table(rows, cols) + + try: + # Call the generate_puzzle function with our table + premises = generate_puzzle( + table, + level=level, + minimal_conditions=minimal_conditions, + max_seconds_for_minimizing=30.0, + tries=3 + ) + + ground_truth = table_to_dict(table) + # zebralogic convention and ground truth in format of n_objects x n_attributes + puzzle_key = f"puzzle_{config['cols']}x{config['rows']}_level{config['level']}_instance{instance_id}" + list_attributes = "" + for row in table: + list_attributes += f"\n{row[0]}: " + ', '.join(sorted(row[1:])) + instruction = INSTRUCTION.replace("{{table}}", list_attributes) + instruction = instruction.replace("{{n_objects}}", str(cols)) + result = { + "worker_id": worker_id, + "puzzle_id": puzzle_key, + "config": config, + "instruction": instruction, + "clues": premises, + "ground_truth": ground_truth, + } + print(f"Worker {worker_id}: Successfully generated puzzle") + return result + + except Exception as e: + print(f"Worker {worker_id}: Error generating puzzle: {e}") + return None + +def run_parallel_generation(total_puzzles = None, num_processes=None): + """Run puzzle generation in parallel using multiple threads""" + num_configs = len(PUZZLE_CONFIGS) + + if total_puzzles is None: + total_puzzles = num_configs + + puzzles_per_config = total_puzzles // num_configs + + extra_puzzles = total_puzzles % num_configs + expanded_configs = [] + for i, config in enumerate(PUZZLE_CONFIGS): + repeats = puzzles_per_config + (1 if i < extra_puzzles else 0) + for j in range(repeats): + # Create a copy with a unique instance ID + config_copy = config.copy() + config_copy['instance_id'] = j + expanded_configs.append(config_copy) + + if num_processes is None: + num_processes = min(len(expanded_configs), + concurrent.futures.cpu_count()) + + print(f"Generating {total_puzzles} puzzles across {num_configs} configurations using {num_processes} processes") + + worker_args = [(config, i) for i, config in enumerate(expanded_configs)] + + with concurrent.futures.ProcessPoolExecutor(max_workers=num_processes) as executor: + results = list(executor.map(worker_function, worker_args)) + + # Collect all results from the queue + results = [r for r in results if r is not None] + + return results + +def save_puzzles(results, output_dir="./data/raw", output_file="zebra_puzzles.json"): + import os + import json + + # Create output directory if it doesn't exist + os.makedirs(output_dir, exist_ok=True) + + output_path = os.path.join(output_dir, output_file) + with open(output_path, "w") as f: + json.dump(results, f, indent=4) + +if __name__ == "__main__": + # Parse command-line arguments + parser = argparse.ArgumentParser(description="Generate puzzles in parallel.") + parser.add_argument("--num_puzzles", type=int, default=100, help="Total number of puzzles to generate") + parser.add_argument("--num_processes", type=int, default=16, help="Number of processes to use for generation") + parser.add_argument("--output_dir", type=str, default="data/raw", help="Directory to save generated puzzles") + parser.add_argument("--output_file", type=str, default="zebra_puzzles.json", help="Filename for the generated puzzles") + args = parser.parse_args() + multiprocessing.freeze_support() + start_time = time.time() + print("Starting parallel puzzle generation...") + + # Run parallel generation with 4 threads (or adjust as needed) + results = run_parallel_generation(total_puzzles = args.num_puzzles, num_processes=args.num_processes) + + # Save the generated puzzles + save_puzzles(results, output_dir=args.output_dir, output_file=args.output_file) + + end_time = time.time() + print(f"Completed in {end_time - start_time:.2f} seconds") + print(f"Generated {len(results)} puzzles") \ No newline at end of file diff --git a/reasoners/data_preprocess/math/bigmath_preview_filtered_mar21.py b/reasoners/data_preprocess/math/bigmath_preview_filtered_mar21.py new file mode 100644 index 000000000..a5a2a6956 --- /dev/null +++ b/reasoners/data_preprocess/math/bigmath_preview_filtered_mar21.py @@ -0,0 +1,159 @@ +"""Downloads, processes, and saves math datasets.""" + +import os +import datasets +from typing import Dict, List, Optional, Any, Union +import enum +import argparse +import pandas as pd +import json +import random +import transformers + +from verl.utils.reward_score.math import remove_boxed, last_boxed_only_string +from verl.utils.data_process.prompt import build_zero_style_prompt +from verl.utils.data_process.utils import add_suffix, set_seed, sample_dataset, save_dataset + +def get_datasets(cache_dir: str): + """ + Load the math datasets from Hugging Face Hub. + + Args: + train_data_source (str): Source for training data + test_data_sources (list): List of sources for test data + + Returns: + tuple: (train_dataset, test_datasets) as Dataset objects + """ + train_data_source = "SDSB/big_math_partial_mar21_filtered_basic" + test_data_sources = [ # hard-coded for now as only bigmath handles all the test datasets + "nanoverl/minerva", + "SDSB/aime_repeated_8x", + "SDSB/amc_repeated_4x", + "nanoverl/olympiad_bench", + "nanoverl/math", + ] + + print(f"Loading the {train_data_source} dataset...") + train_dataset = datasets.load_dataset( + train_data_source, trust_remote_code=True, split="train", cache_dir=cache_dir + ) + + print(f"Loading the test datasets...") + test_datasets = [ + datasets.load_dataset(test_data_source, trust_remote_code=True, split="test", cache_dir=cache_dir) + for test_data_source in test_data_sources + ] + + return train_dataset, test_datasets + + +def make_map_fn(split: str, data_source: str, prompt_style: str="zero_style", reward_metric: str=None) -> callable: + def process_fn(example: Dict[str, Any], idx: int) -> Dict[str, Any]: + question = example.pop("problem") + answer = example.pop("answer") + data = { + "data_source": data_source, + "prompt": [ + {"role": "user", "content": question + " Please output the final answer within \\boxed{}."}, + ], + "ability": "math", + "apply_chat_template": True, + "reward_model": {"style": "rule", "ground_truth": answer}, + "extra_info": { + "split": split, + "index": idx, + "reward_metric": reward_metric, + "original_question": question, + }, + } + + if idx == 0 or idx == 1: + print(f"data_source: {data_source}, split: {split}, idx: {idx}") + print("\n" + "=" * 100 + f"{data_source} {split} {idx}" + "=" * 10) + print(data) + return data + + return process_fn + + +if __name__ == "__main__": + """Main script execution: parse args, load, process, and save datasets.""" + parser = argparse.ArgumentParser() + parser.add_argument( + "--data-dir", + default="data", + help="Directory to save the processed data files. Will be modified based on other parameters.", + ) + parser.add_argument( + "--domain", + default="math", + type=str, + help="Data domain", + ) + parser.add_argument( + "--name", + default="bigmath_preview_filtered_mar21", + type=str, + help="Dataset name", + ) + parser.add_argument( + "--train-sample-size", + type=int, + default=None, + help="Number of samples to use from the training dataset. If None, use all samples.", + ) + parser.add_argument( + "--train-reward-metric", + type=str, + default=None, + help="Reward metric to use for training. If None, use the naive_dapo.compute_score.", + ) + parser.add_argument( + "--test-reward-metric", + type=str, + default=None, + help="Reward metric to use for testing. If None, use the naive_dapo.compute_score.", + ) + parser.add_argument( + "--suffix", + type=str, + default=None, + help="Suffix to add to the dataset name.", + ) + parser.add_argument( + "--seed", + type=int, + default=42, + help="Random seed for reproducibility when sampling data.", + ) + args = parser.parse_args() + set_seed(args.seed) + + # Download the datasets from Hugging Face Hub + cache_dir = datasets.config.HF_DATASETS_CACHE + train_dataset, _ = get_datasets(cache_dir) + data_source = f"{args.domain}__{args.name}" + + # Process train dataset + process_train_fn = make_map_fn("train", data_source, args.train_reward_metric) + train_data = train_dataset.map(function=process_train_fn, with_indices=True) + # Sample + train_data = sample_dataset(train_data, args.train_sample_size) + # Save + train_output_dir = os.path.join(args.data_dir, "train") + args.train_sample_size = len(train_dataset) if args.train_sample_size is None else args.train_sample_size + train_output_path = save_dataset( + dataset=train_data, + output_dir=train_output_dir, + filename_prefix=data_source + args.suffix if args.suffix else data_source, + sample_size=args.train_sample_size + ) + + print(f"Done! \n" + f"Train data saved to {train_output_path}") + + # python bigmath_preview_filtered_mar21.py + + # with llm judge + # python bigmath_preview_filtered_mar21.py --train-reward-metric math_llm_judge --suffix _llm_judge \ No newline at end of file diff --git a/reasoners/data_preprocess/math/dapo_or1_merge_dedup_apr30.py b/reasoners/data_preprocess/math/dapo_or1_merge_dedup_apr30.py new file mode 100644 index 000000000..4adc216d7 --- /dev/null +++ b/reasoners/data_preprocess/math/dapo_or1_merge_dedup_apr30.py @@ -0,0 +1,183 @@ +"""Downloads, processes, and saves math datasets.""" + +import os +import datasets +from typing import Dict, List, Optional, Any, Union +import enum +import argparse +import pandas as pd +import json +import random +import transformers + +from verl.utils.reward_score.math import remove_boxed, last_boxed_only_string +from verl.utils.data_process.prompt import build_zero_style_prompt +from verl.utils.data_process.utils import add_suffix, set_seed, sample_dataset, save_dataset + +def get_datasets(cache_dir: str): + """ + Load the math datasets from Hugging Face Hub. + + Args: + train_data_source (str): Source for training data + test_data_sources (list): List of sources for test data + + Returns: + tuple: (train_dataset, test_datasets) as Dataset objects + """ + train_data_source = "SDSB/merged_deduped_dapo_or1_dataset" + test_data_sources = [ + "nanoverl/minerva", + "SDSB/aime_repeated_8x", + "SDSB/amc_repeated_4x", + "nanoverl/olympiad_bench", + "nanoverl/math", + "nanoverl/aime2025_repeated_8x", + ] + + print(f"Loading the {train_data_source} dataset...") + train_dataset = datasets.load_dataset( + train_data_source, trust_remote_code=True, split="train", cache_dir=cache_dir + ) + + print(f"Loading the test datasets...") + test_datasets = { + os.path.basename(test_data_source.lower()): + datasets.load_dataset(test_data_source, trust_remote_code=True, split="test", cache_dir=cache_dir) + for test_data_source in test_data_sources + } + + return train_dataset, test_datasets + + +def make_map_fn(split: str, data_source: str, reward_metric: str='default') -> callable: + def process_fn(example: Dict[str, Any], idx: int) -> Dict[str, Any]: + # if original_question in example, use it as the question + question = example.pop("problem") + # question = example.pop("problem") + answer = example.pop("answer") + if isinstance(answer, list): + answer = answer[0] + data = { + "data_source": data_source, + "prompt": [ + {"role": "user", "content": question + " Please output the final answer within \\boxed{}."}, + ], + "ability": "math", + "apply_chat_template": True, + "reward_model": {"style": "rule", "ground_truth": answer}, + "extra_info": { + "split": split, + "index": idx, + "reward_metric": reward_metric, + "original_question": question, + }, + } + + if idx == 0 or idx == 1: + print(f"data_source: {data_source}, split: {split}, idx: {idx}") + print("\n" + "=" * 100 + f"{data_source} {split} {idx}" + "=" * 10) + print(data) + return data + + return process_fn + + +if __name__ == "__main__": + """Main script execution: parse args, load, process, and save datasets.""" + parser = argparse.ArgumentParser() + parser.add_argument( + "--data-dir", + default="data", + help="Directory to save the processed data files. Will be modified based on other parameters.", + ) + parser.add_argument( + "--domain", + default="math", + type=str, + help="Data domain", + ) + parser.add_argument( + "--name", + default="merged_deduped_dapo_or1_dataset", + type=str, + help="Dataset name", + ) + parser.add_argument( + "--train-sample-size", + type=int, + default=None, + help="Number of samples to use from the training dataset. If None, use all samples.", + ) + parser.add_argument( + "--train-reward-metric", + type=str, + default="default", + help="Reward metric to use for training. If None, use the naive_dapo.compute_score.", + ) + parser.add_argument( + "--test-reward-metric", + type=str, + default="default", + help="Reward metric to use for testing. If None, use the naive_dapo.compute_score.", + ) + parser.add_argument( + "--suffix", + type=str, + default=None, + help="Suffix to add to the dataset name.", + ) + parser.add_argument( + "--seed", + type=int, + default=42, + help="Random seed for reproducibility when sampling data.", + ) + args = parser.parse_args() + set_seed(args.seed) + + # Download the datasets from Hugging Face Hub + cache_dir = datasets.config.HF_DATASETS_CACHE + train_dataset, test_datasets = get_datasets(cache_dir) + data_source = f"{args.domain}__{args.name}" + + # Process train dataset + process_train_fn = make_map_fn("train", data_source, args.train_reward_metric) + train_data = train_dataset.map(function=process_train_fn, with_indices=True) + # Sample + train_data = sample_dataset(train_data, args.train_sample_size) + # Save + train_output_dir = os.path.join(args.data_dir, "train") + args.train_sample_size = len(train_dataset) if args.train_sample_size is None else args.train_sample_size + train_output_path = save_dataset( + dataset=train_data, + output_dir=train_output_dir, + filename_prefix=data_source + args.suffix if args.suffix else data_source, + sample_size=args.train_sample_size + ) + + # Process test datasets + test_output_dir = os.path.join(args.data_dir, "test") + test_output_paths = [] + for test_data_source, test_data in test_datasets.items(): + test_data_source = f"{args.domain}__{test_data_source}" + process_fn = make_map_fn("test", test_data_source, args.test_reward_metric) + test_data = test_data.map(process_fn, with_indices=True) + test_output_path = save_dataset( + dataset=test_data, + output_dir=test_output_dir, + filename_prefix=test_data_source, + sample_size=None + ) + test_output_paths.append(test_output_path) + print(f"test_data_source: {test_data_source}") + print(f"Test data saved to {test_output_path}") + + print(f"Done! \n" + f"Train data saved to {train_output_path}\n" + f"Test data saved to {test_output_paths}") + + # python dapo_or1_merge_dedup_apr30.py + + # with llm judge + # python dapo_or1_merge_dedup_apr30.py --train-reward-metric math_llm_judge --suffix _llm_judge \ No newline at end of file diff --git a/reasoners/data_preprocess/math/deepscaler_preview.py b/reasoners/data_preprocess/math/deepscaler_preview.py new file mode 100644 index 000000000..ef37ed0bd --- /dev/null +++ b/reasoners/data_preprocess/math/deepscaler_preview.py @@ -0,0 +1,143 @@ +# 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. +""" +Preprocess the https://huggingface.co/datasets/agentica-org/DeepScaleR-Preview-Dataset to parquet format +Code referring to https://github.com/agentica-project/deepscaler/blob/main/scripts/data/deepscaler_dataset.py +""" + +import os +import datasets +from typing import Dict, Any +import argparse + +from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset +from verl.utils.data_process.prompt import build_zero_style_prompt +from verl.utils.reward_score.math import remove_boxed, last_boxed_only_string + +def get_datasets(cache_dir: str): + """ + Load the math datasets from Hugging Face Hub. + + Args: + train_data_source (str): Source for training data + test_data_sources (list): List of sources for test data + + Returns: + tuple: (train_dataset, test_datasets) as Dataset objects + """ + train_data_source = "agentica-org/DeepScaleR-Preview-Dataset" + print(f"Loading the {train_data_source} dataset...") + train_dataset = datasets.load_dataset( + train_data_source, trust_remote_code=True, split="train", cache_dir=cache_dir + ) + + return train_dataset, None + + +def make_map_fn(split: str, data_source: str) -> callable: + """ + Create a mapping function for processing dataset examples. + + Args: + split (str): Dataset split ('train' or 'test') + data_source (str): Source of the data + + Returns: + callable: Function to process individual examples + """ + def process_fn(example: Dict[str, Any], idx: int) -> Dict[str, Any]: + question = example.pop("problem") + answer = example.pop("answer") + + if isinstance(answer, list): + answer = answer[0] # minerva, olympiad_bench + + data = { + "data_source": data_source, + "prompt": [ + {"role": "user", "content": question + " Please output the final answer within \\boxed{}."}, + ], + "ability": "math", + "apply_chat_template": True, + "reward_model": {"style": "rule", "ground_truth": answer}, + "extra_info": {"split": split, "index": idx, "original_question": question}, + } + if idx == 0 or idx == 1: + print(f"data_source: {data_source}, split: {split}, idx: {idx}") + print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) + print(data) + return data + + return process_fn + + +if __name__ == "__main__": + """Main script execution: parse args, load, process, and save datasets.""" + parser = argparse.ArgumentParser() + parser.add_argument( + "--data-dir", + default="data", + help="Directory to save the processed data files. Will be modified based on other parameters.", + ) + parser.add_argument( + "--domain", + default="math", + type=str, + help="Data domain", + ) + parser.add_argument( + "--name", + default="deepscaler_preview", + type=str, + help="Dataset name", + ) + parser.add_argument( + "--train-sample-size", + type=int, + default=None, + help="Number of samples to use from the training dataset. If None, use all samples.", + ) + parser.add_argument( + "--seed", + type=int, + default=42, + help="Random seed for reproducibility when sampling data.", + ) + args = parser.parse_args() + set_seed(args.seed) + data_source = f"{args.domain}__{args.name}" + + # Configure data sources + cache_dir = datasets.config.HF_DATASETS_CACHE + # Download the datasets from Hugging Face Hub + train_dataset, test_datasets = get_datasets(cache_dir) + + # Process train dataset + process_train_fn = make_map_fn("train", data_source) + train_data = train_dataset.map(function=process_train_fn, with_indices=True) + + # Sample + train_data = sample_dataset(train_data, args.train_sample_size) + + # Save + train_output_dir = os.path.join(args.data_dir, "train") + args.train_sample_size = len(train_dataset) if args.train_sample_size is None else args.train_sample_size + train_output_path = save_dataset( + dataset=train_data, + output_dir=train_output_dir, + filename_prefix=data_source, + sample_size=args.train_sample_size + ) + print(f"Done! \n" + f"Train data saved to {train_output_path}") \ No newline at end of file diff --git a/reasoners/data_preprocess/ood/ifeval.py b/reasoners/data_preprocess/ood/ifeval.py new file mode 100644 index 000000000..b02436ccd --- /dev/null +++ b/reasoners/data_preprocess/ood/ifeval.py @@ -0,0 +1,139 @@ +import argparse +import json +import os + +import datasets +import transformers +from datasets import load_dataset +from tqdm import tqdm + +from verl.utils.data_process.filter import LengthFilter +from verl.utils.data_process.utils import (sample_dataset, save_dataset, + set_seed) + +""" +python data_preprocess/ood/ifeval.py +""" + + +def get_datasets(cache_dir: str): + """ + Loads the IFEval dataset. + """ + try: + dataset = load_dataset("google/IFEval", cache_dir=cache_dir)["train"] # they updated the dataset under train, but its still called test + print(f"IFEval dataset: {len(dataset)} examples") + return None, dataset + except Exception as e: + print(f"Error loading dataset: {e}") + return None, None + + +non_numeric_answer = 0 +non_arithmetic_question = 0 + +PromptTemplate = """{{context}}""" + + +def make_map_fn(split: str, data_source: str) -> callable: + + def process_fn(example, idx): + prompt = example["prompt"] + instruction_id_list = example["instruction_id_list"] + kwargs = example["kwargs"] + + data = { + "data_source": data_source, + "prompt": [ + { + "role": "user", + "content": PromptTemplate.replace("{{context}}", prompt) + } + ], + "ability": "ood", + "apply_chat_template": True, + "reward_model": { + "style": "rule", + "ground_truth": kwargs, + }, + "extra_info": { + "instruction_id_list": instruction_id_list, + "prompt": prompt, + } + } + + if idx == 0 or idx == 1: + print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) + print(data) + + return data + + return process_fn + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Download, process, and save OOD datasets." + ) + parser.add_argument( + "--data-dir", + default="data", + help="Base directory to save the processed data files.", + ) + parser.add_argument("--domain", default="ood", help="Domain of the dataset.") + parser.add_argument("--name", default="ifeval", help="Name of the dataset.") + parser.add_argument( + "--sample-size", + type=int, + default=None, + help="Number of samples to use from dataset. If None, use all samples.", + ) + parser.add_argument( + "--seed", type=int, default=42, help="Random seed for reproducibility" + ) + + args = parser.parse_args() + + # Config + set_seed(args.seed) + data_source = f"{args.domain}__{args.name}" + test_output_dir = os.path.join(args.data_dir, "test") + + # Download dataset + cache_dir = datasets.config.HF_DATASETS_CACHE + _, dataset = get_datasets(cache_dir) + + # Process datasets + process_fn = make_map_fn("test", data_source) + + dataset = dataset.map(function=process_fn, with_indices=True) + + # Filter dataset + try: + # length filter + tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B") + length_filter = LengthFilter(tokenizer=tokenizer, max_length=4096) + dataset = dataset.filter(lambda x: length_filter.check(x)) + + # non-numeric answer filter + dataset = dataset.filter(lambda x: x["reward_model"]["ground_truth"] is not None) + except Exception as e: + print(f"Warning: Could not perform length filtering. Error: {e}") + print("Proceeding without length filtering.") + + # Sample the dataset + dataset = sample_dataset(dataset, args.sample_size) + + # Save the dataset to test directory + test_output_path = save_dataset( + dataset=dataset, + output_dir=test_output_dir, + filename_prefix=data_source, + sample_size=len(dataset), + ) + + print( + f"\nDone! \n" + f"Data source: {data_source}\n" + f"Test data saved to {test_output_path} ({len(dataset)} samples)" + ) diff --git a/reasoners/data_preprocess/ood/livebench.py b/reasoners/data_preprocess/ood/livebench.py new file mode 100644 index 000000000..e972c7673 --- /dev/null +++ b/reasoners/data_preprocess/ood/livebench.py @@ -0,0 +1,160 @@ +import argparse +import json +import os +import datetime + +import datasets +import transformers +from datasets import load_dataset +from tqdm import tqdm + +from verl.utils.data_process.filter import LengthFilter +from verl.utils.data_process.utils import (sample_dataset, save_dataset, + set_seed) + +""" +python data_preprocess/ood/livebench.py +""" + + +def get_datasets(cache_dir: str): + """ + Loads the LiveBench dataset. + """ + + test_data_sources = [ + "livebench/reasoning", + "livebench/data_analysis", + # "livebench/coding", + # "livebench/instruction_following", + # "livebench/math", + "livebench/language", + ] + + print(f"Loading the test datasets...") + test_datasets = { + os.path.basename(test_data_source.lower()): + datasets.load_dataset(test_data_source, trust_remote_code=True, split="test", cache_dir=cache_dir) + for test_data_source in test_data_sources + } + + return test_datasets + + + +non_numeric_answer = 0 +non_arithmetic_question = 0 + +PromptTemplate = """{{context}}""" + + +def make_map_fn(split: str, data_source: str) -> callable: + + def process_fn(example, idx): + turns = example["turns"] + if len(turns) != 1: + print(f"⚠️ Warning: {len(turns)} turns found in {example['task']}") + ground_truth = example["ground_truth"] + date = example["livebench_release_date"] + + data = { + "data_source": data_source, + "prompt": [ + { + "role": "user", + "content": PromptTemplate.replace("{{context}}", turns[0]) + } + ], + "ability": "ood", + "apply_chat_template": True, + "reward_model": { + "style": "rule", + "ground_truth": ground_truth, + }, + "extra_info": { + "category": example["category"], + "task": example["task"], + "prompt": turns, + "date": date, + } + } + + if idx == 0 or idx == 1: + print(data["prompt"][0]["content"]) + print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) + print(data) + + return data + + return process_fn + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Download, process, and save OOD datasets." + ) + parser.add_argument( + "--data-dir", + default="data", + help="Base directory to save the processed data files.", + ) + parser.add_argument("--domain", default="ood", help="Domain of the dataset.") + parser.add_argument("--name", default="livebench", help="Name of the dataset.") + parser.add_argument( + "--sample-size", + type=int, + default=None, + help="Number of samples to use from dataset. If None, use all samples.", + ) + parser.add_argument( + "--seed", type=int, default=42, help="Random seed for reproducibility" + ) + + args = parser.parse_args() + + # Config + set_seed(args.seed) + data_source = f"{args.domain}__{args.name}" + test_output_dir = os.path.join(args.data_dir, "test") + + # Download dataset + cache_dir = datasets.config.HF_DATASETS_CACHE + test_datasets = get_datasets(cache_dir) + + # Process datasets + process_fn = make_map_fn("test", data_source) + for test_data_source, test_data in test_datasets.items(): + dataset = test_data.map(function=process_fn, with_indices=True) + # filter date before 2024-6 and after 2024-07 + dataset = dataset.filter(lambda x: x["extra_info"]["date"] >= datetime.datetime(2024, 6, 1) and x["extra_info"]["date"] <= datetime.datetime(2024, 7, 31)) + # Filter dataset + try: + # length filter + tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B") + length_filter = LengthFilter(tokenizer=tokenizer, max_length=4096) + dataset = dataset.filter(lambda x: length_filter.check(x)) + + # null answer filter + dataset = dataset.filter(lambda x: x["reward_model"]["ground_truth"] is not None) + except Exception as e: + print(f"Warning: Could not perform length filtering. Error: {e}") + print("Proceeding without length filtering.") + + # Sample the dataset + dataset = sample_dataset(dataset, args.sample_size) + + # Save the dataset to test directory + print(test_data_source, data_source) + file_prefix = test_data_source.replace("/", "_") + test_output_path = save_dataset( + dataset=dataset, + output_dir=test_output_dir, + filename_prefix=f"{data_source}_{file_prefix}", + sample_size=len(dataset), + ) + + print( + f"\nDone! \n" + f"Data source: {data_source}\n" + f"Test data saved to {test_output_path} ({len(dataset)} samples)" + ) diff --git a/reasoners/data_preprocess/sft_rlhf/full_hh_rlhf.py b/reasoners/data_preprocess/sft_rlhf/full_hh_rlhf.py new file mode 100644 index 000000000..07e0884cb --- /dev/null +++ b/reasoners/data_preprocess/sft_rlhf/full_hh_rlhf.py @@ -0,0 +1,146 @@ +# 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. +""" +- Preprocess data and split the training set into 75% for training RM and 25% for validting RM. +- All the training data is used to train SFT and RL. +- Both chosen and rejected is used to train SFT +""" +import argparse +import os + +import pandas as pd +from datasets import load_dataset + +from tqdm.auto import tqdm + +from verl.utils.fs import copy, makedirs + + +def generate_sft_dataset(target_hdfs_path_dir, local_dir='~/data/full_hh_rlh/sft'): + dataset = load_dataset('Dahoas/full-hh-rlhf') + output = {'prompt': [], 'response': []} + for data in tqdm(dataset['train']): + # add chosen + output['prompt'].append(data['prompt']) + output['response'].append(data['chosen']) + + # add rejection + output['prompt'].append(data['prompt']) + output['response'].append(data['rejected']) + + df = pd.DataFrame(output) + + local_dir = os.path.expanduser(local_dir) + os.makedirs(local_dir, exist_ok=True) + + local_path = os.path.join(local_dir, 'train.parquet') + + df.to_parquet(path=local_path) + + if target_hdfs_path_dir is not None: + hdfs_dir = target_hdfs_path_dir + '/' + 'train.parquet' + makedirs(hdfs_dir) + + copy(local_path, hdfs_dir) + + +def generate_rm_dataset(target_hdfs_path_dir, local_dir='~/data/full_hh_rlh/rm'): + train_dataset = load_dataset('Dahoas/full-hh-rlhf', split='train[:75%]') + test_dataset = load_dataset('Dahoas/full-hh-rlhf', split='train[-25%:]') + + local_dir = os.path.expanduser(local_dir) + os.makedirs(local_dir, exist_ok=True) + + for dataset, name in zip([train_dataset, test_dataset], ['train', 'test']): + output = {'prompt': [], 'chosen': [], 'rejected': []} + for data in tqdm(dataset): + # add chosen + output['prompt'].append(data['prompt']) + output['chosen'].append(data['chosen']) + output['rejected'].append(data['rejected']) + + df = pd.DataFrame(output) + + local_path = os.path.join(local_dir, name + '.parquet') + + df.to_parquet(path=local_path) + + if target_hdfs_path_dir is not None: + hdfs_dir = target_hdfs_path_dir + '/' + name + '.parquet' + makedirs(hdfs_dir) + + copy(local_path, hdfs_dir) + + +def generate_rl_dataset(target_hdfs_path_dir, local_dir='~/data/full_hh_rlhf/rl'): + dataset = load_dataset('Dahoas/full-hh-rlhf') + train_dataset = dataset['train'] + + data_source = 'Dahoas/full-hh-rlhf' + + # add a row to each data item that represents a unique id + def make_map_fn(split): + + def process_fn(example, idx): + prompt = example.pop('prompt') + response = example.pop('response') + + data = { + "data_source": data_source, + "prompt": [{ + "role": "user", + "content": prompt + }], + "ability": "alignment", + "reward_model": { + "style": "model", + "ground_truth": response # should not be used + }, + "extra_info": { + 'split': split, + 'index': idx + } + } + return data + + return process_fn + + train_dataset = train_dataset.map(function=make_map_fn('train'), with_indices=True) + local_dir = os.path.expanduser(local_dir) + local_path = os.path.join(local_dir, 'train.parquet') + train_dataset.to_parquet(local_path) + + if target_hdfs_path_dir is not None: + hdfs_dir = target_hdfs_path_dir + '/' + 'train.parquet' + makedirs(hdfs_dir) + + copy(local_path, hdfs_dir) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--split', type=str, choices=['sft', 'rm', 'rl'], required=True) + parser.add_argument('--local_dir', type=str, default='~/data/full_hh_rlhf') + parser.add_argument('--hdfs_dir', type=str, required=False, default=None) + + args = parser.parse_args() + + if args.split == 'sft': + generate_sft_dataset(args.hdfs_dir, os.path.join(args.local_dir, args.split)) + elif args.split == 'rm': + generate_rm_dataset(args.hdfs_dir, os.path.join(args.local_dir, args.split)) + elif args.split == 'rl': + generate_rl_dataset(args.hdfs_dir, os.path.join(args.local_dir, args.split)) + else: + raise NotImplementedError diff --git a/reasoners/data_preprocess/sft_rlhf/limo.py b/reasoners/data_preprocess/sft_rlhf/limo.py new file mode 100644 index 000000000..b9ee0dd3f --- /dev/null +++ b/reasoners/data_preprocess/sft_rlhf/limo.py @@ -0,0 +1,93 @@ +# 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. +""" +Preprocess the GSM8k dataset to parquet format +""" + +import re +import os +import datasets + +from verl.utils.hdfs_io import copy, makedirs +import argparse + + +def extract_solution(solution_str): + solution = re.search(r"\\boxed\{(.*?)\}", solution_str) + assert solution is not None + final_solution = solution.group(1) + return final_solution + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--local_dir', default='~/data/limo') + parser.add_argument('--hdfs_dir', default=None) + + args = parser.parse_args() + + data_source = 'GAIR/LIMO' + + dataset = datasets.load_dataset(data_source) + + train_dataset = dataset['train'] + test_dataset = dataset['train'] + + instruction_following = "Let's think step by step and output the final answer in \\boxed\{\}." + + # add a row to each data item that represents a unique id + def make_map_fn(split): + + def process_fn(example, idx): + question_raw = example.pop('question') + + question = question_raw + ' ' + instruction_following + + solution = example.pop('solution') + exact_answer = example.pop('answer') + data = { + "data_source": data_source, + "prompt": [{ + "role": "user", + "content": question, + }], + "ability": "math", + "reward_model": { + "style": "rule", + "ground_truth": exact_answer + }, + "extra_info": { + 'split': split, + 'index': idx, + 'answer': solution, + "question": question, # TODO: use prompt \\boxed later + } + } + return data + + return process_fn + + train_dataset = train_dataset.map(function=make_map_fn('train'), with_indices=True) + test_dataset = test_dataset.map(function=make_map_fn('test'), with_indices=True) + + local_dir = args.local_dir + hdfs_dir = args.hdfs_dir + + train_dataset.to_parquet(os.path.join(local_dir, 'train_w_prompt.parquet')) + test_dataset.to_parquet(os.path.join(local_dir, 'test_w_prompt.parquet')) + + if hdfs_dir is not None: + makedirs(hdfs_dir) + + copy(src=local_dir, dst=hdfs_dir) diff --git a/reasoners/data_preprocess/simulation/arcagi.py b/reasoners/data_preprocess/simulation/arcagi.py new file mode 100644 index 000000000..a862d8289 --- /dev/null +++ b/reasoners/data_preprocess/simulation/arcagi.py @@ -0,0 +1,183 @@ +"""Downloads, processes, and saves Abstraction and Reasoning Corpus for Artificial General Intelligence v2 (ARC-AGI-2) datasets.""" + +import os +import argparse +import json +import datasets +from datasets import Dataset +import subprocess + +from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset +from verl.utils.data_process.filter import LengthFilter + +import transformers + + + +RawARCPrompt = """You are a world-class puzzle solver with exceptional pattern recognition skills. Your task is to analyze puzzles, spot patterns, and provide direct solutions. +Given input-output grid pairs as reference examples, carefully observe the patterns to predict the output grid for new test input. Each pair follows the same transformation rule. Grids are 2D arrays. Here are the input and output grids for the reference examples: +---------------------------------------- +{{training_data}} +---------------------------------------- +Now, solve the following puzzle based on its input grid by applying the rules you have learned from the training data. +---------------------------------------- +{{input_test_data}} +---------------------------------------- +What is the output grid? Please put your answer within and tags, your final answer should be only the output grid (2d array). +""" + + +def get_datasets(cache_dir, name, download=False): + """ + Downloads (if specified) and loads the ARC-AGI-1 and ARC-AGI-2 datasets. + """ + # Download the dataset + if name == "arcagi1": + dir_path = os.path.join(cache_dir, "ARC-AGI-1/") + url = "https://github.com/fchollet/ARC-AGI.git" + elif name == "arcagi2": + dir_path = os.path.join(cache_dir, "ARC-AGI-2/") + url = "https://github.com/arcprize/ARC-AGI-2.git" + if download: + if os.path.isdir(dir_path): + pass + else: + try: + subprocess.run(["git", "clone", url, str(dir_path)], check=True) + except subprocess.CalledProcessError as e: + print("Error occurred while cloning the repo:", e) + + # Build the dataset. + train_dataset, test_dataset = [], [] + data_dir = dir_path + train_dir = os.path.join(data_dir, "data/training/") + for data_path in os.listdir(train_dir): + with open(os.path.join(train_dir, data_path), 'r') as f: + data = json.load(f) + train_dataset.append(data) + test_dir = os.path.join(data_dir, "data/evaluation/") + for data_path in os.listdir(test_dir): + with open(os.path.join(test_dir, data_path), 'r') as f: + data = json.load(f) + test_dataset.append(data) + + print(f"Training {len(train_dataset)} samples.") + print(f"Evaluation {len(test_dataset)} samples.") + return Dataset.from_list(train_dataset), Dataset.from_list(test_dataset) + + + +def make_map_fn(split: str, data_source: str, prompt_style: str="zero_style") -> callable: + """ + Creates a mapping function to process individual examples of the dataset. + + Args: + split: The dataset split ('train' or 'test'). + data_source: Identifier for the data source. + prompt_style: The style of prompt to use (e.g., 'zero_style'). + + Returns: + A callable function that takes an example and index, and returns the processed data. + """ + def process_fn(example, idx): + max_examples = 6 # max number of examples to show in the prompt + train_data = example.pop("train") + N_pairs = len(train_data) + training_data_prompt = "" + for i in range(min(N_pairs, max_examples)): + pair = train_data[i] + training_data_prompt += f"Example {i+1}\n" + training_data_prompt += f"Input: {pair['input']}\nOutput: {pair['output']}\n\n" + raw_prompt = RawARCPrompt.replace("{{training_data}}", training_data_prompt) + test_data = example.pop("test") + raw_prompt = raw_prompt.replace("{{input_test_data}}", str(test_data[0]["input"])) + answer = test_data[0]["output"] + data = { + "data_source": data_source, + "prompt": [{ + "role": "user", + "content": raw_prompt, + }], + "ability": "reasoning", + "apply_chat_template": True, + "reward_model": {"style": "rule", "ground_truth": answer}, + "extra_info": {"split": split, + "index": idx, + }, + } + if idx == 0 or idx == 1: + print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) + print(data) + return data + + return process_fn + + + +if __name__ == "__main__": + """Main script execution: parse args, load, process, and save datasets.""" + parser = argparse.ArgumentParser(description="Download, process, and save ARC-AGI-2 datasets.") + parser.add_argument('--data-dir', default='data', + help='Base directory to save the processed data files.') + parser.add_argument('--domain', default="simulation", help='Domain of the dataset.') + parser.add_argument('--name', default="arcagi1", choices=['arcagi1', 'arcagi2'], help='Name of the dataset.') + parser.add_argument('--train-sample-size', type=int, default=None, + help='Number of samples to use from training dataset. If None, use all samples.') + parser.add_argument('--test-sample-size', type=int, default=None, + help='Number of samples to use from test dataset. If None, use all samples.') + parser.add_argument('--prompt-style', type=str, choices=['zero_style'], default='zero_style', + help='Prompt style to use (currently only zero_style supported).') + parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility') + + args = parser.parse_args() + + set_seed(args.seed) + + data_source = f"{args.domain}__{args.name}" + train_output_dir = os.path.join(args.data_dir, 'train') + test_output_dir = os.path.join(args.data_dir, 'test') + os.makedirs(train_output_dir, exist_ok=True) + os.makedirs(test_output_dir, exist_ok=True) + + # Download the dataset from Github, but still saving to HF_datasets_cache + cache_dir = datasets.config.HF_DATASETS_CACHE + train_dataset, test_dataset = get_datasets(cache_dir, args.name, download=True) + + # Process the dataset + process_train_fn = make_map_fn('train', data_source, args.prompt_style) + process_test_fn = make_map_fn('test', data_source, args.prompt_style) + train_dataset = train_dataset.map(function=process_train_fn, with_indices=True) + test_dataset = test_dataset.map(function=process_test_fn, with_indices=True) + + # Length filter + try: + tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B") + length_filter = LengthFilter(tokenizer=tokenizer, max_length=4096) + train_dataset = train_dataset.filter(lambda x: length_filter.check(x), num_proc=64) + test_dataset = test_dataset.filter(lambda x: length_filter.check(x), num_proc=64) + except Exception as e: + print(f"Warning: Could not perform length filtering. Error: {e}") + print("Proceeding without length filtering.") + + # Sample the datasets using utility function + train_dataset = sample_dataset(train_dataset, args.train_sample_size) + test_dataset = sample_dataset(test_dataset, args.test_sample_size) + + # Save the datasets using utility function + train_output_path = save_dataset( + dataset=train_dataset, + output_dir=train_output_dir, + filename_prefix=data_source, + sample_size=len(train_dataset) + ) + test_output_path = save_dataset( + dataset=test_dataset, + output_dir=test_output_dir, + filename_prefix=data_source, + sample_size=len(test_dataset) + ) + + print(f"\nDone! \n" + f"Data source: {data_source}\n" + f"Train data saved to {train_output_path} ({len(train_dataset)} samples)\n" + f"Test data saved to {test_output_path} ({len(test_dataset)} samples)") diff --git a/reasoners/data_preprocess/simulation/barc.py b/reasoners/data_preprocess/simulation/barc.py new file mode 100644 index 000000000..902645730 --- /dev/null +++ b/reasoners/data_preprocess/simulation/barc.py @@ -0,0 +1,166 @@ +"""Downloads, processes, and saves BARC (Bootstrapping ARC: Synthetic Problem Generation for ARC Visual Reasoning Tasks) datasets.""" + +import os +import argparse +import datasets +import numpy as np +from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset +from verl.utils.data_process.filter import LengthFilter + +import transformers + + +RawARCPrompt = """You are a world-class puzzle solver with exceptional pattern recognition skills. Your task is to analyze puzzles, spot patterns, and provide direct solutions. +Given input-output grid pairs as reference examples, carefully observe the patterns to predict the output grid for new test input. Each pair follows the same transformation rule. Grids are 2D arrays. Here are the input and output grids for the reference examples: +---------------------------------------- +{{training_data}} +---------------------------------------- +Now, solve the following puzzle based on its input grid by applying the rules you have learned from the training data. +---------------------------------------- +{{input_test_data}} +---------------------------------------- +What is the output grid? Please put your answer within and tags, your final answer should be only the output grid (2d array). +""" + + +def get_datasets(cache_dir, N_samples=10000): + """ + Downloads (if specified) and loads the ARC-AGI-1 and ARC-AGI-2 datasets. + """ + # Load the dataset + ds = datasets.load_dataset("barc0/200k_HEAVY_gpt4o-description-gpt4omini-code_generated_problems", + split="train", cache_dir=cache_dir) # only have train split, no test split, max 200k samples + + ds = ds.select(range(min(N_samples, len(ds)))) + split_dataset = ds.train_test_split(test_size=0.2) + train_dataset = split_dataset["train"] + test_dataset = split_dataset["test"] + + return train_dataset, test_dataset + + + + +def make_map_fn(split: str, data_source: str, prompt_style: str="zero_style") -> callable: + """ + Creates a mapping function to process individual examples of the dataset. + + Args: + split: The dataset split ('train' or 'test'). + data_source: Identifier for the data source. + prompt_style: The style of prompt to use (e.g., 'zero_style'). + + Returns: + A callable function that takes an example and index, and returns the processed data. + """ + def process_fn(example, idx): + """ + BARC dataset contains multiple examples, we only use the maximum 6 examples as the training examples. + """ + max_examples = 6 # max number of examples to show in the prompt + train_data = example.pop("examples") + N_pairs = min(len(train_data)-1, max_examples) + training_data_prompt = "" + for i in range(N_pairs): + pair = train_data[i] + training_data_prompt += f"Example {i+1}\n" + input_data = [[int(float(x)) for x in row] for row in pair[0]] + output_data = [[int(float(x)) for x in row] for row in pair[1]] + training_data_prompt += f"Input: {input_data}\nOutput: {output_data}\n\n" + raw_prompt = RawARCPrompt.replace("{{training_data}}", training_data_prompt) + test_pair = train_data[N_pairs] + input_test_data = [[int(float(x)) for x in row] for row in test_pair[0]] + output_test_data = [[int(float(x)) for x in row] for row in test_pair[1]] + raw_prompt = raw_prompt.replace("{{input_test_data}}", str(input_test_data)) + answer = np.array(output_test_data) + data = { + "data_source": data_source, + "prompt": [{ + "role": "user", + "content": raw_prompt, + }], + "ability": "reasoning", + "apply_chat_template": True, + "reward_model": {"style": "rule", "ground_truth": answer}, + "extra_info": {"split": split, + "index": idx, + }, + } + if idx == 0 or idx == 1: + print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) + print(data) + return data + + return process_fn + + + + + +if __name__ == "__main__": + """Main script execution: parse args, load, process, and save datasets.""" + parser = argparse.ArgumentParser(description="Download, process, and save ARC-AGI-2 datasets.") + parser.add_argument('--data-dir', default='data', + help='Base directory to save the processed data files.') + parser.add_argument('--domain', default="simulation", help='Domain of the dataset.') + parser.add_argument('--name', default="barc", help='Name of the dataset.') + parser.add_argument('--train-sample-size', type=int, default=None, + help='Number of samples to use from training dataset. If None, use all samples.') + parser.add_argument('--test-sample-size', type=int, default=None, + help='Number of samples to use from test dataset. If None, use all samples.') + parser.add_argument('--prompt-style', type=str, choices=['zero_style'], default='zero_style', + help='Prompt style to use (currently only zero_style supported).') + parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility') + + args = parser.parse_args() + + set_seed(args.seed) + + data_source = f"{args.domain}__{args.name}" + train_output_dir = os.path.join(args.data_dir, 'train') + test_output_dir = os.path.join(args.data_dir, 'test') + os.makedirs(train_output_dir, exist_ok=True) + os.makedirs(test_output_dir, exist_ok=True) + + # Download the dataset from Github, but still saving to HF_datasets_cache + cache_dir = datasets.config.HF_DATASETS_CACHE + train_dataset, test_dataset = get_datasets(cache_dir) + + # Process the dataset + process_train_fn = make_map_fn('train', data_source, args.prompt_style) + process_test_fn = make_map_fn('test', data_source, args.prompt_style) + train_dataset = train_dataset.map(function=process_train_fn, with_indices=True) + test_dataset = test_dataset.map(function=process_test_fn, with_indices=True) + + # Length filter + try: + tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B") + length_filter = LengthFilter(tokenizer=tokenizer, max_length=4096) + train_dataset = train_dataset.filter(lambda x: length_filter.check(x), num_proc=64) + test_dataset = test_dataset.filter(lambda x: length_filter.check(x), num_proc=64) + except Exception as e: + print(f"Warning: Could not perform length filtering. Error: {e}") + print("Proceeding without length filtering.") + + # Sample the datasets using utility function + train_dataset = sample_dataset(train_dataset, args.train_sample_size) + test_dataset = sample_dataset(test_dataset, args.test_sample_size) + + # Save the datasets using utility function + train_output_path = save_dataset( + dataset=train_dataset, + output_dir=train_output_dir, + filename_prefix=data_source, + sample_size=len(train_dataset) + ) + test_output_path = save_dataset( + dataset=test_dataset, + output_dir=test_output_dir, + filename_prefix=data_source, + sample_size=len(test_dataset) + ) + + print(f"\nDone! \n" + f"Data source: {data_source}\n" + f"Train data saved to {train_output_path} ({len(train_dataset)} samples)\n" + f"Test data saved to {test_output_path} ({len(test_dataset)} samples)") diff --git a/reasoners/data_preprocess/simulation/codeio.py b/reasoners/data_preprocess/simulation/codeio.py new file mode 100644 index 000000000..97f5d097b --- /dev/null +++ b/reasoners/data_preprocess/simulation/codeio.py @@ -0,0 +1,193 @@ +"""Downloads, processes, and saves CodeIO PyEdu Reasoning datasets.""" + +import os +import argparse +import json +import random +import time +import transformers +import datasets +from datasets import Dataset + +from verl.utils.data_process.prompt import build_zero_style_prompt +from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset +from verl.utils.data_process.filter import LengthFilter + +InstructionFollow = "Please output the final answer within ```json```" +RawInputPredictionPrompt = """You are given a question that requires some input and output variables as follows: +{{problem_description}} +The input and output requirements are as follows: +{{io_requirements}} +Given the following {{given_type}}: +{{given}} +Can you predict a feasible input without writing any code? Please reason and put your final answer in the following json format: "input": , where should be a dictionary, even if the there is only one input variable, with keys strictly match the input variables' names as specified. Please put your answer in \\boxed{} tags. +Tip: Here is a reference code snippet for this question. You can refer to this code tov guide your reasoning but not copy spans of code directly. +{{refcode}} + +""" +RawOutputPredictionPrompt = """You are given a question that requires some input and output variables as follows: +{{problem_description}} +The input and output requirements are as follows: +{{io_requirements}} +Given the following {{given_type}}: +{{given}} +Can you predict the output without writing any code? Please reason and put your final answer in the following json format: "output": , where should strictly match the the output requirement as specified. Please put your answer in \\boxed{} tags. +Tip: Here is a reference code snippet for this question. You can refer to this code tov guide your reasoning but not copy spans of code directly. +{{refcode}} + +""" +AnswerTemplate = """"{{predict_type}}": {{sol}}""" + + +def get_datasets(cache_dir, download=False): + """ + Downloads (if specified) and loads the CodeIO PyEdu Reasoning dataset. + """ + # Download the dataset + data_path = os.path.join(cache_dir, "PythonEdu-Reasoning.jsonl") + if download: + if os.path.exists(data_path): + pass + else: + url = "https://huggingface.co/datasets/hkust-nlp/CodeIO-PyEdu-Reasoning-Raw/resolve/main/0_368500_filtered_v2_ds25.sced.jsonl" + os.system(f'wget -O {data_path} {url}') + dataset = [] + N_code = 0 + if not os.path.exists(data_path): + time.sleep(5) + + # Build the dataset. Only keep the first "input/output" pair for each code sample for diversity. + max_entry_to_load = 200000 # hard-code to save processing time + with open(data_path, "r") as f: + for line in f: + if N_code >= max_entry_to_load: + break + N_code += 1 + data = json.loads(line) + common_fields = {k: v for k, v in data.items() if k != "ios"} + if data["ios"]: + io = data["ios"][0] + if random.random() < 0.5: + dataset.append({**common_fields, "input": json.dumps(io["input"]), "output": json.dumps(io["output"]), "given_type": "input", "predict_type": "output"}) + else: + dataset.append({**common_fields, "input": json.dumps(io["input"]), "output": json.dumps(io["output"]), "given_type": "output", "predict_type": "input"}) + + N = len(dataset) + N_train = int(N * 0.1) + train_dataset = dataset[:N_train] + test_dataset = dataset[N_train:] + print(f"Total {N_code} code samples, {N} I/O samples.") + return Dataset.from_list(train_dataset), Dataset.from_list(test_dataset) + + +def make_map_fn(split: str, data_source: str) -> callable: + """ + Creates a mapping function to process individual examples of the dataset. + + Args: + split: The dataset split ('train' or 'test'). + data_source: Identifier for the data source. + + Returns: + A callable function that takes an example and index, and returns the processed data. + """ + def process_fn(example, idx): + given_type = example.pop("given_type") + predict_type = example.pop("predict_type") + if predict_type == "input": + PromptTemplate = RawInputPredictionPrompt + else: + PromptTemplate = RawOutputPredictionPrompt + prompt = PromptTemplate.replace("{{given_type}}", given_type) + InstructionFollow + for key in ["problem_description", "io_requirements", given_type, "refcode"]: + feature = example.pop(key) + if key in ["input", "output"]: + prompt = prompt.replace("{{given}}", str(feature)) + else: + prompt = prompt.replace(f"{{{{{key}}}}}", str(feature)) + + sol = example.pop(predict_type) + answer = AnswerTemplate.replace("{{predict_type}}", predict_type) + answer = answer.replace("{{sol}}", sol) + data = { + "data_source": data_source, + "prompt": [{"role": "user", "content": prompt}], + "ability": "coding-inference", + "apply_chat_template": True, + "reward_model": {"style": "rule", "ground_truth": answer}, + "extra_info": {"split": split, + "index": idx, + }, + } + if idx == 0 or idx == 1: + print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) + print(data) + return data + + return process_fn + +if __name__ == '__main__': + """Main script execution: parse args, load, process, and save datasets.""" + parser = argparse.ArgumentParser(description="Download, process, and save CodeIO PyEdu Reasoning datasets.") + parser.add_argument('--data-dir', default='data', + help='Base directory to save the processed data files.') + parser.add_argument('--domain', default="simulation", help='Domain of the dataset.') + parser.add_argument('--name', default="codeio", help='Name of the dataset.') + parser.add_argument('--train-sample-size', type=int, default=None, + help='Number of samples to use from training dataset. If None, use all samples.') + parser.add_argument('--test-sample-size', type=int, default=None, + help='Number of samples to use from test dataset. If None, use all samples.') + parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility') + + args = parser.parse_args() + + set_seed(args.seed) + + data_source = f"{args.domain}__{args.name}" + train_output_dir = os.path.join(args.data_dir, 'train') + test_output_dir = os.path.join(args.data_dir, 'test') + os.makedirs(train_output_dir, exist_ok=True) + os.makedirs(test_output_dir, exist_ok=True) + + # Download the dataset + cache_dir = datasets.config.HF_DATASETS_CACHE + train_dataset, test_dataset = get_datasets(cache_dir, download=True) + + # Process the dataset + process_train_fn = make_map_fn('train', data_source) + process_test_fn = make_map_fn('test', data_source) + train_dataset = train_dataset.map(function=process_train_fn, with_indices=True) + test_dataset = test_dataset.map(function=process_test_fn, with_indices=True) + + # Length filter + try: + tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B") + length_filter = LengthFilter(tokenizer=tokenizer, max_length=2048) + train_dataset = train_dataset.filter(lambda x: length_filter.check(x), num_proc=64) + test_dataset = test_dataset.filter(lambda x: length_filter.check(x), num_proc=64) + except Exception as e: + print(f"Warning: Could not perform length filtering. Error: {e}") + print("Proceeding without length filtering.") + + # Sample the datasets using utility function + train_dataset = sample_dataset(train_dataset, args.train_sample_size) + test_dataset = sample_dataset(test_dataset, args.test_sample_size) + + # Save the datasets using utility function + train_output_path = save_dataset( + dataset=train_dataset, + output_dir=train_output_dir, + filename_prefix=data_source, + sample_size=len(train_dataset) + ) + test_output_path = save_dataset( + dataset=test_dataset, + output_dir=test_output_dir, + filename_prefix=data_source, + sample_size=len(test_dataset) + ) + + print(f"\nDone! \n" + f"Data source: {data_source}\n" + f"Train data saved to {train_output_path} ({len(train_dataset)} samples)\n" + f"Test data saved to {test_output_path} ({len(test_dataset)} samples)") diff --git a/reasoners/data_preprocess/simulation/cruxeval-i.py b/reasoners/data_preprocess/simulation/cruxeval-i.py new file mode 100644 index 000000000..44d54095a --- /dev/null +++ b/reasoners/data_preprocess/simulation/cruxeval-i.py @@ -0,0 +1,136 @@ +"""Downloads, processes, and saves HumanEval dataset.""" + +import os +import argparse +import json +import transformers +import datasets +from datasets import load_dataset, Dataset + +from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset +from verl.utils.data_process.filter import LengthFilter +from verl.utils.reward_score.coder1 import code_exec + + +def get_datasets(cache_dir: str): + """ + Loads the CruxEval dataset. + """ + try: + dataset = load_dataset("cruxeval-org/cruxeval", cache_dir=cache_dir)["test"] + print(f"CruxEval dataset: {len(dataset)} examples") + return None, dataset + except Exception as e: + print(f"Error loading dataset: {e}") + return None, None + + +def make_map_fn(split: str, data_source: str) -> callable: + def process_fn(example, idx): + + # Extract task ID and prompt + task_id = example["id"] + + prompt = ( + "You will be given a function f and an output in the form f(??) == output. Your task is to find any input such that executing f on the input leads to the given output. There may be multiple answers, but only output one. First, think step by step. You MUST surround the answer with ```python and ``` tags. Express your answer as a passing assertion containing the input and the given output.\n\n" + "```python\n" + f"{example['code'].strip()}\n" + f"assert f(??) == {example['output']}\n" + "```\n" + ) + + data = { + "data_source": data_source, + "prompt": [{"role": "user", "content": prompt}], + "ability": "simulation", + "apply_chat_template": True, + "reward_model": { + "style": "rule", + "ground_truth": json.dumps({"functional": example["code"].strip()}), + }, + "extra_info": { + "split": split, + "index": idx, + "reference": example[ + "output" + ], # Include the canonical solution as reference + "original_prompt": prompt, + "dataset": "cruxeval-i", + "task_id": task_id, + }, + } + + if idx == 0 or idx == 1: + print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) + print(data) + + return data + + return process_fn + + +if __name__ == "__main__": + """Main script execution: parse args, load, process, and save datasets.""" + parser = argparse.ArgumentParser(description="Process and save HumanEval dataset.") + parser.add_argument( + "--data-dir", + default="data", + help="Base directory to save the processed data files.", + ) + parser.add_argument("--domain", default="simulation", help="Domain of the dataset.") + parser.add_argument("--name", default="cruxeval-i", help="Name of the dataset.") + parser.add_argument( + "--sample-size", + type=int, + default=None, + help="Number of samples to use from dataset. If None, use all samples.", + ) + parser.add_argument( + "--seed", type=int, default=42, help="Random seed for reproducibility" + ) + + args = parser.parse_args() + + set_seed(args.seed) + + data_source = f"{args.domain}__{args.name}" + test_output_dir = os.path.join(args.data_dir, "test") + os.makedirs(test_output_dir, exist_ok=True) + + # Load the dataset + cache_dir = datasets.config.HF_DATASETS_CACHE + _, dataset = get_datasets(cache_dir=cache_dir) + + # Process the dataset + process_fn = make_map_fn("test", data_source) + + dataset = dataset.map(function=process_fn, with_indices=True) + + # Filter out examples where processing failed + dataset = dataset.filter(lambda x: x["data_source"] == data_source) + + # Length filter + try: + tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B") + length_filter = LengthFilter(tokenizer=tokenizer, max_length=4096) + dataset = dataset.filter(lambda x: length_filter.check(x)) + except Exception as e: + print(f"Warning: Could not perform length filtering. Error: {e}") + print("Proceeding without length filtering.") + + # Sample the dataset + dataset = sample_dataset(dataset, args.sample_size) + + # Save the dataset to test directory + test_output_path = save_dataset( + dataset=dataset, + output_dir=test_output_dir, + filename_prefix=data_source, + sample_size=len(dataset), + ) + + print( + f"\nDone! \n" + f"Data source: {data_source}\n" + f"Test data saved to {test_output_path} ({len(dataset)} samples)" + ) diff --git a/reasoners/data_preprocess/simulation/cruxeval-o.py b/reasoners/data_preprocess/simulation/cruxeval-o.py new file mode 100644 index 000000000..d34618145 --- /dev/null +++ b/reasoners/data_preprocess/simulation/cruxeval-o.py @@ -0,0 +1,136 @@ +"""Downloads, processes, and saves HumanEval dataset.""" + +import os +import argparse +import json +import transformers +import datasets +from datasets import load_dataset, Dataset + +from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset +from verl.utils.data_process.filter import LengthFilter +from verl.utils.reward_score.coder1 import code_exec + + +def get_datasets(cache_dir: str): + """ + Loads the CruxEval dataset. + """ + try: + dataset = load_dataset("cruxeval-org/cruxeval", cache_dir=cache_dir)["test"] + print(f"CruxEval dataset: {len(dataset)} examples") + return None, dataset + except Exception as e: + print(f"Error loading dataset: {e}") + return None, None + + +def make_map_fn(split: str, data_source: str) -> callable: + def process_fn(example, idx): + + # Extract task ID and prompt + task_id = example["id"] + + prompt = ( + "You will be given a Python function and an assertion containing an input to the function. Complete the assertion with a literal (no unsimplified expressions, no function calls) containing the output when executing the provided code on the given input, even if the function is incorrect or incomplete. Do NOT output any extra information. Execute the program step by step before arriving at an answer, and provide the full assertion with the correct output in ```python and ``` tags, following the examples.\n\n" + "```python\n" + f"{example['code'].strip()}\n" + f"assert f({example['input']}) == ??\n" + "```\n" + ) + + data = { + "data_source": data_source, + "prompt": [{"role": "user", "content": prompt}], + "ability": "simulation", + "apply_chat_template": True, + "reward_model": { + "style": "rule", + "ground_truth": json.dumps({"functional": example["code"].strip()}), + }, + "extra_info": { + "split": split, + "index": idx, + "reference": example[ + "output" + ], # Include the canonical solution as reference + "original_prompt": prompt, + "dataset": "cruxeval-o", + "task_id": task_id, + }, + } + + if idx == 0 or idx == 1: + print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) + print(data) + + return data + + return process_fn + + +if __name__ == "__main__": + """Main script execution: parse args, load, process, and save datasets.""" + parser = argparse.ArgumentParser(description="Process and save HumanEval dataset.") + parser.add_argument( + "--data-dir", + default="data", + help="Base directory to save the processed data files.", + ) + parser.add_argument("--domain", default="simulation", help="Domain of the dataset.") + parser.add_argument("--name", default="cruxeval-o", help="Name of the dataset.") + parser.add_argument( + "--sample-size", + type=int, + default=None, + help="Number of samples to use from dataset. If None, use all samples.", + ) + parser.add_argument( + "--seed", type=int, default=42, help="Random seed for reproducibility" + ) + + args = parser.parse_args() + + set_seed(args.seed) + + data_source = f"{args.domain}__{args.name}" + test_output_dir = os.path.join(args.data_dir, "test") + os.makedirs(test_output_dir, exist_ok=True) + + # Load the dataset + cache_dir = datasets.config.HF_DATASETS_CACHE + _, dataset = get_datasets(cache_dir=cache_dir) + + # Process the dataset + process_fn = make_map_fn("test", data_source) + + dataset = dataset.map(function=process_fn, with_indices=True) + + # Filter out examples where processing failed + dataset = dataset.filter(lambda x: x["data_source"] == data_source) + + # Length filter + try: + tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B") + length_filter = LengthFilter(tokenizer=tokenizer, max_length=4096) + dataset = dataset.filter(lambda x: length_filter.check(x)) + except Exception as e: + print(f"Warning: Could not perform length filtering. Error: {e}") + print("Proceeding without length filtering.") + + # Sample the dataset + dataset = sample_dataset(dataset, args.sample_size) + + # Save the dataset to test directory + test_output_path = save_dataset( + dataset=dataset, + output_dir=test_output_dir, + filename_prefix=data_source, + sample_size=len(dataset), + ) + + print( + f"\nDone! \n" + f"Data source: {data_source}\n" + f"Test data saved to {test_output_path} ({len(dataset)} samples)" + ) diff --git a/reasoners/data_preprocess/stem/bbeh.py b/reasoners/data_preprocess/stem/bbeh.py new file mode 100644 index 000000000..e69de29bb diff --git a/reasoners/data_preprocess/stem/gpqa.py b/reasoners/data_preprocess/stem/gpqa.py new file mode 100644 index 000000000..c06607cc1 --- /dev/null +++ b/reasoners/data_preprocess/stem/gpqa.py @@ -0,0 +1,137 @@ +""" +Preprocess the GPQA dataset to parquet format +""" + +import os +import re +import argparse +import random +from datasets import load_dataset, Dataset + +from verl.utils.hdfs_io import copy, makedirs +from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset + + +def get_datasets(): + """ + Loads the GPQA dataset. + """ + try: + dataset = load_dataset("Idavidrein/gpqa", "gpqa_main")["train"] + print(f"GPQA dataset: {len(dataset)} examples") + return None, dataset + except Exception as e: + print(f"Error loading dataset: {e}") + return None, None + + +# adopted from https://github.com/EleutherAI/lm-evaluation-harness/blob/main/lm_eval/tasks/gpqa/zeroshot/utils.py +def preprocess(text): + if text is None: + return " " + text = text.strip() + text = text.replace(" [title]", ". ") + text = re.sub("\\[.*?\\]", "", text) + text = text.replace(" ", " ") + return text + + +def make_map_fn(split: str, data_source: str) -> callable: + def process_fn(example, idx): + # Create a default "skip" response with all required fields + question = example["Question"].strip() + correct = preprocess(example["Correct Answer"]) + incorrect1 = preprocess(example["Incorrect Answer 1"]) + incorrect2 = preprocess(example["Incorrect Answer 2"]) + incorrect3 = preprocess(example["Incorrect Answer 3"]) + + all_choices = [incorrect1, incorrect2, incorrect3, correct] + random.shuffle(all_choices) + + correct_index = all_choices.index(correct) + correct_letter = chr(65 + correct_index) + + formatted_choices = "" + for i, choice in enumerate(all_choices): + letter = chr(65 + i) + formatted_choices += f"{letter}) {choice}\n" + + # deepseek uses OpenAI's simple-eval for GPQA-Diamond, so we adopt prompts from here: https://github.com/openai/simple-evals/blob/main/gpqa_eval.py + prompt = ( + f"Answer the following multiple choice question. The last line of your response should be of the following format: 'Answer: $LETTER' (without quotes) where LETTER is one of ABCD. Think step by step before answering." + f"\n{question}\n" + f"{formatted_choices}" + ) + + + data = { + "data_source": data_source, + "prompt": [ + {"role": "user", "content": prompt} + ], + "ability": "stem", + "apply_chat_template": True, + "reward_model": { + "style": "rule", + "ground_truth": correct_letter, + }, + "extra_info": { + "split": split, + "index": idx, + "original_prompt": prompt, + "dataset": "Idavidrein/gpqa", + }, + } + + if idx == 0 or idx == 1: + print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) + print(data) + print(f'\none prompt example is \n{prompt}') + + return data + + return process_fn + +if __name__ == '__main__': + """Main script execution: parse args, load, process, and save datasets.""" + parser = argparse.ArgumentParser(description="Process and save GPQA dataset.") + parser.add_argument('--data-dir', default='data', + help='Base directory to save the processed data files.') + parser.add_argument('--domain', default="stem", + help='Domain of the dataset.') + parser.add_argument('--name', default="gpqa", + help='Name of the dataset.') + parser.add_argument('--sample-size', type=int, default=None, + help='Number of samples to use from dataset. If None, use all samples.') + parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility') + + args = parser.parse_args() + + set_seed(args.seed) + + data_source = f"{args.domain}__{args.name}" + test_output_dir = os.path.join(args.data_dir, 'test') + os.makedirs(test_output_dir, exist_ok=True) + + # Load the dataset + _, dataset = get_datasets() + + # Process the dataset + process_fn = make_map_fn('test', data_source) + + dataset = dataset.map(function=process_fn, with_indices=True) + + # Sample the dataset + dataset = sample_dataset(dataset, args.sample_size) + + # Save the dataset to test directory + test_output_path = save_dataset( + dataset=dataset, + output_dir=test_output_dir, + filename_prefix=data_source, + sample_size=len(dataset) + ) + + print(f"\nDone! \n" + f"Data source: {data_source}\n" + f"Test data saved to {test_output_path} ({len(dataset)} samples)") \ No newline at end of file diff --git a/reasoners/data_preprocess/stem/gpqa_diamond.py b/reasoners/data_preprocess/stem/gpqa_diamond.py new file mode 100644 index 000000000..c1dbad484 --- /dev/null +++ b/reasoners/data_preprocess/stem/gpqa_diamond.py @@ -0,0 +1,142 @@ +""" +Preprocess the GPQA - Diamond dataset to parquet format +""" + +import os +import re +import argparse +import random +from datasets import load_dataset, Dataset + +from verl.utils.hdfs_io import copy, makedirs +from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset + + +def get_datasets(): + """ + Loads the GPQA - Diamond dataset. + """ + try: + dataset = load_dataset("Idavidrein/gpqa", "gpqa_diamond")["train"] + print(f"GPQA - Diamond dataset: {len(dataset)} examples") + return None, dataset + except Exception as e: + print(f"Error loading dataset: {e}") + return None, None + + +# adopted from https://github.com/EleutherAI/lm-evaluation-harness/blob/main/lm_eval/tasks/gpqa/zeroshot/utils.py +def preprocess(text): + if text is None: + return " " + text = text.strip() + text = text.replace(" [title]", ". ") + text = re.sub("\\[.*?\\]", "", text) + text = text.replace(" ", " ") + return text + + +def make_map_fn(split: str, data_source: str) -> callable: + def process_fn(example, idx): + # Create a default "skip" response with all required fields + question = example["Question"].strip() + correct = preprocess(example["Correct Answer"]) + incorrect1 = preprocess(example["Incorrect Answer 1"]) + incorrect2 = preprocess(example["Incorrect Answer 2"]) + incorrect3 = preprocess(example["Incorrect Answer 3"]) + + all_choices = [incorrect1, incorrect2, incorrect3, correct] + random.shuffle(all_choices) + + correct_index = all_choices.index(correct) + correct_letter = chr(65 + correct_index) + + formatted_choices = "" + for i, choice in enumerate(all_choices): + letter = chr(65 + i) + formatted_choices += f"{letter}) {choice}\n" + + # # deepseek uses OpenAI's simple-eval for GPQA-Diamond, so we adopt prompts from here: https://github.com/openai/simple-evals/blob/main/gpqa_eval.py + # prompt = ( + # f"Answer the following multiple choice question. The last line of your response should be of the following format: 'Answer: $LETTER' (without quotes) where LETTER is one of ABCD. Think step by step before answering." + # f"\n{question}\n" + # f"{formatted_choices}" + # ) + prompt = ( + f"{question}\n" + f"{formatted_choices}" + "Please reason step by step, and put your final answer option within \\boxed{}. Only put the letter in the box, e.g. \\boxed{A}. There is only one correct answer." + ) + + + data = { + "data_source": data_source, + "prompt": [ + {"role": "user", "content": prompt} + ], + "ability": "stem", + "apply_chat_template": True, + "reward_model": { + "style": "rule", + "ground_truth": correct_letter, + }, + "extra_info": { + "split": split, + "index": idx, + "original_prompt": prompt, + "dataset": "Idavidrein/gpqa", + }, + } + + if idx == 0 or idx == 1: + print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) + print(data["prompt"][0]["content"]) + # print(f'\none prompt example is \n{prompt}') + + return data + + return process_fn + +if __name__ == '__main__': + """Main script execution: parse args, load, process, and save datasets.""" + parser = argparse.ArgumentParser(description="Process and save GPQA - Diamond dataset.") + parser.add_argument('--data-dir', default='data', + help='Base directory to save the processed data files.') + parser.add_argument('--domain', default="stem", + help='Domain of the dataset.') + parser.add_argument('--name', default="gpqa_diamond_no_box", + help='Name of the dataset.') + parser.add_argument('--sample-size', type=int, default=None, + help='Number of samples to use from dataset. If None, use all samples.') + parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility') + + args = parser.parse_args() + + set_seed(args.seed) + + data_source = f"{args.domain}__{args.name}" + test_output_dir = os.path.join(args.data_dir, 'test') + os.makedirs(test_output_dir, exist_ok=True) + + # Load the dataset + _, dataset = get_datasets() + + # Process the dataset + process_fn = make_map_fn('test', data_source) + + dataset = dataset.map(function=process_fn, with_indices=True) + + # Sample the dataset + dataset = sample_dataset(dataset, args.sample_size) + + # Save the dataset to test directory + test_output_path = save_dataset( + dataset=dataset, + output_dir=test_output_dir, + filename_prefix=data_source, + sample_size=len(dataset) + ) + + print(f"\nDone! \n" + f"Data source: {data_source}\n" + f"Test data saved to {test_output_path} ({len(dataset)} samples)") \ No newline at end of file diff --git a/reasoners/data_preprocess/stem/supergpqa.py b/reasoners/data_preprocess/stem/supergpqa.py new file mode 100644 index 000000000..e954ef43e --- /dev/null +++ b/reasoners/data_preprocess/stem/supergpqa.py @@ -0,0 +1,118 @@ +""" +Preprocess the GPQA dataset to parquet format +""" + +import os +import argparse +import random +from datasets import load_dataset, Dataset + +from verl.utils.hdfs_io import copy, makedirs +from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset + + +def get_datasets(): + """ + Loads the SuperGPQA dataset. + """ + try: + dataset = load_dataset("m-a-p/SuperGPQA")["train"] + print(f"SuperGPQA dataset: {len(dataset)} examples") + return None, dataset + except Exception as e: + print(f"Error loading dataset: {e}") + return None, None + + + +def make_map_fn(split: str, data_source: str) -> callable: + def process_fn(example, idx): + def form_options(options: list): + option_str = 'Options are:\n' + opts = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J'] + for opt, o in zip(options, opts): + option_str += f'({o}): {opt}\n' + return option_str + + question = example["question"].strip() + options = form_options(example["options"]) + query = question + '\n' + options + '\n' + + # prompt format is adopted from "General Reasoner" in https://github.com/TIGER-AI-Lab/General-Reasoner/blob/main/evaluation/eval_supergpqa.py + prompt = ( + f"{query}" + "Please reason step by step, and put your final answer option within \\boxed{}. Only put the letter in the box, e.g. \\boxed{A}. There is only one correct answer." + ) + + + data = { + "data_source": data_source, + "prompt": [ + {"role": "user", "content": prompt} + ], + "ability": "stem", + "apply_chat_template": True, + "reward_model": { + "style": "rule", + "ground_truth": example["answer_letter"], + }, + "extra_info": { + "split": split, + "index": idx, + "original_prompt": prompt, + "dataset": "m-a-p/SuperGPQA", + }, + } + + if idx == 0 or idx == 1: + print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) + print(data) + print(f'\none prompt example is \n{prompt}') + + return data + + return process_fn + +if __name__ == '__main__': + """Main script execution: parse args, load, process, and save datasets.""" + parser = argparse.ArgumentParser(description="Process and save GPQA dataset.") + parser.add_argument('--data-dir', default='data', + help='Base directory to save the processed data files.') + parser.add_argument('--domain', default="stem", + help='Domain of the dataset.') + parser.add_argument('--name', default="supergpqa", + help='Name of the dataset.') + parser.add_argument('--sample-size', type=int, default=1000, + help='Number of samples to use from dataset. If None, use all samples.') + parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility') + + args = parser.parse_args() + + set_seed(args.seed) + + data_source = f"{args.domain}__{args.name}" + test_output_dir = os.path.join(args.data_dir, 'test') + os.makedirs(test_output_dir, exist_ok=True) + + # Load the dataset + _, dataset = get_datasets() + + # Process the dataset + process_fn = make_map_fn('test', data_source) + + dataset = dataset.map(function=process_fn, with_indices=True) + + # Sample the dataset + dataset = sample_dataset(dataset, args.sample_size) + + # Save the dataset to test directory + test_output_path = save_dataset( + dataset=dataset, + output_dir=test_output_dir, + filename_prefix=data_source, + sample_size=len(dataset) + ) + + print(f"\nDone! \n" + f"Data source: {data_source}\n" + f"Test data saved to {test_output_path} ({len(dataset)} samples)") \ No newline at end of file diff --git a/reasoners/data_preprocess/stem_web/README.md b/reasoners/data_preprocess/stem_web/README.md new file mode 100644 index 000000000..5f2a389a4 --- /dev/null +++ b/reasoners/data_preprocess/stem_web/README.md @@ -0,0 +1,20 @@ +# STEM Web Dataset Preprocessing + +This directory contains scripts for preprocessing the STEM web dataset and running the verifier judge. + +## Running the Verifier Judge + +To run the verifier judge service: + +1. Submit the job using sbatch: + ```bash + sbatch run_verifier_judge.sh + ``` + +2. The script will: + - Detect the node's IP address automatically + - Launch a vLLM server running the TIGER-Lab/general-verifier model + +3. The service will be available at `http://:8000` + +4. User need to set a environment variable using following format: *export STEM_LLM_JUDGE_URL="http://{node_ip}:8000"* \ No newline at end of file diff --git a/reasoners/data_preprocess/stem_web/json2_parquet.py b/reasoners/data_preprocess/stem_web/json2_parquet.py new file mode 100644 index 000000000..0a5528122 --- /dev/null +++ b/reasoners/data_preprocess/stem_web/json2_parquet.py @@ -0,0 +1,109 @@ +#!/usr/bin/env python3 +# make_webinstruct_le30_parquet.py +""" +Convert samples_le30.json into a chat-style template and save as Parquet. +""" + +import argparse, json, os +from typing import Dict, Any, List +import datasets +from transformers import AutoTokenizer +import pprint +import random +from verl.utils.data_process.utils import set_seed, save_dataset + +# ------------------ prompt template definition ------------------ # +SYSTEM_MSG = ( + "You are a knowledgeable assistant. " + "Answer the following questions and think step by step. Please output the final answer within \\boxed{}. " +) + +def make_prompt(question: str) -> List[Dict[str, str]]: + """ + Chat-style message list: + - system: fixed instruction + - user : original question + """ + return [ + {"role": "user", "content": SYSTEM_MSG + question.strip()} + ] + +# --------------------------- main ------------------------------- # +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--in", dest="in_file", + default="samples_le30.json", + help="Input JSON file created earlier.") + parser.add_argument("--out", dest="out_file", + default="webinstruct_le30.parquet", + help="Output Parquet filename.") + parser.add_argument("--tokenizer", + default="Qwen/Qwen3-8B", + help="Name or path of tokenizer used for stats.") + args = parser.parse_args() + + # 1. load json + with open(args.in_file, encoding="utf-8") as f: + raw: List[Dict[str, Any]] = json.load(f) + + # 2. wrap into new template + processed = [] + for ex in raw: + question = ex["question"] + answer = ex["answer"] + processed.append({ + "data_source": "WebInstruct-le30", + "prompt": make_prompt(question), + "raw_prompt": question, + "ability": "QA", + "apply_chat_template": True, + "response": answer, + "reward_model": {"ground_truth": answer}, + "extra_info": { + "category": ex.get("category"), + "difficulty": ex.get("difficulty"), + "answer_type": ex.get("answer_type"), + "token_length": ex["token_length"] + }, + }) + # ---------- NEW: print a few samples ----------------- + print("\n*** Example prompts (3 random rows) ***") + for sample in random.sample(processed, k=min(3, len(processed))): + pprint.pprint( + { + "prompt": sample["prompt"], + "response": sample["response"], + "extra_info": sample["extra_info"] + }, + compact=True, + width=120 + ) + print("-" * 80) + # ------------------------------------------------------ + + # 3. build Dataset and save parquet + ds = datasets.Dataset.from_list(processed) + save_path = save_dataset(dataset=ds, + output_dir="./", + filename_prefix="stem", + sample_size=len(ds)) + + print(f"Saved {len(ds)} rows to {save_path}") + + + # 4. (optional) token-length stats over *prompt* messages + tok = AutoTokenizer.from_pretrained(args.tokenizer, trust_remote_code=True) + lengths = [ + len(tok.encode(msg["content"], add_special_tokens=False)) + for row in ds + for msg in row["prompt"] + if msg["role"] == "user" + ] + import numpy as np + lengths = np.asarray(lengths) + print(f"User-prompt token stats — min:{lengths.min()}, " + f"max:{lengths.max()}, mean:{lengths.mean():.1f}, " + f"std:{lengths.std():.1f}") + +if __name__ == "__main__": + main() diff --git a/reasoners/data_preprocess/stem_web/simple_filter.py b/reasoners/data_preprocess/stem_web/simple_filter.py new file mode 100644 index 000000000..1c2f50075 --- /dev/null +++ b/reasoners/data_preprocess/stem_web/simple_filter.py @@ -0,0 +1,35 @@ +import os +import json +from datasets import load_dataset +from transformers import AutoTokenizer + +dataset = load_dataset("TIGER-Lab/WebInstruct-verified", split="train") +def keep_example(example): + return ( + example.get("category") != "Mathematics" and + example.get("difficulty") in ["University", "PhD"] and + example.get("answer_type") not in ["Boolean", "Multiple Choice"] + ) +filtered = dataset.filter(keep_example) + +tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-8B", trust_remote_code=True) +answers = filtered["answer"] +encodings = tokenizer(answers, padding=False, truncation=False) +token_lengths = [len(ids) for ids in encodings["input_ids"]] + +selected_indices = [i for i, length in enumerate(token_lengths) if length <= 30] +print(f"Number of examples with answer token length ≤ 30: {len(selected_indices)}") + +samples = [] +for idx in selected_indices: + example = filtered[idx] + ex_dict = dict(example) + ex_dict["token_length"] = token_lengths[idx] + samples.append(ex_dict) + +os.makedirs("output", exist_ok=True) +output_path = os.path.join("output", "samples_le30.json") +with open(output_path, "w", encoding="utf-8") as f: + json.dump(samples, f, ensure_ascii=False, indent=2) + +print(f"Saved {len(samples)} examples to {output_path}") \ No newline at end of file diff --git a/reasoners/data_preprocess/stem_web/stem.py b/reasoners/data_preprocess/stem_web/stem.py new file mode 100644 index 000000000..da91e179e --- /dev/null +++ b/reasoners/data_preprocess/stem_web/stem.py @@ -0,0 +1,133 @@ +#!/usr/bin/env python3 +import os +import json +import argparse +import random +import pprint +import re + +from datasets import load_dataset, Dataset +from transformers import AutoTokenizer +import numpy as np + +# ------------------ prompt template ------------------ # +SYSTEM_MSG = ( + "You are a knowledgeable assistant. " + "Answer the following questions and think step by step. Please output the final answer within \\boxed{}. " +) + +def make_prompt(question: str): + return [ + {"role": "user", "content": SYSTEM_MSG + question.strip()} + ] + +# ------------------ filter predicates ------------------ # +def keep_example(example, max_answer_len): + if len(example.get("answer", "")) == 1 and example.get("answer", "").isalpha(): + return False + return ( + example.get("category") != "Mathematics" + and example.get("difficulty") in ["University", "PhD"] + and example.get("answer_type") not in ["Boolean", "Multiple Choice"] + and len(example.get("answer", "").split()) > 0 + ) + +def has_too_many_decimal_places(s: str, max_decimals: int = 3) -> bool: + pattern = re.compile(r"\d+\.(\d{" + str(max_decimals+1) + r",})") + return bool(pattern.search(s)) + +# --------------------------- main ------------------------------- # +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--json_out", default="stem_web.json") + parser.add_argument("--parquet_out", default="stem_web.parquet") + parser.add_argument("--dataset", default="TIGER-Lab/WebInstruct-verified") + parser.add_argument("--split", default="train") + parser.add_argument("--tokenizer", default="Qwen/Qwen3-8B") + parser.add_argument("--max_answer_len", type=int, default=30) + args = parser.parse_args() + + # 1) load and master filter + ds = load_dataset(args.dataset, split=args.split) + + def master_filter(ex): + ans = ex.get("answer", "") + if not keep_example(ex, args.max_answer_len): + return False + if has_too_many_decimal_places(ans, max_decimals=3): + return False + return True + + ds = ds.filter(master_filter) + print(f"Filtered down to {len(ds)} examples after applying all criteria") + + # 2) tokenize answers and select by token length + tok = AutoTokenizer.from_pretrained(args.tokenizer, trust_remote_code=True) + enc = tok(ds["answer"], padding=False, truncation=False) + lengths = [len(ids) for ids in enc["input_ids"]] + keep_idxs = [i for i, L in enumerate(lengths) if L <= args.max_answer_len] + print(f"{len(keep_idxs)} examples with answer ≤ {args.max_answer_len} tokens") + + # 3) build JSON samples + samples_q2idx = {} + samples = [] + for i in keep_idxs: + ex = dict(ds[i]) + ex["token_length"] = lengths[i] + prompt_str = make_prompt(ex["question"])[0]["content"] + samples_q2idx[prompt_str] = ex + samples.append(ex) + + os.makedirs(os.path.dirname(args.json_out) or ".", exist_ok=True) + with open(args.json_out, "w", encoding="utf-8") as f: + json.dump(samples_q2idx, f, ensure_ascii=False, indent=2) + print(f"Saved {len(samples)} samples to {args.json_out}") + + # 4) wrap into chat‐template records + processed = [] + for ex in samples: + processed.append({ + "data_source": "stem_web", + "prompt": make_prompt(ex["question"]), + "raw_prompt": ex["question"], + "ability": "QA", + "apply_chat_template": True, + "response": ex["answer"], + "reward_model": {"ground_truth": ex["answer"]}, + "extra_info": { + "idx": ex.get("id"), + "question": ex["question"], + "category": ex.get("category"), + "difficulty": ex.get("difficulty"), + "answer_type": ex.get("answer_type"), + "token_length": ex["token_length"], + }, + }) + + # 5) show a few examples + print("\n*** Example prompts (3 random rows) ***") + for row in random.sample(processed, k=min(3, len(processed))): + pprint.pprint({ + "prompt": row["prompt"], + "response": row["response"], + "extra_info": row["extra_info"], + }, compact=True, width=120) + print("-" * 80) + + # 6) save to Parquet + hf_ds = Dataset.from_list(processed) + hf_ds.to_parquet(args.parquet_out) + print(f"Saved {len(hf_ds)} rows to {args.parquet_out}") + + # 7) (optional) prompt‐length stats + user_lens = [] + for row in processed: + for msg in row["prompt"]: + if msg["role"] == "user": + user_lens.append(len(tok.encode(msg["content"], add_special_tokens=False))) + arr = np.array(user_lens) + print(f"User‐prompt token stats — " + f"min:{arr.min()}, max:{arr.max()}, mean:{arr.mean():.1f}, std:{arr.std():.1f}") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/reasoners/data_preprocess/table/finqa.py b/reasoners/data_preprocess/table/finqa.py new file mode 100644 index 000000000..16d307a1a --- /dev/null +++ b/reasoners/data_preprocess/table/finqa.py @@ -0,0 +1,170 @@ +import argparse +import json +import os + +import datasets +import transformers +from datasets import load_dataset +from tqdm import tqdm + +from verl.utils.data_process.filter import LengthFilter +from verl.utils.data_process.utils import (sample_dataset, save_dataset, + set_seed) + +""" +python data_preprocess/table/finqa.py +""" + + +def get_datasets(cache_dir: str): + """ + Loads the FinQA dataset. + """ + try: + dataset = load_dataset("nanoverl/finqa", cache_dir=cache_dir)["test"] + print(f"FinQA dataset: {len(dataset)} examples") + return None, dataset + except Exception as e: + print(f"Error loading dataset: {e}") + return None, None + + +non_numeric_answer = 0 +non_arithmetic_question = 0 + +InstructionFollow = "Please output the final numeric answer within \\boxed{}." +PromptTemplate = """You are given some text and tables. Use the information in the given context to answer the following question. +{{context}} +The question is: +{{question}} +""" + + +def make_map_fn(split: str, data_source: str) -> callable: + + def process_fn(example, idx): + pre_text = example["pre_text"] + post_text = example["post_text"] + table = example["table"] + question = example["question"] + answer = example["answer"] + + # custom table + tablerows = [] + for idx, row in enumerate(table): + if idx == 0: + tablerows.append("| " + " | ".join(row) + " |") + tablerows.append("| " + " | ".join(["---"] * len(row)) + " |") + else: + tablerows.append("| " + " | ".join(row) + " |") + + table_str = "\n".join(tablerows) + context_str = ( + "\n".join(pre_text) + "\n" + table_str + "\n" + "\n".join(post_text) + ) + + # filter answer which is not numeric + nanswer = ( + answer.replace(",", "") + .replace("%", " / 100") + .replace("$", "") + .replace(":", "/") + ) + + try: + nanswer = float(eval(nanswer)) + # print(nanswer) + except: + nanswer = None + + data = { + "data_source": data_source, + "prompt": [ + { + "role": "user", + "content": PromptTemplate.replace("{{context}}", context_str).replace("{{question}}", question) + InstructionFollow + } + ], + "ability": "table", + "apply_chat_template": True, + "reward_model": { + "style": "rule", + "ground_truth": nanswer, + }, + } + + if idx == 0 or idx == 1: + print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) + print(data) + + return data + + return process_fn + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Download, process, and save MultiHierTT table reasoning datasets." + ) + parser.add_argument( + "--data-dir", + default="data", + help="Base directory to save the processed data files.", + ) + parser.add_argument("--domain", default="table", help="Domain of the dataset.") + parser.add_argument("--name", default="finqa", help="Name of the dataset.") + parser.add_argument( + "--sample-size", + type=int, + default=None, + help="Number of samples to use from dataset. If None, use all samples.", + ) + parser.add_argument( + "--seed", type=int, default=42, help="Random seed for reproducibility" + ) + + args = parser.parse_args() + + # Config + set_seed(args.seed) + data_source = f"{args.domain}__{args.name}" + test_output_dir = os.path.join(args.data_dir, "test") + + # Download dataset + cache_dir = datasets.config.HF_DATASETS_CACHE + _, dataset = get_datasets(cache_dir) + + # Process datasets + process_fn = make_map_fn("test", data_source) + + dataset = dataset.map(function=process_fn, with_indices=True) + + # Filter dataset + try: + # length filter + tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B") + length_filter = LengthFilter(tokenizer=tokenizer, max_length=4096) + dataset = dataset.filter(lambda x: length_filter.check(x)) + + # non-numeric answer filter + dataset = dataset.filter(lambda x: x["reward_model"]["ground_truth"] is not None) + except Exception as e: + print(f"Warning: Could not perform length filtering. Error: {e}") + print("Proceeding without length filtering.") + + # Sample the dataset + dataset = sample_dataset(dataset, args.sample_size) + + # Save the dataset to test directory + test_output_path = save_dataset( + dataset=dataset, + output_dir=test_output_dir, + filename_prefix=data_source, + sample_size=len(dataset), + ) + + print( + f"\nDone! \n" + f"Data source: {data_source}\n" + f"Test data saved to {test_output_path} ({len(dataset)} samples)" + ) diff --git a/reasoners/data_preprocess/table/hitab.py b/reasoners/data_preprocess/table/hitab.py new file mode 100644 index 000000000..367db29ed --- /dev/null +++ b/reasoners/data_preprocess/table/hitab.py @@ -0,0 +1,263 @@ +"""Downloads, processes, and saves HiTab table reasoning datasets.""" + +import os +import json +import argparse +import random +import transformers +from datasets import Dataset +import datasets +import requests +import zipfile + +from verl.utils.data_process.filter import LengthFilter +from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset +from verl.utils.data_process.prompt import build_zero_style_prompt + + +InstructionFollow = "Please output the final answer within \\boxed{}. If there are multiple answers, please output them separated by |." +PromptTemplate = """You are given one or more tables. Use the information in the table to answer the following question. +{{tables}} +The question is: +{{question}} +""" + + +class HiTabFilter(LengthFilter): + """Filter for HiTab dataset to keep only arithmetic questions without text evidence.""" + + def __init__(self, tokenizer, max_length=2048): + super().__init__(tokenizer=tokenizer, max_length=max_length) + + def check(self, example): + # Only keep questions requiring operations on the table + aggregation = example["aggregation"] + # TODO: Now comment to include non-operation questions to enlarge the dataset + # if 'none' in aggregation: + # return False + + # Ensure the prompt length is within the specified range + length_check = super().check(example) + if not length_check: + return False + + return True + + +def get_datasets(cache_dir, download=False): + """ + Load the HiTab dataset from cache or download it if necessary. + + Args: + cache_dir (str): Directory to cache the dataset + download (bool): Whether to download the dataset if not found + + Returns: + tuple: (train_dataset, test_dataset) as Dataset objects + """ + data_path = os.path.join(cache_dir, "hitab/") + if download: + if os.path.exists(data_path): + pass + else: + # Base URL for raw content + def download_file(url, output_path): + """Download a file from GitHub raw content URL.""" + response = requests.get(url) + if response.status_code == 200: + os.makedirs(os.path.dirname(output_path), exist_ok=True) + with open(output_path, 'wb') as f: + f.write(response.content) + print(f"Successfully downloaded: {output_path}") + else: + print(f"Failed to download: {url}") + base_url = "https://raw.githubusercontent.com/microsoft/HiTab/main/data" + + # Files to download + files = [ + "tables.zip", + "train_samples.jsonl", + "test_samples.jsonl" + ] + + # Download each file + for file in files: + url = f"{base_url}/{file}" + output_path = os.path.join(data_path, file) + download_file(url, output_path) + # Unzip the tables.zip file + zip_path = os.path.join(data_path, "tables.zip") + with zipfile.ZipFile(zip_path, 'r') as zip_ref: + zip_ref.extractall(data_path) + + train_path = os.path.join(data_path, "train_samples.jsonl") + test_path = os.path.join(data_path, "test_samples.jsonl") + + def _preprocess_hitab_entry(entry): + new_entry = {} + new_entry['id'] = entry['id'] + new_entry['question'] = entry['question'] + new_entry['answer'] = '|'.join([str(answer) for answer in entry['answer']]) + new_entry['aggregation'] = entry['aggregation'] + table_id = entry['table_id'] + table_path = os.path.join(data_path, "tables/raw", f"{table_id}.json") + with open(table_path, "r") as f: + table = json.load(f) + new_entry['table'] = table["texts"] + return new_entry + + with open(train_path, "r") as f: + train_dataset = [json.loads(line) for line in f] + train_dataset = [_preprocess_hitab_entry(entry) for entry in train_dataset] + for idx, entry in enumerate(train_dataset): + print(entry) + if idx > 5: + break + + with open(test_path, "r") as f: + test_dataset = [json.loads(line) for line in f] + test_dataset = [_preprocess_hitab_entry(entry) for entry in test_dataset] + + print("Total train samples:", len(train_dataset)) + print("Total test samples:", len(test_dataset)) + + return Dataset.from_list(train_dataset), Dataset.from_list(test_dataset) + + +def table_matrix_to_markdown(table_matrix: list[list[str]]) -> str: + """ + Convert table matrix to markdown format. + + Args: + table_matrix (list[list[str]]): Table matrix + + Returns: + str: Markdown representation of the table + """ + if not table_matrix or not table_matrix[0]: + return "" + + rows = [] + + # Header row + header = "| " + " | ".join(str(cell) for cell in table_matrix[0]) + " |" + rows.append(header) + + # Separator row + separator = "|" + "|".join("-" * 3 for _ in table_matrix[0]) + "|" + rows.append(separator) + + # Data rows + for row in table_matrix[1:]: + formatted_row = "| " + " | ".join(str(cell) for cell in row) + " |" + rows.append(formatted_row) + + return "\n".join(rows) + + +def make_map_fn(split: str, data_source: str): + def process_fn(example, idx): + try: + table = example.pop("table") + table_string = table_matrix_to_markdown(table) + question = example["question"] + answer = example["answer"] + + prompt = PromptTemplate.replace("{{tables}}", table_string).replace("{{question}}", question) + InstructionFollow + + except Exception as e: + print(e) + print(table) + exit() + + data = { + "data_source": data_source, + "prompt": [ + { + "role": "user", + "content": prompt + } + ], + "ability": "table", + "apply_chat_template": True, + "reward_model": {"style": "rule", "ground_truth": answer}, + "extra_info": { + "split": split, + "index": idx, + }, + } + if idx == 0 or idx == 1: + print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) + print(data) + return data + + return process_fn + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Download, process, and save Hitab table reasoning datasets.") + parser.add_argument('--data-dir', default='data', + help='Base directory to save the processed data files.') + parser.add_argument('--domain', default="table", help='Domain of the dataset.') + parser.add_argument('--name', default="hitab", help='Name of the dataset.') + parser.add_argument('--train-sample-size', type=int, default=200, + help='Number of samples to use from training dataset. If None, use all samples.') + parser.add_argument('--test-sample-size', type=int, default=None, + help='Number of samples to use from test dataset. If None, use all samples.') + parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility') + + args = parser.parse_args() + + # Config + set_seed(args.seed) + data_source = f"{args.domain}__{args.name}" + train_output_dir = os.path.join(args.data_dir, 'train') + test_output_dir = os.path.join(args.data_dir, 'test') + + # Download dataset + cache_dir = datasets.config.HF_DATASETS_CACHE + train_dataset, test_dataset = get_datasets(cache_dir, download=True) + + # Process datasets + process_train_fn = make_map_fn('train', data_source) + process_test_fn = make_map_fn('test', data_source) + train_dataset = train_dataset.map(function=process_train_fn, with_indices=True) + test_dataset = test_dataset.map(function=process_test_fn, with_indices=True) + + print("Total train samples:", len(train_dataset)) + print("Total test samples:", len(test_dataset)) + for idx, example in enumerate(train_dataset): + print(example) + if idx > 5: + break + + # Filter dataset + tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B") + table_filter = HiTabFilter(tokenizer=tokenizer, max_length=4096) + train_dataset = train_dataset.filter(lambda x: table_filter.check(x), num_proc=64) + test_dataset = test_dataset.filter(lambda x: table_filter.check(x), num_proc=64) + args.train_sample_size = len(train_dataset) if args.train_sample_size is None else args.train_sample_size + args.test_sample_size = len(test_dataset) if args.test_sample_size is None else args.test_sample_size + + # Sample datasets if specified + train_dataset = sample_dataset(train_dataset, args.train_sample_size) + test_dataset = sample_dataset(test_dataset, args.test_sample_size) + + # Save datasets + train_output_path = save_dataset( + dataset=train_dataset, + output_dir=train_output_dir, + filename_prefix=data_source, + sample_size=args.train_sample_size + ) + test_output_path = save_dataset( + dataset=test_dataset, + output_dir=test_output_dir, + filename_prefix=data_source, + sample_size=args.test_sample_size + ) + + print(f"\nDone! \n" + f"Data source: {data_source}\n" + f"Train data saved to {train_output_path} ({len(train_dataset)} samples)\n" + f"Test data saved to {test_output_path} ({len(test_dataset)} samples)") diff --git a/reasoners/data_preprocess/table/multihier.py b/reasoners/data_preprocess/table/multihier.py new file mode 100644 index 000000000..26a15d32c --- /dev/null +++ b/reasoners/data_preprocess/table/multihier.py @@ -0,0 +1,274 @@ +"""Downloads, processes, and saves MultiHierTT table reasoning datasets.""" + +import os +import json +import argparse +import random +import transformers +from datasets import Dataset +import datasets + +from verl.utils.data_process.filter import LengthFilter +from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset +from verl.utils.data_process.prompt import build_zero_style_prompt + + +InstructionFollow = "Please output the final answer within \\boxed{}." +PromptTemplate = """You are given one or more tables. Use the information in the tables to answer the following question. +{{tables}} +The question is: +{{question}} +""" + + +class MultiHierTTFilter(LengthFilter): + """Filter for MultiHierTT dataset to keep only arithmetic questions without text evidence.""" + + def __init__(self, tokenizer, max_length=2048): + super().__init__(tokenizer=tokenizer, max_length=max_length) + + def check(self, example): + # Only keep arithmetic questions + # TODO: Now comment to include non-operation questions to enlarge the dataset + # question_type = example["qa"]["question_type"] + # if question_type != "arithmetic": + # return False + # Filter out questions that need text evidence + text_evidence = example["qa"]["text_evidence"] + if text_evidence != []: + return False + # Ensure the prompt length is within the specified range + length_check = super().check(example) + if not length_check: + return False + + return True + + +def get_datasets(cache_dir, download=False): + """ + Load the MultiHierTT dataset from cache or download it if necessary. + + Args: + cache_dir (str): Directory to cache the dataset + download (bool): Whether to download the dataset if not found + + Returns: + tuple: (train_dataset, test_dataset) as Dataset objects + """ + data_path = os.path.join(cache_dir, "multihier/") + if download: + if os.path.exists(data_path): + pass + else: + import gdown + + # Replace with your folder ID or full URL + url = "https://drive.google.com/drive/folders/1ituEWZ5F7G9T9AZ0kzZZLrHNhRigHCZJ" + + # Download the folder to the cache directory + gdown.download_folder(url, output=data_path, quiet=False, use_cookies=False) + + train_path = os.path.join(data_path, "MultiHiertt_dataset/train.json") + test_path = os.path.join(data_path, "MultiHiertt_dataset/dev.json") + + def _align_format(dataset): + for entry in dataset: + entry["qa"]["answer"] = str(entry["qa"]["answer"]) + return dataset + + with open(train_path, "r") as f: + train_dataset = json.load(f) + train_dataset = _align_format(train_dataset) + with open(test_path, "r") as f: + test_dataset = json.load(f) + test_dataset = _align_format(test_dataset) + print("Total train samples:", len(train_dataset)) + print("Total test samples:", len(test_dataset)) + + return Dataset.from_list(train_dataset), Dataset.from_list(test_dataset) + + +def html_table_to_markdown(table): + """ + Convert HTML table to markdown format. + Handles hierarchical tables with colspan attributes. + + Args: + table (str): HTML table string + + Returns: + str: Markdown representation of the table + """ + import re + from bs4 import BeautifulSoup + + # Parse HTML + soup = BeautifulSoup(table, 'html.parser') + + # Get all rows + rows = soup.find_all('tr') + if not rows: + return "" + + # Process each row to determine column structure + table_data = [] + max_cols = 0 + + for row in rows: + row_data = [] + col_idx = 0 + + for cell in row.find_all(['td', 'th']): + # Get cell content + content = cell.get_text().strip() + + # Get colspan (default to 1) + colspan = int(cell.get('colspan', 1)) + + # Add cell with colspan info + row_data.append({ + 'content': content, + 'colspan': colspan, + 'col_idx': col_idx + }) + + col_idx += colspan + + max_cols = max(max_cols, col_idx) + table_data.append(row_data) + + # Create markdown table + markdown_rows = [] + + # Process each row + for i, row in enumerate(table_data): + md_row = [''] * max_cols + + # Fill in cells + for cell in row: + content = cell['content'] + col_idx = cell['col_idx'] + colspan = cell.get('colspan', 1) + + # For cells with colspan > 1, center the content + if colspan > 1: + md_row[col_idx] = content + # Fill in empty cells for the span + for j in range(1, colspan): + if col_idx + j < max_cols: + md_row[col_idx + j] = '' + else: + md_row[col_idx] = content + + # Join cells with pipe separator + markdown_rows.append('| ' + ' | '.join(md_row) + ' |') + + # Add header separator after first row + if i == 0: + separator = '| ' + ' | '.join(['---'] * max_cols) + ' |' + markdown_rows.append(separator) + + return '\n'.join(markdown_rows) + + +def make_map_fn(split: str, data_source: str): + def process_fn(example, idx): + try: + tables = example.pop("tables") + paragraphs = example.pop("paragraphs") + question = example["qa"]["question"] + answer = example["qa"]["answer"] + table_string = "\n".join([html_table_to_markdown(table) for table in tables]) + + prompt = PromptTemplate.replace("{{tables}}", table_string).replace("{{question}}", question) + InstructionFollow + + except Exception as e: + print(e) + print(tables) + exit() + + data = { + "data_source": data_source, + "prompt": [ + { + "role": "user", + "content": prompt + } + ], + "ability": "table", + "apply_chat_template": True, + "reward_model": {"style": "rule", "ground_truth": answer}, + "extra_info": { + "split": split, + "index": idx, + }, + } + if idx == 0 or idx == 1: + print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) + print(data) + return data + + return process_fn + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Download, process, and save MultiHierTT table reasoning datasets.") + parser.add_argument('--data-dir', default='data', + help='Base directory to save the processed data files.') + parser.add_argument('--domain', default="table", help='Domain of the dataset.') + parser.add_argument('--name', default="multihier", help='Name of the dataset.') + parser.add_argument('--train-sample-size', type=int, default=None, + help='Number of samples to use from training dataset. If None, use all samples.') + parser.add_argument('--test-sample-size', type=int, default=500, + help='Number of samples to use from test dataset. If None, use all samples.') + parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility') + + args = parser.parse_args() + + # Config + set_seed(args.seed) + data_source = f"{args.domain}__{args.name}" + train_output_dir = os.path.join(args.data_dir, 'train') + test_output_dir = os.path.join(args.data_dir, 'test') + + # Download dataset + cache_dir = datasets.config.HF_DATASETS_CACHE + train_dataset, test_dataset = get_datasets(cache_dir, download=True) + + # Process datasets + process_train_fn = make_map_fn('train', data_source) + process_test_fn = make_map_fn('test', data_source) + train_dataset = train_dataset.map(function=process_train_fn, with_indices=True) + test_dataset = test_dataset.map(function=process_test_fn, with_indices=True) + + # Filter dataset + tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B") + table_filter = MultiHierTTFilter(tokenizer=tokenizer, max_length=4096) + train_dataset = train_dataset.filter(lambda x: table_filter.check(x), num_proc=64) + test_dataset = test_dataset.filter(lambda x: table_filter.check(x), num_proc=64) + args.train_sample_size = len(train_dataset) if args.train_sample_size is None else args.train_sample_size + args.test_sample_size = len(test_dataset) if args.test_sample_size is None else args.test_sample_size + + # Sample datasets if specified + train_dataset = sample_dataset(train_dataset, args.train_sample_size) + test_dataset = sample_dataset(test_dataset, args.test_sample_size) + + # Save datasets + train_output_path = save_dataset( + dataset=train_dataset, + output_dir=train_output_dir, + filename_prefix=data_source, + sample_size=args.train_sample_size + ) + test_output_path = save_dataset( + dataset=test_dataset, + output_dir=test_output_dir, + filename_prefix=data_source, + sample_size=args.test_sample_size + ) + + print(f"\nDone! \n" + f"Data source: {data_source}\n" + f"Train data saved to {train_output_path} ({len(train_dataset)} samples)\n" + f"Test data saved to {test_output_path} ({len(test_dataset)} samples)") diff --git a/reasoners/reward/__init__.py b/reasoners/reward/__init__.py new file mode 100644 index 000000000..0f1dc5991 --- /dev/null +++ b/reasoners/reward/__init__.py @@ -0,0 +1,92 @@ +# 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. +# from . import gsm8k, math, prime_math, prime_code +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["compressed"])) + # 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 000000000..915c96a77 --- /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 000000000..a9f32c17c --- /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 000000000..a1818ac93 --- /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 reward_score.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 .docker_exec import code_exec_docker + + code_exec = code_exec_docker +elif CODER1_EXEC == "firejail": + from .firejail_exec import code_exec_firejail + + code_exec = code_exec_firejail +elif CODER1_EXEC == "ces": + from .ces_exec import remote_code_exec_ces + + code_exec = remote_code_exec_ces +elif CODER1_EXEC == "kira": + from .kira_exec import remote_code_exec_kira + + code_exec = remote_code_exec_kira +elif CODER1_EXEC == "bwrap": + from .bwrap_exec import code_exec_bwrap + + code_exec = code_exec_bwrap +elif CODER1_EXEC == "unsafe_local": + from .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 000000000..b1b2f87ce --- /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 000000000..34f694e05 --- /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 000000000..80ed07b9c --- /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 000000000..872eca102 --- /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 000000000..0d4288a9b --- /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 000000000..8c01d2627 --- /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 000000000..39f5cba12 --- /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 000000000..dc8e95bbd --- /dev/null +++ b/reasoners/reward/cruxeval/__init__.py @@ -0,0 +1,33 @@ +import re +from verl.utils.reward_score.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/cruxeval.py b/reasoners/reward/cruxeval/cruxeval.py new file mode 100644 index 000000000..e69de29bb diff --git a/reasoners/reward/cruxeval/utils.py b/reasoners/reward/cruxeval/utils.py new file mode 100644 index 000000000..9eec77cd6 --- /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/geo3k.py b/reasoners/reward/geo3k.py new file mode 100644 index 000000000..92c455469 --- /dev/null +++ b/reasoners/reward/geo3k.py @@ -0,0 +1,31 @@ +# 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 +from mathruler.grader import extract_boxed_content, grade_answer + + +def format_reward(predict_str: str) -> float: + pattern = re.compile(r'.*.*\\boxed\{.*\}.*', re.DOTALL) + match_result = re.fullmatch(pattern, predict_str) + return 1.0 if match_result else 0.0 + + +def acc_reward(predict_str: str, ground_truth: str) -> float: + answer = extract_boxed_content(predict_str) + return 1.0 if grade_answer(answer, ground_truth) else 0.0 + + +def compute_score(predict_str: str, ground_truth: str) -> float: + return 0.9 * acc_reward(predict_str, ground_truth) + 0.1 * format_reward(predict_str) diff --git a/reasoners/reward/gpqa.py b/reasoners/reward/gpqa.py new file mode 100644 index 000000000..916e4890c --- /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 000000000..4d0ec08c5 --- /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/gsm8k.py b/reasoners/reward/gsm8k.py new file mode 100644 index 000000000..709103764 --- /dev/null +++ b/reasoners/reward/gsm8k.py @@ -0,0 +1,63 @@ +# 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'): + assert method in ['strict', 'flexible'] + + if method == 'strict': + # this also tests the formatting of the model + solution = re.search("#### (\\-?[0-9\\.\\,]+)", solution_str) + if solution is None: + final_answer = None + else: + final_answer = solution.group(0) + final_answer = final_answer.split('#### ')[1].replace(',', '').replace('$', '') + elif method == 'flexible': + answer = re.findall("(\\-?[0-9\\.\\,]+)", solution_str) + final_answer = None + if len(answer) == 0: + # no reward is there is no answer + pass + else: + invalid_str = ['', '.'] + # find the last number that is not '.' + 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.): + """The scoring function for GSM8k. + + 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 + """ + 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/ifeval/__init__.py b/reasoners/reward/ifeval/__init__.py new file mode 100644 index 000000000..3f6c312fb --- /dev/null +++ b/reasoners/reward/ifeval/__init__.py @@ -0,0 +1,47 @@ +from verl.utils.reward_score.ifeval import instructions_registry +import numpy as np +def compute_score(solution_str, ground_truth, extra_info): + """The scoring function for IFEval. + + 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 "" in solution_str: + answer = solution_str.split("")[1] + else: + answer = solution_str + is_following_list = [] + for index, instruction_id in enumerate(extra_info["instruction_id_list"]): + instruction_cls = instructions_registry.INSTRUCTION_DICT[instruction_id] + instruction = instruction_cls(instruction_id) + + # Remove None values from kwargs to avoid unexpected keyword argument errors in build_description method. + # print(ground_truth) + # print(extra_info["instruction_id_list"]) + # print(len(extra_info["instruction_id_list"])) + # if v is , turn it into a list + # if v is float, turn it into an int + kwargs = { + k: int(v) if isinstance(v, float) else v.tolist() if isinstance(v, np.ndarray) else v for k, v in ground_truth[index].items() if v is not None + } + + instruction.build_description(**kwargs) + args = instruction.get_instruction_args() + if args and "prompt" in args: + instruction.build_description(prompt=extra_info["prompt"]) + + if answer.strip() and instruction.check_following(answer): + is_following_list.append(True) + else: + is_following_list.append(False) + + return { + "score": all(is_following_list), + "acc": all(is_following_list), + } diff --git a/reasoners/reward/ifeval/instructions.py b/reasoners/reward/ifeval/instructions.py new file mode 100644 index 000000000..995536ec1 --- /dev/null +++ b/reasoners/reward/ifeval/instructions.py @@ -0,0 +1,1611 @@ +# Copyright 2023 The Google Research Authors. +# +# 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. + +"""Library of instructions.""" + +import collections +import json +import logging +import random +import re +import string +from typing import Dict, Optional, Sequence, Union + +import langdetect + +from verl.utils.reward_score.ifeval import instructions_util + + +logger = logging.getLogger(__name__) + +_InstructionArgsDtype = Optional[Dict[str, Union[int, str, Sequence[str]]]] + +_LANGUAGES = instructions_util.LANGUAGE_CODES + +# The relational operation for comparison. +_COMPARISON_RELATION = ("less than", "at least") + +# The maximum number of sentences. +_MAX_NUM_SENTENCES = 20 + +# The number of placeholders. +_NUM_PLACEHOLDERS = 4 + +# The number of bullet lists. +_NUM_BULLETS = 5 + +# The options of constrained response. +_CONSTRAINED_RESPONSE_OPTIONS = ( + "My answer is yes.", + "My answer is no.", + "My answer is maybe.", +) + +# The options of starter keywords. +_STARTER_OPTIONS = ( + "I would say", + "My answer is", + "I believe", + "In my opinion", + "I think", + "I reckon", + "I feel", + "From my perspective", + "As I see it", + "According to me", + "As far as I'm concerned", + "To my understanding", + "In my view", + "My take on it is", + "As per my perception", +) + +# The options of ending keywords. +# TODO(jeffreyzhou) add more ending options +_ENDING_OPTIONS = ("Any other questions?", "Is there anything else I can help with?") + +# The number of highlighted sections. +_NUM_HIGHLIGHTED_SECTIONS = 4 + +# The section splitter. +_SECTION_SPLITER = ("Section", "SECTION") + +# The number of sections. +_NUM_SECTIONS = 5 + +# The number of paragraphs. +_NUM_PARAGRAPHS = 5 + +# The postscript marker. +_POSTSCRIPT_MARKER = ("P.S.", "P.P.S") + +# The number of keywords. +_NUM_KEYWORDS = 2 + +# The occurrences of a single keyword. +_KEYWORD_FREQUENCY = 3 + +# The occurrences of a single letter. +_LETTER_FREQUENCY = 10 + +# The occurrences of words with all capital letters. +_ALL_CAPITAL_WORD_FREQUENCY = 20 + +# The number of words in the response. +_NUM_WORDS_LOWER_LIMIT = 100 +_NUM_WORDS_UPPER_LIMIT = 500 + + +class Instruction: + """An instruction template.""" + + def __init__(self, instruction_id): + self.id = instruction_id + + def build_description(self, **kwargs): + raise NotImplementedError("`build_description` not implemented.") + + def get_instruction_args(self): + raise NotImplementedError("`get_instruction_args` not implemented.") + + def get_instruction_args_keys(self): + raise NotImplementedError("`get_instruction_args_keys` not implemented.") + + def check_following(self, value): + raise NotImplementedError("`check_following` not implemented.") + + +class ResponseLanguageChecker(Instruction): + """Check the language of the entire response.""" + + def build_description(self, *, language=None): + """Build the instruction description. + + Args: + language: A string representing the expected language of the response. The + language has to comply to the 97 types defined in + `langid.py` (https://pypi.org/project/langid/1.1.5/), which follows + ISO 639-1 codes (https://en.wikipedia.org/wiki/List_of_ISO_639-1_codes); + for example, `en` for English, `zh` for Chinese, `fr` for French. + + Returns: + A string representing the instruction description. + """ + self._language = language + if self._language is None: + self._language = random.choice(list(_LANGUAGES.keys())) + # TODO(tianjianlu): opens the description generation to more choices. + self._description_pattern = ( + "Your ENTIRE response should be in {language} language, no other " + + "language is allowed." + ) + return self._description_pattern.format(language=_LANGUAGES[self._language]) + + def get_instruction_args(self): + """Returns the keyword args of `build_description`.""" + return {"language": self._language} + + def get_instruction_args_keys(self): + """Returns the args keys of `build_description`.""" + return ["language"] + + def check_following(self, value): + """Check if the language of the entire response follows the instruction. + + Args: + value: A string representing the response. + + Returns: + True if the language of `value` follows instruction; otherwise False. + """ + assert isinstance(value, str) + + try: + return langdetect.detect(value) == self._language + except langdetect.LangDetectException as e: + # Count as instruction is followed. + logging.error( + "Unable to detect language for text %s due to %s", value, e + ) # refex: disable=pytotw.037 + return True + + +class NumberOfSentences(Instruction): + """Check the number of sentences.""" + + def build_description(self, *, num_sentences=None, relation=None): + """Build the instruction description. + + Args: + num_sentences: An integer specifying the number of sentences as a + threshold. + relation: A string in (`less than`, `at least`), defining the relational + operator for comparison. + Two relational comparisons are supported for now: + if 'less than', the actual number of sentences < the threshold; + if 'at least', the actual number of sentences >= the threshold. + + Returns: + A string representing the instruction description. + """ + # The number of sentences as a threshold for comparison. + self._num_sentences_threshold = num_sentences + if self._num_sentences_threshold is None or self._num_sentences_threshold < 0: + self._num_sentences_threshold = random.randint(1, _MAX_NUM_SENTENCES) + + if relation is None: + self._comparison_relation = random.choice(_COMPARISON_RELATION) + elif relation not in _COMPARISON_RELATION: + raise ValueError( + "The supported relation for comparison must be in " + f"{_COMPARISON_RELATION}, but {relation} is given." + ) + else: + self._comparison_relation = relation + + self._description_pattern = ( + "Your response should contain {relation} {num_sentences} sentences." + ) + return self._description_pattern.format( + relation=self._comparison_relation, + num_sentences=self._num_sentences_threshold, + ) + + def get_instruction_args(self): + """Returns the keyword args of `build_description`.""" + return { + "num_sentences": self._num_sentences_threshold, + "relation": self._comparison_relation, + } + + def get_instruction_args_keys(self): + """Returns the args keys of `build_description`.""" + return ["num_sentences", "relation"] + + def check_following(self, value): + """Check if the number of sentences follows the instruction. + + Args: + value: A string representing the response. + + Returns: + True if the response follows the instruction. + + Raise: + ValueError if the string in `instruction_args` is not in + [`less_than`, `at_least`]. + """ + num_sentences = instructions_util.count_sentences(value) + if self._comparison_relation == _COMPARISON_RELATION[0]: + return num_sentences < self._num_sentences_threshold + elif self._comparison_relation == _COMPARISON_RELATION[1]: + return num_sentences >= self._num_sentences_threshold + + +class PlaceholderChecker(Instruction): + """Check the placeholders in template writing.""" + + def build_description(self, *, num_placeholders=None): + """Build the instruction description. + + Args: + num_placeholders: An integer denoting the minimum number of + placeholders required in the response. + + Returns: + A string representing the instruction description. + """ + self._num_placeholders = num_placeholders + if self._num_placeholders is None or self._num_placeholders < 0: + self._num_placeholders = random.randint(1, _NUM_PLACEHOLDERS) + self._description_pattern = ( + "The response must contain at least {num_placeholders} placeholders " + + "represented by square brackets, such as [address]." + ) + return self._description_pattern.format(num_placeholders=self._num_placeholders) + + def get_instruction_args(self): + """Returns the keyword args of `build_description`.""" + return {"num_placeholders": self._num_placeholders} + + def get_instruction_args_keys(self): + """Returns the args keys of `build_description`.""" + return ["num_placeholders"] + + def check_following(self, value): + """Check if the number of placeholders follows the instruction. + + Args: + value: A string representing the response. + + Returns: + True if the actual number of placeholders in the response is greater than + or equal to `num_placeholders`; otherwise, False. + """ + placeholders = re.findall(r"\[.*?\]", value) + num_placeholders = len(placeholders) + return num_placeholders >= self._num_placeholders + + +class BulletListChecker(Instruction): + """Checks the bullet list in the prompt.""" + + def build_description(self, *, num_bullets=None): + """Build the instruction description. + + Args: + num_bullets: An integer specifying the exact number of bullet lists + that is required to appear in the response. + + Returns: + A string representing the instruction description. + """ + self._num_bullets = num_bullets + if self._num_bullets is None or self._num_bullets < 0: + self._num_bullets = random.randint(1, _NUM_BULLETS) + self._description_pattern = ( + "Your answer must contain exactly {num_bullets} bullet points. " + + "Use the markdown bullet points such as:\n" + + "* This is point 1. \n" + + "* This is point 2" + ) + return self._description_pattern.format(num_bullets=self._num_bullets) + + def get_instruction_args(self): + """Returns the keyword args of `build_description`.""" + return {"num_bullets": self._num_bullets} + + def get_instruction_args_keys(self): + """Returns the args keys of `build_description`.""" + return ["num_bullets"] + + def check_following(self, value): + r"""Check if the number of bullet lists meets the requirement. + + Args: + value: A string representing the response. The response is expected to + contain some bullet lists that start with `\*`. + + Returns: + True if the actual number of bullet lists in the response meets the + requirement. + """ + bullet_lists = re.findall(r"^\s*\*[^\*].*$", value, flags=re.MULTILINE) + bullet_lists_2 = re.findall(r"^\s*-.*$", value, flags=re.MULTILINE) + num_bullet_lists = len(bullet_lists) + len(bullet_lists_2) + return num_bullet_lists == self._num_bullets + + +class ConstrainedResponseChecker(Instruction): + """Checks the constrained response.""" + + def build_description(self): + """Build the instruction description.""" + # A sequence of string(s) representing the options of the expected response. + self._constrained_responses = _CONSTRAINED_RESPONSE_OPTIONS + self._description_pattern = ( + "Answer with one of the following options: {response_options}" + ) + return self._description_pattern.format( + response_options=self._constrained_responses + ) + + def get_instruction_args(self): + """Returns the keyword args of `build_description`.""" + return None + + def get_instruction_args_keys(self): + """Returns the args keys of `build_description`.""" + return [] + + def check_following(self, value): + """Checks if the response matches the constrained options. + + Args: + value: A string representing the response. + + Returns: + True if the actual response contains one of the options in the constrained + responses; otherwise False. + """ + value = value.strip() + for constrained_response in self._constrained_responses: + if constrained_response in value: + return True + return False + + +class ConstrainedStartChecker(Instruction): + """Checks the response start.""" + + def build_description(self, *, starter=None): + """Build the instruction description. + + Args: + starter: A string representing the keyword that the response should start + with. + + Returns: + A string representing the instruction description. + """ + self._starter = starter.strip() if isinstance(starter, str) else starter + if self._starter is None: + self._starter = random.choice(_STARTER_OPTIONS) + self._description_pattern = ( + "During the conversation, when it is your turn, " + + "please always start with {starter}" + ) + return self._description_pattern.format(starter=self._starter) + + def get_instruction_args(self): + """Returns the keyword args of `build_description`.""" + return {"starter": self._starter} + + def get_instruction_args_keys(self): + """Returns the args keys of `build_description`.""" + return ["starter"] + + def check_following(self, value): + """Checks if the response starts with the constrained keyword or phrase. + + Args: + value: A string representing the response. + + Returns: + True if the response starts with the given phrase or keyword that is + contained in `instruction_args`; otherwise, False. + """ + response_pattern = r"^\s*" + self._starter + r".*$" + response_with_constrained_start = re.search( + response_pattern, value, flags=re.MULTILINE + ) + return True if response_with_constrained_start else False + + +class HighlightSectionChecker(Instruction): + """Checks the highlighted section.""" + + def build_description(self, *, num_highlights=None): + """Build the instruction description. + + Args: + num_highlights: An integer specifying the minimum number of highlighted + sections. + + Returns: + A string representing the instruction description. + """ + self._num_highlights = num_highlights + if self._num_highlights is None or self._num_highlights < 0: + self._num_highlights = random.randint(1, _NUM_HIGHLIGHTED_SECTIONS) + + self._description_pattern = ( + "Highlight at least {num_highlights} sections in your answer with " + + "markdown, i.e. *highlighted section*." + ) + + return self._description_pattern.format(num_highlights=self._num_highlights) + + def get_instruction_args(self): + """Returns the keyword args of `build_description`.""" + return {"num_highlights": self._num_highlights} + + def get_instruction_args_keys(self): + """Returns the args keys of `build_description`.""" + return ["num_highlights"] + + def check_following(self, value): + """Checks if the number of highlighted sections meets the requirement. + + Args: + value: a string representing the response. The response is expected to + contain highlighted sections in the format of *highlighted*. + + Returns: + True if the actual number of highlighted sections in the format of + *highlighted sections* meets the minimum requirement; otherwise False. + """ + num_highlights = 0 + highlights = re.findall(r"\*[^\n\*]*\*", value) + double_highlights = re.findall(r"\*\*[^\n\*]*\*\*", value) + for highlight in highlights: + if highlight.strip("*").strip(): + num_highlights += 1 + for highlight in double_highlights: + if highlight.removeprefix("**").removesuffix("**").strip(): + num_highlights += 1 + + return num_highlights >= self._num_highlights + + +class SectionChecker(Instruction): + """Checks the sections.""" + + def build_description(self, *, section_spliter=None, num_sections=None): + """Build the instruction description. + + Args: + section_spliter: A string represents the section spliter keyword that + marks a new section, i.e., `Section` or `SECTION`. + num_sections: An integer specifying the number of sections. + + Returns: + A string representing the instruction description. + """ + self._section_spliter = ( + section_spliter.strip() + if isinstance(section_spliter, str) + else section_spliter + ) + if self._section_spliter is None: + self._section_spliter = random.choice(_SECTION_SPLITER) + + self._num_sections = num_sections + if self._num_sections is None or self._num_sections < 0: + self._num_sections = random.randint(1, _NUM_SECTIONS) + + self._description_pattern = ( + "Your response must have {num_sections} sections. Mark the beginning " + + "of each section with {section_spliter} X, such as:\n" + + "{section_spliter} 1\n" + + "[content of section 1]\n" + + "{section_spliter} 2\n" + + "[content of section 2]" + ) + + return self._description_pattern.format( + num_sections=self._num_sections, section_spliter=self._section_spliter + ) + + def get_instruction_args(self): + """Returns the keyword args of `build_description`.""" + return { + "section_spliter": self._section_spliter, + "num_sections": self._num_sections, + } + + def get_instruction_args_keys(self): + """Returns the args keys of `build_description`.""" + return ["section_spliter", "num_sections"] + + def check_following(self, value): + """Checks the response contains multiple sections. + + Args: + value: A string representing the response. The response is expected + to contain multiple sections (number of sections is greater than 1). + A new section starts with `Section 1`, where the number denotes the + section index. + + Returns: + True if the number of sections in the response is greater than or equal to + the minimum number of sections; otherwise, False. + """ + section_splitter_patten = r"\s?" + self._section_spliter + r"\s?\d+\s?" + sections = re.split(section_splitter_patten, value) + num_sections = len(sections) - 1 + return num_sections >= self._num_sections + + +class ParagraphChecker(Instruction): + """Checks the paragraphs.""" + + def build_description(self, *, num_paragraphs=None): + """Build the instruction description. + + Args: + num_paragraphs: An integer specifying the number of paragraphs. + + Returns: + A string representing the instruction description. + """ + self._num_paragraphs = num_paragraphs + if self._num_paragraphs is None or self._num_paragraphs < 0: + self._num_paragraphs = random.randint(1, _NUM_PARAGRAPHS) + + self._description_pattern = ( + "There should be {num_paragraphs} paragraphs. " + + "Paragraphs are separated with the markdown divider: ***" + ) + + return self._description_pattern.format(num_paragraphs=self._num_paragraphs) + + def get_instruction_args(self): + """Returns the keyword args of `build_description`.""" + return {"num_paragraphs": self._num_paragraphs} + + def get_instruction_args_keys(self): + """Returns the args keys of `build_description`.""" + return ["num_paragraphs"] + + def check_following(self, value): + """Checks the response contains required number of paragraphs. + + Args: + value: A string representing the response. The response may contain + paragraphs that are separated by the markdown divider: `***`. + + Returns: + True if the actual number of paragraphs is the same as required; + otherwise, False. + """ + paragraphs = re.split(r"\s?\*\*\*\s?", value) + num_paragraphs = len(paragraphs) + + for index, paragraph in enumerate(paragraphs): + if not paragraph.strip(): + if index == 0 or index == len(paragraphs) - 1: + num_paragraphs -= 1 + else: + return False + + return num_paragraphs == self._num_paragraphs + + +class PostscriptChecker(Instruction): + """Checks the postscript.""" + + def build_description(self, *, postscript_marker=None): + """Build the instruction description. + + Args: + postscript_marker: A string containing the keyword that marks the start + of the postscript section. + + Returns: + A string representing the instruction description. + """ + self._postscript_marker = ( + postscript_marker.strip() + if isinstance(postscript_marker, str) + else postscript_marker + ) + if self._postscript_marker is None: + self._postscript_marker = random.choice(_POSTSCRIPT_MARKER) + + self._description_pattern = ( + "At the end of your response, please explicitly add a postscript " + + "starting with {postscript}" + ) + + return self._description_pattern.format(postscript=self._postscript_marker) + + def get_instruction_args(self): + """Returns the keyword args of `build_description`.""" + return {"postscript_marker": self._postscript_marker} + + def get_instruction_args_keys(self): + """Returns the args keys of `build_description`.""" + return ["postscript_marker"] + + def check_following(self, value): + """Checks if the response follows the postscript format. + + Args: + value: a string representing the response. The response is expected to + contain a postscript section. + + Returns: + True if the response contains a postscript section starting with + the keyword containing in the `instruction_args`; otherwise False. + """ + value = value.lower() + if self._postscript_marker == "P.P.S": + postscript_pattern = r"\s*p\.\s?p\.\s?s.*$" + elif self._postscript_marker == "P.S.": + postscript_pattern = r"\s*p\.\s?s\..*$" + else: + postscript_pattern = r"\s*" + self._postscript_marker.lower() + r".*$" + postscript = re.findall(postscript_pattern, value, flags=re.MULTILINE) + return True if postscript else False + + +class RephraseChecker(Instruction): + """Checks the rephrase.""" + + def build_description(self, *, original_message): + """Build the instruction description. + + Args: + original_message: A string representing the original message. The + rephrased response should only change its words/sentences in between + its two asterisks, for example, *change me*. Both original and rephrased + messages should contain the changes in the form of *change me*. + + Returns: + A string representing the instruction description. + """ + if not self.is_change(original_message): + raise ValueError( + f"Message {original_message} does not contain changes " + "in the form of *change me*." + ) + + self._reference_without_change = original_message + self._description = ( + "Rephrasing: Your rephrased response should only" + + "change the words/sentences in between two asterisks" + + "such as *change me*." + ) + return self._description + + def get_instruction_args(self): + """Returns the keyword args of `build_description`.""" + return {"original_message": self._reference_without_change} + + def get_instruction_args_keys(self): + """Returns the args keys of `build_description`.""" + return ["original_message"] + + def check_following(self, value): + r"""Checks if the rephrasing follows the instruction. + + Args: + value: A string representing the response, which is expected to rephras + the string of `instruction_args`. + + Returns: + True if `value` and `instruction_args` only differ by the words/sentences + in between two asterisks such as *change me*; otherwise, False. + """ + + if not self.is_change(value): + raise ValueError( + f"value {value} does not contain changes in the form of *change me*." + ) + + response_without_changes = self.strip_changes(value) + reference_without_changes = self.strip_changes(self._reference_without_change) + + return response_without_changes == reference_without_changes + + def is_change(self, response): + """Check if there is change in the response in the form of *change me*.""" + return re.search(r"\*.*\*", response) + + def strip_changes(self, response): + """Strips off the changes.""" + return re.sub(r"\*.*\*", "", response) + + +class KeywordChecker(Instruction): + """Check the exisitence of certain keywords.""" + + def build_description(self, *, keywords=None): + """Build the instruction description. + + Args: + keywords: A sequence of strings representing the keywords that are + expected in the response. + + Returns: + A string representing the instruction description. + """ + if not keywords: + self._keywords = instructions_util.generate_keywords( + num_keywords=_NUM_KEYWORDS + ) + else: + self._keywords = keywords + self._keywords = sorted(self._keywords) + + self._description_pattern = "Include keywords {keywords} in the response." + + return self._description_pattern.format(keywords=self._keywords) + + def get_instruction_args(self): + """Returns the keyword args of `build_description`.""" + return {"keywords": self._keywords} + + def get_instruction_args_keys(self): + """Returns the args keys of `build_description`.""" + return ["keywords"] + + def check_following(self, value): + """Check if the response contain the expected keywords.""" + for keyword in self._keywords: + if not re.search(keyword, value, flags=re.IGNORECASE): + return False + return True + + +class KeywordFrequencyChecker(Instruction): + """Check the keyword frequency.""" + + def build_description(self, *, keyword=None, frequency=None, relation=None): + """Build the instruction description. + + Args: + keyword: A string representing a keyword that is expected in the response. + frequency: An integer specifying the number of times `keyword` is expected + to appear in the response. + relation: A string in (`less than`, `at least`), defining the relational + operator for comparison. + Two relational comparisons are supported for now: + if 'less than', the actual number of occurrences < frequency; + if 'at least', the actual number of occurrences >= frequency. + + Returns: + A string representing the instruction description. + """ + if not keyword: + self._keyword = instructions_util.generate_keywords(num_keywords=1)[0] + else: + self._keyword = keyword.strip() + + self._frequency = frequency + if self._frequency is None or self._frequency < 0: + self._frequency = random.randint(1, _KEYWORD_FREQUENCY) + + if relation is None: + self._comparison_relation = random.choice(_COMPARISON_RELATION) + elif relation not in _COMPARISON_RELATION: + raise ValueError( + "The supported relation for comparison must be in " + f"{_COMPARISON_RELATION}, but {relation} is given." + ) + else: + self._comparison_relation = relation + + self._description_pattern = ( + "In your response, the word {keyword} should appear {relation} " + + "{frequency} times." + ) + + return self._description_pattern.format( + keyword=self._keyword, + relation=self._comparison_relation, + frequency=self._frequency, + ) + + def get_instruction_args(self): + """Returns the keyword args of `build_description`.""" + return { + "keyword": self._keyword, + "frequency": self._frequency, + "relation": self._comparison_relation, + } + + def get_instruction_args_keys(self): + """Returns the args keys of `build_description`.""" + return ["keyword", "frequency", "relation"] + + def check_following(self, value): + """Checks if the response contain the keyword with required frequency.""" + actual_occurrences = len(re.findall(self._keyword, value, flags=re.IGNORECASE)) + + if self._comparison_relation == _COMPARISON_RELATION[0]: + return actual_occurrences < self._frequency + elif self._comparison_relation == _COMPARISON_RELATION[1]: + return actual_occurrences >= self._frequency + + +class NumberOfWords(Instruction): + """Checks the number of words.""" + + def build_description(self, *, num_words=None, relation=None): + """Build the instruction description. + + Args: + num_words: An integer specifying the number of words contained in the + response. + relation: A string in (`less than`, `at least`), defining the relational + operator for comparison. + Two relational comparisons are supported for now: + if 'less than', the actual number of words < num_words; + if 'at least', the actual number of words >= num_words. + + Returns: + A string representing the instruction description. + """ + + self._num_words = num_words + if self._num_words is None or self._num_words < 0: + self._num_words = random.randint( + _NUM_WORDS_LOWER_LIMIT, _NUM_WORDS_UPPER_LIMIT + ) + + if relation is None: + self._comparison_relation = random.choice(_COMPARISON_RELATION) + elif relation not in _COMPARISON_RELATION: + raise ValueError( + "The supported relation for comparison must be in " + f"{_COMPARISON_RELATION}, but {relation} is given." + ) + else: + self._comparison_relation = relation + + self._description_pattern = "Answer with {relation} {num_words} words." + + return self._description_pattern.format( + relation=self._comparison_relation, num_words=self._num_words + ) + + def get_instruction_args(self): + """Returns the keyword args of `build_description`.""" + return {"num_words": self._num_words, "relation": self._comparison_relation} + + def get_instruction_args_keys(self): + """Returns the args keys of `build_description`.""" + return ["num_words", "relation"] + + def check_following(self, value): + """Checks if the response contains the expected number of words.""" + num_words = instructions_util.count_words(value) + + if self._comparison_relation == _COMPARISON_RELATION[0]: + return num_words < self._num_words + elif self._comparison_relation == _COMPARISON_RELATION[1]: + return num_words >= self._num_words + + +class JsonFormat(Instruction): + """Check the Json format.""" + + def build_description(self): + self._description_pattern = ( + "Entire output should be wrapped in JSON format. You can use markdown" + " ticks such as ```." + ) + return self._description_pattern + + def get_instruction_args(self): + """Returns the keyword args of `build_description`.""" + return None + + def get_instruction_args_keys(self): + """Returns the args keys of `build_description`.""" + return [] + + def check_following(self, value): + value = ( + value.strip() + .removeprefix("```json") + .removeprefix("```Json") + .removeprefix("```JSON") + .removeprefix("```") + .removesuffix("```") + .strip() + ) + try: + json.loads(value) + except ValueError: + return False + return True + + +class ParagraphFirstWordCheck(Instruction): + """Check the paragraph and the first word of the nth paragraph.""" + + def build_description( + self, num_paragraphs=None, nth_paragraph=None, first_word=None + ): + r"""Build the instruction description. + + Args: + num_paragraphs: An integer indicating the number of paragraphs expected + in the response. A paragraph is a subset of the string that is + expected to be separated by '\n\n'. + nth_paragraph: An integer indicating the paragraph number that we look at. + Note that n starts from 1. + first_word: A string that represent the first word of the bth paragraph. + + Returns: + A string representing the instruction description. + """ + self._num_paragraphs = num_paragraphs + if self._num_paragraphs is None or self._num_paragraphs < 0: + self._num_paragraphs = random.randint(1, _NUM_PARAGRAPHS) + + self._nth_paragraph = nth_paragraph + if ( + self._nth_paragraph is None + or self._nth_paragraph <= 0 + or self._nth_paragraph > self._num_paragraphs + ): + self._nth_paragraph = random.randint(1, self._num_paragraphs + 1) + + self._first_word = first_word + if self._first_word is None: + self._first_word = instructions_util.generate_keywords(num_keywords=1)[0] + self._first_word = self._first_word.lower() + + self._description_pattern = ( + "There should be {num_paragraphs} paragraphs. " + + "Paragraphs and only paragraphs are separated with each other by two " + + "new lines as if it was '\\n\\n' in python. " + + "Paragraph {nth_paragraph} must start with word {first_word}." + ) + + return self._description_pattern.format( + num_paragraphs=self._num_paragraphs, + nth_paragraph=self._nth_paragraph, + first_word=self._first_word, + ) + + def get_instruction_args(self): + """Returns the keyword args of `build_description`.""" + return { + "num_paragraphs": self._num_paragraphs, + "nth_paragraph": self._nth_paragraph, + "first_word": self._first_word, + } + + def get_instruction_args_keys(self): + """Returns the args keys of `build_description`.""" + return ["num_paragraphs", "nth_paragraph", "first_word"] + + def check_following(self, value): + """Checks for required number of paragraphs and correct first word. + + Args: + value: a string representing the response. The response may contain + paragraphs that are separated by two new lines and the first word of + the nth paragraph will have to match a specified word. + + Returns: + True if the number of paragraphs is the same as required and the first + word of the specified paragraph is the same as required. Otherwise, false. + """ + + paragraphs = re.split(r"\n\n", value) + num_paragraphs = len(paragraphs) + + for paragraph in paragraphs: + if not paragraph.strip(): + num_paragraphs -= 1 + + # check that index doesn't go out of bounds + if self._nth_paragraph <= num_paragraphs: + paragraph = paragraphs[self._nth_paragraph - 1].strip() + if not paragraph: + return False + else: + return False + + first_word = "" + punctuation = {".", ",", "?", "!", "'", '"'} + + # get first word and remove punctuation + word = paragraph.split()[0].strip() + # TODO(jeffrey): make more complex? + word = word.lstrip("'") + word = word.lstrip('"') + + for letter in word: + if letter in punctuation: + break + first_word += letter.lower() + + return num_paragraphs == self._num_paragraphs and first_word == self._first_word + + +# TODO(jeffrey) add relation - at least/at most? +class KeySentenceChecker(Instruction): + """Check the existence of certain key sentences.""" + + def build_description(self, key_sentences=None, num_sentences=None): + """Build the instruction description. + + Args: + key_sentences: A sequences of strings representing the key sentences that + are expected in the response. + num_sentences: The number of key sentences that are expected to be seen in + the response. + + Returns: + A string representing the instruction description. + """ + + if not key_sentences: + # TODO(jeffrey) make a generate sentences function? wonderwords package + self._key_sentences = set(["For now, this is fine."]) + else: + self._key_sentences = key_sentences + + if not num_sentences: + self._num_sentences = random.randint(1, len(self._key_sentences)) + else: + self._num_sentences = num_sentences + + self._description_pattern = ( + "Include {num_sentences} of the following sentences {key_sentences}" + ) + + return self._description_pattern.format( + num_sentences=self._num_sentences, key_sentences=self._key_sentences + ) + + def get_instruction_args(self): + """Returns the keyword args of `build_description`.""" + return { + "num_sentences": self._num_sentences, + "key_sentences": list(self._key_sentences), + } + + def get_instruction_args_keys(self): + """Returns the args keys of `build_description`.""" + return ["num_sentences", "key_sentences"] + + def check_following(self, value): + """Checks if the response contains the expected key sentences.""" + count = 0 + sentences = instructions_util.split_into_sentences(value) + for sentence in self._key_sentences: + if sentence in sentences: + count += 1 + + return count == self._num_sentences + + +class ForbiddenWords(Instruction): + """Checks that specified words are not used in response.""" + + def build_description(self, forbidden_words=None): + """Build the instruction description. + + Args: + forbidden_words: A sequences of strings representing words that are not + allowed in the response. + + Returns: + A string representing the instruction description. + """ + + if not forbidden_words: + self._forbidden_words = instructions_util.generate_keywords( + num_keywords=_NUM_KEYWORDS + ) + else: + self._forbidden_words = list(set(forbidden_words)) + self._forbidden_words = sorted(self._forbidden_words) + self._description_pattern = ( + "Do not include keywords {forbidden_words} in the response." + ) + + return self._description_pattern.format(forbidden_words=self._forbidden_words) + + def get_instruction_args(self): + """Returns the keyword args of `build_description`.""" + return {"forbidden_words": self._forbidden_words} + + def get_instruction_args_keys(self): + """Returns the args keys of `build_description`.""" + return ["forbidden_words"] + + def check_following(self, value): + """Check if the response does not contain the expected keywords.""" + for word in self._forbidden_words: + if re.search(r"\b" + word + r"\b", value, flags=re.IGNORECASE): + return False + return True + + +class RephraseParagraph(Instruction): + """Checks that the paragraph is rephrased.""" + + def build_description(self, *, original_paragraph, low, high): + """Builds the instruction description. + + Args: + original_paragraph: A string presenting the original paragraph. The + rephrases response should have betweeb low-high words in common. + low: An integer presenting the lower bound of similar words. + high: An integer representing the upper bound of similar words. + + Returns: + A string representing the instruction description. + """ + # TODO(jeffrey) make more encompassing + self._original_paragraph = original_paragraph + self._low = low + self._high = high + + self._description = ( + "Rephrase the following paragraph: " + + "{original_paragraph}\nYour response should have " + + "between {low} and {high} of the same words. " + + "Words are the same if and only if all of the " + + "letters, ignoring cases, are the same. For " + + "example, 'run' is the same as 'Run' but different " + + "to 'ran'." + ) + + return self._description.format( + original_paragraph=original_paragraph, low=self._low, high=self._high + ) + + def get_instruction_args(self): + """Returns the keyword args of `build_description`.""" + return { + "original_paragraph": self._original_paragraph, + "low": self._low, + "high": self._high, + } + + def get_instruction_args_keys(self): + """Returns the args keys of `build_description`.""" + return ["original_paragraph", "low", "high"] + + def check_following(self, value): + val_words = re.findall(r"\w+", value.lower()) + original_words = re.findall(r"\w+", self._original_paragraph.lower()) + similar_words = 0 + + dict_val = collections.Counter(val_words) + dict_original = collections.Counter(original_words) + + for word in dict_original: + similar_words += min(dict_original[word], dict_val[word]) + + return similar_words >= self._low and similar_words <= self._high + + +class TwoResponsesChecker(Instruction): + """Check that two responses were given.""" + + def build_description(self): + """Build the instruction description.""" + self._description_pattern = ( + "Give two different responses. Responses and only responses should" + " be separated by 6 asterisk symbols: ******." + ) + return self._description_pattern + + def get_instruction_args(self): + """Returns the keyword args of `build_description`.""" + return None + + def get_instruction_args_keys(self): + """Returns the args keys of `build_description`.""" + return [] + + def check_following(self, value): + """Checks if the response has two different answers. + + Args: + value: A string representing the response. + + Returns: + True if two responses are detected and false otherwise. + """ + valid_responses = list() + responses = value.split("******") + for index, response in enumerate(responses): + if not response.strip(): + if index != 0 and index != len(responses) - 1: + return False + else: + valid_responses.append(response) + return ( + len(valid_responses) == 2 + and valid_responses[0].strip() != valid_responses[1].strip() + ) + + +class RepeatPromptThenAnswer(Instruction): + """Checks that Prompt is first repeated then answered.""" + + def build_description(self, *, prompt_to_repeat=None): + """Build the instruction description. + + Args: + prompt_to_repeat: The prompt that is meant to be repeated. + + Returns: + A string representing the instruction description. + """ + if not prompt_to_repeat: + raise ValueError("prompt_to_repeat must be set.") + else: + self._prompt_to_repeat = prompt_to_repeat + self._description_pattern = ( + "First repeat the request word for word without change," + " then give your answer (1. do not say any words or characters" + " before repeating the request; 2. the request you need to repeat" + " does not include this sentence)" + ) + return self._description_pattern + + def get_instruction_args(self): + return {"prompt_to_repeat": self._prompt_to_repeat} + + def get_instruction_args_keys(self): + """Returns the args keys of `build_description`.""" + return ["prompt_to_repeat"] + + def check_following(self, value): + if value.strip().lower().startswith(self._prompt_to_repeat.strip().lower()): + return True + return False + + +class EndChecker(Instruction): + """Checks that the prompt ends with a given phrase.""" + + def build_description(self, *, end_phrase=None): + """Build the instruction description. + + Args: + end_phrase: A string representing the phrase the response should end with. + + Returns: + A string representing the instruction description. + """ + self._end_phrase = ( + end_phrase.strip() if isinstance(end_phrase, str) else end_phrase + ) + if self._end_phrase is None: + self._end_phrase = random.choice(_ENDING_OPTIONS) + self._description_pattern = ( + "Finish your response with this exact phrase {ender}. " + "No other words should follow this phrase." + ) + return self._description_pattern.format(ender=self._end_phrase) + + def get_instruction_args(self): + return {"end_phrase": self._end_phrase} + + def get_instruction_args_keys(self): + """Returns the args keys of `build_description`.""" + return ["end_phrase"] + + def check_following(self, value): + """Checks if the response ends with the expected phrase.""" + value = value.strip().strip('"').lower() + self._end_phrase = self._end_phrase.strip().lower() + return value.endswith(self._end_phrase) + + +class TitleChecker(Instruction): + """Checks the response for a title.""" + + def build_description(self): + """Build the instruction description.""" + self._description_pattern = ( + "Your answer must contain a title, wrapped in double angular brackets," + " such as <>." + ) + return self._description_pattern + + def get_instruction_args(self): + return None + + def get_instruction_args_keys(self): + """Returns the args keys of `build_description`.""" + return [] + + def check_following(self, value): + """Checks if the response contains a title.""" + pattern = r"<<[^\n]+>>" + re_pattern = re.compile(pattern) + titles = re.findall(re_pattern, value) + + for title in titles: + if title.lstrip("<").rstrip(">").strip(): + return True + return False + + +class LetterFrequencyChecker(Instruction): + """Checks letter frequency.""" + + def build_description(self, *, letter=None, let_frequency=None, let_relation=None): + """Build the instruction description. + + Args: + letter: A string representing a letter that is expected in the response. + let_frequency: An integer specifying the number of times `keyword` is + expected to appear in the response. + let_relation: A string in (`less than`, `at least`), defining the + relational operator for comparison. Two relational comparisons are + supported for now; if 'less than', the actual number of + occurrences < frequency; if 'at least', the actual number of + occurrences >= frequency. + + Returns: + A string representing the instruction description. + """ + if ( + not letter + or len(letter) > 1 + or ord(letter.lower()) < 97 + or ord(letter.lower()) > 122 + ): + self._letter = random.choice(list(string.ascii_letters)) + else: + self._letter = letter.strip() + self._letter = self._letter.lower() + + self._frequency = let_frequency + if self._frequency is None or self._frequency < 0: + self._frequency = random.randint(1, _LETTER_FREQUENCY) + + if let_relation is None: + self._comparison_relation = random.choice(_COMPARISON_RELATION) + elif let_relation not in _COMPARISON_RELATION: + raise ValueError( + "The supported relation for comparison must be in " + f"{_COMPARISON_RELATION}, but {let_relation} is given." + ) + else: + self._comparison_relation = let_relation + + self._description_pattern = ( + "In your response, the letter {letter} should appear {let_relation}" + " {let_frequency} times." + ) + + return self._description_pattern.format( + letter=self._letter, + let_frequency=self._frequency, + let_relation=self._comparison_relation, + ) + + def get_instruction_args(self): + """Returns the keyword args of build description.""" + return { + "letter": self._letter, + "let_frequency": self._frequency, + "let_relation": self._comparison_relation, + } + + def get_instruction_args_keys(self): + """Returns the args keys of `build_description`.""" + return ["letter", "let_frequency", "let_relation"] + + def check_following(self, value): + """Checks that the response contains the letter at the right frequency.""" + value = value.lower() + letters = collections.Counter(value) + + if self._comparison_relation == _COMPARISON_RELATION[0]: + return letters[self._letter] < self._frequency + else: + return letters[self._letter] >= self._frequency + + +class CapitalLettersEnglishChecker(Instruction): + """Checks that the response is in english and is in all capital letters.""" + + def build_description(self): + """Build the instruction description.""" + self._description_pattern = ( + "Your entire response should be in English, and in all capital letters." + ) + return self._description_pattern + + def get_instruction_args(self): + return None + + def get_instruction_args_keys(self): + """Returns the args keys of `build_description`.""" + return [] + + def check_following(self, value): + """Checks that the response is in English and in all capital letters.""" + assert isinstance(value, str) + + try: + return value.isupper() and langdetect.detect(value) == "en" + except langdetect.LangDetectException as e: + # Count as instruction is followed. + logging.error( + "Unable to detect language for text %s due to %s", value, e + ) # refex: disable=pytotw.037 + return True + + +class LowercaseLettersEnglishChecker(Instruction): + """Checks that the response is in english and is in all lowercase letters.""" + + def build_description(self): + """Build the instruction description.""" + self._description_pattern = ( + "Your entire response should be in English, and in all lowercase" + " letters. No capital letters are allowed." + ) + return self._description_pattern + + def get_instruction_args(self): + return None + + def get_instruction_args_keys(self): + """Returns the args keys of `build_description`.""" + return [] + + def check_following(self, value): + """Checks that the response is in English and in all lowercase letters.""" + assert isinstance(value, str) + + try: + return value.islower() and langdetect.detect(value) == "en" + except langdetect.LangDetectException as e: + # Count as instruction is followed. + logging.error( + "Unable to detect language for text %s due to %s", value, e + ) # refex: disable=pytotw.037 + return True + + +class CommaChecker(Instruction): + """Checks the response for no commas.""" + + def build_description(self): + """Build the instruction description.""" + self._description_pattern = ( + "In your entire response, refrain from the use of any commas." + ) + return self._description_pattern + + def get_instruction_args(self): + return None + + def get_instruction_args_keys(self): + """Returns the args keys of `build_description`.""" + return [] + + def check_following(self, value): + """Checks that the response does not contain commas.""" + return not re.search(r"\,", value) + + +class CapitalWordFrequencyChecker(Instruction): + """Checks frequency of words with all capital letters.""" + + def build_description( + self, + capital_frequency=None, + capital_relation=None, + ): + """Build the instruction description. + + Args: + capital_frequency: An integer that represents the number of words that + should be in all capital letters. + capital_relation: A string that is 'at least' or 'at most' that refers to + the frequency. + + Returns: + A string representing the instruction description. + """ + self._frequency = capital_frequency + if self._frequency is None: + self._frequency = random.randint(1, _ALL_CAPITAL_WORD_FREQUENCY) + + self._comparison_relation = capital_relation + if capital_relation is None: + self._comparison_relation = random.choice(_COMPARISON_RELATION) + elif capital_relation not in _COMPARISON_RELATION: + raise ValueError( + "The supported relation for comparison must be in " + f"{_COMPARISON_RELATION}, but {capital_relation} is given." + ) + + self._description_pattern = ( + "In your response, words with all capital letters should appear" + " {relation} {frequency} times." + ) + + return self._description_pattern.format( + frequency=self._frequency, relation=self._comparison_relation + ) + + def get_instruction_args(self): + """Returns the keyword args of build description.""" + return { + "capital_frequency": self._frequency, + "capital_relation": self._comparison_relation, + } + + def get_instruction_args_keys(self): + """Returns the args keys of `build_description`.""" + return ["capital_frequency", "capital_relation"] + + def check_following(self, value): + """Checks the frequency of words with all capital letters.""" + # Hyphenated words will count as one word + words = instructions_util.nltk.word_tokenize(value) + capital_words = [word for word in words if word.isupper()] + + capital_words = len(capital_words) + + if self._comparison_relation == _COMPARISON_RELATION[0]: + return capital_words < self._frequency + else: + return capital_words >= self._frequency + + +class QuotationChecker(Instruction): + """Checks response is wrapped with double quotation marks.""" + + def build_description(self): + """Build the instruction description.""" + self._description_pattern = ( + "Wrap your entire response with double quotation marks." + ) + return self._description_pattern + + def get_instruction_args(self): + """Returns the keyword args of build description.""" + return None + + def get_instruction_args_keys(self): + """Returns the args keys of `build_description`.""" + return [] + + def check_following(self, value): + """Checks if the response is wrapped with double quotation marks.""" + value = value.strip() + return len(value) > 1 and value[0] == '"' and value[-1] == '"' \ No newline at end of file diff --git a/reasoners/reward/ifeval/instructions_registry.py b/reasoners/reward/ifeval/instructions_registry.py new file mode 100644 index 000000000..cde7fb810 --- /dev/null +++ b/reasoners/reward/ifeval/instructions_registry.py @@ -0,0 +1,168 @@ +# Copyright 2023 The Google Research Authors. +# +# 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. + +"""Registry of all instructions.""" + +from verl.utils.reward_score.ifeval import instructions + + +_KEYWORD = "keywords:" + +_LANGUAGE = "language:" + +_LENGTH = "length_constraints:" + +_CONTENT = "detectable_content:" + +_FORMAT = "detectable_format:" + +_MULTITURN = "multi-turn:" + +_COMBINATION = "combination:" + +_STARTEND = "startend:" + +_CHANGE_CASES = "change_case:" + +_PUNCTUATION = "punctuation:" + +INSTRUCTION_DICT = { + _KEYWORD + "existence": instructions.KeywordChecker, + _KEYWORD + "frequency": instructions.KeywordFrequencyChecker, + # TODO(jeffreyzhou): make a proper set of sentences to choose from + # _KEYWORD + "key_sentences": instructions.KeySentenceChecker, + _KEYWORD + "forbidden_words": instructions.ForbiddenWords, + _KEYWORD + "letter_frequency": instructions.LetterFrequencyChecker, + _LANGUAGE + "response_language": instructions.ResponseLanguageChecker, + _LENGTH + "number_sentences": instructions.NumberOfSentences, + _LENGTH + "number_paragraphs": instructions.ParagraphChecker, + _LENGTH + "number_words": instructions.NumberOfWords, + _LENGTH + "nth_paragraph_first_word": instructions.ParagraphFirstWordCheck, + _CONTENT + "number_placeholders": instructions.PlaceholderChecker, + _CONTENT + "postscript": instructions.PostscriptChecker, + _FORMAT + "number_bullet_lists": instructions.BulletListChecker, + # TODO(jeffreyzhou): Pre-create paragraph or use prompt to replace + # _CONTENT + "rephrase_paragraph": instructions.RephraseParagraph, + _FORMAT + "constrained_response": instructions.ConstrainedResponseChecker, + _FORMAT + "number_highlighted_sections": (instructions.HighlightSectionChecker), + _FORMAT + "multiple_sections": instructions.SectionChecker, + # TODO(tianjianlu): Re-enable rephrasing with preprocessing the message. + # _FORMAT + "rephrase": instructions.RephraseChecker, + _FORMAT + "json_format": instructions.JsonFormat, + _FORMAT + "title": instructions.TitleChecker, + # TODO(tianjianlu): Re-enable with specific prompts. + # _MULTITURN + "constrained_start": instructions.ConstrainedStartChecker, + _COMBINATION + "two_responses": instructions.TwoResponsesChecker, + _COMBINATION + "repeat_prompt": instructions.RepeatPromptThenAnswer, + _STARTEND + "end_checker": instructions.EndChecker, + _CHANGE_CASES + "capital_word_frequency": instructions.CapitalWordFrequencyChecker, + _CHANGE_CASES + "english_capital": instructions.CapitalLettersEnglishChecker, + _CHANGE_CASES + "english_lowercase": instructions.LowercaseLettersEnglishChecker, + _PUNCTUATION + "no_comma": instructions.CommaChecker, + _STARTEND + "quotation": instructions.QuotationChecker, +} + +INSTRUCTION_CONFLICTS = { + _KEYWORD + "existence": {_KEYWORD + "existence"}, + _KEYWORD + "frequency": {_KEYWORD + "frequency"}, + # TODO(jeffreyzhou): make a proper set of sentences to choose from + # _KEYWORD + "key_sentences": instructions.KeySentenceChecker, + _KEYWORD + "forbidden_words": {_KEYWORD + "forbidden_words"}, + _KEYWORD + "letter_frequency": {_KEYWORD + "letter_frequency"}, + _LANGUAGE + "response_language": { + _LANGUAGE + "response_language", + _FORMAT + "multiple_sections", + _KEYWORD + "existence", + _KEYWORD + "frequency", + _KEYWORD + "forbidden_words", + _STARTEND + "end_checker", + _CHANGE_CASES + "english_capital", + _CHANGE_CASES + "english_lowercase", + }, + _LENGTH + "number_sentences": {_LENGTH + "number_sentences"}, + _LENGTH + "number_paragraphs": { + _LENGTH + "number_paragraphs", + _LENGTH + "nth_paragraph_first_word", + _LENGTH + "number_sentences", + _LENGTH + "nth_paragraph_first_word", + }, + _LENGTH + "number_words": {_LENGTH + "number_words"}, + _LENGTH + "nth_paragraph_first_word": { + _LENGTH + "nth_paragraph_first_word", + _LENGTH + "number_paragraphs", + }, + _CONTENT + "number_placeholders": {_CONTENT + "number_placeholders"}, + _CONTENT + "postscript": {_CONTENT + "postscript"}, + _FORMAT + "number_bullet_lists": {_FORMAT + "number_bullet_lists"}, + # TODO(jeffreyzhou): Pre-create paragraph or use prompt to replace + # _CONTENT + "rephrase_paragraph": instructions.RephraseParagraph, + _FORMAT + "constrained_response": set(INSTRUCTION_DICT.keys()), + _FORMAT + "number_highlighted_sections": {_FORMAT + "number_highlighted_sections"}, + _FORMAT + "multiple_sections": { + _FORMAT + "multiple_sections", + _LANGUAGE + "response_language", + _FORMAT + "number_highlighted_sections", + }, + # TODO(tianjianlu): Re-enable rephrasing with preprocessing the message. + # _FORMAT + "rephrase": instructions.RephraseChecker, + _FORMAT + "json_format": set(INSTRUCTION_DICT.keys()).difference( + {_KEYWORD + "forbidden_words", _KEYWORD + "existence"} + ), + _FORMAT + "title": {_FORMAT + "title"}, + # TODO(tianjianlu): Re-enable with specific prompts. + # _MULTITURN + "constrained_start": instructions.ConstrainedStartChecker, + _COMBINATION + "two_responses": set(INSTRUCTION_DICT.keys()).difference( + { + _KEYWORD + "forbidden_words", + _KEYWORD + "existence", + _LANGUAGE + "response_language", + _FORMAT + "title", + _PUNCTUATION + "no_comma", + } + ), + _COMBINATION + "repeat_prompt": set(INSTRUCTION_DICT.keys()).difference( + {_KEYWORD + "existence", _FORMAT + "title", _PUNCTUATION + "no_comma"} + ), + _STARTEND + "end_checker": {_STARTEND + "end_checker"}, + _CHANGE_CASES + "capital_word_frequency": { + _CHANGE_CASES + "capital_word_frequency", + _CHANGE_CASES + "english_lowercase", + _CHANGE_CASES + "english_capital", + }, + _CHANGE_CASES + "english_capital": {_CHANGE_CASES + "english_capital"}, + _CHANGE_CASES + "english_lowercase": { + _CHANGE_CASES + "english_lowercase", + _CHANGE_CASES + "english_capital", + }, + _PUNCTUATION + "no_comma": {_PUNCTUATION + "no_comma"}, + _STARTEND + "quotation": {_STARTEND + "quotation", _FORMAT + "title"}, +} + + +def conflict_make(conflicts): + """Makes sure if A conflicts with B, B will conflict with A. + + Args: + conflicts: Dictionary of potential conflicts where key is instruction id + and value is set of instruction ids that it conflicts with. + + Returns: + Revised version of the dictionary. All instructions conflict with + themselves. If A conflicts with B, B will conflict with A. + """ + for key in conflicts: + for k in conflicts[key]: + conflicts[k].add(key) + conflicts[key].add(key) + return conflicts \ No newline at end of file diff --git a/reasoners/reward/ifeval/instructions_util.py b/reasoners/reward/ifeval/instructions_util.py new file mode 100644 index 000000000..33e0a0a00 --- /dev/null +++ b/reasoners/reward/ifeval/instructions_util.py @@ -0,0 +1,1701 @@ +# Copyright 2023 The Google Research Authors. +# +# 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. + +"""Utility library of instructions.""" + +import functools +import os +import random +import re +from importlib.metadata import version + +import immutabledict +import nltk +from packaging.version import parse as parse_version + + +# Downloading 'punkt' with nltk<3.9 has a remote code vuln. +# see https://github.com/EleutherAI/lm-evaluation-harness/issues/2210 +# and https://github.com/nltk/nltk/issues/3266 +# for more information. +NLTK_MIN_VERSION = "3.9.1" +RANK = os.environ.get("LOCAL_RANK", "0") + + +def download_nltk_resources(): + """Download 'punkt' if not already installed""" + assert (nltk_version := parse_version(version("nltk"))) >= parse_version( + NLTK_MIN_VERSION + ), ( + f"`nltk` version {nltk_version} is not >= {NLTK_MIN_VERSION}. Please update `nltk` before proceeding--older versions are vulnerable to a remote code execution vulnerability." + ) + + try: + nltk.data.find("tokenizers/punkt_tab") + except LookupError: + if RANK == "0": + nltk.download("punkt_tab") + print("Downloaded punkt_tab on rank 0") + + +download_nltk_resources() + +WORD_LIST = [ + "western", + "sentence", + "signal", + "dump", + "spot", + "opposite", + "bottom", + "potato", + "administration", + "working", + "welcome", + "morning", + "good", + "agency", + "primary", + "wish", + "responsibility", + "press", + "problem", + "president", + "steal", + "brush", + "read", + "type", + "beat", + "trainer", + "growth", + "lock", + "bone", + "case", + "equal", + "comfortable", + "region", + "replacement", + "performance", + "mate", + "walk", + "medicine", + "film", + "thing", + "rock", + "tap", + "total", + "competition", + "ease", + "south", + "establishment", + "gather", + "parking", + "world", + "plenty", + "breath", + "claim", + "alcohol", + "trade", + "dear", + "highlight", + "street", + "matter", + "decision", + "mess", + "agreement", + "studio", + "coach", + "assist", + "brain", + "wing", + "style", + "private", + "top", + "brown", + "leg", + "buy", + "procedure", + "method", + "speed", + "high", + "company", + "valuable", + "pie", + "analyst", + "session", + "pattern", + "district", + "pleasure", + "dinner", + "swimming", + "joke", + "order", + "plate", + "department", + "motor", + "cell", + "spend", + "cabinet", + "difference", + "power", + "examination", + "engine", + "horse", + "dimension", + "pay", + "toe", + "curve", + "literature", + "bother", + "fire", + "possibility", + "debate", + "activity", + "passage", + "hello", + "cycle", + "background", + "quiet", + "author", + "effect", + "actor", + "page", + "bicycle", + "error", + "throat", + "attack", + "character", + "phone", + "tea", + "increase", + "outcome", + "file", + "specific", + "inspector", + "internal", + "potential", + "staff", + "building", + "employer", + "shoe", + "hand", + "direction", + "garden", + "purchase", + "interview", + "study", + "recognition", + "member", + "spiritual", + "oven", + "sandwich", + "weird", + "passenger", + "particular", + "response", + "reaction", + "size", + "variation", + "a", + "cancel", + "candy", + "exit", + "guest", + "condition", + "fly", + "price", + "weakness", + "convert", + "hotel", + "great", + "mouth", + "mind", + "song", + "sugar", + "suspect", + "telephone", + "ear", + "roof", + "paint", + "refrigerator", + "organization", + "jury", + "reward", + "engineering", + "day", + "possession", + "crew", + "bar", + "road", + "description", + "celebration", + "score", + "mark", + "letter", + "shower", + "suggestion", + "sir", + "luck", + "national", + "progress", + "hall", + "stroke", + "theory", + "offer", + "story", + "tax", + "definition", + "history", + "ride", + "medium", + "opening", + "glass", + "elevator", + "stomach", + "question", + "ability", + "leading", + "village", + "computer", + "city", + "grand", + "confidence", + "candle", + "priest", + "recommendation", + "point", + "necessary", + "body", + "desk", + "secret", + "horror", + "noise", + "culture", + "warning", + "water", + "round", + "diet", + "flower", + "bus", + "tough", + "permission", + "week", + "prompt", + "connection", + "abuse", + "height", + "save", + "corner", + "border", + "stress", + "drive", + "stop", + "rip", + "meal", + "listen", + "confusion", + "girlfriend", + "living", + "relation", + "significance", + "plan", + "creative", + "atmosphere", + "blame", + "invite", + "housing", + "paper", + "drink", + "roll", + "silver", + "drunk", + "age", + "damage", + "smoke", + "environment", + "pack", + "savings", + "influence", + "tourist", + "rain", + "post", + "sign", + "grandmother", + "run", + "profit", + "push", + "clerk", + "final", + "wine", + "swim", + "pause", + "stuff", + "singer", + "funeral", + "average", + "source", + "scene", + "tradition", + "personal", + "snow", + "nobody", + "distance", + "sort", + "sensitive", + "animal", + "major", + "negotiation", + "click", + "mood", + "period", + "arrival", + "expression", + "holiday", + "repeat", + "dust", + "closet", + "gold", + "bad", + "sail", + "combination", + "clothes", + "emphasis", + "duty", + "black", + "step", + "school", + "jump", + "document", + "professional", + "lip", + "chemical", + "front", + "wake", + "while", + "inside", + "watch", + "row", + "subject", + "penalty", + "balance", + "possible", + "adult", + "aside", + "sample", + "appeal", + "wedding", + "depth", + "king", + "award", + "wife", + "blow", + "site", + "camp", + "music", + "safe", + "gift", + "fault", + "guess", + "act", + "shame", + "drama", + "capital", + "exam", + "stupid", + "record", + "sound", + "swing", + "novel", + "minimum", + "ratio", + "machine", + "shape", + "lead", + "operation", + "salary", + "cloud", + "affair", + "hit", + "chapter", + "stage", + "quantity", + "access", + "army", + "chain", + "traffic", + "kick", + "analysis", + "airport", + "time", + "vacation", + "philosophy", + "ball", + "chest", + "thanks", + "place", + "mountain", + "advertising", + "red", + "past", + "rent", + "return", + "tour", + "house", + "construction", + "net", + "native", + "war", + "figure", + "fee", + "spray", + "user", + "dirt", + "shot", + "task", + "stick", + "friend", + "software", + "promotion", + "interaction", + "surround", + "block", + "purpose", + "practice", + "conflict", + "routine", + "requirement", + "bonus", + "hole", + "state", + "junior", + "sweet", + "catch", + "tear", + "fold", + "wall", + "editor", + "life", + "position", + "pound", + "respect", + "bathroom", + "coat", + "script", + "job", + "teach", + "birth", + "view", + "resolve", + "theme", + "employee", + "doubt", + "market", + "education", + "serve", + "recover", + "tone", + "harm", + "miss", + "union", + "understanding", + "cow", + "river", + "association", + "concept", + "training", + "recipe", + "relationship", + "reserve", + "depression", + "proof", + "hair", + "revenue", + "independent", + "lift", + "assignment", + "temporary", + "amount", + "loss", + "edge", + "track", + "check", + "rope", + "estimate", + "pollution", + "stable", + "message", + "delivery", + "perspective", + "mirror", + "assistant", + "representative", + "witness", + "nature", + "judge", + "fruit", + "tip", + "devil", + "town", + "emergency", + "upper", + "drop", + "stay", + "human", + "neck", + "speaker", + "network", + "sing", + "resist", + "league", + "trip", + "signature", + "lawyer", + "importance", + "gas", + "choice", + "engineer", + "success", + "part", + "external", + "worker", + "simple", + "quarter", + "student", + "heart", + "pass", + "spite", + "shift", + "rough", + "lady", + "grass", + "community", + "garage", + "youth", + "standard", + "skirt", + "promise", + "blind", + "television", + "disease", + "commission", + "positive", + "energy", + "calm", + "presence", + "tune", + "basis", + "preference", + "head", + "common", + "cut", + "somewhere", + "presentation", + "current", + "thought", + "revolution", + "effort", + "master", + "implement", + "republic", + "floor", + "principle", + "stranger", + "shoulder", + "grade", + "button", + "tennis", + "police", + "collection", + "account", + "register", + "glove", + "divide", + "professor", + "chair", + "priority", + "combine", + "peace", + "extension", + "maybe", + "evening", + "frame", + "sister", + "wave", + "code", + "application", + "mouse", + "match", + "counter", + "bottle", + "half", + "cheek", + "resolution", + "back", + "knowledge", + "make", + "discussion", + "screw", + "length", + "accident", + "battle", + "dress", + "knee", + "log", + "package", + "it", + "turn", + "hearing", + "newspaper", + "layer", + "wealth", + "profile", + "imagination", + "answer", + "weekend", + "teacher", + "appearance", + "meet", + "bike", + "rise", + "belt", + "crash", + "bowl", + "equivalent", + "support", + "image", + "poem", + "risk", + "excitement", + "remote", + "secretary", + "public", + "produce", + "plane", + "display", + "money", + "sand", + "situation", + "punch", + "customer", + "title", + "shake", + "mortgage", + "option", + "number", + "pop", + "window", + "extent", + "nothing", + "experience", + "opinion", + "departure", + "dance", + "indication", + "boy", + "material", + "band", + "leader", + "sun", + "beautiful", + "muscle", + "farmer", + "variety", + "fat", + "handle", + "director", + "opportunity", + "calendar", + "outside", + "pace", + "bath", + "fish", + "consequence", + "put", + "owner", + "go", + "doctor", + "information", + "share", + "hurt", + "protection", + "career", + "finance", + "force", + "golf", + "garbage", + "aspect", + "kid", + "food", + "boot", + "milk", + "respond", + "objective", + "reality", + "raw", + "ring", + "mall", + "one", + "impact", + "area", + "news", + "international", + "series", + "impress", + "mother", + "shelter", + "strike", + "loan", + "month", + "seat", + "anything", + "entertainment", + "familiar", + "clue", + "year", + "glad", + "supermarket", + "natural", + "god", + "cost", + "conversation", + "tie", + "ruin", + "comfort", + "earth", + "storm", + "percentage", + "assistance", + "budget", + "strength", + "beginning", + "sleep", + "other", + "young", + "unit", + "fill", + "store", + "desire", + "hide", + "value", + "cup", + "maintenance", + "nurse", + "function", + "tower", + "role", + "class", + "camera", + "database", + "panic", + "nation", + "basket", + "ice", + "art", + "spirit", + "chart", + "exchange", + "feedback", + "statement", + "reputation", + "search", + "hunt", + "exercise", + "nasty", + "notice", + "male", + "yard", + "annual", + "collar", + "date", + "platform", + "plant", + "fortune", + "passion", + "friendship", + "spread", + "cancer", + "ticket", + "attitude", + "island", + "active", + "object", + "service", + "buyer", + "bite", + "card", + "face", + "steak", + "proposal", + "patient", + "heat", + "rule", + "resident", + "broad", + "politics", + "west", + "knife", + "expert", + "girl", + "design", + "salt", + "baseball", + "grab", + "inspection", + "cousin", + "couple", + "magazine", + "cook", + "dependent", + "security", + "chicken", + "version", + "currency", + "ladder", + "scheme", + "kitchen", + "employment", + "local", + "attention", + "manager", + "fact", + "cover", + "sad", + "guard", + "relative", + "county", + "rate", + "lunch", + "program", + "initiative", + "gear", + "bridge", + "breast", + "talk", + "dish", + "guarantee", + "beer", + "vehicle", + "reception", + "woman", + "substance", + "copy", + "lecture", + "advantage", + "park", + "cold", + "death", + "mix", + "hold", + "scale", + "tomorrow", + "blood", + "request", + "green", + "cookie", + "church", + "strip", + "forever", + "beyond", + "debt", + "tackle", + "wash", + "following", + "feel", + "maximum", + "sector", + "sea", + "property", + "economics", + "menu", + "bench", + "try", + "language", + "start", + "call", + "solid", + "address", + "income", + "foot", + "senior", + "honey", + "few", + "mixture", + "cash", + "grocery", + "link", + "map", + "form", + "factor", + "pot", + "model", + "writer", + "farm", + "winter", + "skill", + "anywhere", + "birthday", + "policy", + "release", + "husband", + "lab", + "hurry", + "mail", + "equipment", + "sink", + "pair", + "driver", + "consideration", + "leather", + "skin", + "blue", + "boat", + "sale", + "brick", + "two", + "feed", + "square", + "dot", + "rush", + "dream", + "location", + "afternoon", + "manufacturer", + "control", + "occasion", + "trouble", + "introduction", + "advice", + "bet", + "eat", + "kill", + "category", + "manner", + "office", + "estate", + "pride", + "awareness", + "slip", + "crack", + "client", + "nail", + "shoot", + "membership", + "soft", + "anybody", + "web", + "official", + "individual", + "pizza", + "interest", + "bag", + "spell", + "profession", + "queen", + "deal", + "resource", + "ship", + "guy", + "chocolate", + "joint", + "formal", + "upstairs", + "car", + "resort", + "abroad", + "dealer", + "associate", + "finger", + "surgery", + "comment", + "team", + "detail", + "crazy", + "path", + "tale", + "initial", + "arm", + "radio", + "demand", + "single", + "draw", + "yellow", + "contest", + "piece", + "quote", + "pull", + "commercial", + "shirt", + "contribution", + "cream", + "channel", + "suit", + "discipline", + "instruction", + "concert", + "speech", + "low", + "effective", + "hang", + "scratch", + "industry", + "breakfast", + "lay", + "join", + "metal", + "bedroom", + "minute", + "product", + "rest", + "temperature", + "many", + "give", + "argument", + "print", + "purple", + "laugh", + "health", + "credit", + "investment", + "sell", + "setting", + "lesson", + "egg", + "middle", + "marriage", + "level", + "evidence", + "phrase", + "love", + "self", + "benefit", + "guidance", + "affect", + "you", + "dad", + "anxiety", + "special", + "boyfriend", + "test", + "blank", + "payment", + "soup", + "obligation", + "reply", + "smile", + "deep", + "complaint", + "addition", + "review", + "box", + "towel", + "minor", + "fun", + "soil", + "issue", + "cigarette", + "internet", + "gain", + "tell", + "entry", + "spare", + "incident", + "family", + "refuse", + "branch", + "can", + "pen", + "grandfather", + "constant", + "tank", + "uncle", + "climate", + "ground", + "volume", + "communication", + "kind", + "poet", + "child", + "screen", + "mine", + "quit", + "gene", + "lack", + "charity", + "memory", + "tooth", + "fear", + "mention", + "marketing", + "reveal", + "reason", + "court", + "season", + "freedom", + "land", + "sport", + "audience", + "classroom", + "law", + "hook", + "win", + "carry", + "eye", + "smell", + "distribution", + "research", + "country", + "dare", + "hope", + "whereas", + "stretch", + "library", + "if", + "delay", + "college", + "plastic", + "book", + "present", + "use", + "worry", + "champion", + "goal", + "economy", + "march", + "election", + "reflection", + "midnight", + "slide", + "inflation", + "action", + "challenge", + "guitar", + "coast", + "apple", + "campaign", + "field", + "jacket", + "sense", + "way", + "visual", + "remove", + "weather", + "trash", + "cable", + "regret", + "buddy", + "beach", + "historian", + "courage", + "sympathy", + "truck", + "tension", + "permit", + "nose", + "bed", + "son", + "person", + "base", + "meat", + "usual", + "air", + "meeting", + "worth", + "game", + "independence", + "physical", + "brief", + "play", + "raise", + "board", + "she", + "key", + "writing", + "pick", + "command", + "party", + "yesterday", + "spring", + "candidate", + "physics", + "university", + "concern", + "development", + "change", + "string", + "target", + "instance", + "room", + "bitter", + "bird", + "football", + "normal", + "split", + "impression", + "wood", + "long", + "meaning", + "stock", + "cap", + "leadership", + "media", + "ambition", + "fishing", + "essay", + "salad", + "repair", + "today", + "designer", + "night", + "bank", + "drawing", + "inevitable", + "phase", + "vast", + "chip", + "anger", + "switch", + "cry", + "twist", + "personality", + "attempt", + "storage", + "being", + "preparation", + "bat", + "selection", + "white", + "technology", + "contract", + "side", + "section", + "station", + "till", + "structure", + "tongue", + "taste", + "truth", + "difficulty", + "group", + "limit", + "main", + "move", + "feeling", + "light", + "example", + "mission", + "might", + "wait", + "wheel", + "shop", + "host", + "classic", + "alternative", + "cause", + "agent", + "consist", + "table", + "airline", + "text", + "pool", + "craft", + "range", + "fuel", + "tool", + "partner", + "load", + "entrance", + "deposit", + "hate", + "article", + "video", + "summer", + "feature", + "extreme", + "mobile", + "hospital", + "flight", + "fall", + "pension", + "piano", + "fail", + "result", + "rub", + "gap", + "system", + "report", + "suck", + "ordinary", + "wind", + "nerve", + "ask", + "shine", + "note", + "line", + "mom", + "perception", + "brother", + "reference", + "bend", + "charge", + "treat", + "trick", + "term", + "homework", + "bake", + "bid", + "status", + "project", + "strategy", + "orange", + "let", + "enthusiasm", + "parent", + "concentrate", + "device", + "travel", + "poetry", + "business", + "society", + "kiss", + "end", + "vegetable", + "employ", + "schedule", + "hour", + "brave", + "focus", + "process", + "movie", + "illegal", + "general", + "coffee", + "ad", + "highway", + "chemistry", + "psychology", + "hire", + "bell", + "conference", + "relief", + "show", + "neat", + "funny", + "weight", + "quality", + "club", + "daughter", + "zone", + "touch", + "tonight", + "shock", + "burn", + "excuse", + "name", + "survey", + "landscape", + "advance", + "satisfaction", + "bread", + "disaster", + "item", + "hat", + "prior", + "shopping", + "visit", + "east", + "photo", + "home", + "idea", + "father", + "comparison", + "cat", + "pipe", + "winner", + "count", + "lake", + "fight", + "prize", + "foundation", + "dog", + "keep", + "ideal", + "fan", + "struggle", + "peak", + "safety", + "solution", + "hell", + "conclusion", + "population", + "strain", + "alarm", + "measurement", + "second", + "train", + "race", + "due", + "insurance", + "boss", + "tree", + "monitor", + "sick", + "course", + "drag", + "appointment", + "slice", + "still", + "care", + "patience", + "rich", + "escape", + "emotion", + "royal", + "female", + "childhood", + "government", + "picture", + "will", + "sock", + "big", + "gate", + "oil", + "cross", + "pin", + "improvement", + "championship", + "silly", + "help", + "sky", + "pitch", + "man", + "diamond", + "most", + "transition", + "work", + "science", + "committee", + "moment", + "fix", + "teaching", + "dig", + "specialist", + "complex", + "guide", + "people", + "dead", + "voice", + "original", + "break", + "topic", + "data", + "degree", + "reading", + "recording", + "bunch", + "reach", + "judgment", + "lie", + "regular", + "set", + "painting", + "mode", + "list", + "player", + "bear", + "north", + "wonder", + "carpet", + "heavy", + "officer", + "negative", + "clock", + "unique", + "baby", + "pain", + "assumption", + "disk", + "iron", + "bill", + "drawer", + "look", + "double", + "mistake", + "finish", + "future", + "brilliant", + "contact", + "math", + "rice", + "leave", + "restaurant", + "discount", + "sex", + "virus", + "bit", + "trust", + "event", + "wear", + "juice", + "failure", + "bug", + "context", + "mud", + "whole", + "wrap", + "intention", + "draft", + "pressure", + "cake", + "dark", + "explanation", + "space", + "angle", + "word", + "efficiency", + "management", + "habit", + "star", + "chance", + "finding", + "transportation", + "stand", + "criticism", + "flow", + "door", + "injury", + "insect", + "surprise", + "apartment", +] # pylint: disable=line-too-long + +# ISO 639-1 codes to language names. +LANGUAGE_CODES = immutabledict.immutabledict( + { + "en": "English", + "es": "Spanish", + "pt": "Portuguese", + "ar": "Arabic", + "hi": "Hindi", + "fr": "French", + "ru": "Russian", + "de": "German", + "ja": "Japanese", + "it": "Italian", + "bn": "Bengali", + "uk": "Ukrainian", + "th": "Thai", + "ur": "Urdu", + "ta": "Tamil", + "te": "Telugu", + "bg": "Bulgarian", + "ko": "Korean", + "pl": "Polish", + "he": "Hebrew", + "fa": "Persian", + "vi": "Vietnamese", + "ne": "Nepali", + "sw": "Swahili", + "kn": "Kannada", + "mr": "Marathi", + "gu": "Gujarati", + "pa": "Punjabi", + "ml": "Malayalam", + "fi": "Finnish", + } +) + +_ALPHABETS = "([A-Za-z])" +_PREFIXES = "(Mr|St|Mrs|Ms|Dr)[.]" +_SUFFIXES = "(Inc|Ltd|Jr|Sr|Co)" +_STARTERS = r"(Mr|Mrs|Ms|Dr|Prof|Capt|Cpt|Lt|He\s|She\s|It\s|They\s|Their\s|Our\s|We\s|But\s|However\s|That\s|This\s|Wherever)" +_ACRONYMS = "([A-Z][.][A-Z][.](?:[A-Z][.])?)" +_WEBSITES = "[.](com|net|org|io|gov|edu|me)" +_DIGITS = "([0-9])" +_MULTIPLE_DOTS = r"\.{2,}" + + +def split_into_sentences(text): + """Split the text into sentences. + + Args: + text: A string that consists of more than or equal to one sentences. + + Returns: + A list of strings where each string is a sentence. + """ + text = " " + text + " " + text = text.replace("\n", " ") + text = re.sub(_PREFIXES, "\\1", text) + text = re.sub(_WEBSITES, "\\1", text) + text = re.sub(_DIGITS + "[.]" + _DIGITS, "\\1\\2", text) + text = re.sub( + _MULTIPLE_DOTS, + lambda match: "" * len(match.group(0)) + "", + text, + ) + if "Ph.D" in text: + text = text.replace("Ph.D.", "PhD") + text = re.sub(r"\s" + _ALPHABETS + "[.] ", " \\1 ", text) + text = re.sub(_ACRONYMS + " " + _STARTERS, "\\1 \\2", text) + text = re.sub( + _ALPHABETS + "[.]" + _ALPHABETS + "[.]" + _ALPHABETS + "[.]", + "\\1\\2\\3", + text, + ) + text = re.sub(_ALPHABETS + "[.]" + _ALPHABETS + "[.]", "\\1\\2", text) + text = re.sub(" " + _SUFFIXES + "[.] " + _STARTERS, " \\1 \\2", text) + text = re.sub(" " + _SUFFIXES + "[.]", " \\1", text) + text = re.sub(" " + _ALPHABETS + "[.]", " \\1", text) + if "”" in text: + text = text.replace(".”", "”.") + if '"' in text: + text = text.replace('."', '".') + if "!" in text: + text = text.replace('!"', '"!') + if "?" in text: + text = text.replace('?"', '"?') + text = text.replace(".", ".") + text = text.replace("?", "?") + text = text.replace("!", "!") + text = text.replace("", ".") + sentences = text.split("") + sentences = [s.strip() for s in sentences] + if sentences and not sentences[-1]: + sentences = sentences[:-1] + return sentences + + +def count_words(text): + """Counts the number of words.""" + tokenizer = nltk.tokenize.RegexpTokenizer(r"\w+") + tokens = tokenizer.tokenize(text) + num_words = len(tokens) + return num_words + + +@functools.lru_cache(maxsize=None) +def _get_sentence_tokenizer(): + return nltk.data.load("nltk:tokenizers/punkt/english.pickle") + + +def count_sentences(text): + """Count the number of sentences.""" + tokenizer = _get_sentence_tokenizer() + tokenized_sentences = tokenizer.tokenize(text) + return len(tokenized_sentences) + + +def generate_keywords(num_keywords): + """Randomly generates a few keywords.""" + return random.sample(WORD_LIST, k=num_keywords) diff --git a/reasoners/reward/livebench/__init__.py b/reasoners/reward/livebench/__init__.py new file mode 100644 index 000000000..908100cf0 --- /dev/null +++ b/reasoners/reward/livebench/__init__.py @@ -0,0 +1,68 @@ +import re +import numpy as np +import ast +from ast import literal_eval +import json + +# data analysis split +from verl.utils.reward_score.livebench.data_analysis.tablereformat.utils import table_process_results +from verl.utils.reward_score.livebench.data_analysis.cta.utils import cta_process_results +from verl.utils.reward_score.livebench.data_analysis.tablejoin.utils import joinmap_process_results + +# reasoning split +from verl.utils.reward_score.livebench.reasoning.web_of_lies_v2.utils import web_of_lies_process_results +from verl.utils.reward_score.livebench.reasoning.web_of_lies_v3.utils import web_of_lies_v3_process_results +from verl.utils.reward_score.livebench.reasoning.house_traversal.utils import house_traversal_process_results +from verl.utils.reward_score.livebench.reasoning.zebra_puzzle.utils import zebra_puzzle_process_results_old # get_zebra_puzzle_evaluator +from verl.utils.reward_score.livebench.reasoning.spatial.utils import spatial_process_results + +# language split +from verl.utils.reward_score.livebench.writing.plot_unscrambling.utils import plot_unscrambling_process_results +from verl.utils.reward_score.livebench.writing.typos.utils import typos_process_results +from verl.utils.reward_score.livebench.writing.connections.utils import connections_process_results_old # get_connections_puzzle_evaluator + +# excluded splits +# from livebench.process_results.math.math_competitions.utils import mathcontest_process_results,aime_process_results +# from livebench.process_results.math.olympiad.utils import proof_rearrangement_process_results +# from livebench.process_results.math.AMPS_Hard.utils import amps_hard_process_results +# from livebench.process_results.coding.utils import LCB_generation_process_results, code_generation_process_results +# from livebench.process_results.instruction_following.utils import instruction_following_process_results + + + +def compute_score(solution_str, ground_truth, extra_info): + """ + Compute the score for the solution string. + """ + score = 0 + + # data analysis split + if extra_info["task"] == "cta": + score = cta_process_results(ground_truth, solution_str) + elif extra_info["task"] == "tablejoin": + score = joinmap_process_results(ground_truth, solution_str) + elif extra_info["task"] == "tableformat": + score = table_process_results(ground_truth, solution_str) + # reasoning split + elif extra_info["task"] == "web_of_lies_v2": + score = web_of_lies_process_results(ground_truth, solution_str) + elif extra_info["task"] == "web_of_lies_v3": + score = web_of_lies_v3_process_results(ground_truth, solution_str) + elif extra_info["task"] == "house_traversal": + score = house_traversal_process_results(ground_truth, solution_str) + elif extra_info["task"] == "zebra_puzzle": + score = zebra_puzzle_process_results_old(ground_truth, solution_str) + elif extra_info["task"] == "spatial": + score = spatial_process_results(ground_truth, solution_str) + # language split + elif extra_info["task"] == "plot_unscrambling": + score = plot_unscrambling_process_results(ground_truth, solution_str) + elif extra_info["task"] == "typos": + score = typos_process_results(ground_truth, solution_str) + elif extra_info["task"] == "connections": + score = connections_process_results_old(ground_truth, solution_str) + + return { + "score": score, + "acc": score, + } diff --git a/reasoners/reward/livebench/data_analysis/cta/utils.py b/reasoners/reward/livebench/data_analysis/cta/utils.py new file mode 100644 index 000000000..c96f20bd8 --- /dev/null +++ b/reasoners/reward/livebench/data_analysis/cta/utils.py @@ -0,0 +1,32 @@ +import re +from verl.utils.reward_score.livebench.util import last_boxed_only_string, remove_boxed + + +def clean_text(text): + text = text.lower().strip() + text = re.sub(r'[^\w]', '', text) + return text + + +def cta_process_results(ground_truth: str, llm_answer: str, debug=False) -> int: + + parsed_answer = llm_answer + + if '\\boxed{' in parsed_answer: + parsed_answer = last_boxed_only_string(parsed_answer) + parsed_answer = remove_boxed(parsed_answer) + parsed_answer = parsed_answer.replace('\\text{', '').replace('}', '').replace('\\', '') + + if clean_text(ground_truth) == clean_text(parsed_answer): + return 1 + elif clean_text(ground_truth) == clean_text(parsed_answer)[-len(clean_text(ground_truth)):]: + return 1 + else: + if debug: + print('INCORRECT') + print('GROUND TRUTH', ground_truth) + print('SOLUTION', parsed_answer) + if parsed_answer != llm_answer: + print('END OF OUTPUT', llm_answer[-100:]) + + return 0 diff --git a/reasoners/reward/livebench/data_analysis/tablejoin/utils.py b/reasoners/reward/livebench/data_analysis/tablejoin/utils.py new file mode 100644 index 000000000..2b564086b --- /dev/null +++ b/reasoners/reward/livebench/data_analysis/tablejoin/utils.py @@ -0,0 +1,82 @@ +import re +import numpy as np +from ast import literal_eval +import json +import ast +from verl.utils.reward_score.livebench.util import last_boxed_only_string, remove_boxed + +def clean_llm_output(s): + pattern_solution = r'(.*?)' + matches = re.findall(pattern_solution, s, re.DOTALL) + if len(matches) > 0: + return clean_llm_output(matches[-1].strip()) + try: + match_d = literal_eval(s) + except: + matches = re.findall('%s(.*?)%s' % ("```python", "```"), s.replace("\n",""),re.MULTILINE) + if len(matches) == 0: + matches = re.findall('%s(.*?)%s' % ("```json", "```"), s.replace("\n",""),re.MULTILINE) + if len(matches) == 0: + matches = re.findall('%s(.*?)%s' % ("```", "```"), s.replace("\n",""),re.MULTILINE) + if len(matches) == 0: + if '\\boxed' in s: + boxed = last_boxed_only_string(s.replace('\n', '')) + if boxed: + no_boxed = remove_boxed(boxed) + matches = [re.sub(r"\\text{[\'|\"](.*?)[\'|\"]}", r"'\1'", no_boxed).replace('\\', '')] + if len(matches) == 0: + matches = [s] + if len(matches) >= 1: + matches = matches[-1] + matches = matches.replace('null', 'None') + try: + match_d = literal_eval(matches) + except: + return {} + if not isinstance(match_d, dict): + return {} + else: + keys = list(match_d.keys()) + for k in keys: + if match_d[k] is None: + del match_d[k] + return match_d + +def joinmap_process_results(ground_truth, llm, debug=True): + if type(ground_truth) != dict: + ground_truth = ast.literal_eval(ground_truth) + llm_clean = clean_llm_output(llm) + if len(llm_clean) == 0: + if debug: + print('could not parse output') + print('GROUND TRUTH', ground_truth) + print('END OF OUTPUT', llm[-min(500, len(llm)):]) + print('=' * 100) + return 0.0 + tp = 0 + fp = 0 + fn = 0 + for k, v in llm_clean.items(): + gt = ground_truth.get(k, None) + if not gt: + fp += 1 + elif gt == v: + tp += 1 + else: + fp += 1 + fn += 1 + if isinstance(ground_truth, str): + ground_truth = json.loads(ground_truth) + for k, v in ground_truth.items(): + llm_resp = llm_clean.get(k, None) + if not llm_resp: + fn += 1 + result = np.round(((2.0 * tp) / ((2.0 * tp) + fp + fn)), 2) + if debug and result < 1: + print('INCORRECT') + print('GROUND TRUTH', ground_truth) + print('SOLUTION', llm_clean) + print('END OF OUTPUT', llm[-min(500, len(llm)):]) + print('RESULT', result) + print("=" * 100) + return result diff --git a/reasoners/reward/livebench/data_analysis/tablereformat/utils.py b/reasoners/reward/livebench/data_analysis/tablereformat/utils.py new file mode 100644 index 000000000..958136420 --- /dev/null +++ b/reasoners/reward/livebench/data_analysis/tablereformat/utils.py @@ -0,0 +1,240 @@ +import json +import traceback +import pandas as pd +from io import StringIO +import re +import math +from pandas.api.types import is_string_dtype +from pandas.api.types import is_numeric_dtype + + +def read_df_func(df_type, df_str): + if df_type == "json": + try: + return pd.read_json(StringIO(df_str), orient="index", encoding="utf-8") + except: + pass + try: + return pd.read_json(StringIO(df_str), orient="records", lines=False, encoding="utf-8") + except: + pass + try: + return pd.read_json(StringIO(df_str), orient="records", lines=True, encoding="utf-8") + except: + pass + try: + return pd.read_json(StringIO(df_str), orient="table", encoding="utf-8") + except: + pass + return pd.read_json(StringIO(df_str), orient="values", encoding="utf-8") + elif df_type == "jsonl": + return pd.read_json(StringIO(df_str), orient="records", lines=True, encoding="utf-8") + elif df_type == "html": + return pd.concat(pd.read_html(StringIO(df_str), encoding="utf-8"), axis=0) + elif df_type == "csv": + return pd.read_csv(StringIO(df_str), encoding="utf-8") + elif df_type == "markdown": + return pd.read_table(StringIO(df_str), sep="|", header=0, index_col=1, skipinitialspace=True) + elif df_type == "tsv": + return pd.read_csv(StringIO(df_str), sep='\t', encoding="utf-8") + +def read_df_func_v2(df_type, df_str): + if df_type == "json": + try: + # Try table orientation first as it preserves dtypes + return pd.read_json(StringIO(df_str), orient="table", encoding="utf-8") + except: + try: + return pd.read_json(StringIO(df_str), orient="index", lines=False, encoding="utf-8") + except: + try: + return pd.read_json(StringIO(df_str), orient="records", lines=False, encoding="utf-8") + except: + print("Could not read JSON") + return None + elif df_type == "jsonl": + return pd.read_json(StringIO(df_str), orient="records", lines=True, encoding="utf-8") + elif df_type == "html": + return pd.concat(pd.read_html(StringIO(df_str), encoding="utf-8"), axis=0) + elif df_type == "csv": + return pd.read_csv(StringIO(df_str), encoding="utf-8") + elif df_type == "markdown": + # Process markdown table by removing the separator line + lines = df_str.strip().split("\n") + header = lines[0] + # Skip the separator line (typically line with |:-----|:-----|) + data_lines = lines[2:] if len(lines) > 2 else [] + processed_md = header + "\n" + "\n".join(data_lines) + df = pd.read_table(StringIO(processed_md), sep="|", header=0, skipinitialspace=True).iloc[:, 1:-1] + + # Strip whitespace from all string columns + for col in df.columns: + if df[col].dtype == 'object': + df[col] = df[col].astype(str).str.strip() + + return df + elif df_type == "tsv": + return pd.read_csv(StringIO(df_str), sep='\t', encoding="utf-8") + else: + raise ValueError(f"Unsupported file type: {df_type}") + +def clean_llm_output(s): + pattern_solution = r'(.*?)' + matches = re.findall(pattern_solution, s, re.DOTALL) + if len(matches) > 0: + return clean_llm_output(matches[-1].strip()) + pattern_json = r'```json\n(.*?)```' + matches = re.findall(pattern_json, s, re.DOTALL) + if len(matches) > 0: + return matches[-1].strip() + pattern_html = r'```html\n(.*?)```' + matches = re.findall(pattern_html, s, re.DOTALL) + if len(matches) > 0: + return matches[-1].strip() + pattern_a = r'^```.*\n' + s = re.sub(pattern_a, "", s) + # re.findall(pattern, text, re.MULTILINE) + s = s.replace("&", "&") + return s.replace("```", "").strip() + +def read_sep_table_from_text(text, header, sep=','): + text = text.strip() + # look for the first line that contains the header + header_line = 0 + while header_line < len(text.split('\n')) and text.split('\n')[header_line].strip() != header.strip(): + header_line += 1 + if header_line == len(text.split('\n')) or text.split('\n')[header_line].strip() != header.strip(): + return None + # read the table from the header index + table = text.split('\n')[header_line:] + # read the table as a csv + parsed_table = None + while parsed_table is None: + try: + parsed_table = pd.read_csv(StringIO('\n'.join(table)), sep=sep) + except: + # in case there's extra text after the table + table = table[:-1] + return parsed_table + +def read_jsonl_table_from_text(text, header): + text = text.strip().split('\n') + table = [] + for line in text: + if len(line) < 2 or line[0] != '{' or line[-1] != '}': + continue + if not all(key in line for key in header): + continue + try: + table.append(json.loads(line)) + except: + continue + if len(table) == 0: + return None + return pd.DataFrame(table) + + +def remove_initial_phrase(text): + # remove phrases like "Here is the updated table:" , "Here is the table in a new format:" + pattern = r'^\s*(Here|Input)\b.*?\b(format|table)\s*[:)]\s*' + modified_text = re.sub(pattern, '', text, flags=re.IGNORECASE) + return modified_text.strip() + +def table_process_results(input_command: str, ground_truth: str, llm_answer: str, version: str = "v1", debug=True) -> int: + if version == "v1": + input_format = input_command.split("Please convert the Input Table from ")[1].split(" format")[0].lower() + output_format = input_command.split("Please convert the Input Table from ")[1].split("format to ")[1].split(" format")[0].lower() + else: + command_lines = input_command.split('\n') + input_format_lines = [line for line in command_lines if "Source Format" in line] + input_format = input_format_lines[-1].split("Source Format: ")[1].strip().lower() + output_format_lines = [line for line in command_lines if "Target Format" in line] + output_format = output_format_lines[-1].split("Target Format: ")[1].strip().lower() + + df_read_func = read_df_func if version == "v1" else read_df_func_v2 + + gt_df = df_read_func(output_format, ground_truth) + + llm_clean = clean_llm_output(llm_answer) + # if there's an initial phrase before the table, remove it and try to score again + llm_clean = remove_initial_phrase(llm_clean) + # first check the raw LLM output + llm_df = None + try: + llm_df = df_read_func(output_format, llm_clean) + except: + if output_format == 'csv' or output_format == 'tsv': + header = (',', '\t')[output_format == 'tsv'].join(gt_df.columns) + llm_df = read_sep_table_from_text(llm_clean, header, sep=',' if output_format == 'csv' else '\t') + elif output_format == 'jsonl': + llm_df = read_jsonl_table_from_text(llm_clean, gt_df.columns) + if llm_df is None: + print('Could not read the LLM output') + print('GROUND TRUTH\n', ground_truth) + print('END OF OUTPUT\n', llm_answer[-min(3000, len(llm_answer)):]) + print('=' * 100) + return 0 + score = check_table_reformat(output_format, llm_df, gt_df, debug) + + if score == 0: + # try to read table directly from the text, in case the LLM added some text before/after the table + if output_format == 'csv': + header = ','.join(gt_df.columns) + llm_df = read_sep_table_from_text(llm_clean, header, sep=',' if output_format == 'csv' else '\t') + + elif output_format == 'jsonl': + llm_df = read_jsonl_table_from_text(llm_clean, gt_df.columns) + + if llm_df is not None: + score = check_table_reformat(output_format, llm_df, gt_df, debug) + if debug and score == 0: + print('INCORRECT') + print('GROUND TRUTH\n', 'None' if gt_df is None else gt_df.head()) + print('LLM DF\n', 'None' if llm_df is None else llm_df.head()) + print('LLM ANSWER\n', llm_clean) + print('=' * 100) + return score + +def check_table_reformat(output_format, llm_df, gt_df, debug=False): + try: + gt_df.columns = [s.strip() for s in gt_df.columns] + if 'index' in gt_df.columns: + gt_df = gt_df.drop(columns=['index']) + llm_df.columns = [s.strip() for s in llm_df.columns] + if 'index' in llm_df.columns: + llm_df = llm_df.drop(columns=['index']) + assert len(llm_df) == len(gt_df), f"DataFrame Length does not match, {len(llm_df)} (LLM) vs {len(gt_df)} (Ground Truth)" + assert list(sorted(llm_df.columns)) == list(sorted(gt_df.columns)), f"Columns do not match:\n{sorted(llm_df.columns)} (LLM)\n{sorted(gt_df.columns)} (Ground Truth)" + # for test_col in llm_df.columns: + # assert sorted(llm_df[test_col].tolist()) == sorted(gt_df[test_col].tolist()), f"Column content {test_col} does not match" + for i in range(len(llm_df)): + for key in llm_df.columns: + llm_val = llm_df.iloc[i][key] + gt_val = gt_df.iloc[i][key] + if isinstance(llm_val, str): + llm_val = llm_val.strip() + if isinstance(gt_val, str): + gt_val = gt_val.strip() + + if (isinstance(llm_val, float) or is_numeric_dtype(llm_val)) and (isinstance(gt_val, float) or is_numeric_dtype(gt_val)): + try: + llm_val = float(llm_val) + gt_val = float(gt_val) + except: + assert str(llm_val).strip() == str(gt_val).strip(), f"Mismatched types of values {llm_val} (LLM) vs {gt_val} (Ground Truth) for key {key} in row {i}" + continue + if math.isnan(llm_val) and math.isnan(gt_val): + continue + assert abs(llm_val - gt_val) < 1e-6, f"Unequal float values {llm_val} (LLM) vs {gt_val} (Ground Truth) for key {key} in row {i}" + else: + assert llm_val == gt_val, f"Value {llm_val} (LLM) vs {gt_val} (Ground Truth) for key {key} in row {i}" + except AssertionError as e: + if debug: + print(e) + return 0 + except Exception as e: + if debug: + print(e) + traceback.print_exc() + return 0 + return 1 diff --git a/reasoners/reward/livebench/reasoning/house_traversal/utils.py b/reasoners/reward/livebench/reasoning/house_traversal/utils.py new file mode 100644 index 000000000..af21006cc --- /dev/null +++ b/reasoners/reward/livebench/reasoning/house_traversal/utils.py @@ -0,0 +1,41 @@ + +import re + + +def house_traversal_process_results(ground_truth: str, llm_answer: str, debug=False) -> int: + # pull out words in bold + bold_words = re.findall(r'(\*{2,})(.*?)\1', llm_answer.lower()) + if not len(bold_words): + if debug: + print('INCORRECT') + print('GROUND TRUTH', ground_truth) + print('SOLUTION', llm_answer) + return 0 + + last_bold = bold_words[-1][1] + ground_truth_names = ground_truth.lower().split(" ") + + # check if all the ground truth names are in the last few bolded words, in order + if len(bold_words) >= len(ground_truth_names): + if all([name in bold_words[-1 - i] for i,name in enumerate(ground_truth_names[::-1])]): + return 1 + + score = 1 + # check if all the ground truth names are in the last bolded part, in order + last_index = -1 + for name in ground_truth_names: + index = last_bold.find(name) + if index == -1: + score = 0 + break + elif index <= last_index: + score = 0 + break + else: + last_index = index + + if debug and score == 0: + print('INCORRECT', score) + print('GROUND TRUTH', ground_truth) + print('SOLUTION', llm_answer) + return score \ No newline at end of file diff --git a/reasoners/reward/livebench/reasoning/spatial/utils.py b/reasoners/reward/livebench/reasoning/spatial/utils.py new file mode 100644 index 000000000..f4e82358d --- /dev/null +++ b/reasoners/reward/livebench/reasoning/spatial/utils.py @@ -0,0 +1,61 @@ +import re +from verl.utils.reward_score.livebench.util import last_boxed_only_string, remove_boxed + +def spatial_process_results(ground_truth: str, llm_answer: str, debug=False) -> int: + + if llm_answer == ground_truth: + return 1 + + word_to_number = { + 'zero': '0', 'one': '1', 'two': '2', 'three': '3', 'four': '4', 'five': '5', + 'six': '6', 'seven': '7', 'eight': '8', 'nine': '9', 'ten': '10', + 'eleven': '11', 'twelve': '12', 'thirteen': '13', 'fourteen': '14', 'fifteen': '15', + 'sixteen': '16', 'seventeen': '17', 'eighteen': '18', 'nineteen': '19', 'twenty': '20' + } + + bold_words = re.findall(r'\*\*([^\*]+)\*\*', llm_answer) + score = 0 + + # allow the answer to be within the last 3 bolded words + words_to_check = [] + for i in range(3): + if bold_words and len(bold_words) > i: + words_to_check.append(bold_words[-i-1].strip().lower()) + + for i, word in enumerate(words_to_check): + if word == ground_truth.strip().lower(): + score = 1 + + # allow the answer to be the number spelled out + if word in word_to_number and word_to_number[word] == ground_truth.strip().lower(): + score = 1 + + # allow certain cases like "two tetrahedra" == "tetrahedra" and "equilateral triangle" == "triangle" + # while still disallowing cases like "circle square triangle" == "triangle" + for answer in ["tetrahedra", "tetrahedron", "triangle", "square"]: + if ground_truth.strip().lower() == answer and answer in word and len(word) < (2 * len(answer) + 5): + score = 1 + + allow_boxed = True + if score == 0 and allow_boxed: + llm_answer = llm_answer.replace("\\\\fbox{", "\\\\boxed{") + last_boxed = last_boxed_only_string(llm_answer) + if last_boxed: + parsed_answer = remove_boxed(last_boxed) + parsed_answer = parsed_answer.replace("\\textbf{", "") + parsed_answer = parsed_answer.replace("\\mathbf{", "") + parsed_answer = parsed_answer.replace("}", "") + if parsed_answer == ground_truth: + score = 1 + + if debug and score == 0: + print("INCORRECT") + print("GROUND TRUTH", ground_truth.strip().lower()) + if bold_words: + print("BOLD WORDS:", bold_words[-1].strip().lower()) + if last_boxed: + print("LAST BOXED", last_boxed) + print("PARSED ANSWER", parsed_answer) + print("END OF OUTPUT", llm_answer[-50:]) + + return score diff --git a/reasoners/reward/livebench/reasoning/web_of_lies_v2/utils.py b/reasoners/reward/livebench/reasoning/web_of_lies_v2/utils.py new file mode 100644 index 000000000..07be3af68 --- /dev/null +++ b/reasoners/reward/livebench/reasoning/web_of_lies_v2/utils.py @@ -0,0 +1,79 @@ +import re +import itertools +from verl.utils.reward_score.livebench.util import last_boxed_only_string, remove_boxed + +def web_of_lies_process_results(ground_truth: str, llm_answer: str, debug=False) -> int: + + score = 0 + parsed_answer = None + + # extract text from tags + solution_matches = re.findall(r'(.*?)', llm_answer) + + if len(solution_matches) == 0: + solution_matches = re.findall(r'(.*?)', llm_answer) + + if len(solution_matches) > 0: + parsed_answer = solution_matches[-1] + + + # pull out words in bold + bold_words = re.findall(r'\*\*(.*?)\*\*', llm_answer) + + if parsed_answer is None and bold_words: + bold_words = [word.lower().strip().replace(',', '').replace('.', '')[0:max(len(word), 3)] for match in bold_words for word in match.split()] + parsed_answer = [] + i = len(bold_words) - 1 + while i >= 0 and len(parsed_answer) < 3: + if bold_words[i] in ["yes", "no", "unknown"]: + parsed_answer = [bold_words[i]] + parsed_answer + i -= 1 + if len(parsed_answer) > 0: + parsed_answer = ", ".join(parsed_answer) + else: + parsed_answer = None + + allow_latex = True + if parsed_answer is None or parsed_answer.strip() == '' and allow_latex: + llm_answer = llm_answer.replace("\\\\boxed{\\\\textbf{", "\\\\boxed{") + llm_answer = llm_answer.replace("\\\\fbox{", "\\\\boxed{") + llm_answer = llm_answer.replace("\\textbf{", "\\boxed{") + + last_boxed = last_boxed_only_string(llm_answer) + if last_boxed: + parsed_answer = remove_boxed(last_boxed) + + allow_plain = True + if allow_plain and parsed_answer is None: + combs = itertools.product(['yes', 'no', 'unknown'], repeat=3) + # find all instances of these combinations in the answer, then treat the last one as the actual answer + # to compare to the ground truth + final_comb = None + final_comb_index = -1 + for comb in combs: + index = llm_answer.lower().find(', '.join(comb)) + if index != -1 and index > final_comb_index: + final_comb = comb + final_comb_index = index + if final_comb is not None: + parsed_answer = ', '.join(final_comb) + + # Check if parsed_answer is an exact match of ground_truth + if parsed_answer and parsed_answer == ground_truth.lower(): + score = 1 + + # Check if parsed_answer contains the ground_truth + if parsed_answer and parsed_answer.count("yes") + parsed_answer.count("no") + parsed_answer.count("unknown") == 3 and ground_truth.lower() in parsed_answer: + score = 1 + + if debug and score == 0: + print("INCORRECT") + print('GROUND TRUTH', ground_truth) + if parsed_answer: + print('PARSED ANSWER', parsed_answer) + else: + print('NO PARSED ANSWER') + print('END OF OUTPUT', llm_answer[-50:]) + + return score + diff --git a/reasoners/reward/livebench/reasoning/web_of_lies_v3/utils.py b/reasoners/reward/livebench/reasoning/web_of_lies_v3/utils.py new file mode 100644 index 000000000..ae9313e4e --- /dev/null +++ b/reasoners/reward/livebench/reasoning/web_of_lies_v3/utils.py @@ -0,0 +1,87 @@ +import re +import itertools +from verl.utils.reward_score.livebench.util import last_boxed_only_string, remove_boxed + +def web_of_lies_v3_process_results(ground_truth: str, llm_answer: str, debug=False) -> float: + score = 0 + parsed_answer = None + + # extract text from tags + solution_matches = re.findall(r'(.*?)', llm_answer) + + if len(solution_matches) > 0: + parsed_answer = solution_matches[-1] + + # if parsed_answer is None: + # bold_words = re.findall(r'\*\*(.*?)\*\*', llm_answer) + # if bold_words: + # bold_words = [word.lower().strip().replace(',', '').replace('.', '')[0:max(len(word), 3)] for match in bold_words for word in match.split()] + + # parsed_answer = [] + # i = len(bold_words) - 1 + # while i >= 0 and len(parsed_answer) < 3: + # if bold_words[i] in ["yes", "no", "unknown"]: + # parsed_answer = [bold_words[i]] + parsed_answer + # i -= 1 + + # if len(parsed_answer) > 0: + # parsed_answer = ", ".join(parsed_answer) + # else: + # parsed_answer = None + + allow_latex = True + if parsed_answer is None or parsed_answer.strip() == '' and allow_latex: + llm_answer = llm_answer.replace("\\\\boxed{\\\\textbf{", "\\\\boxed{") + llm_answer = llm_answer.replace("\\\\fbox{", "\\\\boxed{") + llm_answer = llm_answer.replace("\\textbf{", "\\boxed{") + + last_boxed = last_boxed_only_string(llm_answer) + + if last_boxed: + parsed_answer = remove_boxed(last_boxed) + + allow_plain = True + if allow_plain and parsed_answer is None: + combs = itertools.product(['yes', 'no', 'unknown'], repeat=3) + + final_comb = None + final_comb_index = -1 + for comb in combs: + index = llm_answer.lower().find(', '.join(comb)) + if index != -1 and index > final_comb_index: + final_comb = comb + final_comb_index = index + + if final_comb is not None: + parsed_answer = ', '.join(final_comb) + + if parsed_answer is None: + score = 0 + else: + parsed_answer_list = parsed_answer.split(',') + ground_truth_list = ground_truth.split(',') + + num_correct = 0 + total = len(ground_truth_list) + for i in range(total): + if i >= len(parsed_answer_list): + break + gt_word = ground_truth_list[i] + llm_word = parsed_answer_list[i] + if llm_word == gt_word: + num_correct += 1 + + score = ((num_correct == total) + num_correct / total) / 2 + + # score = 1 if parsed_answer.strip().replace(' ', '') == ground_truth.strip().replace(' ', '') else 0 + + if debug and score < 1: + print("INCORRECT") + print('GROUND TRUTH', ground_truth) + if parsed_answer: + print('PARSED ANSWER', parsed_answer) + else: + print('PARSED ANSWER', parsed_answer) + print('END OF OUTPUT', llm_answer[-50:]) + + return score diff --git a/reasoners/reward/livebench/reasoning/zebra_puzzle/utils.py b/reasoners/reward/livebench/reasoning/zebra_puzzle/utils.py new file mode 100644 index 000000000..92d29fd28 --- /dev/null +++ b/reasoners/reward/livebench/reasoning/zebra_puzzle/utils.py @@ -0,0 +1,128 @@ + +import re +from verl.utils.reward_score.livebench.util import last_boxed_only_string, remove_boxed + +def zebra_puzzle_process_results_old(ground_truth: str, llm_answer: str, debug=False) -> int: + # Mapping of numbers to words for 1 to 9 + number_to_word = { + '1': 'one', '2': 'two', '3': 'three', '4': 'four', + '5': 'five', '6': 'six', '7': 'seven', '8': 'eight', '9': 'nine' + } + + # Pull out words in bold + bold_words = re.findall(r'\*\*\*(\w+)\*\*\*', llm_answer) + + score = 0 + + # Remove any trailing punctuation from the last bold word if exists + if bold_words: + if (bold_words[-1].lower() == ground_truth.lower() or + (bold_words[-1] in number_to_word and number_to_word[bold_words[-1]].lower() == ground_truth.lower()) + or bold_words[-1].lower() + ' movies' == ground_truth.lower()): + score = 1 + else: + # Split the text into words and remove punctuation. + words = re.findall(r'\b\w+\b', llm_answer) + last_word = words[-1] if words else '' + # Check if the last bold word is a number and matches the word representation of the ground_truth + if (last_word.lower() == ground_truth.lower() or + (last_word in number_to_word and number_to_word[last_word].lower() == ground_truth.lower()) + or last_word.lower() + ' movies' == ground_truth.lower()): + score = 1 + + if debug and score == 0: + print('INCORRECT') + print('GROUND TRUTH', ground_truth.lower()) + if bold_words: + print('LAST BOLD WORD', bold_words[-1].lower()) + else: + print('LAST WORD', last_word.lower()) + print('END OF OUTPUT', llm_answer[-50:]) + return score + + +def zebra_puzzle_process_results(ground_truth: str, llm_answer: str, debug=False) -> float: + # Mapping of numbers to words for 1 to 9 + word_to_num = { + 'one': '1', + 'two': '2', + 'three': '3', + 'four': '4', + 'five': '5', + 'six': '6', + 'seven': '7', + 'eight': '8', + 'nine': '9' + } + + ground_truth = ground_truth.split(',') + + # extract text from tags + solution_matches = re.findall(r'(.*?)', llm_answer) + + if len(solution_matches) == 0: + solution_matches = re.findall(r'(.*?)', llm_answer) + + + allow_boxed = True + if len(solution_matches) == 0 and allow_boxed: + llm_answer = llm_answer.replace("\\\\fbox{", "\\\\boxed{") + last_boxed = last_boxed_only_string(llm_answer) + if last_boxed: + boxed_removed = remove_boxed(last_boxed) + boxed_removed = boxed_removed.replace("\\text{", "").replace("}", "").replace('\\', '') + solution_matches.append(boxed_removed) + + if len(solution_matches) == 0: + last_line = llm_answer.strip().split('\n')[-1] + if last_line.count(',') == len(ground_truth) - 1: + solution_matches.append(last_line) + + + if len(solution_matches) == 0: + if debug: + print('No solution text found for zebra puzzle') + print('GROUND TRUTH', ground_truth) + print('END OF OUTPUT', llm_answer[-100:]) + return 0 + + if len(solution_matches) > 1: + if debug: + print('WARNING: Multiple solution texts found for zebra puzzle, combining starting from last') + all_solution_text = [] + for match in solution_matches: + all_solution_text += match.split(',') + # get last len(ground_truth) elements + solution_text = all_solution_text[-len(ground_truth):] + else: + solution_text = solution_matches[-1].split(',') + + if debug and len(solution_text) < len(ground_truth): + print(f'WARNING: LLM does not have an answer for all zebra puzzle questions (expected {len(ground_truth)}, got {len(solution_text)})') + + num_correct = 0 + total = len(ground_truth) + for i in range(total): + gt_word = ground_truth[i].strip().lower().replace('-', ' ') + if i >= len(solution_text): + continue + llm_word = solution_text[i].strip().lower().replace('-', ' ').replace('position', '') + if llm_word in word_to_num: + # llm_word = word_to_num[llm_word] + if debug: + print('WARNING: LLM answer contains numbers in word form') + + if i < len(solution_text) and (gt_word == llm_word or gt_word in llm_word): + num_correct += 1 + if debug and num_correct < total: + print('INCORRECT', num_correct / total) + print('GROUND TRUTH', ground_truth) + print('SOLUTION', solution_text) + print('END OF OUTPUT', llm_answer[-50:]) + return ((num_correct == total) + num_correct / total) / 2 + + +def get_zebra_puzzle_evaluator(release_date: str): + if release_date < '2024-11-25': + return zebra_puzzle_process_results_old + return zebra_puzzle_process_results diff --git a/reasoners/reward/livebench/util.py b/reasoners/reward/livebench/util.py new file mode 100644 index 000000000..084a3e17b --- /dev/null +++ b/reasoners/reward/livebench/util.py @@ -0,0 +1,46 @@ +from typing import Optional + + +def last_boxed_only_string(string: str) -> Optional[str]: + 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].replace("$", "").replace("fbox","boxed") + + return retval + + +def remove_boxed(s: str) -> str: + 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] \ No newline at end of file diff --git a/reasoners/reward/livebench/writing/connections/utils.py b/reasoners/reward/livebench/writing/connections/utils.py new file mode 100644 index 000000000..7f4877ebd --- /dev/null +++ b/reasoners/reward/livebench/writing/connections/utils.py @@ -0,0 +1,108 @@ +import re +from verl.utils.reward_score.livebench.util import last_boxed_only_string, remove_boxed + + +def group_words(words): + groups = [set()] + words = [w.strip().lower() for w in words] + for word in words: + if len(groups[-1]) == 4: + groups.append(set()) + groups[-1].add(word) + return groups + + +def connections_process_results_old(ground_truth: str, llm_answer: str, debug=False) -> int: + + # pull out words in bold + bold_words = re.findall(r'\*\*(.*?)\*\*', llm_answer.replace('\n', '')) + + if not bold_words: + if debug: + print('No bold words found for connections') + print('END OF OUTPUT', llm_answer[-50:]) + return 0 + + bold_words = [words.split(',') for words in bold_words] + + ground_truth_groups = group_words(ground_truth.split(',')) + max_score = 0 + for output_groups in list(map(group_words, bold_words)): + + correct_groups = 0 + for ground_truth_group in ground_truth_groups: + for output_group in output_groups: + if all([word in output_group for word in ground_truth_group]): + correct_groups += 1 + break + + max_score = max(max_score, correct_groups / len(ground_truth_groups)) + if debug and max_score < 1: + print('INCORRECT', max_score) + print('GROUND TRUTH', ground_truth_groups) + print('SOLUTION', output_groups) + print('END OF OUTPUT', llm_answer[-50:]) + return max_score + + +def connections_process_results(ground_truth: str, llm_answer: str, debug=False) -> int: + + # extract text from tags + solution_matches = re.findall(r'(.*?)<\/solution>', llm_answer) + if len(solution_matches) == 0: + solution_matches = re.findall(r'(.*?)<\/solution>', llm_answer.replace('\n', '')) + if len(solution_matches) == 0: + solution_matches = re.findall(r'(.*?)<\/solution>', llm_answer) + if len(solution_matches) == 0: + solution_matches = re.findall(r'(.*?)<\/solution>', llm_answer.replace('\n', '')) + + ground_truth_words = ground_truth.split(',') + + if len(solution_matches) == 0 and '\\boxed' in llm_answer: + boxed = last_boxed_only_string(llm_answer) + no_box = remove_boxed(boxed) + solution_matches = [no_box.replace('\\text{', '').replace('}', '').replace('\\', '')] + + solution_matches = [match.replace('\n', '') for match in solution_matches] + + if len(solution_matches) == 0: + if debug: + print('No solution text found for connections') + print(llm_answer[-500:]) + return 0 + elif len(solution_matches) > 1: + if debug: + print('Multiple solution texts found for connections. Combining starting from last') + all_words = [] + num_words = len(ground_truth_words) + for match in solution_matches: + all_words.extend(match.split(',')) + solution_words = all_words[-num_words:] + else: + solution_words = solution_matches[-1].split(',') + + if len(solution_words) != len(ground_truth_words): + if debug: + print(f'Number of words in solution does not match number of words in ground truth ({len(solution_words)} vs expected {len(ground_truth_words)})') + + llm_groups = group_words(solution_words) + ground_truth_groups = group_words(ground_truth_words) + + correct_groups = 0 + for llm_group in llm_groups: + if llm_group in ground_truth_groups: + correct_groups += 1 + + if debug and correct_groups / len(ground_truth_groups) < 1: + print('INCORRECT', correct_groups / len(ground_truth_groups)) + print('GROUND TRUTH', sorted([sorted(list(group)) for group in ground_truth_groups])) + print('SOLUTION', sorted([sorted(list(group)) for group in llm_groups])) + print('END OF OUTPUT', llm_answer[-500:]) + + return correct_groups / len(ground_truth_groups) + + +def get_connections_puzzle_evaluator(release_date: str): + if release_date < '2024-11-25': + return connections_process_results_old + return connections_process_results diff --git a/reasoners/reward/livebench/writing/plot_unscrambling/utils.py b/reasoners/reward/livebench/writing/plot_unscrambling/utils.py new file mode 100644 index 000000000..084fc26f4 --- /dev/null +++ b/reasoners/reward/livebench/writing/plot_unscrambling/utils.py @@ -0,0 +1,66 @@ +import difflib +import re + + +def levenshtein_distance(A, B): + N, M = len(A), len(B) + # Create an array of size NxM + dp = [[0 for i in range(M + 1)] for j in range(N + 1)] + + # Base Case: When N = 0 + for j in range(M + 1): + dp[0][j] = j + # Base Case: When M = 0 + for i in range(N + 1): + dp[i][0] = i + # Transitions + for i in range(1, N + 1): + for j in range(1, M + 1): + if A[i - 1] == B[j - 1]: + dp[i][j] = dp[i-1][j-1] + else: + dp[i][j] = 1 + min( + dp[i-1][j], # Insertion + dp[i][j-1], # Deletion + dp[i-1][j-1] # Replacement + ) + + return dp[N][M] + + +def extract_plot_summary(text: str) -> str: + pattern = r'(.*)' + match = re.search(pattern, text, re.DOTALL) # re.DOTALL allows . to match newline characters + if not match: + pattern = r'(.*)' + match = re.search(pattern, text, re.DOTALL) + return match.group(1) if match else text + + +def plot_unscrambling_process_results(ground_truth: str, llm_answer: str, debug=False) -> float: + # Split the ground truth and answer into sentences based on full stops + llm_answer = extract_plot_summary(llm_answer) + + gt_sentences = [s.strip() for s in ground_truth.split('.')] + ans_sentences = [s.strip() for s in llm_answer.split('.') if s.strip() != '' and s.strip() != '**End of Plot Summary**'] + + # Remove empty sentences resulting from trailing or double full stops. + gt_sentences = [s for s in gt_sentences if s] + ans_sentences = [s for s in ans_sentences if s] + + ans_ordering = [] + for x in gt_sentences: + best_match = difflib.get_close_matches(x, ans_sentences, n=1, cutoff=0.0) + if best_match: + ans_ordering.append(ans_sentences.index(best_match[0])) + + n_sentences_gt = len(gt_sentences) + raw_distance = levenshtein_distance(list(range(len(gt_sentences))), ans_ordering) + score = 1 - (raw_distance / n_sentences_gt) + + if debug and score < 1: + print('INCORRECT', score) + print('GROUND TRUTH', gt_sentences) + print('SOLUTION', ans_sentences) + print('END OF OUTPUT', llm_answer[-50:]) + return score diff --git a/reasoners/reward/livebench/writing/typos/utils.py b/reasoners/reward/livebench/writing/typos/utils.py new file mode 100644 index 000000000..1b7715da2 --- /dev/null +++ b/reasoners/reward/livebench/writing/typos/utils.py @@ -0,0 +1,59 @@ +import difflib +import re + +def extract_answer(llm_answer): + pattern = r'.* --- (.*?) --- .*' + match = re.search(pattern, llm_answer) + return match.group(1) if match else llm_answer + +def typos_process_results(ground_truth: str, llm_answer: str, debug=False) -> int: + + parsed_answer = None + + # extract text from tags + solution_matches = re.findall(r'(.*?)', llm_answer) + if len(solution_matches) > 0: + match = solution_matches[-1] + parsed_answer = match + else: + parsed_answer = llm_answer.replace('', '').replace('', '') + parsed_answer = extract_answer(parsed_answer) + + parsed_answer = ' '.join(list(filter(None, parsed_answer.strip().split('\n')))) + + if int(ground_truth in parsed_answer): + return 1 + + score = 0 + + if debug and score == 0: + + a = ground_truth + b = parsed_answer + m = difflib.SequenceMatcher(a=a, b=b) + pad = 10 + + for tag, i1, i2, j1, j2 in m.get_opcodes(): + length = min(len(parsed_answer), len(ground_truth)) + mi1, mi2, mj1, mj2 = max(0,i1-pad), min(length, i2+pad), max(0, j1-pad), min(length, j2+pad) + + mistake_length = 0 + if tag == 'replace': + print("", a[i1:i2], b[j1:j2], "::::", a[mi1:mi2], "-->", b[mj1:mj2]) + mistake_length = i2 - i1 + if tag == 'delete': + print("", a[i1:i2], "::::", a[mi1:mi2], "-->", b[mj1:mj2]) + mistake_length = i2 - i1 + if tag == 'insert': + print("", b[j1:j2], "::::", a[mi1:mi2], "-->", b[mj1:mj2]) + mistake_length = j2 - j1 + #score = score - (mistake_length / len(ground_truth)) * 100 + + + if debug and score == 0: + print('INCORRECT') + print('GROUND TRUTH', ground_truth) + print('SOLUTION', parsed_answer) + print('END OF OUTPUT', llm_answer[-len(parsed_answer):]) + + return score diff --git a/reasoners/reward/math.py b/reasoners/reward/math.py new file mode 100644 index 000000000..50792aa6e --- /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_batch.py b/reasoners/reward/math_batch.py new file mode 100644 index 000000000..46a7cfec6 --- /dev/null +++ b/reasoners/reward/math_batch.py @@ -0,0 +1,25 @@ +# Copyright 2025 Individual Contributor: Mert Unsal + # + # 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. + + from .math import compute_score + + + def compute_score_batched(data_sources, solution_strs, ground_truths, extra_infos): + """ + This is a demonstration of how the batched reward function should look like. + Typically, you want to use batched reward to speed up the process with parallelization + """ + return [ + compute_score(solution_str, ground_truth) for solution_str, ground_truth in zip(solution_strs, ground_truths) + ] \ No newline at end of file diff --git a/reasoners/reward/math_dapo.py b/reasoners/reward/math_dapo.py new file mode 100644 index 000000000..5ce35811f --- /dev/null +++ b/reasoners/reward/math_dapo.py @@ -0,0 +1,294 @@ +# 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 + + +def last_boxed_only_string(string: str) -> Optional[str]: + """Extract the last LaTeX boxed expression from a string. + + Args: + string: Input string containing LaTeX code + + Returns: + The last boxed expression or None if not found + """ + idx = string.rfind("\\boxed{") + 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 + + return string[idx:right_brace_idx + 1] if right_brace_idx is not None else None + + +def remove_boxed(s: str) -> str: + """Remove the LaTeX boxed command from a string. + + Args: + s: String with format "\\boxed{content}" + + Returns: + The content inside the boxed command + """ + left = "\\boxed{" + assert s[:len(left)] == left, f"box error: {s}" + assert s[-1] == "}", f"box error: {s}" + return s[len(left):-1] + + +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() + + +def is_correct_minerva(solution_str: str, + gt: str, + gt_need_extract: bool = False, + answer_pattern: str = r"(?i)Answer\s*:\s*([^\n]+)") -> tuple[bool, str]: + """Check if the solution is correct according to Minerva criteria. + + Args: + solution_str: The solution string to check + gt: The ground truth answer + gt_need_extract: Whether the ground truth needs extraction + answer_pattern: Regex pattern to extract the answer + + Returns: + Tuple of (is_correct, normalized_prediction) + """ + print(f"[DAPO compute_score: gt]: {gt}, type: {type(gt)}") + print(f"[DAPO compute_score: solution_str]: {solution_str}") + # Extract answer from solution + match = re.findall(answer_pattern, solution_str) + print(f"[DAPO compute_score: match]: {match}") + extracted_answer = match[-1] if match else "[INVALID]" + pred = normalize_final_answer(extracted_answer) + print(f"[DAPO compute_score: extracted_answer]: {extracted_answer}") + print(f"[DAPO compute_score: pred]: {pred}") + + # Process ground truth + if gt_need_extract: + gt = normalize_final_answer(remove_boxed(last_boxed_only_string(gt))) + else: + gt = normalize_final_answer(gt) + + return (pred == gt), pred + + +def is_correct_strict_box(pred: str, + gt: str, + pause_tokens_index: Optional[list[int]] = None) -> tuple[int, Optional[str]]: + """Check if the prediction is correct using strict boxed answer criteria. + + Args: + pred: The prediction string + gt: The ground truth answer + pause_tokens_index: Indices of pause tokens + + Returns: + Tuple of (score, extracted_prediction) + """ + # Extract the relevant part of the prediction + if pause_tokens_index is not None: + assert len(pause_tokens_index) == 4 + pred = pred[pause_tokens_index[-1] - 100:] + else: + pred = pred[-100:] + + # Extract and check the boxed answer + boxed_pred = last_boxed_only_string(pred) + extracted_pred = remove_boxed(boxed_pred) if boxed_pred is not None else None + + return 1 if (extracted_pred == gt) else -1, extracted_pred + + +def verify(solution_str: str, + answer: str, + strict_box_verify: bool = False, + pause_tokens_index: Optional[list[int]] = None) -> bool: + """Verify if the solution is correct. + + Args: + solution_str: The solution string to verify + answer: The ground truth answer + strict_box_verify: Whether to use strict box verification + pause_tokens_index: Indices of pause tokens + + Returns: + True if the solution is correct, False otherwise + """ + if strict_box_verify: + correct, pred = is_correct_strict_box(solution_str, answer, pause_tokens_index) + return correct == 1, pred + + correct, pred = is_correct_minerva(solution_str, answer) + return correct, pred + + +def compute_score(solution_str: str, + ground_truth: str, + strict_box_verify: bool = False, + pause_tokens_index: Optional[list[int]] = None) -> float: + """Compute the reward score for a solution. + + Args: + solution_str: The solution string + ground_truth: The ground truth answer + config: Configuration object containing reward model settings + pause_tokens_index: Indices of pause tokens + + Returns: + Reward score (1.0 for correct, -1.0 for incorrect) + """ + # Limit solution length for efficiency + solution_str = solution_str[-300:] # The longest answer in MATH-500 has 159 characters + + # Verify the solution + correct, pred = verify(solution_str, ground_truth, strict_box_verify, pause_tokens_index) + + reward = 1.0 if correct else -1.0 + acc = correct + + return { + "score": reward, + "acc": acc, + } \ No newline at end of file diff --git a/reasoners/reward/math_llm_judge/__init__.py b/reasoners/reward/math_llm_judge/__init__.py new file mode 100644 index 000000000..8ac4ec51e --- /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 000000000..34f0c7f52 --- /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 000000000..e9921a5ad --- /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/math_verify.py b/reasoners/reward/math_verify.py new file mode 100644 index 000000000..66669ef94 --- /dev/null +++ b/reasoners/reward/math_verify.py @@ -0,0 +1,33 @@ +# 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. + +from math_verify.metric import math_metric +from math_verify.parser import LatexExtractionConfig, ExprExtractionConfig + + +def compute_score(model_output: str, ground_truth: str) -> bool: + verify_func = math_metric( + gold_extraction_target=(LatexExtractionConfig(),), + pred_extraction_target=(ExprExtractionConfig(), LatexExtractionConfig()), + ) + ret_score = 0. + + # Wrap the ground truth in \boxed{} format for verification + ground_truth_boxed = "\\boxed{" + ground_truth + "}" + try: + ret_score, _ = verify_func([ground_truth_boxed], [model_output]) + except Exception as e: + print(e) + + return ret_score diff --git a/reasoners/reward/naive_dapo.py b/reasoners/reward/naive_dapo.py new file mode 100644 index 000000000..a4f5c30a6 --- /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/orz/__init__.py b/reasoners/reward/orz/__init__.py new file mode 100644 index 000000000..0a513de80 --- /dev/null +++ b/reasoners/reward/orz/__init__.py @@ -0,0 +1,3 @@ +""" +Adapted from https://github.com/Open-Reasoner-Zero/Open-Reasoner-Zero +""" diff --git a/reasoners/reward/orz/math_utils.py b/reasoners/reward/orz/math_utils.py new file mode 100644 index 000000000..859a2d4f0 --- /dev/null +++ b/reasoners/reward/orz/math_utils.py @@ -0,0 +1,448 @@ +import asyncio +import re +from itertools import islice, zip_longest + +from sympy.parsing.latex import parse_latex + +try: + from math_verify import parse, verify +except ImportError: + print("math_verify is not installed in this environment") + parse = None + verify = None + + +def repeatness(s: str): + def ranks(l): + index = {v: i for i, v in enumerate(sorted(set(l)))} + return [index[v] for v in l] + + def suffixArray(s): + line = ranks(s) + n, k, ans, sa = len(s), 1, line, [0] * len(s) + while k < n - 1: + line = ranks(list(zip_longest(line, islice(line, k, None), fillvalue=-1))) + ans, k = line, k << 1 + for i, k in enumerate(ans): + sa[k] = i + return ans, sa + + def lcp(arr, suffixArr, inv_suff): + n, ans, k = len(arr), [0] * len(arr), 0 + + for i in range(n): + if inv_suff[i] == n - 1: + k = 0 + continue + + j = suffixArr[inv_suff[i] + 1] + while i + k < n and j + k < n and arr[i + k] == arr[j + k]: + k += 1 + + ans[inv_suff[i]] = k + if k > 0: + k -= 1 + + return ans + + arr = [ord(i) for i in s] + n = len(arr) + if n <= 1: + return 0 + c, sa = suffixArray(arr) + cnt = sum(lcp(arr, sa, c)) + + return (cnt * 2 / (n * (n + 1))) > 0.2 + + +SUBSTITUTIONS = [ + ("an ", ""), + ("a ", ""), + (".$", "$"), + ("\\$", ""), + (r"\ ", ""), + (" ", ""), + ("mbox", "text"), + (",\\text{and}", ","), + ("\\text{and}", ","), + ("\\text{m}", "\\text{}"), +] + + +REMOVED_EXPRESSIONS = [ + "square", + "ways", + "integers", + "dollars", + "mph", + "inches", + "ft", + "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. + This code comes from https://arxiv.org/pdf/2206.14858.pdf, page18. + """ + # final_answer = final_answer.split("=")[-1] + + for before, after in SUBSTITUTIONS: + final_answer = final_answer.replace(before, after) + for expr in REMOVED_EXPRESSIONS: + final_answer = final_answer.replace(expr, "") + + # Extract answer that is in LaTeX math, is bold, + # is surrounded by a box, etc. + 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 100,000 -> 100000 + if final_answer.replace(",", "").isdigit(): + final_answer = final_answer.replace(",", "") + + return final_answer + + +def latex_eval(latex): + sym = parse_latex(latex) + val = sym.evalf() + return sym, val + + +def _is_latex_equal(str1, str2): + try: + sym1, val1 = latex_eval(str1) + sym2, val2 = latex_eval(str2) + if sym1 == sym2 or val1 == val2: + return True + else: + raise ValueError + except Exception: # noqa + try: + norm1, norm2 = normalize_final_answer(str1), normalize_final_answer(str2) + sym1, val1 = latex_eval(norm1) + sym2, val2 = latex_eval(norm2) + if sym1 == sym2 or val1 == val2: + return True + except Exception: # noqa + return norm1 == norm2 + return False + + +async def is_latex_equal(str1, str2, executor, math_mode="legacy"): + if math_mode == "legacy": + if (len(str1) > 128 and repeatness(str1)) or ( + len(str2) > 128 and repeatness(str2) + ): + return False + + try: + loop = asyncio.get_event_loop() + task = loop.run_in_executor(executor, _is_latex_equal, str1, str2) + result = await asyncio.wait_for(task, timeout=1.0) + return result + except asyncio.exceptions.TimeoutError: + return False + elif math_mode == "math_verify": + try: + loop = asyncio.get_event_loop() + task = loop.run_in_executor(executor, verify, parse(str1), parse(str2)) + result = await asyncio.wait_for(task, timeout=1.0) + return result + except asyncio.exceptions.TimeoutError: + return False + else: + raise NotImplementedError(f"Math mode {math_mode} is not implemented") + + +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 Exception: # noqa + 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 Exception: # noqa + 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", "") + # print(string) + + # remove inverse spaces + string = string.replace("\\!", "") + # print(string) + + # replace \\ with \ + string = string.replace("\\\\", "\\") + # print(string) + + # replace tfrac and dfrac with frac + string = string.replace("tfrac", "frac") + string = string.replace("dfrac", "frac") + # print(string) + + # remove \left and \right + string = string.replace("\\left", "") + string = string.replace("\\right", "") + # print(string) + + # Remove circ (degrees) + string = string.replace("^{\\circ}", "") + string = string.replace("^\\circ", "") + + # remove dollar signs + string = string.replace("\\$", "") + string = string.replace("$", "") + 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 + + +def is_equiv(str1, str2, verbose=False) -> bool: + 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) + try: + return float(ss1) == (float(ss2)) + except Exception: # noqa + return ss1 == ss2 + except Exception: # noqa + return str1 == str2 + + +def last_boxed_only_string(string): + idx = string.rfind("\\boxed") + 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 remove_boxed(s): + left = "\\boxed{" + try: + assert s[: len(left)] == left + assert s[-1] == "}" + return s[len(left) : -1] + except Exception: + return None + + +def get_answer_str(s: str) -> str: + res = remove_boxed(last_boxed_only_string(s)) + if res is not None: + return res + return s + + +async def is_equal(str1, str2, executor, math_mode="legacy"): + first_equal = is_equiv(str1, str2) + if first_equal: + return True + return await is_latex_equal(str1, str2, executor, math_mode) + + +def solution2answer(solution: str, math_mode="eval_peeking") -> str: + answer = solution + if math_mode == "eval_peeking": + answer = get_answer_str(solution) + else: + raise ValueError(f"Invalid math_mode: {math_mode}") + return answer + + +def get_final_answer(output: str) -> str: + output = output.replace("is:", "is").replace("answer:", "answer is").strip() + if output.endswith("."): + output = output[:-1] + if ".$" in output: + output = output.replace(".$", "$") + pattern_list = [ + r"answer is (-?\d+\.?\d*)$", + r"answer is (.+?)$", + ] + matches = [] + for pat in pattern_list: + matches = re.findall(pat, output, re.S) + if matches: + return get_answer_str(matches[0]) + + return get_answer_str(output) diff --git a/reasoners/reward/orz/math_utils_sync.py b/reasoners/reward/orz/math_utils_sync.py new file mode 100644 index 000000000..5d453f930 --- /dev/null +++ b/reasoners/reward/orz/math_utils_sync.py @@ -0,0 +1,442 @@ +import asyncio +import re +from itertools import islice, zip_longest + +from sympy.parsing.latex import parse_latex + +try: + from math_verify import parse, verify +except ImportError: + print("math_verify is not installed in this environment") + parse = None + verify = None + + +def repeatness(s: str): + def ranks(l): + index = {v: i for i, v in enumerate(sorted(set(l)))} + return [index[v] for v in l] + + def suffixArray(s): + line = ranks(s) + n, k, ans, sa = len(s), 1, line, [0] * len(s) + while k < n - 1: + line = ranks(list(zip_longest(line, islice(line, k, None), fillvalue=-1))) + ans, k = line, k << 1 + for i, k in enumerate(ans): + sa[k] = i + return ans, sa + + def lcp(arr, suffixArr, inv_suff): + n, ans, k = len(arr), [0] * len(arr), 0 + + for i in range(n): + if inv_suff[i] == n - 1: + k = 0 + continue + + j = suffixArr[inv_suff[i] + 1] + while i + k < n and j + k < n and arr[i + k] == arr[j + k]: + k += 1 + + ans[inv_suff[i]] = k + if k > 0: + k -= 1 + + return ans + + arr = [ord(i) for i in s] + n = len(arr) + if n <= 1: + return 0 + c, sa = suffixArray(arr) + cnt = sum(lcp(arr, sa, c)) + + return (cnt * 2 / (n * (n + 1))) > 0.2 + + +SUBSTITUTIONS = [ + ("an ", ""), + ("a ", ""), + (".$", "$"), + ("\\$", ""), + (r"\ ", ""), + (" ", ""), + ("mbox", "text"), + (",\\text{and}", ","), + ("\\text{and}", ","), + ("\\text{m}", "\\text{}"), +] + + +REMOVED_EXPRESSIONS = [ + "square", + "ways", + "integers", + "dollars", + "mph", + "inches", + "ft", + "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. + This code comes from https://arxiv.org/pdf/2206.14858.pdf, page18. + """ + # final_answer = final_answer.split("=")[-1] + + for before, after in SUBSTITUTIONS: + final_answer = final_answer.replace(before, after) + for expr in REMOVED_EXPRESSIONS: + final_answer = final_answer.replace(expr, "") + + # Extract answer that is in LaTeX math, is bold, + # is surrounded by a box, etc. + 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 100,000 -> 100000 + if final_answer.replace(",", "").isdigit(): + final_answer = final_answer.replace(",", "") + + return final_answer + + +def latex_eval(latex): + sym = parse_latex(latex) + val = sym.evalf() + return sym, val + + +def _is_latex_equal(str1, str2): + try: + sym1, val1 = latex_eval(str1) + sym2, val2 = latex_eval(str2) + if sym1 == sym2 or val1 == val2: + return True + else: + raise ValueError + except Exception: # noqa + try: + norm1, norm2 = normalize_final_answer(str1), normalize_final_answer(str2) + sym1, val1 = latex_eval(norm1) + sym2, val2 = latex_eval(norm2) + if sym1 == sym2 or val1 == val2: + return True + except Exception: # noqa + return norm1 == norm2 + return False + + +def is_latex_equal(str1, str2, math_mode="legacy"): + if math_mode == "legacy": + if (len(str1) > 128 and repeatness(str1)) or ( + len(str2) > 128 and repeatness(str2) + ): + return False + + try: + return _is_latex_equal(str1, str2) + except Exception: # noqa + return False + elif math_mode == "math_verify": + try: + return verify(parse(str1), parse(str2)) + except Exception: # noqa + return False + else: + raise NotImplementedError(f"Math mode {math_mode} is not implemented") + + +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 Exception: # noqa + 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 Exception: # noqa + 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", "") + # print(string) + + # remove inverse spaces + string = string.replace("\\!", "") + # print(string) + + # replace \\ with \ + string = string.replace("\\\\", "\\") + # print(string) + + # replace tfrac and dfrac with frac + string = string.replace("tfrac", "frac") + string = string.replace("dfrac", "frac") + # print(string) + + # remove \left and \right + string = string.replace("\\left", "") + string = string.replace("\\right", "") + # print(string) + + # Remove circ (degrees) + string = string.replace("^{\\circ}", "") + string = string.replace("^\\circ", "") + + # remove dollar signs + string = string.replace("\\$", "") + string = string.replace("$", "") + 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 + + +def is_equiv(str1, str2, verbose=False) -> bool: + 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) + try: + return float(ss1) == (float(ss2)) + except Exception: # noqa + return ss1 == ss2 + except Exception: # noqa + return str1 == str2 + + +def last_boxed_only_string(string): + idx = string.rfind("\\boxed") + 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 remove_boxed(s): + left = "\\boxed{" + try: + assert s[: len(left)] == left + assert s[-1] == "}" + return s[len(left) : -1] + except Exception: + return None + + +def get_answer_str(s: str) -> str: + res = remove_boxed(last_boxed_only_string(s)) + if res is not None: + return res + return s + + +def is_equal(str1, str2, math_mode="legacy"): + first_equal = is_equiv(str1, str2) + if first_equal: + return True + return is_latex_equal(str1, str2, math_mode) + + +def solution2answer(solution: str, math_mode="eval_peeking") -> str: + answer = solution + if math_mode == "eval_peeking": + answer = get_answer_str(solution) + else: + raise ValueError(f"Invalid math_mode: {math_mode}") + return answer + + +def get_final_answer(output: str) -> str: + output = output.replace("is:", "is").replace("answer:", "answer is").strip() + if output.endswith("."): + output = output[:-1] + if ".$" in output: + output = output.replace(".$", "$") + pattern_list = [ + r"answer is (-?\d+\.?\d*)$", + r"answer is (.+?)$", + ] + matches = [] + for pat in pattern_list: + matches = re.findall(pat, output, re.S) + if matches: + return get_answer_str(matches[0]) + + return get_answer_str(output) diff --git a/reasoners/reward/prime_code/__init__.py b/reasoners/reward/prime_code/__init__.py new file mode 100644 index 000000000..49b0d30f5 --- /dev/null +++ b/reasoners/reward/prime_code/__init__.py @@ -0,0 +1,73 @@ +# 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. + +from .utils import check_correctness as apps_check_correctness +import json +import re +import traceback + + +def compute_score(completion, test_cases, continuous=False): + # try to get code solution from completion. if the completion is pure code, this will not take effect. + solution = completion.split('```python')[-1].split('```')[0] + try: + try: + if not isinstance(test_cases, dict): + test_cases = json.loads(test_cases) + except Exception as e: + print(f"Error:{e}") + + # Complete check on all in-out pairs first. If there is no failure, per-sample test can be skipped. + try: + res, metadata = apps_check_correctness(in_outs=test_cases, generation=solution, timeout=5, debug=False) + metadata = dict(enumerate(metadata))[0] + success = all(map(lambda x: x == True, res)) + if success: + return success, metadata + except Exception as e: + pass + + test_cases_list = [] + inputs = test_cases["inputs"] + outputs = test_cases["outputs"] + for i in range(len(inputs)): + test_cases_list.append({"inputs": [inputs[i]], "outputs": [outputs[i]]}) + + if continuous: + # per sample test: if continuous score is needed, test first 10 samples regardless of failures + # do not test all samples cuz some problems have enormous test cases + metadata_list = [] + res_list = [] + for test_case_id, test_case in enumerate(test_cases_list): + res, metadata = apps_check_correctness(in_outs=test_case, generation=solution, timeout=5, debug=False) + try: + metadata = dict(enumerate(metadata))[0] # metadata can be empty occasionally + except Exception as e: + metadata = {} + metadata["test_case"] = {} + metadata["test_case"]["input"] = str(test_case["inputs"][0]) + metadata["test_case"]["output"] = str(test_case["outputs"][0]) + metadata["test_case"]["res"] = str(res) + metadata_list.append(metadata) + res_list.extend(res) + + if test_case_id >= 9: + break + res_count = len(res_list) if len(res_list) > 0 else 1 + success = sum(map(lambda x: x == True, res_list)) / res_count + except Exception as e: + traceback.print_exc(10) + success = False + metadata_list = None + return success, metadata_list diff --git a/reasoners/reward/prime_code/testing_util.py b/reasoners/reward/prime_code/testing_util.py new file mode 100644 index 000000000..34845d07e --- /dev/null +++ b/reasoners/reward/prime_code/testing_util.py @@ -0,0 +1,721 @@ +# 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. + +import ast +import json +import sys +import faulthandler +import platform + +# used for debugging to time steps +from datetime import datetime + +# to run the solution files we're using a timing based approach +import signal + +import numpy as np + +# for capturing the stdout +from io import StringIO + +# used for testing the code that reads from input +from unittest.mock import patch, mock_open + +from pyext import RuntimeModule + +from enum import Enum + +import traceback + + +def truncatefn(s, length=300): + assert isinstance(s, str) + if len(s) <= length: + return s + + return s[:length // 2] + "...(truncated) ..." + s[-length // 2:] + + +class CODE_TYPE(Enum): + call_based = 0 + standard_input = 1 + + +# stuff for setting up signal timer +class TimeoutException(Exception): + pass + + +def timeout_handler(signum, frame): + print("alarm went off") + return + # raise TimeoutException # this is an unhandled exception. just return None is OK + + +signal.signal(signal.SIGALRM, timeout_handler) + +# timeout = 6 # seconds + + +# used to capture stdout as a list +# from https://stackoverflow.com/a/16571630/6416660 +# alternative use redirect_stdout() from contextlib +class Capturing(list): + + def __enter__(self): + self._stdout = sys.stdout + sys.stdout = self._stringio = StringIO() + # Make closing the StringIO a no-op + self._stringio.close = lambda x: 1 + return self + + def __exit__(self, *args): + self.append(self._stringio.getvalue()) + del self._stringio # free up some memory + sys.stdout = self._stdout + + +def only_int_check(val): + return isinstance(val, int) + + +def string_int_check(val): + return isinstance(val, str) and val.isdigit() + + +def combined_int_check(val): + return only_int_check(val) or string_int_check(val) + + +def clean_traceback(error_traceback): + file_start = error_traceback.find('File \"\"') + # print(file_start) + error_traceback = "Traceback (most recent call last):\n " + error_traceback[file_start:] + return error_traceback + + +def run_test(in_outs, test=None, debug=False, timeout=15): + """ + if test(generated_code) is not None it'll try to run the code. + otherwise it'll just return an input and output pair. + """ + # Disable functionalities that can make destructive changes to the test. + reliability_guard() + + if debug: + print(f"start = {datetime.now().time()}") + + if in_outs: + if in_outs.get("fn_name") is None: + which_type = CODE_TYPE.standard_input # Standard input + method_name = None + else: + which_type = CODE_TYPE.call_based # Call-based + method_name = in_outs["fn_name"] + + if debug: + print(f"loaded input_output = {datetime.now().time()}") + + if test is None: + assert False, "should not happen: test code is none" + return in_outs, {"error": "no test code provided"} + elif test is not None: + results = [] + sol = "from string import *\nfrom re import *\nfrom datetime import *\nfrom collections import *\nfrom heapq import *\nfrom bisect import *\nfrom copy import *\nfrom math import *\nfrom random import *\nfrom statistics import *\nfrom itertools import *\nfrom functools import *\nfrom operator import *\nfrom io import *\nfrom sys import *\nfrom json import *\nfrom builtins import *\nfrom typing import *\nimport string\nimport re\nimport datetime\nimport collections\nimport heapq\nimport bisect\nimport copy\nimport math\nimport random\nimport statistics\nimport itertools\nimport functools\nimport operator\nimport io\nimport sys\nimport json\nsys.setrecursionlimit(6*10**5)\n" + if debug: + print(f"loading test code = {datetime.now().time()}") + + if which_type == CODE_TYPE.call_based: + + sol += test + if debug: + print(f"sol = {sol}") + signal.alarm(timeout) + try: + tmp_sol = RuntimeModule.from_string("tmp_sol", "", sol) + if "class Solution" not in test: + tmp = tmp_sol + else: + tmp = tmp_sol.Solution() + signal.alarm(0) + except Exception as e: + signal.alarm(0) + error_traceback = traceback.format_exc() + if debug: + print(f"type 0 compilation error = {e}") + results.append(-2) + return results, { + "error": repr(e), + # "error_code": -1, + # "error_message": "Compilation Error", + "traceback": clean_traceback(error_traceback), + } + signal.alarm(0) + + elif which_type == CODE_TYPE.standard_input: + # sol + # if code has if __name__ == "__main__": then remove it + try: + astree = ast.parse(test) + last_block = astree.body[-1] + if isinstance(last_block, ast.If): + condition = last_block.test + if ast.unparse(condition).strip() == "__name__ == '__main__'": + test = (ast.unparse(astree.body[:-1]) + "\n" + ast.unparse(last_block.body)) + except: + pass + + tmp_test = test.split("\n") + + new_test = [] + for x in tmp_test: + if (not x.startswith("from ")) and (not x.startswith("import ")): + new_test.append("\t" + x + "\n") + else: + new_test.append(x + "\n") + tmp_test = new_test + + new_test = "" + started = False + for i in tmp_test: + if i.startswith("\t") and not started: + new_test += "stdin = sys.stdin\nstdout = sys.stdout\n" + new_test += "def code():\n" + new_test += i + started = True + elif started and ((i.startswith("from ")) or (i.startswith("import "))): + new_test += "\t" + i + else: + new_test += i + tmp_test = new_test + + sol += tmp_test + if debug: + print(f"sol = {sol}") + method_name = "code" + signal.alarm(timeout) + try: + tmp_sol = RuntimeModule.from_string("tmp_sol", "", sol) + tmp = tmp_sol + signal.alarm(0) + except Exception as e: + signal.alarm(0) + error_traceback = traceback.format_exc() + if debug: + print(f"type 1 compilation error = {e}") + results.append(-2) + return results, { + "error": repr(e), + # "error_code": -1, + # "error_message": "Compilation Error", + "traceback": clean_traceback(error_traceback), + } + signal.alarm(0) + if debug: + print(f"get method = {datetime.now().time()}") + + try: + method = getattr(tmp, method_name) # get_attr second arg must be str + except: + signal.alarm(0) + error_traceback = traceback.format_exc() + e = sys.exc_info() + print(f"unable to get function error = {e}") + results.append(-2) + return results, { + "error": repr(e), + # "error_code": -1, + # "error_message": "Unable to extract code", + "traceback": clean_traceback(error_traceback), + } + + for index, inputs in enumerate(in_outs["inputs"]): + raw_inputs = inputs + raw_outputs = in_outs["outputs"][index] + if which_type == CODE_TYPE.call_based: + inputs = [json.loads(line) for line in inputs.split("\n")] + in_outs["outputs"][index] = json.loads(in_outs["outputs"][index]) + + truncate_line_size = 300 // (raw_inputs.count("\n") + 1) + raw_inputs = "\n".join( + [truncatefn(line, truncate_line_size) for line in raw_inputs.strip().split("\n")]) + raw_outputs = truncatefn(raw_outputs, 200) + else: + raw_inputs = truncatefn(raw_inputs) + raw_outputs = truncatefn(raw_outputs, 200) + # JSON forces dictionaries to have string keys; this undoes this (assuming a singleton list) + try: + if isinstance(inputs[0], dict): + inputs = [{int(k): v for k, v in inputs[0].items()}] + except: + True + try: + if isinstance(in_outs["outputs"][index], dict): + in_outs["outputs"][index] = [{int(k): v for k, v in in_outs["outputs"][index].items()}] + except: + True + try: + if isinstance(in_outs["outputs"][index][0], dict): + in_outs["outputs"][index] = [{int(k): v for k, v in in_outs["outputs"][index][0].items()}] + except: + True + + if debug: + print( + f"time: {datetime.now().time()} testing index = {index} inputs = {inputs}, {type(inputs)}. type = {which_type}" + ) + if which_type == CODE_TYPE.call_based: # Call-based + signal.alarm(timeout) + faulthandler.enable() + try: + output = method(*inputs) + raw_true_output = output + + raw_true_output_copy = json.dumps(output) + raw_true_output_copy = truncatefn(raw_true_output_copy, 200) + + # ground truth sequences are not tuples + if isinstance(output, tuple): + output = list(output) + + tmp_result = output == in_outs["outputs"][index] + if (isinstance(in_outs["outputs"][index], list) and in_outs["outputs"][index]): + tmp_result = tmp_result or (output == in_outs["outputs"][index][0]) + + # ground truth sequences are not tuples + try: + if isinstance(output[0], tuple): + tmp_result = tmp_result or ([list(x) for x in output] == in_outs["outputs"][index][0]) + except: + True + results.append(tmp_result) + if tmp_result != True: + return results, { + "output": raw_true_output_copy, + "expected": raw_outputs, + "inputs": raw_inputs, + # "error_code": -2, + "error_message": "Wrong Answer", + } + # reset the alarm + signal.alarm(0) + except Exception as e: + signal.alarm(0) + error_traceback = traceback.format_exc() + faulthandler.disable() + if debug: + print(f"Standard input runtime error or time limit exceeded error = {e}") + results.append(-1) + if "timeoutexception" in repr(e).lower(): + return results, { + "error": repr(e), + # "error_code": -3, + # "error_message": "Time Limit Exceeded", + # "inputs": raw_inputs, + # "expected": raw_outputs, + "traceback": clean_traceback(error_traceback), + } + else: + return results, { + "error": repr(e), + # "error_code": -4, + # "error_message": "Runtime Error", + # "inputs": raw_inputs, + # "expected": raw_outputs, + "traceback": clean_traceback(error_traceback), + } + faulthandler.disable() + signal.alarm(0) + if debug: + print( + f"outputs = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs}, {type(inputs)}, {output == [in_outs['outputs'][index]]}" + ) + elif which_type == CODE_TYPE.standard_input: # Standard input + faulthandler.enable() + passed = False + + if isinstance(inputs, list): + inputs = "\n".join(inputs) + if isinstance(in_outs["outputs"][index], list): + in_outs["outputs"][index] = "\n".join(in_outs["outputs"][index]) + + signal.alarm(timeout) + with Capturing() as output: + try: + call_method(method, inputs) + # reset the alarm + signal.alarm(0) + passed = True + except Exception as e: + # runtime error or took too long + signal.alarm(0) + error_traceback = traceback.format_exc() + print(f"Call-based runtime error or time limit exceeded error = {repr(e)}{e}") + results.append(-1) + if "timeoutexception" in repr(e).lower(): + return results, { + "error": repr(e), + # "error_code": -3, + # "error_message": "Time Limit Exceeded", + # "inputs": raw_inputs, + # "expected": raw_outputs, + "traceback": clean_traceback(error_traceback), + } + else: + return results, { + "error": repr(e), + # "error_code": -4, + # "error_message": "Runtime Error", + # "inputs": raw_inputs, + # "expected": raw_outputs, + "traceback": clean_traceback(error_traceback), + } + signal.alarm(0) + raw_true_output = output[0] + raw_true_output_copy = truncatefn(raw_true_output, 200) + output = raw_true_output.splitlines() + if not passed: + if debug: + nl = "\n" + if not isinstance(inputs, list): + print( + f"not passed output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs.replace(nl,' new-line ')}, {type(inputs)}, {output == [in_outs['outputs'][index]]}" + ) + else: + print( + f"not passed output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs}, {type(inputs)}, {output == [in_outs['outputs'][index]]}" + ) + continue + + if passed and debug: + print(f"==> output = {output}, test outputs = {in_outs['outputs'][index]}") + + if custom_compare_(output, in_outs["outputs"][index]): + tmp_result = True + results.append(tmp_result) + continue + + # ground truth sequences are expressed as lists not tuples + if isinstance(output, tuple): + output = list(output) + + tmp_result = False + try: + tmp_result = output == [in_outs["outputs"][index]] + if isinstance(in_outs["outputs"][index], list): + tmp_result = tmp_result or (output == in_outs["outputs"][index]) + if isinstance(output[0], str): + tmp_result = tmp_result or ([e.strip() for e in output] == in_outs["outputs"][index]) + except Exception as e: + if debug: + print(f"Failed check1 exception = {e}") + pass + + if tmp_result == True: + results.append(tmp_result) + continue + + # try one more time without \n + if isinstance(in_outs["outputs"][index], list): + for tmp_index, i in enumerate(in_outs["outputs"][index]): + in_outs["outputs"][index][tmp_index] = i.split("\n") + in_outs["outputs"][index][tmp_index] = [ + x.strip() for x in in_outs["outputs"][index][tmp_index] if x + ] + else: + in_outs["outputs"][index] = in_outs["outputs"][index].split("\n") + in_outs["outputs"][index] = list(filter(len, in_outs["outputs"][index])) + in_outs["outputs"][index] = list(map(lambda x: x.strip(), in_outs["outputs"][index])) + + try: + tmp_result = output == [in_outs["outputs"][index]] + if isinstance(in_outs["outputs"][index], list): + tmp_result = tmp_result or (output == in_outs["outputs"][index]) + except Exception as e: + if debug: + print(f"Failed check2 exception = {e}") + pass + + if tmp_result == True: + results.append(tmp_result) + continue + + # try by converting the output into a split up list too + if isinstance(output, list): + output = list(filter(len, output)) + + if debug: + nl = "\n" + if not isinstance(inputs, list): + print( + f"@1 output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs.replace(nl,' new-line ')}, {type(inputs)}, {output == [in_outs['outputs'][index]]} {tmp_result=}" + ) + else: + print( + f"@1 output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs}, {type(inputs)}, {output == [in_outs['outputs'][index]]} {tmp_result=}" + ) + + if tmp_result == True: + results.append(tmp_result) + continue + + if debug: + print(f"{tmp_result=} @a") + + try: + tmp_result = output == [in_outs["outputs"][index]] + if isinstance(in_outs["outputs"][index], list): + tmp_result = tmp_result or (output == in_outs["outputs"][index]) + except Exception as e: + if debug: + print(f"Failed check3 exception = {e}") + pass + + if debug: + print(f"{tmp_result=} @b") + + try: + all_ints = all( + combined_int_check(e1) and combined_int_check(e2) + for e1, e2 in zip(output, in_outs["outputs"][index])) + if not all_ints: + if debug: + print([ + combined_int_check(e1) and combined_int_check(e2) + for e1, e2 in zip(output, in_outs["outputs"][index]) + ]) + output_float = [float(e) for e in output] + gt_float = [float(e) for e in in_outs["outputs"][index]] + tmp_result = tmp_result or ((len(output_float) == len(gt_float)) and + np.allclose(output_float, gt_float)) + except Exception as e: + pass + + if debug: + print(f"{tmp_result=} @c") + + try: + if isinstance(output[0], list): + all_ints = all( + combined_int_check(e1) and combined_int_check(e2) + for e1, e2 in zip(output[0], in_outs["outputs"][index])) + if not all_ints: + output_float = [float(e) for e in output[0]] + gt_float = [float(e) for e in in_outs["outputs"][index][0]] + tmp_result = tmp_result or ((len(output_float) == len(gt_float)) and + np.allclose(output_float, gt_float)) + except Exception as e: + pass + + if tmp_result == True: + results.append(tmp_result) + continue + + if debug: + print(f"{tmp_result=} @d") + # try by converting the stuff into split up list + if isinstance(in_outs["outputs"][index], list): + for tmp_index, i in enumerate(in_outs["outputs"][index]): + in_outs["outputs"][index][tmp_index] = set(i.split()) + else: + in_outs["outputs"][index] = set(in_outs["outputs"][index].split()) + + if debug: + print(f"{tmp_result=} @e") + + try: + tmp_result = output == in_outs["outputs"][index] + except Exception as e: + if debug: + print(f"Failed check4 exception = {e}") + continue + + if tmp_result == True: + results.append(tmp_result) + continue + + if debug: + print(f"{tmp_result=} @f") + + # try by converting the output into a split up list too + if isinstance(output, list): + for tmp_index, i in enumerate(output): + output[tmp_index] = i.split() + output = list(filter(len, output)) + for tmp_index, i in enumerate(output): + output[tmp_index] = set(i) + else: + output = output.split() + output = list(filter(len, output)) + output = set(output) + + if debug: + print(f"{tmp_result=} @g") + + if tmp_result == True and debug: + print("PASSED") + + results.append(tmp_result) + if tmp_result != True: + return results, { + "output": raw_true_output_copy, + "expected": raw_outputs, + "inputs": raw_inputs, + # "error_code": -2, + "error_message": "Wrong Answer", + } + + if debug: + nl = "\n" + if not isinstance(inputs, list): + print( + f"@2 output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs.replace(nl,' new-line ')}, {type(inputs)}, {output == [in_outs['outputs'][index]]}" + ) + else: + print( + f"@2 output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs}, {type(inputs)}, {output == [in_outs['outputs'][index]]}" + ) + + print(f"results = {results}") + + return results, {} + + +def custom_compare_(output, ground_truth): + + if isinstance(output, list): + output_1 = "\n".join(output) + if stripped_string_compare(output_1, ground_truth): + return True + + if isinstance(output, list): + output_2 = [o.lstrip().rstrip() for o in output] + output_2 = "\n".join(output_2) + if stripped_string_compare(output_2, ground_truth): + return True + + return False + + +def stripped_string_compare(s1, s2): + s1 = s1.lstrip().rstrip() + s2 = s2.lstrip().rstrip() + return s1 == s2 + + +def call_method(method, inputs): + + if isinstance(inputs, list): + inputs = "\n".join(inputs) + + inputs_line_iterator = iter(inputs.split("\n")) + + # sys.setrecursionlimit(10000) + + # @patch('builtins.input', side_effect=inputs.split("\n")) + @patch("builtins.open", mock_open(read_data=inputs)) + @patch("sys.stdin", StringIO(inputs)) + @patch("sys.stdin.readline", lambda *args: next(inputs_line_iterator)) + @patch("sys.stdin.readlines", lambda *args: inputs.split("\n")) + @patch("sys.stdin.read", lambda *args: inputs) + # @patch('sys.stdout.write', print) + def _inner_call_method(_method): + try: + return _method() + except SystemExit as e: + pass + finally: + pass + + return _inner_call_method(method) + + +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 # 防止干扰repl评测 + 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 diff --git a/reasoners/reward/prime_code/utils.py b/reasoners/reward/prime_code/utils.py new file mode 100644 index 000000000..07f3db88e --- /dev/null +++ b/reasoners/reward/prime_code/utils.py @@ -0,0 +1,59 @@ +# 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. + +# Borrowed from: https://huggingface.co/spaces/codeparrot/apps_metric/blob/main/utils.py + +import multiprocessing +from typing import Dict, Optional +from datasets import load_dataset +from .testing_util import run_test +import traceback +import os, sys + + +def _temp_run(sample, generation, debug, result, metadata_list, timeout): + with open(os.devnull, 'w') as devnull: + sys.stdout = devnull + sys.stderr = devnull + try: + res, metadata = run_test(in_outs=sample, test=generation, debug=debug, timeout=timeout) + result.append(res) + metadata_list.append(metadata) + except Exception as e: + # print(e) # some tracebacks are extremely long. + traceback.print_exc(10) + result.append([-1 for i in range(len(sample['inputs']))]) + metadata_list.append({}) + + +def check_correctness(in_outs: Optional[dict], generation, timeout=10, debug=True): + """Check correctness of code generation with a global timeout. + The global timeout is to catch some extreme/rare cases not handled by the timeouts + inside `run_test`""" + + manager = multiprocessing.Manager() + result = manager.list() + metadata_list = manager.list() + p = multiprocessing.Process(target=_temp_run, args=(in_outs, generation, debug, result, metadata_list, timeout)) + p.start() + p.join(timeout=timeout + 1) + if p.is_alive(): + p.kill() + # p.terminate() + if not result: + # consider that all tests failed + result = [[-1 for i in range(len(in_outs["inputs"]))]] + if debug: + print(f"global timeout") + return result[0], metadata_list diff --git a/reasoners/reward/prime_math/__init__.py b/reasoners/reward/prime_math/__init__.py new file mode 100644 index 000000000..01390e5b6 --- /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 000000000..e3e1686f1 --- /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 000000000..e9921a5ad --- /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 000000000..135013ead --- /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 000000000..dcc4e79ba --- /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 000000000..6c3ef69c9 --- /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 000000000..187d68f0a --- /dev/null +++ b/reasoners/reward/tablereason.py @@ -0,0 +1,62 @@ +from verl.utils.reward_score.prime_math.grader import math_equal +from verl.utils.reward_score 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 000000000..aa86dd2ad --- /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 000000000..ed42cb609 --- /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/README.md b/reasoners/tools/model_filtering/README.md new file mode 100644 index 000000000..040a15edc --- /dev/null +++ b/reasoners/tools/model_filtering/README.md @@ -0,0 +1,194 @@ +# Difficulty Filter Pipeline + +A pipeline for difficulty filtering using reward functions, supporting data-parallel and tensor-parallel processing. The pipeline now operates in two stages: +1. **Inference stage**: Runs the model to generate responses (GPU-intensive) +2. **Reward stage**: Evaluates responses with reward functions (CPU-intensive) + +## Resource Requirements + +- 8x H200 (140GB) GPUs + +## Example Usage + +### Stage 1: Inference + +> [!IMPORTANT] +> To run the model, you have to install verl and upgrade vllm to 0.8.5, please first follow the instructions in the main README.md, and then run the following command: + +```bash +pip uninstall vllm +pip install vllm==0.8.5 +``` + +#### Qwen2.5-7B-Instruct (~0.09s 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 +``` + +#### Qwen3-30B-A3B (~12s 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 +``` +### Stage 2: Reward Calculation + +After the inference stage is complete, run the 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 +``` + +## Checkpoint & Resumption + +The pipeline automatically saves batch results to JSON files in the output directory as it processes data. If a run is interrupted, the pipeline will automatically resume from where it left off by: + +1. Scanning the output directory for existing batch files +2. Loading results from these files +3. Continuing from the next unprocessed batch + +### Advanced Checkpoint Options + +The pipeline provides flags to control checkpoint behavior: + +- **Inference stage**: + - `--force_regenerate`: Ignores all existing checkpoint files and starts processing from the beginning, overwriting previous results + +- **Reward stage**: + - `--recalculate_rewards`: Recalculates all reward scores using the previously generated model responses (useful when implementing new evaluation metrics without re-running the expensive model inference) + +## Utility Functions + +After running models with data-parallel processing, you can use the utility functions to concatenate and analyze results. + +### Generating Pass Rate Mappings + +Generate idx to pass_rate mapping for a specific model and dataset: + +```bash +python -m model_filtering.utils map \ + --output_dir "./diff_filter_output" \ + --dataset "codegen__leetcode2k_2.4k" \ + --model "Qwen2.5-7B-Instruct" +``` + +This command will create an `idx_to_passrate.json` file in the model's directory that can be used for analysis. + +### Analyzing Results + +Analyze difficulty distributions of one or more datasets using a specific model: + +```bash +python -m model_filtering.utils analyze \ + --output_dir "./diff_filter_output" \ + --datasets "codegen__leetcode2k_2.4k" \ + --model "Qwen2.5-7B-Instruct" +``` + +Analyze multiple dataset chunks with the 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" +``` + +Some advanced options: +- `--regenerate`: Force regeneration of idx_to_passrate mapping +- `--save_combined`: Save the combined mapping when analyzing multiple datasets + +The analysis output categorizes problems into seven difficulty levels: +- Impossible (pass rate exactly 0.0) +- Very Hard (pass rate 0.0-0.2, exclusive) +- Hard (pass rate 0.2-0.4) +- Medium (pass rate 0.4-0.6) +- Easy (pass rate 0.6-0.8) +- Very Easy (pass rate 0.8-1.0, exclusive) +- Perfect (pass rate exactly 1.0) + + +# Annotate Pass Rates Script *run_add_pr.py* + +This script adds two new columns to each row of a Parquet dataset, reflecting pass-rate values loaded from two JSON files: + +* **`low_pass_rate`**: pass rate from the `--low_json` mapping (default is 0.0 if missing) +* **`high_pass_rate`**: pass rate from the `--high_json` mapping (default is 0.0 if missing) + +## Usage + +```bash +python run_add_pr.py \ + --parquet_in .parquet \ + --low_json low_pass_rates.json \ + --high_json high_pass_rates.json \ + --parquet_out_dir directory for saving processed parquet +``` + +## Command-Line Arguments + +| Argument | Description | Required | +| --------------- | -------------------------------------------------------- | -------- | +| `--parquet_in` | Path to the input Parquet file to annotate. | Yes | +| `--low_json` | JSON file mapping `idx → pass_rate` for low thresholds. | Yes | +| `--high_json` | JSON file mapping `idx → pass_rate` for high thresholds. | Yes | +| `--parquet_out_dir` | Directory where the annotated Parquet file will be saved. | Yes | + +# Filter by Pass-Rate Thresholds *run_filter.py* + +This script filters a Parquet dataset based on pass-rate thresholds defined in two JSON files. It uses **Rich** for a colored, tabular console summary. + +## Usage + +```bash +python run_filter.py \ + --parquet_in .parquet \ + --parquet_out_dir \ + --low_json low_pass_rates.json \ + --high_json high_pass_rates.json \ + --low_thresh 0.1 \ + --high_thresh 0.9 +``` + +## Arguments + +| Flag | Description | +| ------------------- | ----------------------------------------------------------- | +| `--parquet_in` | Path to the source Parquet file | +| `--parquet_out_dir` | Directory where filtered Parquet files will be saved | +| `--low_json` | JSON mapping `idx → pass_rate` for low-threshold filtering | +| `--high_json` | JSON mapping `idx → pass_rate` for high-threshold filtering | +| `--low_thresh` | Drop rows with pass rate **≤** this value in low json | +| `--high_thresh` | Drop rows with pass rate **≥** this value in high json | diff --git a/reasoners/tools/model_filtering/__init__.py b/reasoners/tools/model_filtering/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/reasoners/tools/model_filtering/pipeline.py b/reasoners/tools/model_filtering/pipeline.py new file mode 100644 index 000000000..299d969f1 --- /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/run_add_pr.py b/reasoners/tools/model_filtering/run_add_pr.py new file mode 100644 index 000000000..cbcdfa3b2 --- /dev/null +++ b/reasoners/tools/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/reasoners/tools/model_filtering/run_filter.py b/reasoners/tools/model_filtering/run_filter.py new file mode 100644 index 000000000..214c06946 --- /dev/null +++ b/reasoners/tools/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/reasoners/tools/model_filtering/run_inference.py b/reasoners/tools/model_filtering/run_inference.py new file mode 100644 index 000000000..b2dbe78d3 --- /dev/null +++ b/reasoners/tools/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 model_filtering.utils import console +from 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/reasoners/tools/model_filtering/run_reward.py b/reasoners/tools/model_filtering/run_reward.py new file mode 100644 index 000000000..48298d441 --- /dev/null +++ b/reasoners/tools/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 verl.utils.reward_score import _default_compute_score +from 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/reasoners/tools/model_filtering/run_sample_and_postprocess.py b/reasoners/tools/model_filtering/run_sample_and_postprocess.py new file mode 100644 index 000000000..3240c6fed --- /dev/null +++ b/reasoners/tools/model_filtering/run_sample_and_postprocess.py @@ -0,0 +1,186 @@ +import argparse +import os +import datasets +import json +import transformers +import numpy as np +from verl.utils.data_process.utils import save_dataset +from verl.utils.data_process.filter import LengthFilter + + +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, default="/mnt/weka/home/zhuojun.cheng/leo/Reasoning360/data/train_filtered") + parser.add_argument("--input_data_names", type=str, nargs="+") + parser.add_argument("--output_data_dir", type=str, default="/mnt/weka/home/zhuojun.cheng/leo/Reasoning360/data/train_guru15k") + 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/tools/model_filtering/test_parquet.py b/reasoners/tools/model_filtering/test_parquet.py new file mode 100644 index 000000000..2f7a0c194 --- /dev/null +++ b/reasoners/tools/model_filtering/test_parquet.py @@ -0,0 +1,113 @@ +#!/usr/bin/env python +# coding: utf-8 + +# In[1]: + + +import polars as pl +import pandas as pd +import glob +import os + +paths_and_dfs = {} +for file in glob.glob("/lustrefs/users/zhuojun.cheng/Reasoning360/data/train_guru15k/code*.parquet"): + if "livecodebench" in file.lower(): + df = pl.read_parquet(file).to_pandas() + else: + df = pd.read_parquet(file) + paths_and_dfs[file] = df + + +# In[3]: + + +import json +import matplotlib.pyplot as plt + +for key, df in paths_and_dfs.items(): + print(key) + print(f"Number of rows: {len(df)}") + unit_test_counts, all_unit_test_counts = [], [] + for i in range(len(df)): + ground_truth = json.loads(df['reward_model'][i]['ground_truth']) + if 'inputs' in ground_truth: + assert len(ground_truth['inputs']) == len(ground_truth['outputs']) + unit_test_counts.append(len(ground_truth['inputs'])) + all_unit_test_counts.append(len(ground_truth['inputs'])) + + # distribution of the unit test counts + plt.figure(figsize=(10, 6)) + plt.hist(unit_test_counts, bins=30, edgecolor='black') + plt.title('Distribution of Unit Test Counts') + plt.xlabel('Number of Unit Tests') + plt.ylabel('Frequency') + plt.grid(True, alpha=0.3) + plt.show() + + +# In[4]: + + +import numpy as np + +for key, df in paths_and_dfs.items(): + print(f"\nProcessing: {os.path.basename(key)}") + + # Set seed for reproducibility + np.random.seed(42) + + for idx in range(len(df)): + ground_truth = json.loads(df['reward_model'][idx]['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'] + + # Get 8 random indices from this test case's inputs/outputs + n_tests = len(complete_inputs) + 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 + ground_truth['complete_outputs'] = complete_outputs + ground_truth['inputs'] = [complete_inputs[i] for i in sample_indices] + ground_truth['outputs'] = [complete_outputs[i] for i in sample_indices] + + df.at[idx, 'reward_model']['ground_truth'] = json.dumps(ground_truth) + + # Update the dictionary with modified dataframe + paths_and_dfs[key] = df + + +# In[ ]: + + +import json +import matplotlib.pyplot as plt + +for key, df in paths_and_dfs.items(): + print(key) + print(f"Number of rows: {len(df)}") + unit_test_counts = [] + for i in range(len(df)): + ground_truth = json.loads(df['reward_model'][i]['ground_truth']) + if 'inputs' in ground_truth: + unit_test_counts.append(len(ground_truth['inputs'])) + + # distribution of the unit test counts + plt.figure(figsize=(10, 6)) + plt.hist(unit_test_counts, bins=30, edgecolor='black') + plt.title('Distribution of Unit Test Counts') + plt.xlabel('Number of Unit Tests') + plt.ylabel('Frequency') + plt.grid(True, alpha=0.3) + plt.show() + + +# In[ ]: + + +# save the df back to parquet +for key, df in paths_and_dfs.items(): + df.to_parquet(key) + diff --git a/reasoners/tools/model_filtering/utils.py b/reasoners/tools/model_filtering/utils.py new file mode 100644 index 000000000..0179d6b93 --- /dev/null +++ b/reasoners/tools/model_filtering/utils.py @@ -0,0 +1,464 @@ +#!/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 + +# --------------------------------------------------------------------------- # +# 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 From 4cc0fcd9473ea053c5ce2f8a0d6690299547b900 Mon Sep 17 00:00:00 2001 From: Shibo Hao Date: Thu, 29 May 2025 08:17:35 +0000 Subject: [PATCH 2/5] organize data normalization files --- reasoners/data_preprocess/download_hf.py | 41 +++++ reasoners/data_preprocess/json2parquet.py | 57 ++++++ reasoners/data_preprocess/norm_json.py | 203 ++++++++++++++++++++++ reasoners/data_preprocess/parquet2json.py | 53 ++++++ reasoners/reward/README.md | 1 + reasoners/reward/__init__.py | 12 +- 6 files changed, 361 insertions(+), 6 deletions(-) create mode 100644 reasoners/data_preprocess/download_hf.py create mode 100644 reasoners/data_preprocess/json2parquet.py create mode 100644 reasoners/data_preprocess/norm_json.py create mode 100644 reasoners/data_preprocess/parquet2json.py create mode 100644 reasoners/reward/README.md diff --git a/reasoners/data_preprocess/download_hf.py b/reasoners/data_preprocess/download_hf.py new file mode 100644 index 000000000..1e10ea137 --- /dev/null +++ b/reasoners/data_preprocess/download_hf.py @@ -0,0 +1,41 @@ +from huggingface_hub import hf_hub_download + +def download_all_files(repo_id, file_names, cache_dir=None): + """ + Download all specified files from a Hugging Face dataset repository. + + Args: + repo_id (str): The "namespace/repo_name" of the HF dataset (e.g., "your_user/dataset_name"). + file_names (List[str]): List of filenames to download. + cache_dir (str, optional): Custom cache directory where files will be stored. + """ + for fname in file_names: + local_path = hf_hub_download(repo_id=repo_id, filename=fname, cache_dir=cache_dir, repo_type="dataset") + print(f"✅ Downloaded {fname} → {local_path}") + +if __name__ == "__main__": + # Replace with your actual dataset repo ID + repo_id = "LLM360/guru_preview" + # List all files you want to download + file_names = [ + "codegen__deduped_leetcode2k_1.3k.parquet", + "codegen__deduped_livecodebench_440.parquet", + "codegen__deduped_primeintellect_7.5k.parquet", + "codegen__deduped_taco_8.8k.parquet", + "logic__arcagi1_111.parquet", + "logic__arcagi2_190.parquet", + "logic__barc_1.6k.parquet", + "logic__graph_logical_dataset_1.2k.parquet", + "logic__ordering_puzzle_dataset_1.9k.parquet", + "logic__zebra_puzzle_dataset_1.3k.parquet", + "math__combined_54.4k.parquet", + "simulation__codeio_fixed_12.1k_3.7k.parquet", + "stem__web_3.6k.parquet", + "table__hitab_4.3k.parquet", + "table__multihier_1.5k.parquet" + ] + + # Optionally specify a cache directory, or leave as None to use the default HF cache + cache_dir = "./hf_cache" + + download_all_files(repo_id, file_names, cache_dir) diff --git a/reasoners/data_preprocess/json2parquet.py b/reasoners/data_preprocess/json2parquet.py new file mode 100644 index 000000000..ebac077bd --- /dev/null +++ b/reasoners/data_preprocess/json2parquet.py @@ -0,0 +1,57 @@ +#!/usr/bin/env python3 +""" +Convert all JSON files in a directory into Parquet using Polars, +flattening only the `extra_info` field to top-level columns to avoid +nested chunked arrays. +""" +import sys +from pathlib import Path +import json +import os + +try: + import polars as pl # pip install polars +except ImportError: + print("Error: polars is required. Install via `pip install polars`.") + sys.exit(1) + + +def main(): + # Input and output directories + input_dir = Path( + "/mnt/weka/home/shibo.hao/yutao/Reasoning360/release/data/norm_json" + ).expanduser().resolve() + out_parquet = Path( + "/mnt/weka/home/shibo.hao/yutao/Reasoning360/release/data/parquet" + ).expanduser().resolve() + os.makedirs(out_parquet, exist_ok=True) + + for path in input_dir.glob("*.json"): + if "livecodebench" in path.stem: + continue + # Load normalized JSON list + with path.open("r", encoding="utf-8") as f: + records = json.load(f) + + if not records: + print("No records found in", path) + continue + + # Create Polars DataFrame + df = pl.DataFrame(records) + + print(df.head()) + + # Write out to Parquet with PyArrow backend + out_path = out_parquet / f"{path.stem}.parquet" + df.write_parquet( + out_path, + use_pyarrow=True, + row_group_size=len(records) + ) + + print(f"Saved {df.height} rows to {out_path}") + + +if __name__ == "__main__": + main() diff --git a/reasoners/data_preprocess/norm_json.py b/reasoners/data_preprocess/norm_json.py new file mode 100644 index 000000000..fd43c2902 --- /dev/null +++ b/reasoners/data_preprocess/norm_json.py @@ -0,0 +1,203 @@ +#!/usr/bin/env python3 +""" +Normalise every .json file in a directory so that each record contains +only the allowed keys; anything else is moved under `extra_info` as a +compressed Base64‐encoded JSON string split into 10 parts. + +If a record already has `extra_info` or `reward_model` from a previous run, +we first de-serialise, merge with any new stray keys, then re-serialise, +compress, and Base64-encode. + +Canonical keys: + ability + apply_chat_template + data_source + extra_info (dict → JSON → gzip → Base64 split into parts) + prompt + qwen2.5_7b_pass_rate + qwen3_30b_pass_rate + reward_model (dict → JSON → gzip → Base64) +""" + +from pathlib import Path +from typing import Any, Dict, List, Union +import json +import gzip +import base64 +import sys +import argparse +# --------------------------------------------------------------------------- +# Configuration +# --------------------------------------------------------------------------- + +KEEP_KEYS = { + "ability", + "apply_chat_template", + "data_source", + "extra_info", # handled specially + "prompt", + "qwen2.5_7b_pass_rate", + "qwen3_30b_pass_rate", + "reward_model", +} + +Json = Union[Dict[str, Any], List[Any]] + +# --------------------------------------------------------------------------- +# Serialisation & Compression helpers +# --------------------------------------------------------------------------- + +def _serialise_extra(extra: Dict[str, Any]) -> str: + """Return a compact JSON string for `extra`.""" + return json.dumps(extra, ensure_ascii=False, separators=(',', ':')) + +def _deserialise_extra(extra_str: str) -> Dict[str, Any]: + """Load JSON string if valid, else return empty dict.""" + try: + val = json.loads(extra_str) + return val if isinstance(val, dict) else {} + except Exception: + return {} + +def _compress_str(s: str) -> str: + """Gzip-compress then Base64-encode a JSON string.""" + gz = gzip.compress(s.encode('utf-8')) + return base64.b64encode(gz).decode('ascii') + +def _decompress_str(b64: str) -> str: + """Base64-decode then gzip-decompress to original JSON string.""" + gz = base64.b64decode(b64.encode('ascii')) + return gzip.decompress(gz).decode('utf-8') + +# --------------------------------------------------------------------------- +# JSON/NDJSON loader +# --------------------------------------------------------------------------- + +def load_json_or_ndjson(path: Path) -> Json: + text = path.read_text(encoding='utf-8').strip() + try: + return json.loads(text) + except json.JSONDecodeError: + pass + records: List[Dict[str, Any]] = [] + for line in text.splitlines(): + ln = line.strip() + if ln: + records.append(json.loads(ln)) + return records + +# --------------------------------------------------------------------------- +# Record normalisation (with 10-way split for extra_info) +# --------------------------------------------------------------------------- + +def normalise_record(record: Dict[str, Any]) -> Dict[str, Any]: + # 1) extract existing extra_info dict (if any) + extra_raw = record.get('extra_info', {}) + if isinstance(extra_raw, str): + nested = json.loads(extra_raw).get('compressed', '') + extra = _deserialise_extra(_decompress_str(nested)) + elif isinstance(extra_raw, dict) and 'compressed' in extra_raw: + extra = _deserialise_extra(_decompress_str(extra_raw['compressed'])) + elif isinstance(extra_raw, dict): + extra = extra_raw.copy() + else: + extra = {} + + # 2) extract existing reward_model dict (if any) + rm_raw = record.get('reward_model', {}) + if isinstance(rm_raw, str): + reward_model = _deserialise_extra(rm_raw) + elif isinstance(rm_raw, dict) and 'compressed' in rm_raw: + reward_model = _deserialise_extra(_decompress_str(rm_raw['compressed'])) + elif isinstance(rm_raw, dict): + reward_model = rm_raw.copy() + else: + reward_model = {} + + # 3) copy canonical keys, gather others into extra + out: Dict[str, Any] = {} + for k, v in record.items(): + if k in KEEP_KEYS and k not in ('extra_info', 'reward_model'): + out[k] = v + elif k not in KEEP_KEYS: + extra[k] = v + + # 4) serialise + compress + split extra_info into 10 parts + if extra: + ser = _serialise_extra(extra) + comp = _compress_str(ser) + total_len = len(comp) + chunk_size = total_len // 10 + parts: Dict[str, str] = {} + for i in range(10): + start = i * chunk_size + end = start + chunk_size if i < 9 else total_len + parts[f'part_{i+1}'] = comp[start:end] + out['extra_info'] = parts + + # 5) serialise + compress reward_model if present + if reward_model: + ser_rm = _serialise_extra(reward_model) + comp_rm = _compress_str(ser_rm) + out['reward_model'] = {"ground_truth": {"compressed": comp_rm}} + + return out + +# --------------------------------------------------------------------------- +# Directory processing (parallel) +# --------------------------------------------------------------------------- + +from multiprocessing import Pool, cpu_count + +def _process_file(args): + """ + Worker to normalise one JSON file and write out result. + args: tuple(src_path_str, dst_dir_str) + """ + path_str, dst_dir_str = args + src_path = Path(path_str) + dst_dir = Path(dst_dir_str) + + try: + data = load_json_or_ndjson(src_path) + if isinstance(data, list): + out_data = [normalise_record(obj) for obj in data] + else: + out_data = normalise_record(data) + + out_path = dst_dir / src_path.name + with out_path.open('w', encoding='utf-8') as f: + json.dump(out_data, f, ensure_ascii=False, indent=2) + return f"✔ {src_path.name} → {out_path.name}" + except Exception as e: + return f"✖ {src_path.name} failed: {e}" + +def process_dir_parallel(src_dir: Path, dst_dir: Path, workers: int = None) -> None: + """ + Process all *.json files in src_dir in parallel, writing to dst_dir. + workers: number of parallel worker processes. Defaults to number of CPU cores minus one. + """ + dst_dir.mkdir(parents=True, exist_ok=True) + if workers is None: + workers = max(1, cpu_count() - 1) + + args_list = [(str(path), str(dst_dir)) for path in src_dir.glob('*.json')] + with Pool(processes=workers) as pool: + for result in pool.imap_unordered(_process_file, args_list): + print(result) + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description="Normalize JSON files by moving non-canonical keys to extra_info") + parser.add_argument("src", type=str, help="Source directory containing JSON files") + parser.add_argument("dst", type=str, help="Destination directory for normalized JSON files") + parser.add_argument("-w", "--workers", type=int, help="Number of worker processes (default: CPU count - 1)", default=None) + args = parser.parse_args() + + src = Path(args.src).expanduser().resolve() + dst = Path(args.dst).expanduser().resolve() + + if not src.is_dir(): + print(f"Error: {src} is not a directory") + sys.exit(1) + + process_dir_parallel(src, dst, workers=args.workers) diff --git a/reasoners/data_preprocess/parquet2json.py b/reasoners/data_preprocess/parquet2json.py new file mode 100644 index 000000000..89f969cc7 --- /dev/null +++ b/reasoners/data_preprocess/parquet2json.py @@ -0,0 +1,53 @@ +#!/usr/bin/env python3 +""" +Convert every .parquet in a folder to a single JSON array per file, +using Polars + the standard json module. + +Usage: + python parquet2json_array.py +""" + +from pathlib import Path +import sys +import polars as pl # pip install -U polars +import json + +def parquet_to_json_array(src_dir: Path, dst_dir: Path) -> None: + dst_dir.mkdir(parents=True, exist_ok=True) + + for pq_path in src_dir.glob("*.parquet"): + # 1) Read the Parquet file into a Polars DataFrame + df = pl.read_parquet(pq_path) + + # 2) Convert to a list of dicts (one dict per row) + records = df.to_dicts() + + # 3) Write out as one big JSON array, letting json.dumps call str() on datetimes + out_path = dst_dir / f"{pq_path.stem}.json" + with out_path.open("w", encoding="utf-8") as f: + json.dump( + records, + f, + ensure_ascii=False, + indent=2, + default=str # ← this makes datetime → ISO string + ) + + print(f"✔ {pq_path.name} → {out_path.name}") + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser(description="Convert Parquet files to JSON arrays") + parser.add_argument("src", type=str, help="Source directory containing .parquet files") + parser.add_argument("dst", type=str, help="Destination directory for output .json files") + args = parser.parse_args() + + src = Path(args.src).expanduser().resolve() + dst = Path(args.dst).expanduser().resolve() + + if not src.is_dir(): + print(f"Error: {src} is not a directory") + sys.exit(1) + + parquet_to_json_array(src, dst) diff --git a/reasoners/reward/README.md b/reasoners/reward/README.md new file mode 100644 index 000000000..6422b56e0 --- /dev/null +++ b/reasoners/reward/README.md @@ -0,0 +1 @@ +to be done \ No newline at end of file diff --git a/reasoners/reward/__init__.py b/reasoners/reward/__init__.py index 0f1dc5991..11ed8844a 100644 --- a/reasoners/reward/__init__.py +++ b/reasoners/reward/__init__.py @@ -19,13 +19,13 @@ 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(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["compressed"])) - # reward_metric = extra_info.get("reward_metric", None) - # compressed = False + 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"): From b39ccf749d142fd40637154531391febd7412e6e Mon Sep 17 00:00:00 2001 From: AndreasXie Date: Tue, 3 Jun 2025 20:31:18 +0000 Subject: [PATCH 3/5] remove data preparation --- reasoners/data_preprocess/README.md | 104 - reasoners/data_preprocess/__init__.py | 0 ...iew_filtered_mar21_r1_instruction_style.py | 187 - ...th_preview_filtered_mar21_r1_zero_style.py | 190 - .../archived_for_reference/code_utils.py | 241 - .../archived_for_reference/codegen.py | 977 --- ...deepscaler_preview_r1_zero_style_filter.py | 192 - ...eepscaler_preview_r1_zero_style_sampled.py | 201 - .../archived_for_reference/geo3k.py | 86 - .../archived_for_reference/orz_data.py | 163 - .../orz_data_r1_zero_style.py | 192 - .../data_preprocess/codegen/humaneval.py | 156 - .../data_preprocess/codegen/humanevalplus.py | 157 - .../data_preprocess/codegen/leetcode2k.py | 174 - .../data_preprocess/codegen/livecodebench.py | 230 - reasoners/data_preprocess/codegen/mbpp.py | 169 - reasoners/data_preprocess/codegen/mbppplus.py | 158 - .../data_preprocess/codegen/primeintellect.py | 236 - reasoners/data_preprocess/codegen/taco.py | 246 - reasoners/data_preprocess/download_hf.py | 41 - reasoners/data_preprocess/json2parquet.py | 57 - reasoners/data_preprocess/logic/gen_proof.py | 781 --- .../logic/graph_dataset_gen/core/LICENSE | 201 - .../logic/graph_dataset_gen/core/README.md | 21 - .../logic/graph_dataset_gen/core/array.h | 2429 ------- .../logic/graph_dataset_gen/core/core.h | 479 -- .../logic/graph_dataset_gen/core/io.h | 1496 ----- .../logic/graph_dataset_gen/core/lex.h | 504 -- .../logic/graph_dataset_gen/core/map.h | 2508 -------- .../logic/graph_dataset_gen/core/random.h | 321 - .../logic/graph_dataset_gen/core/timer.h | 130 - .../logic/graph_dataset_gen/core/utility.h | 799 --- .../logic/graph_dataset_gen/core/xxhash.h | 5580 ----------------- .../logic/graph_dataset_gen/generator.cpp | 2151 ------- .../logic/graph_dataset_gen/jobfile.sub | 11 - .../logic/graph_dataset_gen/logic.py | 359 -- .../logic/process_graph_dataset.py | 159 - .../logic/process_puzzles_dataset.py | 198 - .../logic/process_zebrapuzzle_dataset.py | 143 - reasoners/data_preprocess/logic/puzzle_gen.py | 318 - .../logic/zebrapuzzle_gen/generator.py | 488 -- .../logic/zebrapuzzle_gen/puzzle_generator.py | 202 - .../math/bigmath_preview_filtered_mar21.py | 159 - .../math/dapo_or1_merge_dedup_apr30.py | 183 - .../math/deepscaler_preview.py | 143 - reasoners/data_preprocess/norm_json.py | 203 - reasoners/data_preprocess/ood/ifeval.py | 139 - reasoners/data_preprocess/ood/livebench.py | 160 - reasoners/data_preprocess/parquet2json.py | 53 - .../data_preprocess/sft_rlhf/full_hh_rlhf.py | 146 - reasoners/data_preprocess/sft_rlhf/limo.py | 93 - .../data_preprocess/simulation/arcagi.py | 183 - reasoners/data_preprocess/simulation/barc.py | 166 - .../data_preprocess/simulation/codeio.py | 193 - .../data_preprocess/simulation/cruxeval-i.py | 136 - .../data_preprocess/simulation/cruxeval-o.py | 136 - reasoners/data_preprocess/stem/bbeh.py | 0 reasoners/data_preprocess/stem/gpqa.py | 137 - .../data_preprocess/stem/gpqa_diamond.py | 142 - reasoners/data_preprocess/stem/supergpqa.py | 118 - reasoners/data_preprocess/stem_web/README.md | 20 - .../data_preprocess/stem_web/json2_parquet.py | 109 - .../data_preprocess/stem_web/simple_filter.py | 35 - reasoners/data_preprocess/stem_web/stem.py | 133 - reasoners/data_preprocess/table/finqa.py | 170 - reasoners/data_preprocess/table/hitab.py | 263 - reasoners/data_preprocess/table/multihier.py | 274 - 67 files changed, 26929 deletions(-) delete mode 100644 reasoners/data_preprocess/README.md delete mode 100644 reasoners/data_preprocess/__init__.py delete mode 100644 reasoners/data_preprocess/archived_for_reference/bigmath_preview_filtered_mar21_r1_instruction_style.py delete mode 100644 reasoners/data_preprocess/archived_for_reference/bigmath_preview_filtered_mar21_r1_zero_style.py delete mode 100644 reasoners/data_preprocess/archived_for_reference/code_utils.py delete mode 100644 reasoners/data_preprocess/archived_for_reference/codegen.py delete mode 100644 reasoners/data_preprocess/archived_for_reference/deepscaler_preview_r1_zero_style_filter.py delete mode 100644 reasoners/data_preprocess/archived_for_reference/deepscaler_preview_r1_zero_style_sampled.py delete mode 100644 reasoners/data_preprocess/archived_for_reference/geo3k.py delete mode 100644 reasoners/data_preprocess/archived_for_reference/orz_data.py delete mode 100644 reasoners/data_preprocess/archived_for_reference/orz_data_r1_zero_style.py delete mode 100644 reasoners/data_preprocess/codegen/humaneval.py delete mode 100644 reasoners/data_preprocess/codegen/humanevalplus.py delete mode 100644 reasoners/data_preprocess/codegen/leetcode2k.py delete mode 100644 reasoners/data_preprocess/codegen/livecodebench.py delete mode 100644 reasoners/data_preprocess/codegen/mbpp.py delete mode 100644 reasoners/data_preprocess/codegen/mbppplus.py delete mode 100644 reasoners/data_preprocess/codegen/primeintellect.py delete mode 100644 reasoners/data_preprocess/codegen/taco.py delete mode 100644 reasoners/data_preprocess/download_hf.py delete mode 100644 reasoners/data_preprocess/json2parquet.py delete mode 100644 reasoners/data_preprocess/logic/gen_proof.py delete mode 100644 reasoners/data_preprocess/logic/graph_dataset_gen/core/LICENSE delete mode 100644 reasoners/data_preprocess/logic/graph_dataset_gen/core/README.md delete mode 100644 reasoners/data_preprocess/logic/graph_dataset_gen/core/array.h delete mode 100644 reasoners/data_preprocess/logic/graph_dataset_gen/core/core.h delete mode 100644 reasoners/data_preprocess/logic/graph_dataset_gen/core/io.h delete mode 100644 reasoners/data_preprocess/logic/graph_dataset_gen/core/lex.h delete mode 100644 reasoners/data_preprocess/logic/graph_dataset_gen/core/map.h delete mode 100644 reasoners/data_preprocess/logic/graph_dataset_gen/core/random.h delete mode 100644 reasoners/data_preprocess/logic/graph_dataset_gen/core/timer.h delete mode 100644 reasoners/data_preprocess/logic/graph_dataset_gen/core/utility.h delete mode 100644 reasoners/data_preprocess/logic/graph_dataset_gen/core/xxhash.h delete mode 100644 reasoners/data_preprocess/logic/graph_dataset_gen/generator.cpp delete mode 100644 reasoners/data_preprocess/logic/graph_dataset_gen/jobfile.sub delete mode 100644 reasoners/data_preprocess/logic/graph_dataset_gen/logic.py delete mode 100644 reasoners/data_preprocess/logic/process_graph_dataset.py delete mode 100644 reasoners/data_preprocess/logic/process_puzzles_dataset.py delete mode 100644 reasoners/data_preprocess/logic/process_zebrapuzzle_dataset.py delete mode 100644 reasoners/data_preprocess/logic/puzzle_gen.py delete mode 100644 reasoners/data_preprocess/logic/zebrapuzzle_gen/generator.py delete mode 100644 reasoners/data_preprocess/logic/zebrapuzzle_gen/puzzle_generator.py delete mode 100644 reasoners/data_preprocess/math/bigmath_preview_filtered_mar21.py delete mode 100644 reasoners/data_preprocess/math/dapo_or1_merge_dedup_apr30.py delete mode 100644 reasoners/data_preprocess/math/deepscaler_preview.py delete mode 100644 reasoners/data_preprocess/norm_json.py delete mode 100644 reasoners/data_preprocess/ood/ifeval.py delete mode 100644 reasoners/data_preprocess/ood/livebench.py delete mode 100644 reasoners/data_preprocess/parquet2json.py delete mode 100644 reasoners/data_preprocess/sft_rlhf/full_hh_rlhf.py delete mode 100644 reasoners/data_preprocess/sft_rlhf/limo.py delete mode 100644 reasoners/data_preprocess/simulation/arcagi.py delete mode 100644 reasoners/data_preprocess/simulation/barc.py delete mode 100644 reasoners/data_preprocess/simulation/codeio.py delete mode 100644 reasoners/data_preprocess/simulation/cruxeval-i.py delete mode 100644 reasoners/data_preprocess/simulation/cruxeval-o.py delete mode 100644 reasoners/data_preprocess/stem/bbeh.py delete mode 100644 reasoners/data_preprocess/stem/gpqa.py delete mode 100644 reasoners/data_preprocess/stem/gpqa_diamond.py delete mode 100644 reasoners/data_preprocess/stem/supergpqa.py delete mode 100644 reasoners/data_preprocess/stem_web/README.md delete mode 100644 reasoners/data_preprocess/stem_web/json2_parquet.py delete mode 100644 reasoners/data_preprocess/stem_web/simple_filter.py delete mode 100644 reasoners/data_preprocess/stem_web/stem.py delete mode 100644 reasoners/data_preprocess/table/finqa.py delete mode 100644 reasoners/data_preprocess/table/hitab.py delete mode 100644 reasoners/data_preprocess/table/multihier.py diff --git a/reasoners/data_preprocess/README.md b/reasoners/data_preprocess/README.md deleted file mode 100644 index c249a1559..000000000 --- a/reasoners/data_preprocess/README.md +++ /dev/null @@ -1,104 +0,0 @@ -# Guideline - -The processed data will be saved in the `data/train` and `data/test` directories. -The naming convention is `___`, where `` is one of `math`, `codegen`, `logic`, `simulation`, `table` for now. - -# Usage -## Math - -**DAPO+OR1** -```bash -python data_preprocess/math/dapo_or1_merge_dedup_apr30.py -``` -(Please see the comment if want to apply llm as judge) -> Note: -> -> Original: OR1 (105055) + DAPO (17917) -> -> After [dedup](../scripts/tools/cross_dedup.py): 117192 -> -> After merging and removing instance with too long answer >100 char ([script](../scripts/tools/merge_and_upload_dataset.ipynb)): 116632 - -**BigMath** -```bash -python data_preprocess/math/bigmath_preview_filtered_mar21.py --train-sample-size -``` - -**DeepScaler** -```bash -python data_preprocess/math/deepscaler_preview.py --train-sample-size -``` - - -## Code -**leetcode2k** -```bash -python data_preprocess/codegen/leetcode2k.py --train-sample-size -``` -**taco** -```bash -python data_preprocess/codegen/taco.py --train-sample-size -``` -**primeintellect** -```bash -python data_preprocess/codegen/primeintellect.py --train-sample-size -``` -**humaneval** -```bash -python data_preprocess/codegen/humaneval.py -``` -**mbpp** -```bash -python data_preprocess/codegen/mbpp.py -``` -**livecodebench** -```bash -python data_preprocess/codegen/livecodebench.py -``` - -## Logic -**zebra_puzzle_dataset** -```bash -python data_preprocess/logic/zebrapuzzle_gen/puzzle_generator.py --output_dir data/raw --num_puzzles --num_processes -cd .. -python data_preprocess/logic/process_zebrapuzzle_dataset.py -``` - -**graph_logical_dataset** -```bash -uv pip install pybind11 -uv pip install Faker==37.1.0 -cd data_preprocess/logic/graph_dataset_gen/ -python logic.py --num_samples -cd ../../.. # return to Reasoning360 -python data_preprocess/logic/process_graph_dataset.py -``` - -**ordering_puzzle_dataset** -```bash -uv pip install Faker==37.1.0 -python data_preprocess/logic/puzzle_gen.py --test True --num_puzzles -python data_preprocess/logic/process_puzzles_dataset.py -``` - -## Simulation -```bash -python data_preprocess/simulation/codeio.py --train-sample-size --test-sample-size -``` - -## Table -```bash -uv pip install gdown -python data_preprocess/table/multihier.py -``` - -## ARC-AGI -```bash -python data_preprocess/simulation/arcagi.py --name arcagi1 -python data_preprocess/simulation/arcagi.py --name arcagi2 -``` - -# Add a new dataset -1. Add a new script in `data_preprocess//.py` -2. Add a new entry in `tests/data_process/test_data_preprocess.py`. -3. Run `pytest tests/data_process` to check the functionality of the data preprocessing scripts. diff --git a/reasoners/data_preprocess/__init__.py b/reasoners/data_preprocess/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/reasoners/data_preprocess/archived_for_reference/bigmath_preview_filtered_mar21_r1_instruction_style.py b/reasoners/data_preprocess/archived_for_reference/bigmath_preview_filtered_mar21_r1_instruction_style.py deleted file mode 100644 index ad0cb66a5..000000000 --- a/reasoners/data_preprocess/archived_for_reference/bigmath_preview_filtered_mar21_r1_instruction_style.py +++ /dev/null @@ -1,187 +0,0 @@ -# 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. -""" -Preprocess the https://huggingface.co/datasets/agentica-org/DeepScaleR-Preview-Dataset to parquet format -Code referring to https://github.com/agentica-project/deepscaler/blob/main/scripts/data/deepscaler_dataset.py -""" - -import os -import datasets -from typing import Dict, List, Optional, Any, Union -import enum -import argparse -import pandas as pd -import json - -from verl.utils.hdfs_io import copy, makedirs -from verl.utils.reward_score.math import remove_boxed, last_boxed_only_string - - -class TrainDataset(enum.Enum): - """Enum for training datasets. - - Contains identifiers for various math problem datasets used during training. - """ - - AIME = "AIME" # American Invitational Mathematics Examination - AMC = "AMC" # American Mathematics Competition - OMNI_MATH = "OMNI_MATH" # Omni Math - NUMINA_OLYMPIAD = "OLYMPIAD" # Unique Olympiad problems from NUMINA - MATH = "MATH" # Dan Hendrycks Math Problems - STILL = "STILL" # STILL dataset - DEEPSCALER = "DEEPSCALER" # DeepScaler (AIME, AMC, OMNI_MATH, MATH, STILL) - - -class TestDataset(enum.Enum): - """Enum for testing/evaluation datasets. - - Contains identifiers for datasets used to evaluate model performance. - """ - - AIME = "AIME" # American Invitational Mathematics Examination - AMC = "AMC" # American Mathematics Competition - MATH = "MATH" # Math 500 problems - MINERVA = "MINERVA" # Minerva dataset - OLYMPIAD_BENCH = "OLYMPIAD_BENCH" # Olympiad benchmark problems - - -"""Type alias for either training or testing dataset types.""" -Dataset = Union[TrainDataset, TestDataset] - - -def load_dataset(dataset: Dataset) -> List[Dict[str, Any]]: - """Load a dataset from a JSON file. - - Loads and parses a JSON dataset file based on the provided dataset enum. - The file path is constructed based on whether it's a training or testing dataset. - - Args: - dataset: A Dataset enum value specifying which dataset to load. - - Returns: - List[Dict[str, Any]]: A list of dictionaries containing the dataset records. - Each dictionary represents one example in the dataset. - - Raises: - ValueError: If the dataset file cannot be found, contains invalid JSON, - or encounters other file access errors. - - Example: - >>> load_dataset(TrainDataset.AIME) - [{'problem': 'Find x...', 'solution': '42', ...}, ...] - """ - dataset_name = dataset.value.lower() - data_dir = "train" if isinstance(dataset, TrainDataset) else "test" - - current_dir = os.path.dirname(os.path.realpath(__file__)) - file_path = os.path.join(data_dir, f"{dataset_name}.json") - file_path = os.path.join(current_dir, file_path) - - if not os.path.exists(file_path): - raise ValueError(f"Dataset file not found: {file_path}") - - try: - with open(file_path, "r", encoding="utf-8") as file: - data = json.load(file) - return data - except json.JSONDecodeError: - raise ValueError(f"Invalid JSON format in {file_path}") - except Exception as exc: # pylint: disable=broad-except - raise ValueError(f"Error loading dataset: {exc}") from exc - - -def extract_solution(solution_str): - return remove_boxed(last_boxed_only_string(solution_str)) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("--local_dir", default="data/bigmath_preview_filtered_mar21_r1_instruction_style") - parser.add_argument("--hdfs_dir", default=None) - - args = parser.parse_args() - - train_data_source = "SDSB/big_math_partial_mar21_filtered_basic" - test_data_sources = [ - "nanoverl/minerva", - "SDSB/aime_repeated_8x", - "SDSB/amc_repeated_4x", - "nanoverl/olympiad_bench", - "nanoverl/math", - ] - train_dataset = datasets.load_dataset( - train_data_source, trust_remote_code=True, split="train" - ) - print(f"Loading the {train_data_source} dataset from huggingface...", flush=True) - test_datasets = [ - datasets.load_dataset(test_data_source, trust_remote_code=True, split="test") - for test_data_source in test_data_sources - ] - - print(f"Loading the {test_data_sources} dataset from huggingface...", flush=True) - instruction_following = ( - "Please output the final answer within \\boxed{}." - ) - - - # add a row to each data item that represents a unique id - def make_map_fn(split, data_source): - - def process_fn(example, idx): - question = example.pop("problem") - - question = question + " " + instruction_following - - answer = example.pop("answer") - data = { - "data_source": data_source, - "prompt": [{ - "role": "user", - "content": question - }], - "ability": "math", - "apply_chat_template": True, - "reward_model": {"style": "rule", "ground_truth": answer}, - "extra_info": {"split": split, - "index": idx, - "question": question}, - } - if idx == 0: - print("=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) - print(data) - return data - - return process_fn - - local_dir = args.local_dir - hdfs_dir = args.hdfs_dir - - train_data = train_dataset.map( - function=make_map_fn("train", train_data_source), with_indices=True - ) - train_data.to_parquet(os.path.join(local_dir, "train.parquet")) - print(f"train data size:", len(train_data)) - print(train_data[0]) - for test_data_source, test_data in zip(test_data_sources, test_datasets): - process_fn = make_map_fn("test", test_data_source) - test_data = test_data.map(process_fn, with_indices=True) - dataset_name = os.path.basename(test_data_source.lower()) - test_df = pd.DataFrame(test_data) - test_df.to_parquet(os.path.join(local_dir, f"{dataset_name}.parquet")) - print(f"test data size: ({dataset_name})", len(test_df)) - - if hdfs_dir is not None: - makedirs(hdfs_dir) - - copy(src=local_dir, dst=hdfs_dir) diff --git a/reasoners/data_preprocess/archived_for_reference/bigmath_preview_filtered_mar21_r1_zero_style.py b/reasoners/data_preprocess/archived_for_reference/bigmath_preview_filtered_mar21_r1_zero_style.py deleted file mode 100644 index 535b405e5..000000000 --- a/reasoners/data_preprocess/archived_for_reference/bigmath_preview_filtered_mar21_r1_zero_style.py +++ /dev/null @@ -1,190 +0,0 @@ -# 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. -""" -Preprocess the https://huggingface.co/datasets/agentica-org/DeepScaleR-Preview-Dataset to parquet format -Code referring to https://github.com/agentica-project/deepscaler/blob/main/scripts/data/deepscaler_dataset.py -""" - -import os -import datasets -from typing import Dict, List, Optional, Any, Union -import enum -import argparse -import pandas as pd -import json - -from verl.utils.hdfs_io import copy, makedirs -from verl.utils.reward_score.math import remove_boxed, last_boxed_only_string - - -class TrainDataset(enum.Enum): - """Enum for training datasets. - - Contains identifiers for various math problem datasets used during training. - """ - - AIME = "AIME" # American Invitational Mathematics Examination - AMC = "AMC" # American Mathematics Competition - OMNI_MATH = "OMNI_MATH" # Omni Math - NUMINA_OLYMPIAD = "OLYMPIAD" # Unique Olympiad problems from NUMINA - MATH = "MATH" # Dan Hendrycks Math Problems - STILL = "STILL" # STILL dataset - DEEPSCALER = "DEEPSCALER" # DeepScaler (AIME, AMC, OMNI_MATH, MATH, STILL) - - -class TestDataset(enum.Enum): - """Enum for testing/evaluation datasets. - - Contains identifiers for datasets used to evaluate model performance. - """ - - AIME = "AIME" # American Invitational Mathematics Examination - AMC = "AMC" # American Mathematics Competition - MATH = "MATH" # Math 500 problems - MINERVA = "MINERVA" # Minerva dataset - OLYMPIAD_BENCH = "OLYMPIAD_BENCH" # Olympiad benchmark problems - - -"""Type alias for either training or testing dataset types.""" -Dataset = Union[TrainDataset, TestDataset] - - -def load_dataset(dataset: Dataset) -> List[Dict[str, Any]]: - """Load a dataset from a JSON file. - - Loads and parses a JSON dataset file based on the provided dataset enum. - The file path is constructed based on whether it's a training or testing dataset. - - Args: - dataset: A Dataset enum value specifying which dataset to load. - - Returns: - List[Dict[str, Any]]: A list of dictionaries containing the dataset records. - Each dictionary represents one example in the dataset. - - Raises: - ValueError: If the dataset file cannot be found, contains invalid JSON, - or encounters other file access errors. - - Example: - >>> load_dataset(TrainDataset.AIME) - [{'problem': 'Find x...', 'solution': '42', ...}, ...] - """ - dataset_name = dataset.value.lower() - data_dir = "train" if isinstance(dataset, TrainDataset) else "test" - - current_dir = os.path.dirname(os.path.realpath(__file__)) - file_path = os.path.join(data_dir, f"{dataset_name}.json") - file_path = os.path.join(current_dir, file_path) - - if not os.path.exists(file_path): - raise ValueError(f"Dataset file not found: {file_path}") - - try: - with open(file_path, "r", encoding="utf-8") as file: - data = json.load(file) - return data - except json.JSONDecodeError: - raise ValueError(f"Invalid JSON format in {file_path}") - except Exception as exc: # pylint: disable=broad-except - raise ValueError(f"Error loading dataset: {exc}") from exc - - -def extract_solution(solution_str): - return remove_boxed(last_boxed_only_string(solution_str)) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument( - "--local_dir", default="data/bigmath_preview_filtered_mar21_r1_zero_style" - ) - parser.add_argument("--hdfs_dir", default=None) - - args = parser.parse_args() - - train_data_source = "SDSB/big_math_partial_mar21_filtered_basic" - test_data_sources = [ - "nanoverl/minerva", - "SDSB/aime_repeated_8x", - "SDSB/amc_repeated_4x", - "nanoverl/olympiad_bench", - "nanoverl/math", - ] - train_dataset = datasets.load_dataset( - train_data_source, trust_remote_code=True, split="train" - ) - print(f"Loading the {train_data_source} dataset from huggingface...", flush=True) - test_datasets = [ - datasets.load_dataset(test_data_source, trust_remote_code=True, split="test") - for test_data_source in test_data_sources - ] - - print(f"Loading the {test_data_sources} dataset from huggingface...", flush=True) - instruction_following = ( - "Let's think step by step and output the final answer within \\boxed{}." - ) - - prompt = """A conversation between User and Assistant. The user asks a question, and the Assistant solves it. The assistant first thinks about the reasoning process in the mind and then provides the user with the response. The reasoning process is enclosed within i.e., reasoning process here respond to the user's question here. - -User: {{prompt}} Please put your answer in \\boxed{} tags. -Assistant: -""" - - # add a row to each data item that represents a unique id - def make_map_fn(split, data_source): - - def process_fn(example, idx): - question = example.pop("problem") - - # question = question + " " + instruction_following - - answer = example.pop("answer") - data = { - "data_source": data_source, - "prompt": [], - "raw_prompt": prompt.replace("{{prompt}}", question), - "ability": "math", - "apply_chat_template": False, - "reward_model": {"style": "rule", "ground_truth": answer}, - "extra_info": {"split": split, "index": idx, "question": question}, - } - if idx == 0: - print("=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) - print(data) - return data - - return process_fn - - local_dir = args.local_dir - hdfs_dir = args.hdfs_dir - - train_data = train_dataset.map( - function=make_map_fn("train", train_data_source), with_indices=True - ) - train_data.to_parquet(os.path.join(local_dir, "train.parquet")) - print(f"train data size:", len(train_data)) - print(train_data[0]) - for test_data_source, test_data in zip(test_data_sources, test_datasets): - process_fn = make_map_fn("test", test_data_source) - test_data = test_data.map(process_fn, with_indices=True) - dataset_name = os.path.basename(test_data_source.lower()) - test_df = pd.DataFrame(test_data) - test_df.to_parquet(os.path.join(local_dir, f"{dataset_name}.parquet")) - print(f"test data size: ({dataset_name})", len(test_df)) - - if hdfs_dir is not None: - makedirs(hdfs_dir) - - copy(src=local_dir, dst=hdfs_dir) diff --git a/reasoners/data_preprocess/archived_for_reference/code_utils.py b/reasoners/data_preprocess/archived_for_reference/code_utils.py deleted file mode 100644 index 8774e39e2..000000000 --- a/reasoners/data_preprocess/archived_for_reference/code_utils.py +++ /dev/null @@ -1,241 +0,0 @@ -""" -Code utils for data preprocessing -""" - -import sys -import pandas as pd -import matplotlib.pyplot as plt - -N_TESTSET_PER_DATASET = 512 # per dataset -MAX_N_TRAINSET_PER_DATASET = 50000 -EMPTY_RETURN = { - "data_source": None, - "prompt": None, - "ability": None, - "reward_model": None, - "extra_info": None -} -SYSTEM_PROMPT = """You are a helpful programming assistant. \ -The user will ask you a question and you as the assistant solve it. \ -The assistant first thinks how to solve the task through reasoning and then provides the user with the final answer. \ -The reasoning process and answer are enclosed within ... and ... tags, respectively.""" -PY_IMPORTS = "import heapq\nfrom math import floor, gcd\nimport random\nimport sys\nfrom typing import *\nfrom functools import *\nimport collections\nfrom collections import *\nfrom itertools import *\nfrom heapq import *\nfrom bisect import *\nfrom string import *\nimport math\nimport datetime\ninf = float('inf')\n" -BLOCK_LIBS = [ - # "requests", - # "subprocess", - # "multiprocessing", - # "threading", - # "os", - # "sys", - "fake-useragent", - "keras", - "socket", - "torch", - "scipy", - "sklearn", - "cv2", - "scipy", - "imageio", - "sphinx-pyproject", - "xgboost", - "tweepy", - "flask", - "matplotlib", - "pillow", - "seaborn", - "smtpd", - "docker", - "psutil", - "paramiko", - "pickle", - "shutil", - "telnetlib", - "ftplib", - "pyautogui", - "pexpect", - "asyncssh", - "cryptography", - "pycrypto", - "scapy", - "nmap", - "pwntools", - "webbrowser", - "ctypes", - "tempfile", -] -BLOCK_USAGES = [ - "open(", - "exec(", - "eval(", - "system(", - "popen(", - "importlib", - "globals(", - "locals(", - "breakpoint(", - "__builtins__", - "compile(", - "fork(", - "chmod(", -] -MAX_PROMPT_LENGTH = 2048 - - -def minimize_stdio(inputs, outputs, max_n_tests=8): - stdin_list = [] - stdout_list = [] - for stdin, stdout in zip(inputs, outputs): - if isinstance(stdin, list): - stdin = "\n".join(stdin) - if isinstance(stdout, list): - stdout = "\n".join(stdout) - if sys.getsizeof(stdin) > 4 * 1024: - continue - stdout.replace("\r\n", "\n") - stdin_list.append(stdin) - stdout_list.append(stdout) - - zipped = sorted(zip(stdin_list, stdout_list), key=lambda x: sys.getsizeof(x[0])) - - if not zipped: - print("No tests found!") - return [], [] - - sorted_stdin, sorted_stdout = zip(*zipped) - return list(sorted_stdin[:max_n_tests]), list(sorted_stdout[:max_n_tests]) - - -# Plot the filter counts -def plot_hist( - data_dict, - file_path, - title, - xlabel, - ylabel, - figsize=(10, 6), - plot_type="bar", - color=None, - sort_by=None, - max_categories=None, - rotation=45, - dpi=300, - log_scale=False, - grid=True, - font_size=10, -): - """ - Plot the histogram of filter counts from a dictionary. - - Args: - data_dict (dict): Dictionary mapping categories to counts - file_path (str): Path to save the output figure - title (str): Title of the plot - xlabel (str): Label for x-axis - ylabel (str): Label for y-axis - figsize (tuple): Figure size in inches (width, height) - plot_type (str): Type of plot ('bar', 'barh', 'pie') - color (str or list): Color or color map for the plot - sort_by (str): How to sort categories ('key', 'value', 'value_desc', None) - max_categories (int): Maximum number of categories to display (None for all) - rotation (int): Rotation angle for x-tick labels - dpi (int): DPI for the output figure - log_scale (bool): Use log scale for y-axis - grid (bool): Show grid lines - font_size (int): Base font size for the plot - - Returns: - None: The figure is saved to file_path - """ - try: - plt.figure(figsize=figsize) - - # Handle empty dictionary - if not data_dict: - plt.text( - 0.5, - 0.5, - "No data available", - ha="center", - va="center", - fontsize=font_size + 4, - ) - plt.savefig(file_path, dpi=dpi, bbox_inches="tight") - plt.close() - return - - # Process data - items = list(data_dict.items()) - - # Sort data if requested - if sort_by == "key": - items.sort(key=lambda x: x[0]) - elif sort_by == "value": - items.sort(key=lambda x: x[1]) - elif sort_by == "value_desc": - items.sort(key=lambda x: x[1], reverse=True) - - # Limit number of categories if requested - if max_categories and len(items) > max_categories: - items = items[: max_categories - 1] - # Add "Other" category if needed - if max_categories < len(data_dict): - other_sum = sum(v for k, v in data_dict.items() if k not in dict(items)) - items.append(("Other", other_sum)) - - labels, counts = zip(*items) if items else ([], []) - - # Create the appropriate plot - if plot_type == "bar": - plt.bar(labels, counts, color=color) - if log_scale and any(counts): - plt.yscale("log") - plt.xticks(rotation=rotation, ha="right", fontsize=font_size) - plt.yticks(fontsize=font_size) - elif plot_type == "barh": - plt.barh(labels, counts, color=color) - if log_scale and any(counts): - plt.xscale("log") - plt.yticks(fontsize=font_size) - plt.xticks(fontsize=font_size) - elif plot_type == "pie": - plt.pie( - counts, labels=labels, autopct="%1.1f%%", startangle=90, colors=color - ) - plt.axis("equal") - - # Set labels and title - plt.xlabel(xlabel, fontsize=font_size + 2) - plt.ylabel(ylabel, fontsize=font_size + 2) - plt.title(title, fontsize=font_size + 4) - - if grid and plot_type != "pie": - plt.grid(alpha=0.3) - - plt.tight_layout() - - # Ensure directory exists - import os - - os.makedirs(os.path.dirname(os.path.abspath(file_path)), exist_ok=True) - - # Save figure - plt.savefig(file_path, dpi=dpi, bbox_inches="tight") - plt.close() - - except Exception as e: - import logging - - logging.error(f"Error plotting histogram: {str(e)}") - # Create a minimal error plot - plt.figure(figsize=(8, 6)) - plt.text( - 0.5, - 0.5, - f"Error creating plot: {str(e)}", - ha="center", - va="center", - color="red", - ) - plt.tight_layout() - plt.savefig(file_path, dpi=dpi) - plt.close() diff --git a/reasoners/data_preprocess/archived_for_reference/codegen.py b/reasoners/data_preprocess/archived_for_reference/codegen.py deleted file mode 100644 index 0e6fe4c50..000000000 --- a/reasoners/data_preprocess/archived_for_reference/codegen.py +++ /dev/null @@ -1,977 +0,0 @@ -""" -Preprocess LeetCode problems (newfacade/LeetCodeDataset) to parquet format. -Thanks to https://github.com/ganler/code-r1/blob/main/examples/data_preprocess/coder1.py -""" - -import os -import json -import sys -from concurrent.futures import ThreadPoolExecutor, as_completed -from transformers import AutoTokenizer -from datasets import load_dataset, concatenate_datasets -from rich.rule import Rule - -import rich -import matplotlib.pyplot as plt -import datasets - -from verl.utils.hdfs_io import copy, makedirs -from verl.utils.reward_score.coder1 import ( - code_exec, - remote_check_stdio, - _ERROR_MSG_PREFIX, - extract_code_from_string, - fuzzy_equal -) - -# from examples.data_preprocess.code.code_utils import * - -WORKDING_DIR = os.path.join(os.environ.get("HOME"), "Reasoning360") - -SYSTEM_PROMPT = """You are a helpful programming assistant. \ -The user will ask you a question and you as the assistant solve it. \ -The assistant first thinks how to solve the task through reasoning and then provides the user with the final answer. \ -The reasoning process and answer are enclosed within ... and ... tags, respectively.""" -EMPTY_RETURN = { - "data_source": None, - "prompt": None, - "ability": None, - "reward_model": None, - "extra_info": None -} -N_TESTSET_PER_DATASET = 0 - -def minimize_stdio(inputs, outputs, max_n_tests=8): - stdin_list = [] - stdout_list = [] - for stdin, stdout in zip(inputs, outputs): - if isinstance(stdin, list): - stdin = "\n".join(stdin) - if isinstance(stdout, list): - stdout = "\n".join(stdout) - if sys.getsizeof(stdin) > 4 * 1024: - continue - stdout.replace("\r\n", "\n") - stdin_list.append(stdin) - stdout_list.append(stdout) - - zipped = sorted(zip(stdin_list, stdout_list), key=lambda x: sys.getsizeof(x[0])) - - if not zipped: - print("No tests found!") - return [], [] - - sorted_stdin, sorted_stdout = zip(*zipped) - return list(sorted_stdin[:max_n_tests]), list(sorted_stdout[:max_n_tests]) - - -def kodcode(): # Thanks!!! to Zhangchen and Yueqin - # library requirements? - rich.print(Rule("Loading KodCode/KodCode-Light-RL-10K...")) - dataset = load_dataset("KodCode/KodCode-Light-RL-10K") - tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-32B-Instruct") - - def make_map_fn(split): - - def process_fn(example, idx): - reference_solution = example["solution"] - test_code = "from solution import *\n" + example["test"].strip() - # Filtering... - # + block libs are used in reference solution and test code - # + block usages are used in reference solution or test code - # + filter out too long prompts - # + filter out easy problems - # + filter out failed unittests - filter_info = {} - for lib in BLOCK_LIBS: - if ( - f"import {lib}" - in reference_solution # naive import detection; ast then detect would be better - or f"from {lib}" in reference_solution - ): - print("===========Blocked lib in solution===========") - print(f"reference_solution:") - print(reference_solution) - print(f"lib: {lib}") - print(f"question_id: {example['question_id']}") - return { - **EMPTY_RETURN, - "filter_info": { - "type": f"blocked_lib", - "detail": lib, - }, - } - for lib in BLOCK_LIBS: - if f"import {lib}" in test_code or f"from {lib}" in test_code: - print("===========Blocked lib in test===========") - print(f"test_code:") - print(test_code) - print(f"lib: {lib}") - print(f"question_id: {example['question_id']}") - return { - **EMPTY_RETURN, - "filter_info": { - "type": f"blocked_lib", - "detail": lib, - }, - } - for usage in BLOCK_USAGES: - if usage in reference_solution: - return { - **EMPTY_RETURN, - "filter_info": { - "type": f"blocked_usage", - "detail": usage, - }, - } - for usage in BLOCK_USAGES: - if usage in test_code: - return { - **EMPTY_RETURN, - "filter_info": { - "type": f"blocked_usage", - "detail": usage, - }, - } - if ( - len(tokenizer.encode(example["question"])) > MAX_PROMPT_LENGTH - 200 - ): # -200 for (approximately) the prompt template extra tokens - return { - **EMPTY_RETURN, - "filter_info": {"type": "prompt_too_long", "detail": None}, - } - if example["gpt_difficulty"] == "easy": - return { - **EMPTY_RETURN, - "filter_info": {"type": "easy_problem", "detail": None}, - } - succ, err = code_exec(code=reference_solution, pytest=test_code) - if not succ: - # The above code is using the `rich` library in Python to print a formatted message in the console. - # The message is in red color and includes the value of `example['conversation_id']`. - # rich.print( - # f"[bold red]Test code failed for {example['question_id']}" - # ) - print("===========Unittest failed===========") - print(f"reference_solution:") - print(reference_solution) - print(f"test_code:") - print(test_code) - print(f"err:") - print(err) - return { - **EMPTY_RETURN, - "filter_info": {"type": "failed_unittests", "detail": None}, - } - - prompt = f"Please solve the programming task below in Python. Code should wrapped in a markdown code block.\n\n{example['question'].strip()}" - if example["test_info"]: - prompt += f"\n\nNote that the output function should be {str(example['test_info']).strip()}." - - return { - "data_source": "codegen-kodcode", - "prompt": [ - {"role": "system", "content": SYSTEM_PROMPT}, - {"role": "user", "content": prompt}, - ], - "ability": "codegen", - "reward_model": { - "style": "rule", - "ground_truth": json.dumps({"pytest": test_code}), - }, - "extra_info": { - "split": split, - "index": idx, - "reference": reference_solution, - "prompt": prompt, - "dataset": "KodCode/KodCode-Light-RL-10K", - "question_subset": example["subset"], - "question_id": example["question_id"], - "gpt_difficulty": example["gpt_difficulty"], - }, - "filter_info": None, - } - - return process_fn - - dataset = dataset["train"].shuffle(seed=666) - - # Preprocess the dataset - print("Executing tests to ensure correctness...") - dataset = dataset.map( - function=make_map_fn("train"), - with_indices=True, - num_proc=64, - load_from_cache_file=False, - ) - # Analyze the filter reasons - filter_counts, filter_counts_fine_block_libs = {}, {} - for entry in dataset: - if entry["filter_info"] is None: - continue - filter_type = entry["filter_info"]["type"] - if filter_type is not None: - # Filter type distribution - if filter_type not in filter_counts: - filter_counts[filter_type] = 0 - filter_counts[filter_type] += 1 - # Filter detail distribution - filter_detail = entry["filter_info"].get("detail", None) - if filter_detail is not None: - if filter_detail not in filter_counts_fine_block_libs: - filter_counts_fine_block_libs[filter_detail] = 0 - filter_counts_fine_block_libs[filter_detail] += 1 - # entry["filter_info"] = None - - print(f"Filtered samples from KodCode: {filter_counts}") - - plot_hist( - filter_counts, - file_path=os.path.join(WORKDING_DIR, "artifacts", "filter_counts.png"), - title="Filter Sample Distribution from KodCode", - xlabel="Filter Reason", - ylabel="Count", - ) - plot_hist( - filter_counts_fine_block_libs, - file_path=os.path.join( - WORKDING_DIR, "artifacts", "filter_counts_fine_block_libs.png" - ), - title="Blocked Library Distribution from KodCode", - xlabel="Blocked Library", - ylabel="Count", - ) - - print(f"Before filtering, KodCode dataset size: {len(dataset)}") - - dataset = dataset.filter(lambda x: x["data_source"] is not None) - print(f"Remaining samples from KodCode: {len(dataset)}") - - # pick random 50k examples for RL, otherwise it's too large - # dataset = dataset.select(range(50000 + N_TESTSET_PER_DATASET)) - - # Split into train and test - # splits = dataset["train"].train_test_split( - # test_size=N_TESTSET_PER_DATASET, seed=666 - # ) - splits = dataset.train_test_split(test_size=N_TESTSET_PER_DATASET, seed=666) - train_dataset = splits["train"].shuffle(seed=666) - test_dataset = splits["test"] - return train_dataset, test_dataset - - -# this dataset is super noisy and needs code execution to verify the tasks -def taco(): - rich.print(Rule("Loading likaixin/TACO-verified...")) - dataset = load_dataset("likaixin/TACO-verified")["train"] - - - # add a row to each data item that represents a unique id - def make_map_fn(split): - - def process_fn(example, idx): - oracle = json.loads(example["input_output"]) - source = example["source"] - - # skip poorly formatted examples - if source in ["geeksforgeeks", "leetcode"]: - return EMPTY_RETURN - - # too short description - if len("".join([c for c in example["question"] if c.isalnum()])) < 100: - return EMPTY_RETURN - - # no image - if "image" in example["question"].lower() or "\n![" in example["question"]: - return EMPTY_RETURN - - prompt_pieces = [ - "Solve the programming task below in a Python markdown code block.", - example["question"].strip(), - ] - if example["starter_code"].strip(): - prompt_pieces.append( - "Also feel free to reuse/extend the following starter code:" - ) - prompt_pieces.append( - f"```python\n{example['starter_code'].strip()}\n```" - ) - - ## - ## Customization - ## - if "fn_name" in oracle: # the dataset is too noisy - fn_name = oracle["fn_name"] - if source == "leetcode": - fn_name = "Solution()." + fn_name - - test_code = f"""\ -_inputs = {oracle["inputs"]} -_outputs = {oracle["outputs"]} -import math -def _deep_eq(a, b, tol=1e-5): - if isinstance(a, float) or isinstance(b, float): - return math.isclose(a, b, rel_tol=tol, abs_tol=tol) - if isinstance(a, (list, tuple)): - if len(a) != len(b): return False - return all(_deep_eq(x, y, tol) for x, y in zip(a, b)) - return a == b - -for i, o in zip(_inputs, _outputs): -""" - - if source in ["leetcode", "hackerrank"]: - test_code += f" assert _deep_eq({fn_name}(*i), o)" - elif source == "codewars": - test_code += f" assert _deep_eq({fn_name}(*i), o[0])" - else: - raise ValueError(f"Unknown source: {source}") - - _check_test = example["solutions"][-1] + "\n" + test_code - if source in ["leetcode"]: - _check_test = PY_IMPORTS + _check_test - - succ, err = code_exec(_check_test) - if not succ: - rich.print(f"[bold red]Test code failed for {source}") - print(_check_test) - print(err) - return EMPTY_RETURN - oracle = json.dumps({"functional": test_code}) - assert example["starter_code"].strip() != "" - elif "inputs" in oracle and "outputs" in oracle: - stdin_list, stdout_list = minimize_stdio( - oracle["inputs"], oracle["outputs"] - ) - if len(stdin_list) == 0: - return EMPTY_RETURN - - with ThreadPoolExecutor( - max_workers=min(len(stdin_list), 8) - ) as executor: - futures = [] - for stdin, stdout in zip(stdin_list, stdout_list): - futures.append( - executor.submit( - remote_check_stdio, - example["solutions"][-1], - stdin, - stdout, - ) - ) - for future in as_completed(futures): - exec_succ, output, stdin, stdout = future.result() - pass_test = exec_succ and fuzzy_equal(output.strip(), stdout.strip()) - if not pass_test: - rich.print(f"[bold red]Test code failed for {source}") - print(example["solutions"][-1]) - print(f"{exec_succ = }") - print(f"{stdin = }", f"{stdout = }") - if output.startswith(_ERROR_MSG_PREFIX): - print("output = \n", output) - else: - print(f"{output = }") - return EMPTY_RETURN - - oracle = json.dumps({"inputs": stdin_list, "outputs": stdout_list}) - else: - raise ValueError(f"Unknown ground truth format: {oracle}") - - prompt = "\n".join(prompt_pieces) - return { - "data_source": "codegen-taco", - "prompt": [ - {"role": "system", "content": SYSTEM_PROMPT}, - {"role": "user", "content": prompt}, - ], - "ability": "codegen", - "reward_model": { - "style": "rule", - "ground_truth": oracle, - }, - "extra_info": { - "split": split, - "index": idx, - "prompt": prompt, - "reference": ( - example["solutions"][0] if example["solutions"] else "" - ), - "dataset": "likaixin/TACO-verified", - }, - } - - return process_fn - - dataset = dataset.map( - function=make_map_fn("train"), - with_indices=True, - num_proc=64, - remove_columns=dataset.column_names, - load_from_cache_file=False, - ).filter(lambda x: x['reward_model'] is not None) - - if N_TESTSET_PER_DATASET == 0: - # Return no test samples if N_TESTSET_PER_DATASET is 0 - test_dataset = datasets.Dataset.from_dict({ - "data_source": [], - "prompt": [], - "ability": [], - "reward_model": [], - "extra_info": [] - }) - train_dataset = dataset - else: - splits = dataset.train_test_split( - test_size=min(N_TESTSET_PER_DATASET, len(dataset) * 0.1), seed=666 - ) - train_dataset = splits["train"] - test_dataset = splits["test"] - - print(f"Taco train set: {train_dataset}") - print(f"Taco test set: {test_dataset}") - - return train_dataset, test_dataset - - -def leetcode2k(): - rich.print(Rule("Loading LeetCodeDataset...")) - commit_hash = "34803eb64eab1979069ba1f80e7ea474282e28f3" - train_dataset = load_dataset("newfacade/LeetCodeDataset", - revision=commit_hash, - cache_dir=cache_dir)["train"] - test_dataset = load_dataset("newfacade/LeetCodeDataset", - revision=commit_hash, - cache_dir=cache_dir)["test"] - print("Train set:", train_dataset) - print("Test set:", test_dataset) - - - # add a row to each data item that represents a unique id - def make_map_fn(split): - - def process_fn(example, idx): - prefix = example["prompt"] - - prompt = example["query"] - # remove the "### Answer: (use the provided format with backticks)" part - prompt = prompt.replace("### Answer: (use the provided format with backticks)", "").strip() - # adjust the "### Format: " part to be more readable - prompt = prompt.replace("### Format: ", "### Format:\n") - - # Build test code (as before) - test_code = f"{example['test']}\n\ncheck({example['entry_point'].strip()})" - # Extract the candidate solution (the original completion) - solution = example["completion"] - - # Combine all code pieces into a single file to execute. - full_code = f"{prefix}\n{solution}\n{test_code}" - - # Validate that the candidate solution passes the tests - # 20s timeout as some leetcode tests are slow - succ, err = code_exec(full_code, timeout=20) - - if not succ: - print("===========Test code failed for LeetCodeDataset===========") - print("Question:", example["meta"]["question_title"]) - print("Error:", err) - # Skip the example if the test code fails - return EMPTY_RETURN - - return { - "data_source": "codegen-leetcode2k", - "prompt": [ - {"role": "system", "content": SYSTEM_PROMPT}, - {"role": "user", "content": prompt}, - ], - "ability": "codegen", - "reward_model": { - "style": "rule", - "ground_truth": json.dumps({ - "functional": test_code - }), - }, - "extra_info": { - "split": split, - "index": idx, - "reference": solution, - "prompt": prompt, - "prefix": prefix, - "dataset": "LeetCodeDataset", - }, - } - - return process_fn - - # filter out empty examples ("reward_model" is None) - train_dataset = train_dataset.map( - function=make_map_fn("train"), - with_indices=True, - load_from_cache_file=False, - ).filter(lambda x: x["reward_model"] is not None) - test_dataset = test_dataset.map( - function=make_map_fn("test"), - with_indices=True, - load_from_cache_file=False, - ).filter(lambda x: x["reward_model"] is not None) - print(f"Leetcode2k train set: {train_dataset}") - print(f"Leetcode2k test set: {test_dataset}") - return train_dataset, test_dataset - - -def humaneval(): - rich.print(Rule("Loading OpenAI HumanEval...")) - dataset = load_dataset("openai_humaneval")["test"] - print("HumanEval dataset:", dataset) - - def process_fn(example, idx): - # HumanEval's prompt already contains the function signature and docstring - prompt = ( - "Write a complete, self-contained Python solution to the following problem. " - "Your solution must include all necessary imports and the full function definition including " - "the signature exactly as specified. Do not modify the function signature or docstring.\n\n" - f"```python\n{example['prompt'].strip()}\n```" - ) - - # Extract test code - test_code = example['test'] - entry_point = example['entry_point'] - - # Validate that the canonical solution passes the tests - solution = example['canonical_solution'] - - # Combine the prompt code + solution + test code to verify it works - full_code = f"{example['prompt']}\n{solution}\n{test_code}\n\ncheck({entry_point})" - - succ, err = code_exec(full_code) - if not succ: - print(f"Error in canonical solution for task {example['task_id']}: {err}") - return EMPTY_RETURN - - return { - "data_source": "codegen-humaneval", - "prompt": [ - {"role": "system", "content": SYSTEM_PROMPT}, - {"role": "user", "content": prompt}, - ], - "ability": "codegen", - "reward_model": { - "style": "rule", - "ground_truth": json.dumps( - {"functional": f"{test_code}\n\ncheck({entry_point})"} - ), - }, - "extra_info": { - "split": "test", - "index": idx, - "reference": solution, - "prompt": prompt, - "dataset": "openai_humaneval", - "task_id": str(example["task_id"]), - }, - } - - test_dataset = dataset.map( - function=process_fn, - with_indices=True, - load_from_cache_file=False, - ).filter(lambda x: x["reward_model"] is not None) - - # Return empty train dataset and test dataset - empty_train = datasets.Dataset.from_dict({ - "data_source": [], - "prompt": [], - "ability": [], - "reward_model": [], - "extra_info": [] - }) if len(test_dataset) > 0 else datasets.Dataset.from_dict({}) - - print(f"HumanEval test set: {test_dataset}") - return empty_train, test_dataset - -def mbpp(): - rich.print(Rule("Loading MBPP dataset...")) - dataset = load_dataset("google-research-datasets/mbpp") - - def make_map_fn(split): - def process_fn(example, idx): - # rewrite the task_id as it is int - example["task_id"] = "MBPP/" + str(example["task_id"]) - - # Create prompt - prompt = ( - f"{example['text']}\n\n" - f"Your solution should be a complete, self-contained function in a markdown code block. " - f"Make sure your solution passes the following test cases:\n" - ) - - # Construct test code - test_code = "" - if example.get('test_setup_code'): - test_code += example['test_setup_code'] + "\n\n" - - # Add all test assertions - for assertion in example['test_list'] + example.get('challenge_test_list', []): - test_code += assertion + "\n" - - # Add test cases to prompt - prompt += f"```python\n{test_code}```" - prompt += "\n\nPlease do not include the test cases in your solution." - - # Validate that the canonical solution passes the tests - solution = example['code'] - full_code = f"{solution}\n\n{test_code}" - - succ, err = code_exec(full_code) - if not succ: - print(f"Error in canonical solution for task {example['task_id']}: {err}") - return EMPTY_RETURN - - return { - "data_source": "codegen-mbpp", - "prompt": [ - {"role": "system", "content": SYSTEM_PROMPT}, - {"role": "user", "content": prompt}, - ], - "ability": "codegen", - "reward_model": { - "style": "rule", - "ground_truth": json.dumps( - {"functional": test_code} - ), - }, - "extra_info": { - "split": split, - "index": idx, - "reference": solution, - "prompt": prompt, - "dataset": "mbpp", - "task_id": str(example["task_id"]), - }, - } - - return process_fn - - # Process train and test splits - train_dataset = dataset["train"].map( - function=make_map_fn("train"), - with_indices=True, - load_from_cache_file=False, - ).filter(lambda x: x['reward_model'] is not None) - - test_dataset = dataset["test"].map( - function=make_map_fn("test"), - with_indices=True, - load_from_cache_file=False, - ).filter(lambda x: x['reward_model'] is not None) - - print(f"MBPP train set: {train_dataset}") - print(f"MBPP test set: {test_dataset}") - return train_dataset, test_dataset - -def primeintellect(): - rich.print(Rule("Loading PrimeIntellect dataset...")) - dataset = load_dataset("agentica-org/DeepCoder-Preview-Dataset", "primeintellect", split="train") - - def make_map_fn(split): - def process_fn(example, idx): - # Get the problem description - prompt = example["problem"] - - # Get the solution code - solution = example["solutions"][0] if example["solutions"] else "" - solution = extract_code_from_string(solution) - - # Process tests - tests = json.loads(example["tests"]) - if not tests: - return EMPTY_RETURN - - # Handle different test types - if tests[0]["type"] == "function_call": - # Function call tests - fn_name = tests[0]["fn_name"] - test_code = f"""\ -def check_{fn_name}(): -""" - for test in tests: - input_args = ", ".join([ - repr(arg) if isinstance(arg, (str, list, tuple, dict)) else str(arg) - for arg in test["input"] - ]) - expected_output = repr(test["output"][0]) if isinstance(test["output"][0], (str, list, tuple, dict)) else test["output"][0] - test_code += f""" assert {fn_name}({input_args}) == {expected_output} -""" - test_code += f""" -check_{fn_name}() -""" - - # Validate the solution - full_code = f"{solution}\n{test_code}" - succ, err = code_exec(full_code) - if not succ: - rich.print(f"[bold red]Test code failed for PrimeIntellect example {idx}") - print(f"===========Full code===========") - print(full_code) - print(f"===========Error===========") - print(err) - return EMPTY_RETURN - - oracle = json.dumps({"functional": test_code}) - - elif tests[0]["type"] == "stdin_stdout": - # STDIN/STDOUT tests - stdin_list = [] - stdout_list = [] - for test in tests: - stdin_list.append(test["input"]) - stdout_list.append(test["output"]) - - # Validate the solution - with ThreadPoolExecutor(max_workers=min(len(stdin_list), 8)) as executor: - futures = [] - for stdin, stdout in zip(stdin_list, stdout_list): - futures.append( - executor.submit( - remote_check_stdio, - solution, - stdin, - stdout, - ) - ) - for future in as_completed(futures): - exec_succ, output, stdin, stdout = future.result() - pass_test = exec_succ and fuzzy_equal(output.strip(), stdout.strip()) - if not pass_test: - rich.print(f"[bold red]Test code failed for PrimeIntellect example {idx}") - print(f"===========Solution===========") - print(solution) - print(f"===========Input===========") - print(stdin) - print(f"===========Expected output===========") - print(stdout) - print(f"===========Actual output===========") - print(output) - return EMPTY_RETURN - - oracle = json.dumps({"inputs": stdin_list, "outputs": stdout_list}) - else: - raise ValueError(f"Unknown test type: {tests[0]['type']}") - - return { - "data_source": "codegen-primeintellect", - "prompt": [ - {"role": "system", "content": SYSTEM_PROMPT}, - {"role": "user", "content": prompt}, - ], - "ability": "codegen", - "reward_model": { - "style": "rule", - "ground_truth": oracle, - }, - "extra_info": { - "split": split, - "index": idx, - "prompt": prompt, - "reference": solution, - "dataset": "PrimeIntellect", - "function_name": fn_name if tests[0]["type"] == "function_call" else None, - }, - } - - return process_fn - - # Process train and test splits - train_dataset = dataset.map( - function=make_map_fn("train"), - with_indices=True, - num_proc=64, - remove_columns=dataset.column_names, - load_from_cache_file=False, - ).filter(lambda x: x['reward_model'] is not None) - - # Split into train and test sets - if N_TESTSET_PER_DATASET == 0: - # Return no test samples if N_TESTSET_PER_DATASET is 0 - test_dataset = datasets.Dataset.from_dict({ - "data_source": [], - "prompt": [], - "ability": [], - "reward_model": [], - "extra_info": [] - }) - else: - splits = train_dataset.train_test_split( - test_size=min(N_TESTSET_PER_DATASET, len(train_dataset) * 0.1), seed=666 - ) - train_dataset = splits["train"].shuffle(seed=666) - test_dataset = splits["test"] - - print(f"PrimeIntellect train set: {train_dataset}") - print(f"PrimeIntellect test set: {test_dataset}") - return train_dataset, test_dataset - -def livecodebench(): - rich.print(Rule("Loading LiveCodeBench dataset...")) - # Load both train and test splits directly - train_dataset = load_dataset("agentica-org/DeepCoder-Preview-Dataset", "lcbv5", split="train") - test_dataset = load_dataset("agentica-org/DeepCoder-Preview-Dataset", "lcbv5", split="test") - - def make_map_fn(split): - def process_fn(example, idx): - # Get the problem description - problem_desc = example["problem"] - starter_code = example["starter_code"] - - # Create the prompt with the starter code - prompt = f"{problem_desc}\n\nComplete the implementation using the provided starter code:\n```python\n{starter_code}\n```\n\nYour solution should implement the method(s) in the Solution class." - - # Process tests - tests = json.loads(example["tests"]) - if not tests: - return EMPTY_RETURN - - # Process metadata to get function name - metadata = example["metadata"] - function_name = metadata.get("func_name") - - # Handle different test types - if tests[0]["testtype"] == "functional": - if not function_name: - return EMPTY_RETURN - - # Function call tests - test_code = f"""\ -def check_{function_name}(): -""" - for test in tests: - # Parse input string by splitting on '\n' - input_parts = test["input"].split('\n') - # Create proper comma-separated arguments for the function call - input_args = ', '.join(input_parts) - # Get the output value - output_val = test["output"] - - test_code += f""" assert Solution().{function_name}({input_args}) == {output_val} -""" - - test_code += f""" -check_{function_name}() -""" - - # For debugging - print a few examples - if idx < 20: - print(f"Generated test code for example {idx}:") - print(test_code) - - oracle = json.dumps({"functional": test_code}) - - elif tests[0]["testtype"] == "stdin": - # STDIN/STDOUT tests - stdin_list = [] - stdout_list = [] - for test in tests: - stdin_list.append(test["input"]) - stdout_list.append(test["output"]) - - # For debugging - print a few examples - if idx < 20: - print(f"Generated test code for example {idx}:") - for i in range(len(stdin_list)): - print(f"Test {i+1}:") - print(f"Input: {stdin_list[i]}") - print(f"Output: {stdout_list[i]}") - - oracle = json.dumps({"inputs": stdin_list, "outputs": stdout_list}) - else: - raise ValueError(f"Unknown test type: {tests[0]['testtype']}") - - return { - "data_source": "codegen-livecodebench", - "prompt": [ - {"role": "system", "content": SYSTEM_PROMPT}, - {"role": "user", "content": prompt}, - ], - "ability": "codegen", - "reward_model": { - "style": "rule", - "ground_truth": oracle, - }, - "extra_info": { - "split": split, - "index": idx, - "prompt": prompt, - "reference": "", # No solution data - "dataset": "LiveCodeBench", - "function_name": function_name, - }, - } - - return process_fn - - # Process train and test datasets - train_dataset = train_dataset.map( - function=make_map_fn("train"), - with_indices=True, - num_proc=64, - load_from_cache_file=False, - ).filter(lambda x: x['reward_model'] is not None) - - test_dataset = test_dataset.map( - function=make_map_fn("test"), - with_indices=True, - num_proc=64, - load_from_cache_file=False, - ).filter(lambda x: x['reward_model'] is not None) - - print(f"LiveCodeBench train set: {train_dataset}") - print(f"LiveCodeBench test set: {test_dataset}") - return train_dataset, test_dataset - - -if __name__ == "__main__": - import argparse - - parser = argparse.ArgumentParser() - parser.add_argument("--root_dir", default=os.path.join(WORKDING_DIR, "data")) - parser.add_argument( - "--dataset_names", default="kodcode", help="comma separated dataset names" - ) - - args = parser.parse_args() - - root_dir = args.root_dir - dataset_names = args.dataset_names.split(",") - - train_datasets = [] - test_datasets = [] - - dataset_map = { - "kodcode": kodcode, - "taco": taco, - "leetcode2k": leetcode2k, - "humaneval": humaneval, - "mbpp": mbpp, - "primeintellect": primeintellect, - "livecodebench": livecodebench, - } - dataset_makes = [dataset_map[name] for name in dataset_names] - names = "-".join([make.__name__ for make in dataset_makes]) - - for make in dataset_makes: - train, test = make() - train_datasets.append(train) - test_datasets.append(test) - - train_dataset = concatenate_datasets(train_datasets).shuffle(seed=666) - test_dataset = concatenate_datasets(test_datasets) - - rich.print(Rule("Saving the final dataset")) - print("Train set:", train_dataset) - print("Test set:", test_dataset) - - local_dir = os.path.join( - root_dir, f"codegen-{round(len(train_dataset) / 1000)}k-{names}" - ) - rich.print(f"[bold green]Saving to {local_dir}...") - train_dataset.to_parquet(os.path.join(local_dir, "train.parquet")) - test_dataset.to_parquet(os.path.join(local_dir, "test.parquet")) \ No newline at end of file diff --git a/reasoners/data_preprocess/archived_for_reference/deepscaler_preview_r1_zero_style_filter.py b/reasoners/data_preprocess/archived_for_reference/deepscaler_preview_r1_zero_style_filter.py deleted file mode 100644 index 942540697..000000000 --- a/reasoners/data_preprocess/archived_for_reference/deepscaler_preview_r1_zero_style_filter.py +++ /dev/null @@ -1,192 +0,0 @@ -# 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. -""" -Preprocess the https://huggingface.co/datasets/agentica-org/DeepScaleR-Preview-Dataset to parquet format -Code referring to https://github.com/agentica-project/deepscaler/blob/main/scripts/data/deepscaler_dataset.py -""" - -import os -import datasets -from typing import Dict, List, Optional, Any, Union -import enum -import argparse -import pandas as pd -import json - -from verl.utils.hdfs_io import copy, makedirs -from verl.utils.reward_score.math import remove_boxed, last_boxed_only_string - - -class TrainDataset(enum.Enum): - """Enum for training datasets. - - Contains identifiers for various math problem datasets used during training. - """ - - AIME = "AIME" # American Invitational Mathematics Examination - AMC = "AMC" # American Mathematics Competition - OMNI_MATH = "OMNI_MATH" # Omni Math - NUMINA_OLYMPIAD = "OLYMPIAD" # Unique Olympiad problems from NUMINA - MATH = "MATH" # Dan Hendrycks Math Problems - STILL = "STILL" # STILL dataset - DEEPSCALER = "DEEPSCALER" # DeepScaler (AIME, AMC, OMNI_MATH, MATH, STILL) - - -class TestDataset(enum.Enum): - """Enum for testing/evaluation datasets. - - Contains identifiers for datasets used to evaluate model performance. - """ - - AIME = "AIME" # American Invitational Mathematics Examination - AMC = "AMC" # American Mathematics Competition - MATH = "MATH" # Math 500 problems - MINERVA = "MINERVA" # Minerva dataset - OLYMPIAD_BENCH = "OLYMPIAD_BENCH" # Olympiad benchmark problems - - -"""Type alias for either training or testing dataset types.""" -Dataset = Union[TrainDataset, TestDataset] - - -def load_dataset(dataset: Dataset) -> List[Dict[str, Any]]: - """Load a dataset from a JSON file. - - Loads and parses a JSON dataset file based on the provided dataset enum. - The file path is constructed based on whether it's a training or testing dataset. - - Args: - dataset: A Dataset enum value specifying which dataset to load. - - Returns: - List[Dict[str, Any]]: A list of dictionaries containing the dataset records. - Each dictionary represents one example in the dataset. - - Raises: - ValueError: If the dataset file cannot be found, contains invalid JSON, - or encounters other file access errors. - - Example: - >>> load_dataset(TrainDataset.AIME) - [{'problem': 'Find x...', 'solution': '42', ...}, ...] - """ - dataset_name = dataset.value.lower() - data_dir = "train" if isinstance(dataset, TrainDataset) else "test" - - current_dir = os.path.dirname(os.path.realpath(__file__)) - file_path = os.path.join(data_dir, f"{dataset_name}.json") - file_path = os.path.join(current_dir, file_path) - - if not os.path.exists(file_path): - raise ValueError(f"Dataset file not found: {file_path}") - - try: - with open(file_path, "r", encoding="utf-8") as file: - data = json.load(file) - return data - except json.JSONDecodeError: - raise ValueError(f"Invalid JSON format in {file_path}") - except Exception as exc: # pylint: disable=broad-except - raise ValueError(f"Error loading dataset: {exc}") from exc - - -def extract_solution(solution_str): - return remove_boxed(last_boxed_only_string(solution_str)) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("--local_dir", default="data/deepscaler_preview_zero_style_mar21_filter") - parser.add_argument("--hdfs_dir", default=None) - - args = parser.parse_args() - - train_data_source = "SDSB/deepscale_partial_mar21_filtered_basic" - - print(f"Loading the {train_data_source} dataset from huggingface...", flush=True) - test_data_sources = [ - "nanoverl/minerva", - "SDSB/aime_repeated_8x", - "SDSB/amc_repeated_4x", - "nanoverl/olympiad_bench", - "nanoverl/math", - ] - print(f"Loading the {test_data_sources} dataset from huggingface...", flush=True) - train_dataset = datasets.load_dataset( - train_data_source, trust_remote_code=True, split="train" - ) - test_datasets = [ - datasets.load_dataset(test_data_source, trust_remote_code=True, split="test") - for test_data_source in test_data_sources - ] - - instruction_following = ( - "Let's think step by step and output the final answer within \\boxed{}." - ) - - - prompt = """A conversation between User and Assistant. The user asks a question, and the Assistant solves it. The assistant first thinks about the reasoning process in the mind and then provides the user with the response. The reasoning process is enclosed within i.e., reasoning process here respond to the user's question here. - -User: {{prompt}} Please put your answer in \\boxed{} tags. -Assistant: -""" - - # add a row to each data item that represents a unique id - def make_map_fn(split, data_source): - - def process_fn(example, idx): - question = example.pop("problem") - - # question = question + " " + instruction_following - - answer = example.pop("answer") - data = { - "data_source": data_source, - "prompt": [], - "raw_prompt": prompt.replace("{{prompt}}", question), - "ability": "math", - "apply_chat_template": False, - "reward_model": {"style": "rule", "ground_truth": answer}, - "extra_info": {"split": split, - "index": idx, - "question": question}, - } - if idx == 0: - print("=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) - print(data) - return data - - return process_fn - - local_dir = args.local_dir - hdfs_dir = args.hdfs_dir - - train_data = train_dataset.map( - function=make_map_fn("train", train_data_source), with_indices=True - ) - train_data.to_parquet(os.path.join(local_dir, "train.parquet")) - print(f"train data size:", len(train_data)) - print(train_data[0]) - for test_data_source, test_data in zip(test_data_sources, test_datasets): - process_fn = make_map_fn("test", test_data_source) - test_data = test_data.map(process_fn, with_indices=True) - dataset_name = os.path.basename(test_data_source.lower()) - test_df = pd.DataFrame(test_data) - test_df.to_parquet(os.path.join(local_dir, f"{dataset_name}.parquet")) - print(f"test data size: ({dataset_name})", len(test_df)) - - if hdfs_dir is not None: - makedirs(hdfs_dir) - - copy(src=local_dir, dst=hdfs_dir) diff --git a/reasoners/data_preprocess/archived_for_reference/deepscaler_preview_r1_zero_style_sampled.py b/reasoners/data_preprocess/archived_for_reference/deepscaler_preview_r1_zero_style_sampled.py deleted file mode 100644 index f278701f1..000000000 --- a/reasoners/data_preprocess/archived_for_reference/deepscaler_preview_r1_zero_style_sampled.py +++ /dev/null @@ -1,201 +0,0 @@ -# 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. -""" -Preprocess the https://huggingface.co/datasets/agentica-org/DeepScaleR-Preview-Dataset to parquet format -Code referring to https://github.com/agentica-project/deepscaler/blob/main/scripts/data/deepscaler_dataset.py -""" - -import os -import datasets -from typing import Dict, List, Optional, Any, Union -import enum -import argparse -import pandas as pd -import json - -from verl.utils.hdfs_io import copy, makedirs -from verl.utils.reward_score.math import remove_boxed, last_boxed_only_string - - -class TrainDataset(enum.Enum): - """Enum for training datasets. - - Contains identifiers for various math problem datasets used during training. - """ - - AIME = "AIME" # American Invitational Mathematics Examination - AMC = "AMC" # American Mathematics Competition - OMNI_MATH = "OMNI_MATH" # Omni Math - NUMINA_OLYMPIAD = "OLYMPIAD" # Unique Olympiad problems from NUMINA - MATH = "MATH" # Dan Hendrycks Math Problems - STILL = "STILL" # STILL dataset - DEEPSCALER = "DEEPSCALER" # DeepScaler (AIME, AMC, OMNI_MATH, MATH, STILL) - - -class TestDataset(enum.Enum): - """Enum for testing/evaluation datasets. - - Contains identifiers for datasets used to evaluate model performance. - """ - - AIME = "AIME" # American Invitational Mathematics Examination - AMC = "AMC" # American Mathematics Competition - MATH = "MATH" # Math 500 problems - MINERVA = "MINERVA" # Minerva dataset - OLYMPIAD_BENCH = "OLYMPIAD_BENCH" # Olympiad benchmark problems - - -"""Type alias for either training or testing dataset types.""" -Dataset = Union[TrainDataset, TestDataset] - - -def load_dataset(dataset: Dataset) -> List[Dict[str, Any]]: - """Load a dataset from a JSON file. - - Loads and parses a JSON dataset file based on the provided dataset enum. - The file path is constructed based on whether it's a training or testing dataset. - - Args: - dataset: A Dataset enum value specifying which dataset to load. - - Returns: - List[Dict[str, Any]]: A list of dictionaries containing the dataset records. - Each dictionary represents one example in the dataset. - - Raises: - ValueError: If the dataset file cannot be found, contains invalid JSON, - or encounters other file access errors. - - Example: - >>> load_dataset(TrainDataset.AIME) - [{'problem': 'Find x...', 'solution': '42', ...}, ...] - """ - dataset_name = dataset.value.lower() - data_dir = "train" if isinstance(dataset, TrainDataset) else "test" - - current_dir = os.path.dirname(os.path.realpath(__file__)) - file_path = os.path.join(data_dir, f"{dataset_name}.json") - file_path = os.path.join(current_dir, file_path) - - if not os.path.exists(file_path): - raise ValueError(f"Dataset file not found: {file_path}") - - try: - with open(file_path, "r", encoding="utf-8") as file: - data = json.load(file) - return data - except json.JSONDecodeError: - raise ValueError(f"Invalid JSON format in {file_path}") - except Exception as exc: # pylint: disable=broad-except - raise ValueError(f"Error loading dataset: {exc}") from exc - - -def extract_solution(solution_str): - return remove_boxed(last_boxed_only_string(solution_str)) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("--local_dir", default="data/deepscaler_default") - parser.add_argument("--hdfs_dir", default=None) - parser.add_argument( - "--train-sample-size", - type=int, - default=None, - help="Number of samples to use from the training dataset. If None, use all samples.", - ) - - args = parser.parse_args() - - train_data_source = "agentica-org/DeepScaleR-Preview-Dataset" - print(f"Loading the {train_data_source} dataset from huggingface...", flush=True) - test_data_sources = [ - "nanoverl/minerva", - "nanoverl/aime", - "nanoverl/amc", - "nanoverl/olympiad_bench", - "nanoverl/math", - ] - print(f"Loading the {test_data_sources} dataset from huggingface...", flush=True) - train_dataset = datasets.load_dataset( - train_data_source, trust_remote_code=True, split="train" - ) - test_datasets = [ - datasets.load_dataset(test_data_source, trust_remote_code=True, split="test") - for test_data_source in test_data_sources - ] - - instruction_following = ( - "Let's think step by step and output the final answer within \\boxed{}." - ) - - prompt = """A conversation between User and Assistant. The user asks a question, and the Assistant solves it. The assistant first thinks about the reasoning process in the mind and then provides the user with the response. The reasoning process is enclosed within i.e., reasoning process here respond to the user's question here. - -User: {{prompt}} Please put your answer in \\boxed{} tags. -Assistant: -""" - - # add a row to each data item that represents a unique id - def make_map_fn(split, data_source): - - def process_fn(example, idx): - question = example.pop("problem") - - # question = question + " " + instruction_following - - answer = example.pop("answer") - data = { - "data_source": data_source, - "prompt": [], - "raw_prompt": prompt.replace("{{prompt}}", question), - "ability": "math", - "apply_chat_template": False, - "reward_model": {"style": "rule", "ground_truth": answer}, - "extra_info": {"split": split, "index": idx}, - } - if idx == 0: - print("=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) - print(data) - return data - - return process_fn - - local_dir = args.local_dir - hdfs_dir = args.hdfs_dir - - train_data = train_dataset.map( - function=make_map_fn("train", train_data_source), with_indices=True - ) - - # Sample the training data if sample size is specified - if args.train_sample_size is not None: - train_data = train_data.select( - range(min(args.train_sample_size, len(train_data))) - ) - - train_data.to_parquet(os.path.join(local_dir, "train.parquet")) - print(f"train data size:", len(train_data)) - print(train_data[0]) - for test_data_source, test_data in zip(test_data_sources, test_datasets): - process_fn = make_map_fn("test", test_data_source) - test_data = test_data.map(process_fn, with_indices=True) - dataset_name = os.path.basename(test_data_source.lower()) - test_df = pd.DataFrame(test_data) - test_df.to_parquet(os.path.join(local_dir, f"{dataset_name}.parquet")) - print(f"test data size: ({dataset_name})", len(test_df)) - - if hdfs_dir is not None: - makedirs(hdfs_dir) - - copy(src=local_dir, dst=hdfs_dir) diff --git a/reasoners/data_preprocess/archived_for_reference/geo3k.py b/reasoners/data_preprocess/archived_for_reference/geo3k.py deleted file mode 100644 index 3a0c77673..000000000 --- a/reasoners/data_preprocess/archived_for_reference/geo3k.py +++ /dev/null @@ -1,86 +0,0 @@ -# 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. -""" -Preprocess the Geometry3k dataset to parquet format -""" - -import os -import datasets - -from verl.utils.hdfs_io import copy, makedirs -import argparse - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument('--local_dir', default='~/data/geo3k') - parser.add_argument('--hdfs_dir', default=None) - - args = parser.parse_args() - - data_source = 'hiyouga/geometry3k' - - dataset = datasets.load_dataset(data_source) - - train_dataset = dataset['train'] - test_dataset = dataset['test'] - - instruction_following = ( - r'You FIRST think about the reasoning process as an internal monologue and then provide the final answer. ' - r'The reasoning process MUST BE enclosed within tags. The final answer MUST BE put in \boxed{}.' - ) - - # add a row to each data item that represents a unique id - def make_map_fn(split): - - def process_fn(example, idx): - problem = example.pop('problem') - prompt = problem + ' ' + instruction_following - answer = example.pop('answer') - images = example.pop('images') - - data = { - "data_source": data_source, - "prompt": [{ - "role": "user", - "content": prompt, - }], - "images": images, - "ability": "math", - "reward_model": { - "style": "rule", - "ground_truth": answer - }, - "extra_info": { - 'split': split, - 'index': idx, - 'answer': answer, - "question": problem, - } - } - return data - - return process_fn - - train_dataset = train_dataset.map(function=make_map_fn('train'), with_indices=True, num_proc=8) - test_dataset = test_dataset.map(function=make_map_fn('test'), with_indices=True, num_proc=8) - - local_dir = args.local_dir - hdfs_dir = args.hdfs_dir - - train_dataset.to_parquet(os.path.join(local_dir, 'train.parquet')) - test_dataset.to_parquet(os.path.join(local_dir, 'test.parquet')) - - if hdfs_dir is not None: - makedirs(hdfs_dir) - copy(src=local_dir, dst=hdfs_dir) diff --git a/reasoners/data_preprocess/archived_for_reference/orz_data.py b/reasoners/data_preprocess/archived_for_reference/orz_data.py deleted file mode 100644 index 98b1b5f3c..000000000 --- a/reasoners/data_preprocess/archived_for_reference/orz_data.py +++ /dev/null @@ -1,163 +0,0 @@ -# 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 os -import datasets -from typing import Dict, List, Optional, Any, Union -import enum -import argparse -import pandas as pd -import json - -from verl.utils.hdfs_io import copy, makedirs -from verl.utils.reward_score.math import remove_boxed, last_boxed_only_string - - -class TrainDataset(enum.Enum): - """Enum for training datasets. - - Contains identifiers for various math problem datasets used during training. - """ - - AIME = "AIME" # American Invitational Mathematics Examination - AMC = "AMC" # American Mathematics Competition - OMNI_MATH = "OMNI_MATH" # Omni Math - NUMINA_OLYMPIAD = "OLYMPIAD" # Unique Olympiad problems from NUMINA - MATH = "MATH" # Dan Hendrycks Math Problems - STILL = "STILL" # STILL dataset - DEEPSCALER = "DEEPSCALER" # DeepScaler (AIME, AMC, OMNI_MATH, MATH, STILL) - - -class TestDataset(enum.Enum): - """Enum for testing/evaluation datasets. - - Contains identifiers for datasets used to evaluate model performance. - """ - - AIME = "AIME" # American Invitational Mathematics Examination - AMC = "AMC" # American Mathematics Competition - MATH = "MATH" # Math 500 problems - MINERVA = "MINERVA" # Minerva dataset - OLYMPIAD_BENCH = "OLYMPIAD_BENCH" # Olympiad benchmark problems - - -"""Type alias for either training or testing dataset types.""" -Dataset = Union[TrainDataset, TestDataset] - - -def load_dataset(dataset: Dataset) -> List[Dict[str, Any]]: - """Load a dataset from a JSON file. - - Loads and parses a JSON dataset file based on the provided dataset enum. - The file path is constructed based on whether it's a training or testing dataset. - - Args: - dataset: A Dataset enum value specifying which dataset to load. - - Returns: - List[Dict[str, Any]]: A list of dictionaries containing the dataset records. - Each dictionary represents one example in the dataset. - - Raises: - ValueError: If the dataset file cannot be found, contains invalid JSON, - or encounters other file access errors. - - Example: - >>> load_dataset(TrainDataset.AIME) - [{'problem': 'Find x...', 'solution': '42', ...}, ...] - """ - dataset_name = dataset.value.lower() - data_dir = "train" if isinstance(dataset, TrainDataset) else "test" - - current_dir = os.path.dirname(os.path.realpath(__file__)) - file_path = os.path.join(data_dir, f"{dataset_name}.json") - file_path = os.path.join(current_dir, file_path) - - if not os.path.exists(file_path): - raise ValueError(f"Dataset file not found: {file_path}") - - try: - with open(file_path, "r", encoding="utf-8") as file: - data = json.load(file) - return data - except json.JSONDecodeError: - raise ValueError(f"Invalid JSON format in {file_path}") - except Exception as exc: # pylint: disable=broad-except - raise ValueError(f"Error loading dataset: {exc}") from exc - - -def extract_solution(solution_str): - return remove_boxed(last_boxed_only_string(solution_str)) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("--local_dir", default="data/orz") - parser.add_argument("--hdfs_dir", default=None) - args = parser.parse_args() - - # Load train data from local JSON file - - # https://raw.githubusercontent.com/Open-Reasoner-Zero/Open-Reasoner-Zero/refs/heads/main/data/orz_math_57k_collected.json - train_data_source = "orz_math_57k_collected.json" - print(f"Loading training data from {train_data_source}...", flush=True) - - train_dataset = json.load(open(train_data_source, "r")) - - # Convert to Dataset format that's compatible with the rest of the code - train_data = datasets.Dataset.from_list([{ - "problem": item[0]["value"], - "answer": item[1]["ground_truth"]["value"] - } for item in train_dataset]) - - print(train_data[0]) - - instruction_following = ( - "Let's think step by step and output the final answer within \\boxed{}." - ) - - # Rest of the processing remains the same - def make_map_fn(split, data_source): - def process_fn(example, idx): - question = example.pop("problem") - question = question + " " + instruction_following - answer = example.pop("answer") - data = { - "data_source": data_source, - "prompt": [{"role": "user", "content": question}], - "ability": "math", - "reward_model": {"style": "rule", "ground_truth": answer}, - "extra_info": {"split": split, "index": idx}, - } - if idx == 0: - print("=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) - print(data) - return data - return process_fn - - local_dir = args.local_dir - hdfs_dir = args.hdfs_dir - - # Process train data - train_data = train_data.map( - function=make_map_fn("train", train_data_source), with_indices=True - ) - print(train_data[0]) - train_data.to_parquet(os.path.join(local_dir, "train.parquet")) - print(f"train data size:", len(train_data)) - - # Remove test data processing since we're only working with the training data - if hdfs_dir is not None: - makedirs(hdfs_dir) - copy(src=local_dir, dst=hdfs_dir) diff --git a/reasoners/data_preprocess/archived_for_reference/orz_data_r1_zero_style.py b/reasoners/data_preprocess/archived_for_reference/orz_data_r1_zero_style.py deleted file mode 100644 index 95a8d972c..000000000 --- a/reasoners/data_preprocess/archived_for_reference/orz_data_r1_zero_style.py +++ /dev/null @@ -1,192 +0,0 @@ -# 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 os -import datasets -from typing import Dict, List, Optional, Any, Union -import enum -import argparse -import pandas as pd -import json - -from verl.utils.hdfs_io import copy, makedirs -from verl.utils.reward_score.math import remove_boxed, last_boxed_only_string - - -class TrainDataset(enum.Enum): - """Enum for training datasets. - - Contains identifiers for various math problem datasets used during training. - """ - - AIME = "AIME" # American Invitational Mathematics Examination - AMC = "AMC" # American Mathematics Competition - OMNI_MATH = "OMNI_MATH" # Omni Math - NUMINA_OLYMPIAD = "OLYMPIAD" # Unique Olympiad problems from NUMINA - MATH = "MATH" # Dan Hendrycks Math Problems - STILL = "STILL" # STILL dataset - DEEPSCALER = "DEEPSCALER" # DeepScaler (AIME, AMC, OMNI_MATH, MATH, STILL) - - -class TestDataset(enum.Enum): - """Enum for testing/evaluation datasets. - - Contains identifiers for datasets used to evaluate model performance. - """ - - AIME = "AIME" # American Invitational Mathematics Examination - AMC = "AMC" # American Mathematics Competition - MATH = "MATH" # Math 500 problems - MINERVA = "MINERVA" # Minerva dataset - OLYMPIAD_BENCH = "OLYMPIAD_BENCH" # Olympiad benchmark problems - - -"""Type alias for either training or testing dataset types.""" -Dataset = Union[TrainDataset, TestDataset] - - -def load_dataset(dataset: Dataset) -> List[Dict[str, Any]]: - """Load a dataset from a JSON file. - - Loads and parses a JSON dataset file based on the provided dataset enum. - The file path is constructed based on whether it's a training or testing dataset. - - Args: - dataset: A Dataset enum value specifying which dataset to load. - - Returns: - List[Dict[str, Any]]: A list of dictionaries containing the dataset records. - Each dictionary represents one example in the dataset. - - Raises: - ValueError: If the dataset file cannot be found, contains invalid JSON, - or encounters other file access errors. - - Example: - >>> load_dataset(TrainDataset.AIME) - [{'problem': 'Find x...', 'solution': '42', ...}, ...] - """ - dataset_name = dataset.value.lower() - data_dir = "train" if isinstance(dataset, TrainDataset) else "test" - - current_dir = os.path.dirname(os.path.realpath(__file__)) - file_path = os.path.join(data_dir, f"{dataset_name}.json") - file_path = os.path.join(current_dir, file_path) - - if not os.path.exists(file_path): - raise ValueError(f"Dataset file not found: {file_path}") - - try: - with open(file_path, "r", encoding="utf-8") as file: - data = json.load(file) - return data - except json.JSONDecodeError: - raise ValueError(f"Invalid JSON format in {file_path}") - except Exception as exc: # pylint: disable=broad-except - raise ValueError(f"Error loading dataset: {exc}") from exc - - -def extract_solution(solution_str): - return remove_boxed(last_boxed_only_string(solution_str)) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("--local_dir", default="data/orz_zero_style_v2") - parser.add_argument("--hdfs_dir", default=None) - args = parser.parse_args() - - # Load train data from local JSON file - - # https://raw.githubusercontent.com/Open-Reasoner-Zero/Open-Reasoner-Zero/refs/heads/main/data/orz_math_57k_collected.json - train_data_source = "examples/data_preprocess/orz_math_57k_collected.json" - print(f"Loading training data from {train_data_source}...", flush=True) - - train_dataset = json.load(open(train_data_source, "r")) - - # Convert to Dataset format that's compatible with the rest of the code - train_data = datasets.Dataset.from_list([{ - "problem": item[0]["value"], - "answer": item[1]["ground_truth"]["value"] - } for item in train_dataset]) - - print(train_data[0]) - - test_data_sources = [ - "nanoverl/minerva", - "nanoverl/aime", - "nanoverl/amc", - "nanoverl/olympiad_bench", - "nanoverl/math", - ] - print(f"Loading the {test_data_sources} dataset from huggingface...", flush=True) - - test_datasets = [ - datasets.load_dataset(test_data_source, trust_remote_code=True, split="test") - for test_data_source in test_data_sources - ] - - prompt = """A conversation between User and Assistant. The user asks a question, and the Assistant solves it. The assistant first thinks about the reasoning process in the mind and then provides the user with the response. The reasoning process is enclosed within i.e., reasoning process here respond to the user's question here. - -User: {{prompt}} Please put your answer in \\boxed{} tags. -Assistant: -""" - # Rest of the processing remains the same - def make_map_fn(split, data_source): - def process_fn(example, idx): - question = example.pop("problem") - answer = example.pop("answer") - data = { - "data_source": data_source, - "question": question, - "prompt": [], - "raw_prompt": prompt.replace("{{prompt}}", question), - "ability": "math", - "apply_chat_template": False, - "reward_model": {"style": "rule", "ground_truth": answer}, - "extra_info": { - "split": split, - "index": idx, - "question": question, - }, - } - if idx == 0: - print("=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) - print(data) - return data - return process_fn - - local_dir = args.local_dir - hdfs_dir = args.hdfs_dir - - # Process train data - train_data = train_data.map( - function=make_map_fn("train", train_data_source), with_indices=True - ) - print(train_data[0]) - train_data.to_parquet(os.path.join(local_dir, "train.parquet")) - print(f"train data size:", len(train_data)) - - for test_data_source, test_data in zip(test_data_sources, test_datasets): - process_fn = make_map_fn("test", test_data_source) - test_data = test_data.map(process_fn, with_indices=True) - dataset_name = os.path.basename(test_data_source.lower()) - test_df = pd.DataFrame(test_data) - test_df.to_parquet(os.path.join(local_dir, f"{dataset_name}.parquet")) - print(f"test data size: ({dataset_name})", len(test_df)) - - # Remove test data processing since we're only working with the training data - if hdfs_dir is not None: - makedirs(hdfs_dir) - copy(src=local_dir, dst=hdfs_dir) diff --git a/reasoners/data_preprocess/codegen/humaneval.py b/reasoners/data_preprocess/codegen/humaneval.py deleted file mode 100644 index 6557b843a..000000000 --- a/reasoners/data_preprocess/codegen/humaneval.py +++ /dev/null @@ -1,156 +0,0 @@ -"""Downloads, processes, and saves HumanEval dataset.""" - -import os -import argparse -import json -import transformers -import datasets -from datasets import load_dataset, Dataset - -from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset -from verl.utils.data_process.filter import LengthFilter -from verl.utils.reward_score.coder1 import code_exec - - -def get_datasets(cache_dir: str): - """ - Loads the HumanEval dataset. - """ - try: - dataset = load_dataset("openai_humaneval", cache_dir=cache_dir)["test"] - print(f"HumanEval dataset: {len(dataset)} examples") - return None, dataset - except Exception as e: - print(f"Error loading dataset: {e}") - return None, None - - -def make_map_fn(split: str, data_source: str) -> callable: - def process_fn(example, idx): - # Create a default "skip" response with all required fields - skip_response = { - "data_source": None, - "prompt": None, - "ability": None, - "reward_model": None, - "extra_info": None - } - - # Extract task ID and prompt - task_id = example["task_id"] - - prompt = ( - "Write a complete, self-contained Python solution to the following problem. " - "Your solution must include all necessary imports and the full function definition including " - "the signature exactly as specified. Do not modify the function signature or docstring.\n\n" - f"```python\n{example['prompt'].strip()}\n```" - ) - - # Extract test function, entry point, and canonical solution - test_code = example["test"] - entry_point = example["entry_point"] - solution = example["canonical_solution"] - - # Build test code that calls the entry point - test_code_with_check = f"{test_code}\n\ncheck({entry_point})" - - # Verify the canonical solution passes the tests - # full_code = f"{prompt}\n{solution}\n{test_code}\n\ncheck({entry_point})" - full_code = f"{example['prompt']}\n{solution}\n{test_code}\n\ncheck({entry_point})" - succ, err = code_exec(full_code, timeout=5) - print(f"[DEBUG] succ: {succ}, err: {err}") - if not succ: - print(f"Error in canonical solution for task {task_id}: {err}") - return skip_response - - - data = { - "data_source": data_source, - "prompt": [ - {"role": "user", "content": prompt} - ], - "ability": "codegen", - "apply_chat_template": True, - "reward_model": { - "style": "rule", - "ground_truth": json.dumps({ - "functional": test_code_with_check - }), - }, - "extra_info": { - "split": split, - "index": idx, - "reference": solution, # Include the canonical solution as reference - "original_prompt": prompt, - "dataset": "openai_humaneval", - "task_id": task_id, - }, - } - - if idx == 0 or idx == 1: - print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) - print(data) - - return data - - return process_fn - - -if __name__ == '__main__': - """Main script execution: parse args, load, process, and save datasets.""" - parser = argparse.ArgumentParser(description="Process and save HumanEval dataset.") - parser.add_argument('--data-dir', default='data', - help='Base directory to save the processed data files.') - parser.add_argument('--domain', default="codegen", - help='Domain of the dataset.') - parser.add_argument('--name', default="humaneval", - help='Name of the dataset.') - parser.add_argument('--sample-size', type=int, default=None, - help='Number of samples to use from dataset. If None, use all samples.') - parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility') - - args = parser.parse_args() - - set_seed(args.seed) - - data_source = f"{args.domain}__{args.name}" - test_output_dir = os.path.join(args.data_dir, 'test') - os.makedirs(test_output_dir, exist_ok=True) - - # Load the dataset - cache_dir = datasets.config.HF_DATASETS_CACHE - _, dataset = get_datasets(cache_dir=cache_dir) - print(f"[DEBUG] dataset: {len(dataset)}") - - # Process the dataset - process_fn = make_map_fn('test', data_source) - - dataset = dataset.map(function=process_fn, with_indices=True) - print(f"[DEBUG] processed dataset: {len(dataset)}") - - # Filter out examples where processing failed - dataset = dataset.filter(lambda x: x["data_source"] == data_source) - print(f"[DEBUG] filtered dataset: {len(dataset)}") - # Length filter - try: - tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B") - length_filter = LengthFilter(tokenizer=tokenizer, max_length=4096) - dataset = dataset.filter(lambda x: length_filter.check(x)) - except Exception as e: - print(f"Warning: Could not perform length filtering. Error: {e}") - print("Proceeding without length filtering.") - - # Sample the dataset - dataset = sample_dataset(dataset, args.sample_size) - - # Save the dataset to test directory - test_output_path = save_dataset( - dataset=dataset, - output_dir=test_output_dir, - filename_prefix=data_source, - sample_size=len(dataset) - ) - - print(f"\nDone! \n" - f"Data source: {data_source}\n" - f"Test data saved to {test_output_path} ({len(dataset)} samples)") \ No newline at end of file diff --git a/reasoners/data_preprocess/codegen/humanevalplus.py b/reasoners/data_preprocess/codegen/humanevalplus.py deleted file mode 100644 index a03d13938..000000000 --- a/reasoners/data_preprocess/codegen/humanevalplus.py +++ /dev/null @@ -1,157 +0,0 @@ -"""Downloads, processes, and saves HumanEval dataset.""" - -import os -import argparse -import json -import transformers -import datasets -from datasets import load_dataset, Dataset - -from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset -from verl.utils.data_process.filter import LengthFilter -from verl.utils.reward_score.coder1 import code_exec - - -def get_datasets(cache_dir: str): - """ - Loads the HumanEvalPlus dataset. - """ - try: - dataset = load_dataset("autoprogrammer/humanevalplus_corrected", cache_dir=cache_dir)["test"] - print(f"HumanEvalPlus dataset: {len(dataset)} examples") - return None, dataset - except Exception as e: - print(f"Error loading dataset: {e}") - return None, None - - -def make_map_fn(split: str, data_source: str) -> callable: - def process_fn(example, idx): - # Create a default "skip" response with all required fields - skip_response = { - "data_source": None, - "prompt": None, - "ability": None, - "apply_chat_template": None, - "reward_model": None, - "extra_info": None - } - - # Extract task ID and prompt - task_id = example["task_id"] - - prompt = ( - "Write a complete, self-contained Python solution to the following problem. " - "Your solution must include all necessary imports and the full function definition including " - "the signature exactly as specified. Do not modify the function signature or docstring.\n\n" - f"```python\n{example['prompt'].strip()}\n```" - ) - - # Extract test function, entry point, and canonical solution - test_code = example["test"] - entry_point = example["entry_point"] - solution = example["canonical_solution"] - - # Build test code that calls the entry point - test_code_with_check = f"{test_code}\n\ncheck({entry_point})" - - # Verify the canonical solution passes the tests - # full_code = f"{prompt}\n{solution}\n{test_code}\n\ncheck({entry_point})" - full_code = f"{example['prompt']}\n{solution}\n{test_code}\n\ncheck({entry_point})" - succ, err = code_exec(full_code, timeout=30) - print(f"[DEBUG] succ: {succ}, err: {err}") - if not succ: - print(f"Error in canonical solution for task {task_id}: {err}") - return skip_response - - - data = { - "data_source": data_source, - "prompt": [ - {"role": "user", "content": prompt} - ], - "ability": "codegen", - "apply_chat_template": True, - "reward_model": { - "style": "rule", - "ground_truth": json.dumps({ - "functional": test_code_with_check - }), - }, - "extra_info": { - "split": split, - "index": idx, - "reference": solution, # Include the canonical solution as reference - "original_prompt": prompt, - "dataset": "evalplus_humanevalplus", - "task_id": task_id, - }, - } - - if idx == 0 or idx == 1: - print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) - print(data) - - return data - - return process_fn - - -if __name__ == '__main__': - """Main script execution: parse args, load, process, and save datasets.""" - parser = argparse.ArgumentParser(description="Process and save HumanEvalPlus dataset.") - parser.add_argument('--data-dir', default='data', - help='Base directory to save the processed data files.') - parser.add_argument('--domain', default="codegen", - help='Domain of the dataset.') - parser.add_argument('--name', default="humanevalplus", - help='Name of the dataset.') - parser.add_argument('--sample-size', type=int, default=None, - help='Number of samples to use from dataset. If None, use all samples.') - parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility') - - args = parser.parse_args() - - set_seed(args.seed) - - data_source = f"{args.domain}__{args.name}" - test_output_dir = os.path.join(args.data_dir, 'test') - os.makedirs(test_output_dir, exist_ok=True) - - # Load the dataset - cache_dir = datasets.config.HF_DATASETS_CACHE - _, dataset = get_datasets(cache_dir=cache_dir) - print(f"[DEBUG] dataset: {len(dataset)}") - - # Process the dataset - process_fn = make_map_fn('test', data_source) - - dataset = dataset.map(function=process_fn, with_indices=True,load_from_cache_file=False) - print(f"[DEBUG] processed dataset: {len(dataset)}") - - # Filter out examples where processing failed - dataset = dataset.filter(lambda x: x["data_source"] == data_source) - print(f"[DEBUG] filtered dataset: {len(dataset)}") - # Length filter - try: - tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B") - length_filter = LengthFilter(tokenizer=tokenizer, max_length=4096) - dataset = dataset.filter(lambda x: length_filter.check(x)) - except Exception as e: - print(f"Warning: Could not perform length filtering. Error: {e}") - print("Proceeding without length filtering.") - - # Sample the dataset - dataset = sample_dataset(dataset, args.sample_size) - - # Save the dataset to test directory - test_output_path = save_dataset( - dataset=dataset, - output_dir=test_output_dir, - filename_prefix=data_source, - sample_size=len(dataset) - ) - - print(f"\nDone! \n" - f"Data source: {data_source}\n" - f"Test data saved to {test_output_path} ({len(dataset)} samples)") \ No newline at end of file diff --git a/reasoners/data_preprocess/codegen/leetcode2k.py b/reasoners/data_preprocess/codegen/leetcode2k.py deleted file mode 100644 index 22d48fc24..000000000 --- a/reasoners/data_preprocess/codegen/leetcode2k.py +++ /dev/null @@ -1,174 +0,0 @@ -"""Downloads, processes, and saves LeetCode2K datasets.""" - -import os -import argparse -import json -import transformers -import datasets -from datasets import load_dataset, Dataset - -from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset -from verl.utils.data_process.filter import LengthFilter -from verl.utils.reward_score.coder1 import code_exec - -EMPTY_EXAMPLE = { - "data_source": None, - "prompt": None, - "apply_chat_template": False, - "ability": None, - "reward_model": None, - "extra_info": None -} - - -def get_datasets(cache_dir: str): - """ - Loads the LeetCode2K dataset. - """ - try: - commit_hash = "34803eb64eab1979069ba1f80e7ea474282e28f3" - train_dataset = load_dataset("newfacade/LeetCodeDataset", - revision=commit_hash, - cache_dir=cache_dir)["train"] - test_dataset = load_dataset("newfacade/LeetCodeDataset", - revision=commit_hash, - cache_dir=cache_dir)["test"] - print(f"Train set: {len(train_dataset)} examples") - print(f"Test set: {len(test_dataset)} examples") - return train_dataset, test_dataset - except Exception as e: - print(f"Error loading dataset: {e}") - return None, None - - -def make_map_fn(split: str, data_source: str) -> callable: - def process_fn(example, idx): - prefix = example["prompt"] - - # Clean up the prompt - prompt = example["query"] - prompt = prompt.replace("### Answer: (use the provided format with backticks)", "").strip() - prompt = prompt.replace("### Format: ", "### Format:\n") - - # Build test code - test_code = f"{example['test']}\n\ncheck({example['entry_point'].strip()})" - - # Extract the candidate solution - solution = example["completion"] - - # Combine all code pieces into a single file to execute - full_code = f"{prefix}\n{solution}\n{test_code}" - - # Validate that the candidate solution passes the tests - # Skip examples where the test code fails - succ, err = code_exec(full_code, timeout=20) - if not succ: - print(f"Test code failed for example {idx}: {example['meta']['question_title']}") - return EMPTY_EXAMPLE - - data = { - "data_source": data_source, - "prompt": [ - {"role": "user", "content": prompt} - ], - "ability": "codegen", - "apply_chat_template": True, - "reward_model": { - "style": "rule", - "ground_truth": json.dumps({ - "functional": test_code - }), - }, - "extra_info": { - "split": split, - "index": idx, - "reference": solution, - "original_prompt": prompt, - "prefix": prefix, - "dataset": "LeetCodeDataset", - "question_title": example["meta"]["question_title"], - "difficulty": example["meta"]["difficulty"], - }, - } - - if idx == 0 or idx == 1: - print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) - print(data) - - return data - - return process_fn - - -if __name__ == '__main__': - """Main script execution: parse args, load, process, and save datasets.""" - parser = argparse.ArgumentParser(description="Process and save LeetCode2K datasets.") - parser.add_argument('--data-dir', default='data', - help='Base directory to save the processed data files.') - parser.add_argument('--domain', default="codegen", - help='Domain of the dataset.') - parser.add_argument('--name', default="leetcode2k", - help='Name of the dataset.') - parser.add_argument('--train-sample-size', type=int, default=None, - help='Number of samples to use from training dataset. If None, use all samples.') - parser.add_argument('--test-sample-size', type=int, default=None, - help='Number of samples to use from test dataset. If None, use all samples.') - parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility') - - args = parser.parse_args() - - set_seed(args.seed) - - data_source = f"{args.domain}__{args.name}" - train_output_dir = os.path.join(args.data_dir, 'train') - test_output_dir = os.path.join(args.data_dir, 'test') - os.makedirs(train_output_dir, exist_ok=True) - os.makedirs(test_output_dir, exist_ok=True) - - # Load the dataset - cache_dir = datasets.config.HF_DATASETS_CACHE - train_dataset, test_dataset = get_datasets(cache_dir=cache_dir) - - # Process the dataset - process_train_fn = make_map_fn('train', data_source) - process_test_fn = make_map_fn('test', data_source) - - train_dataset = train_dataset.map(function=process_train_fn, with_indices=True, num_proc=64) - test_dataset = test_dataset.map(function=process_test_fn, with_indices=True, num_proc=64) - - # Filter out examples where processing failed - train_dataset = train_dataset.filter(lambda x: x["data_source"] == data_source) - test_dataset = test_dataset.filter(lambda x: x["data_source"] == data_source) - - # Length filter - try: - tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B") - length_filter = LengthFilter(tokenizer=tokenizer, max_length=4096) - train_dataset = train_dataset.filter(lambda x: length_filter.check(x)) - test_dataset = test_dataset.filter(lambda x: length_filter.check(x)) - except Exception as e: - print(f"Warning: Could not perform length filtering. Error: {e}") - print("Proceeding without length filtering.") - - # Sample the datasets using utility function - train_dataset = sample_dataset(train_dataset, args.train_sample_size) - test_dataset = sample_dataset(test_dataset, args.test_sample_size) - - # Save the datasets using utility function - train_output_path = save_dataset( - dataset=train_dataset, - output_dir=train_output_dir, - filename_prefix=data_source, - sample_size=len(train_dataset) - ) - test_output_path = save_dataset( - dataset=test_dataset, - output_dir=test_output_dir, - filename_prefix=data_source, - sample_size=len(test_dataset) - ) - - print(f"\nDone! \n" - f"Data source: {data_source}\n" - f"Train data saved to {train_output_path} ({len(train_dataset)} samples)\n" - f"Test data saved to {test_output_path} ({len(test_dataset)} samples)") \ No newline at end of file diff --git a/reasoners/data_preprocess/codegen/livecodebench.py b/reasoners/data_preprocess/codegen/livecodebench.py deleted file mode 100644 index 090ea9345..000000000 --- a/reasoners/data_preprocess/codegen/livecodebench.py +++ /dev/null @@ -1,230 +0,0 @@ -"""Downloads, processes, and saves MBPP datasets.""" - -import os -import argparse -import json -import transformers -import datasets -from datasets import load_dataset, Dataset - -from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset -from verl.utils.data_process.filter import LengthFilter - - -def get_datasets(cache_dir: str): - """ - Loads the LiveCodeBench dataset. - """ - try: - dataset = load_dataset("agentica-org/DeepCoder-Preview-Dataset", "lcbv5", cache_dir=cache_dir) - print(f"Train set: {len(dataset['train'])} examples") - print(f"Test set: {len(dataset['test'])} examples") - return dataset["train"], dataset["test"] - except Exception as e: - print(f"Error loading dataset: {e}") - return None, None - - -def make_map_fn(split: str, data_source: str) -> callable: - def process_fn(example, idx): - # Get the problem description - problem_desc = example["problem"] - starter_code = example["starter_code"] - - # Process tests - try: - tests = json.loads(example["tests"]) - if not tests: - return { - "data_source": None, - "prompt": None, - "ability": None, - "reward_model": None, - "extra_info": None - } - - # Process metadata to get function name - metadata = example["metadata"] - function_name = metadata.get("func_name") - - # Handle different test types - if tests[0]["testtype"] == "functional": - if not function_name: - return { - "data_source": None, - "prompt": None, - "raw_prompt": None, - "ability": None, - "reward_model": None, - "extra_info": None - } - - # create the prompt with the starter code - # prompt adopted from livecodebench official implementation - # https://github.com/LiveCodeBench/LiveCodeBench/blob/main/lcb_runner/prompts/code_generation.py - prompt = f"""You are an expert Python programmer. You will be given a question (problem specification) and will generate a correct Python program that matches the specification and passes all tests. - -Below is the question: - -{problem_desc} - -You will use the following starter code to write the solution to the problem and enclose your code within ```python delimiters. -```python -{starter_code} -```""" - - # Function call tests - test_code = f"""\ -def check_{function_name}(): -""" - for test in tests: - # Parse input string by splitting on '\n' - input_parts = test["input"].split('\n') - # Create proper comma-separated arguments for the function call - input_args = ', '.join(input_parts) - # Get the output value - output_val = test["output"] - - test_code += f""" assert Solution().{function_name}({input_args}) == {output_val} -""" - - test_code += f""" -check_{function_name}() -""" - - oracle = json.dumps({"functional": test_code}) - - elif tests[0]["testtype"] == "stdin": - # create the prompt, and there will be no starter code - # prompt adopted from livecodebench official implementation - # https://github.com/LiveCodeBench/LiveCodeBench/blob/main/lcb_runner/prompts/code_generation.py - prompt = f"""You are an expert Python programmer. You will be given a question (problem specification) and will generate a correct Python program that matches the specification and passes all tests. - -Below is the question: - -{problem_desc} - -Read the inputs from stdin solve the problem and write the answer to stdout (do not directly test on the sample inputs). Enclose your code within delimiters. Ensure that when the python program runs, it reads the inputs, runs the algorithm and writes output to STDOUT.""" - - # STDIN/STDOUT tests - stdin_list = [] - stdout_list = [] - for test in tests: - stdin_list.append(test["input"]) - stdout_list.append(test["output"]) - - oracle = json.dumps({"inputs": stdin_list, "outputs": stdout_list}) - else: - raise ValueError(f"Unknown test type: {tests[0]['testtype']}") - except Exception as e: - print(f"Error processing tests for example {idx}: {e}") - return { - "data_source": None, - "prompt": None, - "ability": None, - "reward_model": None, - "extra_info": None - } - - data = { - "data_source": data_source, - "prompt": [ - {"role": "user", "content": prompt} - ], - "ability": "codegen", - "apply_chat_template": True, - "reward_model": { - "style": "rule", - "ground_truth": oracle, - }, - "extra_info": { - "split": split, - "index": idx, - "reference": "", # No solution data in LiveCodeBench - "dataset": "LiveCodeBench", - "function_name": function_name if function_name else "", - "original_prompt": prompt, - }, - } - - if idx == 0 or idx == 1: - print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) - print(data) - - return data - - return process_fn - - -if __name__ == '__main__': - """Main script execution: parse args, load, process, and save datasets.""" - parser = argparse.ArgumentParser(description="Process and save LiveCodeBench datasets.") - parser.add_argument('--data-dir', default='data', - help='Base directory to save the processed data files.') - parser.add_argument('--domain', type=str, default='codegen', - help='Domain of the dataset.') - parser.add_argument('--name', type=str, default='livecodebench', - help='Name of the dataset.') - parser.add_argument('--train-sample-size', type=int, default=None, - help='Number of samples to use from training dataset. If None, use all samples.') - parser.add_argument('--test-sample-size', type=int, default=None, - help='Number of samples to use from test dataset. If None, use all samples.') - parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility') - - args = parser.parse_args() - - set_seed(args.seed) - - data_source = f"{args.domain}__{args.name}" - train_output_dir = os.path.join(args.data_dir, 'train') - test_output_dir = os.path.join(args.data_dir, 'test') - os.makedirs(train_output_dir, exist_ok=True) - os.makedirs(test_output_dir, exist_ok=True) - - # Load the dataset - cache_dir = datasets.config.HF_DATASETS_CACHE - train_dataset, test_dataset = get_datasets(cache_dir) - - # Process the dataset - process_train_fn = make_map_fn('train', data_source) - process_test_fn = make_map_fn('test', data_source) - - train_dataset = train_dataset.map(function=process_train_fn, with_indices=True, num_proc=64) - test_dataset = test_dataset.map(function=process_test_fn, with_indices=True, num_proc=64) - - # Filter out examples where processing failed - train_dataset = train_dataset.filter(lambda x: x["data_source"] == data_source) - test_dataset = test_dataset.filter(lambda x: x["data_source"] == data_source) - - # Length filter - try: - tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B") - length_filter = LengthFilter(tokenizer=tokenizer, max_length=4096) - train_dataset = train_dataset.filter(lambda x: length_filter.check(x), num_proc=64) - test_dataset = test_dataset.filter(lambda x: length_filter.check(x), num_proc=64) - except Exception as e: - print(f"Warning: Could not perform length filtering. Error: {e}") - print("Proceeding without length filtering.") - - # Sample the datasets using utility function - train_dataset = sample_dataset(train_dataset, args.train_sample_size) - test_dataset = sample_dataset(test_dataset, args.test_sample_size) - - # Save the datasets using utility function - train_output_path = save_dataset( - dataset=train_dataset, - output_dir=train_output_dir, - filename_prefix=data_source, - sample_size=len(train_dataset) - ) - test_output_path = save_dataset( - dataset=test_dataset, - output_dir=test_output_dir, - filename_prefix=data_source, - sample_size=len(test_dataset) - ) - - print(f"\nDone! \n" - f"Data source: {data_source}\n" - f"Train data saved to {train_output_path} ({len(train_dataset)} samples)\n" - f"Test data saved to {test_output_path} ({len(test_dataset)} samples)") \ No newline at end of file diff --git a/reasoners/data_preprocess/codegen/mbpp.py b/reasoners/data_preprocess/codegen/mbpp.py deleted file mode 100644 index d77d5efac..000000000 --- a/reasoners/data_preprocess/codegen/mbpp.py +++ /dev/null @@ -1,169 +0,0 @@ -"""Downloads, processes, and saves MBPP datasets.""" - -import os -import argparse -import json -import transformers -import datasets -from datasets import load_dataset, Dataset - -from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset -from verl.utils.data_process.filter import LengthFilter -from verl.utils.reward_score.coder1 import code_exec - - -def get_datasets(cache_dir: str): - """ - Loads the MBPP dataset. - """ - try: - dataset = load_dataset("google-research-datasets/mbpp", cache_dir=cache_dir) - print(f"Train set: {len(dataset['train'])} examples") - print(f"Test set: {len(dataset['test'])} examples") - return dataset["train"], dataset["test"] - except Exception as e: - print(f"Error loading dataset: {e}") - return None, None - - -def make_map_fn(split: str, data_source: str) -> callable: - def process_fn(example, idx): - # Create prompt - prompt = ( - f"{example['text']}\n\n" - f"Your solution should be a complete, self-contained function in a markdown code block. " - f"Make sure your solution passes the following test cases:\n" - ) - - # Construct test code - test_code = "" - if example.get('test_setup_code'): - test_code += example['test_setup_code'] + "\n\n" - - # Add all test assertions - for assertion in example['test_list'] + example.get('challenge_test_list', []): - test_code += assertion + "\n" - - # Add test cases to prompt - prompt += f"```python\n{test_code}```" - prompt += "\n\nPlease do not include the test cases in your solution." - - # Validate that the canonical solution passes the tests - solution = example['code'] - full_code = f"{solution}\n\n{test_code}" - - succ, err = code_exec(full_code) - if not succ: - print(f"Test code failed for example {idx}: {example.get('task_id', 'unknown')}") - return { - "data_source": None, - "prompt": None, - "ability": None, - "reward_model": None, - "extra_info": None - } - - data = { - "data_source": data_source, - "prompt": [ - {"role": "user", "content": prompt} - ], - "ability": "codegen", - "apply_chat_template": True, - "reward_model": { - "style": "rule", - "ground_truth": json.dumps({ - "functional": test_code - }), - }, - "extra_info": { - "split": split, - "index": idx, - "reference": solution, - "original_prompt": prompt, - "dataset": "mbpp", - "task_id": str(example.get("task_id", "")), - }, - } - - if idx == 0 or idx == 1: - print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) - print(data) - - return data - - return process_fn - - -if __name__ == '__main__': - """Main script execution: parse args, load, process, and save datasets.""" - parser = argparse.ArgumentParser(description="Process and save MBPP datasets.") - parser.add_argument('--data-dir', default='data', - help='Base directory to save the processed data files.') - parser.add_argument('--domain', default="codegen", - help='Domain of the dataset.') - parser.add_argument('--name', default="mbpp", - help='Name of the dataset.') - parser.add_argument('--train-sample-size', type=int, default=None, - help='Number of samples to use from training dataset. If None, use all samples.') - parser.add_argument('--test-sample-size', type=int, default=None, - help='Number of samples to use from test dataset. If None, use all samples.') - parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility') - - args = parser.parse_args() - - set_seed(args.seed) - - data_source = f"{args.domain}__{args.name}" - train_output_dir = os.path.join(args.data_dir, 'train') - test_output_dir = os.path.join(args.data_dir, 'test') - os.makedirs(train_output_dir, exist_ok=True) - os.makedirs(test_output_dir, exist_ok=True) - - # Load the dataset - cache_dir = datasets.config.HF_DATASETS_CACHE - train_dataset, test_dataset = get_datasets(cache_dir=cache_dir) - - # Process the dataset - process_train_fn = make_map_fn('train', data_source) - process_test_fn = make_map_fn('test', data_source) - - train_dataset = train_dataset.map(function=process_train_fn, with_indices=True, num_proc=64) - test_dataset = test_dataset.map(function=process_test_fn, with_indices=True, num_proc=64) - - # Filter out examples where processing failed - train_dataset = train_dataset.filter(lambda x: x["data_source"] == data_source) - test_dataset = test_dataset.filter(lambda x: x["data_source"] == data_source) - - # Length filter - try: - tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B") - length_filter = LengthFilter(tokenizer=tokenizer, max_length=4096) - train_dataset = train_dataset.filter(lambda x: length_filter.check(x)) - test_dataset = test_dataset.filter(lambda x: length_filter.check(x)) - except Exception as e: - print(f"Warning: Could not perform length filtering. Error: {e}") - print("Proceeding without length filtering.") - - # Sample the datasets using utility function - train_dataset = sample_dataset(train_dataset, args.train_sample_size) - test_dataset = sample_dataset(test_dataset, args.test_sample_size) - - # Save the datasets using utility function - train_output_path = save_dataset( - dataset=train_dataset, - output_dir=train_output_dir, - filename_prefix=data_source, - sample_size=len(train_dataset) - ) - test_output_path = save_dataset( - dataset=test_dataset, - output_dir=test_output_dir, - filename_prefix=data_source, - sample_size=len(test_dataset) - ) - - print(f"\nDone! \n" - f"Data source: {data_source}\n" - f"Train data saved to {train_output_path} ({len(train_dataset)} samples)\n" - f"Test data saved to {test_output_path} ({len(test_dataset)} samples)") \ No newline at end of file diff --git a/reasoners/data_preprocess/codegen/mbppplus.py b/reasoners/data_preprocess/codegen/mbppplus.py deleted file mode 100644 index 603e7e367..000000000 --- a/reasoners/data_preprocess/codegen/mbppplus.py +++ /dev/null @@ -1,158 +0,0 @@ -"""Downloads, processes, and saves MBPP+ datasets.""" - -import os -import argparse -import json -import transformers -import datasets -from datasets import load_dataset, Dataset - -from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset -from verl.utils.data_process.filter import LengthFilter -from verl.utils.reward_score.coder1 import code_exec - - -def get_datasets(cache_dir: str): - """ - Loads the MBPP+ dataset. - """ - try: - dataset = load_dataset("evalplus/mbppplus", cache_dir=cache_dir) - print(f"Test set: {len(dataset['test'])} examples") - return None, dataset["test"] - except Exception as e: - print(f"Error loading dataset: {e}") - return None, None - - -def make_map_fn(split: str, data_source: str) -> callable: - def process_fn(example, idx): - # Create prompt - prompt = ( - f"{example['prompt']}\n\n" - f"Your solution should be a complete, self-contained function in a markdown code block. " - f"Make sure your solution passes the following test cases:\n" - ) - # Construct test code: imports + assertions - test_code = "" - if example.get('test_imports'): - for imp in example['test_imports']: - test_code += imp + "\n" - test_code += "\n" - # Add all test assertions - test_code += example['test'] - example_code="" - for assertion in example['test_list'] + example.get('challenge_test_list', []): - example_code += assertion + "\n" - - # Add test cases to prompt - prompt += f"```python\n{example_code}```" - prompt += "\n\nPlease do not include the test cases in your solution." - - # Validate that the canonical solution passes the tests - solution = example['code'] - full_code = f"{solution}\n\n{test_code}" - succ, err = code_exec(full_code) - if not succ: - print(f"Test code failed for example {idx}: {example.get('task_id', 'unknown')}") - return { - "data_source": None, - "prompt": None, - "ability": None, - "reward_model": None, - "extra_info": None - } - - data = { - "data_source": data_source, - "prompt": [ - {"role": "user", "content": prompt} - ], - "ability": "codegen", - "apply_chat_template": True, - "reward_model": { - "style": "rule", - "ground_truth": json.dumps({ - "functional": test_code - }), - }, - "extra_info": { - "split": split, - "index": idx, - "reference": solution, - "original_prompt": prompt, - "dataset": "mbpp", - "task_id": str(example.get("task_id", "")), - }, - } - - if idx == 0 or idx == 1: - print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) - print(data) - - return data - - return process_fn - - -if __name__ == '__main__': - """Main script execution: parse args, load, process, and save datasets.""" - parser = argparse.ArgumentParser(description="Process and save MBPP datasets.") - parser.add_argument('--data-dir', default='data', - help='Base directory to save the processed data files.') - parser.add_argument('--domain', default="codegen", - help='Domain of the dataset.') - parser.add_argument('--name', default="mbppplus", - help='Name of the dataset.') - parser.add_argument('--train-sample-size', type=int, default=None, - help='Number of samples to use from training dataset. If None, use all samples.') - parser.add_argument('--test-sample-size', type=int, default=None, - help='Number of samples to use from test dataset. If None, use all samples.') - parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility') - - args = parser.parse_args() - - set_seed(args.seed) - - data_source = f"{args.domain}__{args.name}" - train_output_dir = os.path.join(args.data_dir, 'train') - test_output_dir = os.path.join(args.data_dir, 'test') - os.makedirs(train_output_dir, exist_ok=True) - os.makedirs(test_output_dir, exist_ok=True) - - # Load the dataset - cache_dir = datasets.config.HF_DATASETS_CACHE - - _, test_dataset = get_datasets(cache_dir=cache_dir) - - # Process the dataset - process_test_fn = make_map_fn('test', data_source) - - test_dataset = test_dataset.map(function=process_test_fn, with_indices=True, num_proc=64, load_from_cache_file=False) - - # Filter out examples where processing failed - test_dataset = test_dataset.filter(lambda x: x["data_source"] == data_source) - - # # Length filter - # try: - # tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B") - # length_filter = LengthFilter(tokenizer=tokenizer, max_length=4096) - # test_dataset = test_dataset.filter(lambda x: length_filter.check(x)) - # except Exception as e: - # print(f"Warning: Could not perform length filtering. Error: {e}") - # print("Proceeding without length filtering.") - - # Sample the datasets using utility function - test_dataset = sample_dataset(test_dataset, args.test_sample_size) - - # Save the datasets using utility function - test_output_path = save_dataset( - dataset=test_dataset, - output_dir=test_output_dir, - filename_prefix=data_source, - sample_size=len(test_dataset) - ) - - print(f"\nDone! \n" - f"Data source: {data_source}\n" - f"Test data saved to {test_output_path} ({len(test_dataset)} samples)") \ No newline at end of file diff --git a/reasoners/data_preprocess/codegen/primeintellect.py b/reasoners/data_preprocess/codegen/primeintellect.py deleted file mode 100644 index f122a4e31..000000000 --- a/reasoners/data_preprocess/codegen/primeintellect.py +++ /dev/null @@ -1,236 +0,0 @@ -"""Downloads, processes, and saves PrimeIntellect datasets.""" - -import os -import argparse -import json -import transformers -import datasets -from datasets import load_dataset, Dataset -from concurrent.futures import ThreadPoolExecutor, as_completed - -from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset -from verl.utils.data_process.filter import LengthFilter -from verl.utils.reward_score.coder1 import code_exec, remote_check_stdio, fuzzy_equal, extract_code_from_string - -# Define a constant for the empty/failed example -EMPTY_EXAMPLE = { - "data_source": None, - "prompt": None, - "apply_chat_template": False, - "ability": None, - "reward_model": None, - "extra_info": None -} - -def get_datasets(cache_dir: str): - """ - Loads the PrimeIntellect dataset. - """ - try: - dataset = load_dataset("agentica-org/DeepCoder-Preview-Dataset", "primeintellect", split="train", cache_dir=cache_dir) - - # TODO: Remove this line before production - only selecting 100 examples for debugging - # dataset = dataset.select(range(100)) - - print(f"Dataset: {len(dataset)} examples") - return dataset, None - except Exception as e: - print(f"Error loading dataset: {e}") - return None, None - - -def make_map_fn(split: str, data_source: str, verbose: bool) -> callable: - def process_fn(example, idx): - # Get the problem description - prompt = example["problem"] - - # Get the solution code - solution = example["solutions"][0] if example["solutions"] else "" - solution = extract_code_from_string(solution) - - # Process tests - tests = json.loads(example["tests"]) - - # Now let's do some filtering - # 1. Remove examples with no tests - if not tests: - print(f"No tests found for example {idx}") - return EMPTY_EXAMPLE - - # 2. Remove examples with "image" in prompt, image typically will be in the following format: - # or [image] - if "" in prompt.lower() or "[image]" in prompt.lower(): - print(f"Image found in prompt for example {idx}") - return EMPTY_EXAMPLE - - # 3. Remove examples with no problem description - # Check if prompt starts with unwanted patterns after removing common prefix - check_prompt = prompt - if "Solve the following coding problem using the programming language python:" in check_prompt: - check_prompt = check_prompt.split("Solve the following coding problem using the programming language python:")[1].lstrip() - - if (check_prompt.lower().startswith("example") or - check_prompt.lower().startswith("input") or - check_prompt.lower().startswith("-----input-----")): - print(f"Example starts with unwanted pattern for example {idx}") - return EMPTY_EXAMPLE - - # Handle different test types - if tests[0]["type"] == "function_call": - # Function call tests - fn_name = tests[0]["fn_name"] - - inputs = [] - outputs = [] - for test in tests: - inputs.append(test["input"]) - outputs.append(test["output"]) - - test_code = f"""\ -_inputs = {inputs} -_outputs = {outputs} -import math -def _deep_eq(a, b, tol=1e-5): - if isinstance(a, float) or isinstance(b, float): - return math.isclose(a, b, rel_tol=tol, abs_tol=tol) - if isinstance(a, (list, tuple)): - if len(a) != len(b): return False - return all(_deep_eq(x, y, tol) for x, y in zip(a, b)) - return a == b - -for i, o in zip(_inputs, _outputs): - assert _deep_eq({fn_name}(*i), o[0] if len(o) == 1 else tuple(o)) -""" - - # Validate the solution - full_code = f"{solution}\n{test_code}" - succ, err = code_exec(full_code) - if not succ: - print(f"Test code failed for example {idx}") - print(f"Error: {err}") - return EMPTY_EXAMPLE - - oracle = json.dumps({"functional": test_code}) - - elif tests[0]["type"] == "stdin_stdout": - # STDIN/STDOUT tests - stdin_list = [] - stdout_list = [] - for test in tests: - stdin_list.append(test["input"]) - stdout_list.append(test["output"]) - - # Validate the solution - with ThreadPoolExecutor(max_workers=min(len(stdin_list), 8)) as executor: - futures = [] - for stdin, stdout in zip(stdin_list, stdout_list): - futures.append( - executor.submit( - remote_check_stdio, - solution, - stdin, - stdout, - ) - ) - for future in as_completed(futures): - exec_succ, output, stdin, stdout = future.result() - pass_test = exec_succ and fuzzy_equal(output.strip(), stdout.strip(), verbose=verbose) - if not pass_test: - print(f"Test code failed for example {idx}") - if verbose: - print(f"Input: {stdin}") - print(f"Expected output: {stdout}") - print(f"Actual output: {output}") - return EMPTY_EXAMPLE - - oracle = json.dumps({"inputs": stdin_list, "outputs": stdout_list}) - else: - print(f"Unknown test type: {tests[0]['type']} for example {idx}") - return EMPTY_EXAMPLE - - data = { - "data_source": data_source, - "prompt": [ - {"role": "user", "content": prompt} - ], - "ability": "codegen", - "apply_chat_template": True, - "reward_model": { - "style": "rule", - "ground_truth": oracle, - }, - "extra_info": { - "split": split, - "index": idx, - "reference": solution, - "original_prompt": prompt, - "dataset": "PrimeIntellect", - "function_name": fn_name if tests[0]["type"] == "function_call" else "" - }, - } - - if idx == 0 or idx == 1: - print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) - print(data) - - return data - - return process_fn - - -if __name__ == '__main__': - """Main script execution: parse args, load, process, and save datasets.""" - parser = argparse.ArgumentParser(description="Process and save PrimeIntellect datasets.") - parser.add_argument('--data-dir', type=str, default='data', help='Directory to save the processed datasets.') - parser.add_argument('--domain', default="codegen", - help='Domain of the dataset.') - parser.add_argument('--name', default="primeintellect", - help='Name of the dataset.') - parser.add_argument('--train-sample-size', type=int, default=None, - help='Number of samples to use from training dataset. If None, use all samples.') - parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility') - parser.add_argument('--verbose', type=bool, default=False, help='Whether to print verbose output.') - - args = parser.parse_args() - - set_seed(args.seed) - - data_source = f"{args.domain}__{args.name}" - train_output_dir = os.path.join(args.data_dir, 'train') - os.makedirs(train_output_dir, exist_ok=True) - - # Load the dataset - cache_dir = datasets.config.HF_DATASETS_CACHE - dataset, _ = get_datasets(cache_dir) - - # Process the dataset - process_fn = make_map_fn('train', data_source, args.verbose) - - dataset = dataset.map(function=process_fn, with_indices=True, num_proc=64) - - # Filter out examples where processing failed - dataset = dataset.filter(lambda x: x["data_source"] == data_source) - - # Length filter - try: - tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B") - length_filter = LengthFilter(tokenizer=tokenizer, max_length=4096) - dataset = dataset.filter(lambda x: length_filter.check(x),) - except Exception as e: - print(f"Warning: Could not perform length filtering. Error: {e}") - print("Proceeding without length filtering.") - - # Sample the dataset using utility function - train_dataset = sample_dataset(dataset, args.train_sample_size) - - # Save the dataset using utility function - train_output_path = save_dataset( - dataset=train_dataset, - output_dir=train_output_dir, - filename_prefix=data_source, - sample_size=len(train_dataset) - ) - - print(f"\nDone! \n" - f"Data source: {data_source}\n" - f"Train data saved to {train_output_path} ({len(train_dataset)} samples)") \ No newline at end of file diff --git a/reasoners/data_preprocess/codegen/taco.py b/reasoners/data_preprocess/codegen/taco.py deleted file mode 100644 index 31eb8be70..000000000 --- a/reasoners/data_preprocess/codegen/taco.py +++ /dev/null @@ -1,246 +0,0 @@ -"""Downloads, processes, and saves TACO dataset.""" - -import os -import argparse -import json -import sys -from concurrent.futures import ThreadPoolExecutor, as_completed -import datasets -import transformers -from datasets import load_dataset, Dataset - -from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset -from verl.utils.data_process.filter import LengthFilter -from verl.utils.reward_score.coder1 import ( - code_exec, - remote_check_stdio, - fuzzy_equal -) - -EMPTY_EXAMPLE = { - "data_source": None, - "prompt": None, - "apply_chat_template": False, - "ability": None, - "reward_model": None, - "extra_info": None -} - - - -def get_datasets(cache_dir: str): - """ - Loads the TACO dataset. - """ - try: - dataset = load_dataset("likaixin/TACO-verified", trust_remote_code=True, split="train", cache_dir=cache_dir) - print(f"TACO dataset: {len(dataset)} examples") - return dataset, None - except Exception as e: - print(f"Error loading dataset: {e}") - return Dataset.from_list([]), None - - -def make_map_fn(split: str, data_source: str) -> callable: - def process_fn(example, idx): - # Create a default "skip" response with all required fields - oracle = json.loads(example["input_output"]) - source = example["source"] - - # Skip poorly formatted examples - if source in ["geeksforgeeks", "leetcode"]: - return EMPTY_EXAMPLE - - # Skip examples with too short descriptions - if len("".join([c for c in example["question"] if c.isalnum()])) < 100: - return EMPTY_EXAMPLE - - # Skip examples with images - if "image" in example["question"].lower() or "\n![" in example["question"]: - return EMPTY_EXAMPLE - - # Build prompt - prompt_pieces = [ - "Solve the programming task below in a Python markdown code block.", - example["question"].strip(), - ] - if example["starter_code"].strip(): - prompt_pieces.append( - "You will use the following starter code to write the solution to the problem and enclose your code within ```python delimiters." - ) - prompt_pieces.append( - f"```python\n{example['starter_code'].strip()}\n```" - ) - else: - prompt_pieces.append( - "Read the inputs from stdin solve the problem and write the answer to stdout (do not directly test on the sample inputs). Enclose your code within ```python delimiters." - ) - - # Process oracle based on format - if "fn_name" in oracle: # Function-based tests - fn_name = oracle["fn_name"] - if source == "leetcode": - fn_name = "Solution()." + fn_name - - test_code = f"""\ -_inputs = {oracle["inputs"]} -_outputs = {oracle["outputs"]} -import math -def _deep_eq(a, b, tol=1e-5): - if isinstance(a, float) or isinstance(b, float): - return math.isclose(a, b, rel_tol=tol, abs_tol=tol) - if isinstance(a, (list, tuple)): - if len(a) != len(b): return False - return all(_deep_eq(x, y, tol) for x, y in zip(a, b)) - return a == b - -for i, o in zip(_inputs, _outputs): -""" - - if source in ["leetcode", "hackerrank"]: - test_code += f" assert _deep_eq({fn_name}(*i), o)" - elif source == "codewars": - test_code += f" assert _deep_eq({fn_name}(*i), o[0])" - else: - print(f"Unknown source: {source}") - return EMPTY_EXAMPLE - - # Verify the solution passes tests - _check_test = example["solutions"][-1] + "\n" + test_code - succ, err = code_exec(_check_test) - if not succ: - print(f"Test code failed for {source}") - return EMPTY_EXAMPLE - - oracle_json = json.dumps({"functional": test_code}) - - elif "inputs" in oracle and "outputs" in oracle: # STDIN/STDOUT tests - stdin_list, stdout_list = oracle["inputs"], oracle["outputs"] - if len(stdin_list) == 0: - return EMPTY_EXAMPLE - - # handle list inputs and normalize line endings - stdin_list = [ - "\n".join(stdin) if isinstance(stdin, list) else stdin - for stdin in stdin_list - ] - stdout_list = [ - ("\n".join(stdout) if isinstance(stdout, list) else stdout).replace("\r\n", "\n") - for stdout in stdout_list - ] - - # Verify the solution passes tests - with ThreadPoolExecutor(max_workers=min(len(stdin_list), 8)) as executor: - futures = [] - for stdin, stdout in zip(stdin_list, stdout_list): - futures.append( - executor.submit( - remote_check_stdio, - example["solutions"][-1], - stdin, - stdout, - ) - ) - for future in as_completed(futures): - exec_succ, output, stdin, stdout = future.result() - pass_test = exec_succ and fuzzy_equal(output.strip(), stdout.strip()) - if not pass_test: - print(f"Test code failed for {source}") - return EMPTY_EXAMPLE - - oracle_json = json.dumps({"inputs": stdin_list, "outputs": stdout_list}) - else: - print(f"Unknown ground truth format: {oracle}") - return EMPTY_EXAMPLE - - # Format the final prompt - prompt = "\n".join(prompt_pieces) - - data = { - "data_source": data_source, - "prompt": [ - {"role": "user", "content": prompt} - ], - "ability": "codegen", - "apply_chat_template": True, - "reward_model": { - "style": "rule", - "ground_truth": oracle_json, - }, - "extra_info": { - "split": split, - "index": idx, - "reference": example["solutions"][0] if example["solutions"] else "", - "original_prompt": prompt, - "dataset": "likaixin/TACO-verified", - "source": source, - }, - } - - if idx == 0 or idx == 1: - print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) - print(data) - - return data - - return process_fn - - -if __name__ == '__main__': - """Main script execution: parse args, load, process, and save datasets.""" - parser = argparse.ArgumentParser(description="Process and save TACO dataset.") - parser.add_argument('--data-dir', default='data', - help='Base directory to save the processed data files.') - parser.add_argument('--name', default="taco", - help='Name of the dataset.') - parser.add_argument('--domain', default="codegen", - help='Domain of the dataset.') - parser.add_argument('--train-sample-size', type=int, default=None, - help='Number of samples to use for training. If None, use all samples.') - parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility') - - args = parser.parse_args() - - set_seed(args.seed) - - data_source = f"{args.domain}__{args.name}" - train_output_dir = os.path.join(args.data_dir, 'train') - test_output_dir = os.path.join(args.data_dir, 'test') - os.makedirs(train_output_dir, exist_ok=True) - os.makedirs(test_output_dir, exist_ok=True) - - # Load the dataset - cache_dir = datasets.config.HF_DATASETS_CACHE - dataset, _ = get_datasets(cache_dir=cache_dir) - - # Process the dataset - process_fn = make_map_fn('train', data_source) - - dataset = dataset.map(function=process_fn, with_indices=True, num_proc=64) - - # Filter out examples where processing failed - dataset = dataset.filter(lambda x: x["data_source"] == data_source) - - # Length filter - try: - tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B") - length_filter = LengthFilter(tokenizer=tokenizer, max_length=4096) - dataset = dataset.filter(lambda x: length_filter.check(x), num_proc=64) - except Exception as e: - print(f"Warning: Could not perform length filtering. Error: {e}") - print("Proceeding without length filtering.") - - # Sample the dataset - dataset = sample_dataset(dataset, args.train_sample_size) - - # Save the datasets - train_output_path = save_dataset( - dataset=dataset, - output_dir=train_output_dir, - filename_prefix=data_source, - sample_size=args.train_sample_size - ) - - print(f"\nDone! \n" - f"Data source: {data_source}\n" - f"Train data saved to {train_output_path} ({len(dataset)} samples)\n") \ No newline at end of file diff --git a/reasoners/data_preprocess/download_hf.py b/reasoners/data_preprocess/download_hf.py deleted file mode 100644 index 1e10ea137..000000000 --- a/reasoners/data_preprocess/download_hf.py +++ /dev/null @@ -1,41 +0,0 @@ -from huggingface_hub import hf_hub_download - -def download_all_files(repo_id, file_names, cache_dir=None): - """ - Download all specified files from a Hugging Face dataset repository. - - Args: - repo_id (str): The "namespace/repo_name" of the HF dataset (e.g., "your_user/dataset_name"). - file_names (List[str]): List of filenames to download. - cache_dir (str, optional): Custom cache directory where files will be stored. - """ - for fname in file_names: - local_path = hf_hub_download(repo_id=repo_id, filename=fname, cache_dir=cache_dir, repo_type="dataset") - print(f"✅ Downloaded {fname} → {local_path}") - -if __name__ == "__main__": - # Replace with your actual dataset repo ID - repo_id = "LLM360/guru_preview" - # List all files you want to download - file_names = [ - "codegen__deduped_leetcode2k_1.3k.parquet", - "codegen__deduped_livecodebench_440.parquet", - "codegen__deduped_primeintellect_7.5k.parquet", - "codegen__deduped_taco_8.8k.parquet", - "logic__arcagi1_111.parquet", - "logic__arcagi2_190.parquet", - "logic__barc_1.6k.parquet", - "logic__graph_logical_dataset_1.2k.parquet", - "logic__ordering_puzzle_dataset_1.9k.parquet", - "logic__zebra_puzzle_dataset_1.3k.parquet", - "math__combined_54.4k.parquet", - "simulation__codeio_fixed_12.1k_3.7k.parquet", - "stem__web_3.6k.parquet", - "table__hitab_4.3k.parquet", - "table__multihier_1.5k.parquet" - ] - - # Optionally specify a cache directory, or leave as None to use the default HF cache - cache_dir = "./hf_cache" - - download_all_files(repo_id, file_names, cache_dir) diff --git a/reasoners/data_preprocess/json2parquet.py b/reasoners/data_preprocess/json2parquet.py deleted file mode 100644 index ebac077bd..000000000 --- a/reasoners/data_preprocess/json2parquet.py +++ /dev/null @@ -1,57 +0,0 @@ -#!/usr/bin/env python3 -""" -Convert all JSON files in a directory into Parquet using Polars, -flattening only the `extra_info` field to top-level columns to avoid -nested chunked arrays. -""" -import sys -from pathlib import Path -import json -import os - -try: - import polars as pl # pip install polars -except ImportError: - print("Error: polars is required. Install via `pip install polars`.") - sys.exit(1) - - -def main(): - # Input and output directories - input_dir = Path( - "/mnt/weka/home/shibo.hao/yutao/Reasoning360/release/data/norm_json" - ).expanduser().resolve() - out_parquet = Path( - "/mnt/weka/home/shibo.hao/yutao/Reasoning360/release/data/parquet" - ).expanduser().resolve() - os.makedirs(out_parquet, exist_ok=True) - - for path in input_dir.glob("*.json"): - if "livecodebench" in path.stem: - continue - # Load normalized JSON list - with path.open("r", encoding="utf-8") as f: - records = json.load(f) - - if not records: - print("No records found in", path) - continue - - # Create Polars DataFrame - df = pl.DataFrame(records) - - print(df.head()) - - # Write out to Parquet with PyArrow backend - out_path = out_parquet / f"{path.stem}.parquet" - df.write_parquet( - out_path, - use_pyarrow=True, - row_group_size=len(records) - ) - - print(f"Saved {df.height} rows to {out_path}") - - -if __name__ == "__main__": - main() diff --git a/reasoners/data_preprocess/logic/gen_proof.py b/reasoners/data_preprocess/logic/gen_proof.py deleted file mode 100644 index 96782ba36..000000000 --- a/reasoners/data_preprocess/logic/gen_proof.py +++ /dev/null @@ -1,781 +0,0 @@ - -import os -from typing import List, Set, Dict, Tuple, Optional -import re -from dataclasses import dataclass, field -from copy import deepcopy -import random -import sys -import pdb -import signal -import time -from contextlib import contextmanager - - -class Constraint: - text: str # Human-readable constraint - object1: Optional[int] # First object's encoded value (if applicable) - object2: Optional[int] # Second object's encoded value (if applicable) - constraint_type: str # Type of mathematical constraint - value: Optional[int] # Additional value needed for constraint (e.g., number of birds between) - - def evaluate(self, positions: List[int]) -> bool: - """Evaluate if the constraint is satisfied for given positions.""" - if self.constraint_type == "left_of": - return positions.index(self.object1) < positions.index(self.object2) - elif self.constraint_type == "right_of": - return positions.index(self.object1) > positions.index(self.object2) - elif self.constraint_type == "adjacent_left": - return positions.index(self.object1) + 1 == positions.index(self.object2) - elif self.constraint_type == "adjacent_right": - return positions.index(self.object1) - 1 == positions.index(self.object2) - elif self.constraint_type == "birds_between": - return abs(positions.index(self.object1) - positions.index(self.object2)) - 1 == self.value - elif self.constraint_type == "absolute_position_left": - return positions.index(self.object1) == self.value - 1 - elif self.constraint_type == "absolute_position_right": - return positions.index(self.object1) == len(positions) - self.value - return False - -DEBUG = False - -@dataclass -class SearchNode: - """Data class to represent a search node in the tree.""" - depth: int - arrangement: Dict[str, int] = field(default_factory=dict) - satisfied_constraints: Set[int] = field(default_factory=set) - children: List['SearchNode'] = field(default_factory=list) - is_leaf: bool = False - parent: Optional['SearchNode'] = None - - - def __str__(self): - positions = [-1] * len(self.arrangement) - for bird, pos in self.arrangement.items(): - if pos != -1: - positions[pos] = bird - return f"Arrangement: {positions}, Satisfied: {self.satisfied_constraints}" - -class PuzzleSolver: - def __init__(self, birds: List[str], puzzle_constraints: List[str]=None): - self.birds = birds - self.num_positions = len(birds) - self.raw_constraints = puzzle_constraints - self.constraints: List[Constraint] = [] - self.absolute_constraints: List[Constraint] = [] - self.relative_constraints: List[Constraint] = [] - self.steps: List[str] = [] - self.ground_truth: List[str] = [] - - - def add_step(self, step: str): - """Add a step to the list of reasoning steps.""" - self.steps.append(step) - # print(step) - - - - def parse_constraints(self): - """Parse text constraints into structured Constraint objects.""" - for text in self.raw_constraints: - constraint = self._parse_single_constraint(text) - self.constraints.append(constraint) - - # Separate absolute and relative constraints - if constraint.constraint_type in ['absolute_position_left', 'absolute_position_right']: - self.absolute_constraints.append(constraint) - else: - self.relative_constraints.append(constraint) - - def _parse_single_constraint(self, text:str)->Constraint: - """Parse a single text constraint into a Constraint object.""" - # Pattern for extracting bird names - bird_pattern = '|'.join(self.birds) - # Absolute position patterns - if "from the left" in text or "from the right" in text: - match = re.search(rf"The ({bird_pattern}) is the (\d+)[a-z]+ from the (left|right)", text) - if match: - bird, position, direction = match.groups() - return Constraint( - text=text, - bird1=bird, - bird2=None, - constraint_type=f"absolute_position_{direction}", - value=int(position) - ) - - # Birds between pattern - if "bird" in text and "between" in text: - match = re.search(rf"There (?:is|are) (\d+) birds? between the ({bird_pattern}) and the ({bird_pattern})", text) - if match: - num_birds, bird1, bird2 = match.groups() - return Constraint( - text=text, - bird1=bird1, - bird2=bird2, - constraint_type="birds_between", - value=int(num_birds) - ) - - # Immediately adjacent patterns - if "immediately" in text: - if "to the left of" in text: - match = re.search(f"The ({bird_pattern}) is immediately to the left of the ({bird_pattern})", text) - bird1, bird2 = match.groups() - return Constraint( - text=text, - bird1=bird1, - bird2=bird2, - constraint_type="adjacent_left", - value=None - ) - elif "to the right of" in text: - match = re.search(f"The ({bird_pattern}) is immediately to the right of the ({bird_pattern})", text) - bird1, bird2 = match.groups() - return Constraint( - text=text, - bird1=bird1, - bird2=bird2, - constraint_type="adjacent_right", - value=None - ) - - # General left/right patterns - if "to the left of" in text: - match = re.search(f"The ({bird_pattern}) is to the left of the ({bird_pattern})", text) - bird1, bird2 = match.groups() - return Constraint( - text=text, - bird1=bird1, - bird2=bird2, - constraint_type="left_of", - value=None - ) - elif "to the right of" in text: - match = re.search(f"The ({bird_pattern}) is to the right of the ({bird_pattern})", text) - bird1, bird2 = match.groups() - return Constraint( - text=text, - bird1=bird1, - bird2=bird2, - constraint_type="right_of", - value=None - ) - - raise ValueError(f"Could not parse constraint: {text}") - - - - - def init_helper(self, file_path: str): - """Initialize the solver with the puzzle file.""" - self.ground_truth = self.read_file(file_path) - self.birds = random.shuffle(self.ground_truth)# Shuffle the birds - self.parse_constraints() - - def verify_constraint(self, arrangement: Dict[str, int], constraint: Constraint) -> bool: - """Check if a single constraint is satisfied by current arrangement""" - """Return true only if the constraint is satisfied else false""" - if constraint.constraint_type == "absolute_position_left": - pos = arrangement.get(constraint.bird1, -1) - return pos == constraint.value - 1 - - elif constraint.constraint_type == "absolute_position_right": - pos = arrangement.get(constraint.bird1, -1) - return pos == self.num_positions - constraint.value if pos != -1 else False - - elif constraint.constraint_type == "adjacent_left": - pos1 = arrangement.get(constraint.bird1, -1) - pos2 = arrangement.get(constraint.bird2, -1) - return pos1 != -1 and pos2 != -1 and pos1 + 1 == pos2 - - elif constraint.constraint_type == "adjacent_right": - pos1 = arrangement.get(constraint.bird1, -1) - pos2 = arrangement.get(constraint.bird2, -1) - return pos1 != -1 and pos2 != -1 and pos1 - 1 == pos2 - - elif constraint.constraint_type == "birds_between": - pos1 = arrangement.get(constraint.bird1, -1) - pos2 = arrangement.get(constraint.bird2, -1) - if pos1 != -1 and pos2 != -1: - return abs(pos1 - pos2) - 1 == constraint.value - return False # Can't verify yet - - elif constraint.constraint_type == "left_of": - pos1 = arrangement.get(constraint.bird1, -1) - pos2 = arrangement.get(constraint.bird2, -1) - return pos1 < pos2 if pos1 != -1 and pos2 != -1 else False - - elif constraint.constraint_type == "right_of": - pos1 = arrangement.get(constraint.bird1, -1) - pos2 = arrangement.get(constraint.bird2, -1) - return pos1 > pos2 if pos1 != -1 and pos2 != -1 else False - - return False - - def solve(self) -> Optional[Dict[str, int]]: - """Main solving function""" - # Parse constraints if not done already - if not self.constraints: - self.parse_constraints() - - # Start with empty arrangement - initial_arrangement = {bird: -1 for bird in self.birds} - # First apply absolute constraints - self.add_step("Starting with absolute position constraints:") - for i, constraint in enumerate(self.absolute_constraints): - bird = constraint.bird1 - if constraint.constraint_type == "absolute_position_left": - initial_arrangement[bird] = constraint.value - 1 - self.add_step(f"Placing {bird} at position {constraint.value} from left") - else: # absolute_position_right - pos = self.num_positions - constraint.value - initial_arrangement[bird] = pos - self.add_step(f"Placing {bird} at position {pos + 1} from left ({constraint.value} from right)") - - # Create initial node - initial_node = SearchNode( - arrangement=initial_arrangement, - satisfied_constraints=set(), - depth=0 - ) - - # Start Tree Search - self.add_step("Starting search to find complete arrangement...") - result = self.dfs(initial_node) - - if result: - positions = [-1] * self.num_positions - for bird, pos in result.arrangement.items(): - positions[pos] = bird - if -1 in positions: - index = positions.index(-1) - positions[index] = [bird for bird in self.birds if bird not in positions][0] - self.add_step(f"{positions[index]} is placed at position {index} since only one spot was left") - self.add_step("Found solution!") - self.add_step(f"Final arrangement: {positions}") - return result.arrangement - else: - self.add_step("\nNo solution found!") - return None - - def dfs(self, node: SearchNode) -> Optional[SearchNode]: - """Depth-first search implementation""" - # First verify all constraints are satisfied for current arrangement - if DEBUG: - pdb.set_trace() - print(f"\nDEBUG: Current depth: {node.depth}") - print(f"DEBUG: Current arrangement: {node.arrangement}") - print(f"DEBUG: Satisfied constraints: {node.satisfied_constraints}") - violated_constraint = None - current_constraint = None - current_constraint_index = None - for i, constraint in enumerate(self.relative_constraints): - if i not in node.satisfied_constraints: - consistent = self.verify_constraint(node.arrangement, constraint) - if consistent: - node.satisfied_constraints.add(i) - else: - current_constraint = constraint - current_constraint_index = i - break - - # return if solution found - if not current_constraint: - if DEBUG: print("DEBUG: No more constraints to satisfy!") - node.is_leaf = True - return node - - filled = set(pos for pos in node.arrangement.values() if pos != -1) - empty_positions = sorted(set(range(self.num_positions)) - filled) - if DEBUG: print(f"DEBUG: Empty positions: {empty_positions}") - # empty_positions = [i for bird, i in node.arrangement.items() if i == -1] - # import pdb; pdb.set_trace() - k = len(empty_positions) - if k==0 and current_constraint: - if DEBUG: print("DEBUG: No empty positions left but constraints remain!") - return None - current_arrangement = ["_"] * self.num_positions - for bird, pos in node.arrangement.items(): - if pos != -1: - current_arrangement[pos] = bird - - if current_constraint: - if current_constraint.constraint_type == "left_of": - if DEBUG: pdb.set_trace() - bird1 = current_constraint.bird1 - bird2 = current_constraint.bird2 - pos1 = node.arrangement.get(bird1, -1) - pos2 = node.arrangement.get(bird2, -1) - if pos1 != -1 and pos2 != -1 and pos1 >= pos2: - violated_constraint = current_constraint - self.add_step(f"Current arrangement: {current_arrangement}; Constraint violated: {current_constraint.text}, cannot proceed on this path") - return None - elif pos1 ==-1 and pos2 == -1: - for i in range(k-1): - for j in range(i+1, k): - new_node = deepcopy(node) - new_node.arrangement[bird1] = empty_positions[i] - new_node.arrangement[bird2] = empty_positions[j] - new_node.satisfied_constraints.add(current_constraint_index) - new_node.depth += 1 - new_node.parent = node - node.children.append(new_node) - self.add_step(f"Current arrangement: {current_arrangement}; Trying to place {bird1} to the left of {bird2} at positions {empty_positions[i]} and {empty_positions[j]}") - result = self.dfs(new_node) - if result: - return result - else: - self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird1} to the left of {bird2} at positions {empty_positions[i]} and {empty_positions[j]}") - else: - if pos1 == -1: - for i in range(k): - if empty_positions[i] < pos2: - new_node = deepcopy(node) - new_node.arrangement[bird1] = empty_positions[i] - new_node.satisfied_constraints.add(current_constraint_index) - new_node.depth += 1 - new_node.parent = node - node.children.append(new_node) - self.add_step(f"Current arrangement: {current_arrangement}; Trying to place {bird1} to the left of {bird2} at positions {empty_positions[i]}") - result = self.dfs(new_node) - if result: - return result - else: - self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird1} to the left of {bird2} at positions {empty_positions[i]}") - elif pos2 == -1: - for i in range(k): - if empty_positions[i] > pos1: - new_node = deepcopy(node) - new_node.arrangement[bird2] = empty_positions[i] - new_node.satisfied_constraints.add(current_constraint_index) - new_node.depth += 1 - new_node.parent = node - node.children.append(new_node) - self.add_step(f"Current arrangement: {current_arrangement}; Trying to place {bird2} to the right of {bird1} at positions {empty_positions[i]}") - result = self.dfs(new_node) - if result: - return result - else: - self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird2} to the right of {bird1} at positions {empty_positions[i]}") - elif current_constraint.constraint_type == "right_of": - if DEBUG: pdb.set_trace() - bird1 = current_constraint.bird1 - bird2 = current_constraint.bird2 - pos1 = node.arrangement.get(bird1, -1) - pos2 = node.arrangement.get(bird2, -1) - if pos1 != -1 and pos2 != -1 and pos1 <= pos2: - violated_constraint = current_constraint - self.add_step(f"Current arrangement: {current_arrangement}; Constraint violated: {current_constraint.text}, cannot proceed on this path") - return None - elif pos1 == -1 and pos2 == -1: - for i in range(k-1): - for j in range(i+1, k): - new_node = deepcopy(node) - new_node.arrangement[bird1] = empty_positions[j] - new_node.arrangement[bird2] = empty_positions[i] - new_node.satisfied_constraints.add(current_constraint_index) - new_node.depth += 1 - new_node.parent = node - node.children.append(new_node) - self.add_step(f"Current arrangement: {current_arrangement}; Trying to place {bird1} to the right of {bird2} at positions {empty_positions[j]} and {empty_positions[i]}") - result = self.dfs(new_node) - if result: - return result - else: - self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird1} to the right of {bird2} at positions {empty_positions[j]} and {empty_positions[i]}") - else: - if pos1 == -1: - for i in range(k): - if empty_positions[i] > pos2: - new_node = deepcopy(node) - new_node.arrangement[bird1] = empty_positions[i] - new_node.satisfied_constraints.add(current_constraint_index) - new_node.depth += 1 - new_node.parent = node - node.children.append(new_node) - self.add_step(f"Current arrangement: {current_arrangement}; Trying to place {bird1} to the right of {bird2} at positions {empty_positions[i]}") - result = self.dfs(new_node) - if result: - return result - else: - self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird1} to the right of {bird2} at positions {empty_positions[i]}") - elif pos2 == -1: - for i in range(k): - if empty_positions[i] < pos1: - new_node = deepcopy(node) - new_node.arrangement[bird2] = empty_positions[i] - new_node.satisfied_constraints.add(current_constraint_index) - new_node.depth += 1 - new_node.parent = node - node.children.append(new_node) - self.add_step(f"Current arrangement: {current_arrangement}; Trying to place {bird2} to the left of {bird1} at positions {empty_positions[i]}") - result = self.dfs(new_node) - if result: - return result - else: - self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird2} to the left of {bird1} at positions {empty_positions[i]}") - elif current_constraint.constraint_type == "adjacent_left": - if DEBUG: pdb.set_trace() - bird1 = current_constraint.bird1 - bird2 = current_constraint.bird2 - pos1 = node.arrangement.get(bird1, -1) - pos2 = node.arrangement.get(bird2, -1) - if pos1 != -1 and pos2 != -1 and pos1 + 1 != pos2: - violated_constraint = current_constraint - self.add_step(f"Current arrangement: {current_arrangement}; Constraint violated: {current_constraint.text}, cannot proceed on this path") - return None - elif pos1 == -1 and pos2 == -1: - for i in range(k-1): - for j in range(i+1, k): - if empty_positions[i] + 1 == empty_positions[j]: - new_node = deepcopy(node) - new_node.arrangement[bird1] = empty_positions[i] - new_node.arrangement[bird2] = empty_positions[j] - new_node.satisfied_constraints.add(current_constraint_index) - new_node.depth += 1 - new_node.parent = node - node.children.append(new_node) - self.add_step(f"Current arrangement: {current_arrangement}; Trying to place {bird1} immediately to the left of {bird2} at positions {empty_positions[i]} and {empty_positions[j]}") - result = self.dfs(new_node) - if result: - return result - else: - self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird1} immediately to the left of {bird2} at positions {empty_positions[i]} and {empty_positions[j]}") - else: - if pos1 == -1: - if pos2-1 not in empty_positions: - self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird1} immediately to the left of {bird2} at positions {pos2-1}") - return None - new_node = deepcopy(node) - new_node.arrangement[bird1] = pos2 - 1 - new_node.satisfied_constraints.add(current_constraint_index) - new_node.depth += 1 - new_node.parent = node - node.children.append(new_node) - self.add_step(f"Current arrangement: {current_arrangement}; Trying to place {bird1} immediately to the left of {bird2} at positions {pos2-1}") - result = self.dfs(new_node) - if result: - return result - else: - self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird1} immediately to the left of {bird2} at positions {pos2-1}") - elif pos2 == -1: - if pos1+1 not in empty_positions: - self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird1} immediately to the left of {bird2} at positions {pos1+1}") - return None - new_node = deepcopy(node) - new_node.arrangement[bird2] = pos1 + 1 - new_node.satisfied_constraints.add(current_constraint_index) - new_node.depth += 1 - new_node.parent = node - node.children.append(new_node) - self.add_step(f"Current arrangement: {current_arrangement}; Trying to place {bird2} immediately to the right of {bird1} at positions {pos1+1}") - result = self.dfs(new_node) - if result: - return result - else: - self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird2} immediately to the right of {bird1} at positions {pos1+1}") - elif current_constraint.constraint_type == "adjacent_right": - if DEBUG: pdb.set_trace() - bird1 = current_constraint.bird1 - bird2 = current_constraint.bird2 - pos1 = node.arrangement.get(bird1, -1) - pos2 = node.arrangement.get(bird2, -1) - if pos1 != -1 and pos2 != -1 and pos1 - 1 != pos2: - violated_constraint = current_constraint - self.add_step(f"Current arrangement: {current_arrangement}; Constraint violated: {current_constraint.text}, cannot proceed on this path ") - return None - elif pos1 == -1 and pos2 == -1: - for i in range(k-1): - for j in range(i+1, k): - if empty_positions[i] + 1 == empty_positions[j]: - new_node = deepcopy(node) - new_node.arrangement[bird1] = empty_positions[j] - new_node.arrangement[bird2] = empty_positions[i] - new_node.satisfied_constraints.add(current_constraint_index) - new_node.depth += 1 - new_node.parent = node - node.children.append(new_node) - self.add_step(f"Current arrangement: {current_arrangement}; Trying to place {bird1} immediately to the right of {bird2} at positions {empty_positions[j]} and {empty_positions[i]}") - result = self.dfs(new_node) - if result: - return result - else: - self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird1} immediately to the right of {bird2} at positions {empty_positions[j]} and {empty_positions[i]}") - else: - if pos1 == -1: - if pos2+1 not in empty_positions: - self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird1} immediately to the right of {bird2} at positions {pos2+1}") - return None - new_node = deepcopy(node) - new_node.arrangement[bird1] = pos2 + 1 - new_node.satisfied_constraints.add(current_constraint_index) - new_node.depth += 1 - new_node.parent = node - node.children.append(new_node) - self.add_step(f"Current arrangement: {current_arrangement}; Trying to place {bird1} immediately to the right of {bird2} at positions {pos2+1}") - result = self.dfs(new_node) - if result: - return result - else: - self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird1} immediately to the right of {bird2} at positions {pos2+1}") - elif pos2 == -1: - if pos1-1 not in empty_positions: - self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird1} immediately to the right of {bird2} at positions {pos1-1}") - return None - new_node = deepcopy(node) - new_node.arrangement[bird2] = pos1 - 1 - new_node.satisfied_constraints.add(current_constraint_index) - new_node.depth += 1 - new_node.parent = node - node.children.append(new_node) - self.add_step(f"Current arrangement: {current_arrangement}; Trying to place {bird2} immediately to the left of {bird1} at positions {pos1-1}") - result = self.dfs(new_node) - if result: - return result - else: - self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird2} immediately to the left of {bird1} at positions {pos1-1}") - elif current_constraint.constraint_type == "birds_between": - if DEBUG: pdb.set_trace() - bird1 = current_constraint.bird1 - bird2 = current_constraint.bird2 - pos1 = node.arrangement.get(bird1, -1) - pos2 = node.arrangement.get(bird2, -1) - value = current_constraint.value - if pos1 != -1 and pos2 != -1 and abs(pos1 - pos2) - 1 != value: - violated_constraint = current_constraint - self.add_step(f"Current arrangement: {current_arrangement}; Constraint violated: {current_constraint.text}, cannot proceed on this path") - return None - elif pos1 == -1 and pos2 == -1: - for p1 in empty_positions: - for p2 in empty_positions: - if p1 != p2 and abs(p1-p2) - 1 == value: - # 2 paths bird 1 left of bird 2 and bird 1 right of bird 2 - new_node = deepcopy(node) - new_node.arrangement[bird1] = p1 - new_node.arrangement[bird2] = p2 - new_node.satisfied_constraints.add(current_constraint_index) - new_node.depth += 1 - new_node.parent = node - node.children.append(new_node) - self.add_step(f"Current arrangement: {current_arrangement}; Trying to place {bird1} to the left of {bird2} at positions {p1} and {p2}") - result = self.dfs(new_node) - if result: - return result - else: - self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird1} to the left of {bird2} at positions {p1} and {p2}") - new_node = deepcopy(node) - new_node.arrangement[bird1] = p2 - new_node.arrangement[bird2] = p1 - new_node.satisfied_constraints.add(current_constraint_index) - new_node.depth += 1 - new_node.parent = node - node.children.append(new_node) - self.add_step(f"Current arrangement: {current_arrangement}; Trying to place {bird1} to the right of {bird2} at positions {p2} and {p1}") - result = self.dfs(new_node) - if result: - return result - else: - self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird1} to the right of {bird2} at positions {p2} and {p1}") - else: - if pos1 == -1: - if pos2 - value - 1 in empty_positions: - new_node = deepcopy(node) - new_node.arrangement[bird1] = pos2 - value - 1 - new_node.satisfied_constraints.add(current_constraint_index) - new_node.depth += 1 - new_node.parent = node - node.children.append(new_node) - self.add_step(f"Current arrangement: {current_arrangement}; Trying to place {bird1} to the left of {bird2} at positions {pos2 - value - 1}") - result = self.dfs(new_node) - if result: - return result - else: - self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird1} to the left of {bird2} at positions {pos2 - value - 1}") - if pos2 + value +1 in empty_positions: - new_node = deepcopy(node) - new_node.arrangement[bird1] = pos2 + value + 1 - new_node.satisfied_constraints.add(current_constraint_index) - new_node.depth += 1 - new_node.parent = node - node.children.append(new_node) - self.add_step(f"Current arrangement: {current_arrangement}; Trying to place {bird1} to the right of {bird2} at positions {pos2 + value + 1}") - result = self.dfs(new_node) - if result: - return result - else: - self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird1} to the right of {bird2} at positions {pos2 + value + 1}") - else: - return None - elif pos2 == -1: - if DEBUG: pdb.set_trace() - if pos1 - value - 1 in empty_positions: - new_node = deepcopy(node) - new_node.arrangement[bird2] = pos1 - value - 1 - new_node.satisfied_constraints.add(current_constraint_index) - new_node.depth += 1 - new_node.parent = node - node.children.append(new_node) - self.add_step(f"Current arrangement: {current_arrangement}; Trying to place {bird2} to the left of {bird1} at positions {pos1 - value - 1}") - result = self.dfs(new_node) - if result: - return result - else: - self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird2} to the left of {bird1} at positions {pos1 - value - 1}") - if pos1 + value + 1 in empty_positions: - new_node = deepcopy(node) - new_node.arrangement[bird2] = pos1 + value + 1 - new_node.satisfied_constraints.add(current_constraint_index) - new_node.depth += 1 - new_node.parent = node - node.children.append(new_node) - self.add_step(f"Current arrangement: {current_arrangement}; Trying to place {bird2} to the right of {bird1} at positions {pos1 + value + 1}") - result = self.dfs(new_node) - if result: - return result - else: - self.add_step(f"Current arrangement: {current_arrangement}; Cannot place {bird2} to the right of {bird1} at positions {pos1 + value + 1}") - else: - return None - return None - - - -def read_file(file_path: str) -> List[str]: - """Read the puzzle file and extract the ground truth and statements.""" - with open(file_path, 'r') as file: - lines = file.readlines() - - constraints = [] - - for line in lines: - if line[0].isdigit(): - match = re.match(r"\d+\.\s*(.+)", line.strip()) - constraints.append(match.group(1)) - # Extract ground truth and statements - start = lines[0].index("[") - ground_truth = lines[0][start:].strip().strip("[]").split(", ") - ground_truth = [item.strip("'") for item in ground_truth] - raw_constraints = constraints - return ground_truth, raw_constraints - -class TimeoutException(Exception): - pass - -@contextmanager -def time_limit(seconds): - """Context manager for timeout""" - def signal_handler(signum, frame): - raise TimeoutException("Solution took too long!") - - # Register signal handler - signal.signal(signal.SIGALRM, signal_handler) - signal.alarm(seconds) # Start timer - - try: - yield - finally: - signal.alarm(0) # Disable timer - -def solve_with_timeout(solver: PuzzleSolver, timeout_seconds: int) -> Optional[Dict[str, int]]: - """Wrapper to call solve() with timeout""" - try: - with time_limit(timeout_seconds): - return solver.solve() - except TimeoutException: - print(f"Solver timed out after {timeout_seconds} seconds!") - return None - -def solve_all_puzzles(num_puzzles: int): - random.seed(0) - timeout_seconds = 30 # Set timeout to 30 seconds - - base_dir = os.path.dirname(os.path.abspath(__file__)) - output_folder = os.path.join(base_dir, "QA_sols") - if not os.path.exists(output_folder): - os.makedirs(output_folder) - - num_sols = 0 - num_timeouts = 0 - - for i in range(num_puzzles): - # filepath = f"generated_QAs/birds_{i}.txt" - - try: - # ground_truth, raw_constraints = read_file(filepath) - - # if len(ground_truth) <= 7: - print(f"Solving puzzle {i}, num_sols: {num_sols}") - birds = deepcopy(ground_truth) - random.shuffle(birds) - sol_file = os.path.join(output_folder, f"sol_{i}.txt") - - try: - with open(sol_file, 'w') as file: - sys.stdout = file - solver = PuzzleSolver(birds=birds, puzzle_constraints=raw_constraints) - solver.parse_constraints() - - start_time = time.time() - final_arrangement = solve_with_timeout(solver, timeout_seconds) - solve_time = time.time() - start_time - - if final_arrangement: - # print(f"Solution found in {solve_time:.2f} seconds.") - num_sols += 1 - else: - # print(f"No solution found or timed out after {solve_time:.2f} seconds.") - if solve_time >= timeout_seconds: - num_timeouts += 1 - # Delete the solution file for timed out puzzles - sys.stdout = sys.__stdout__ - os.remove(sol_file) - continue - - except Exception as e: - print(f"Error solving puzzle {i}: {str(e)}") - finally: - sys.stdout = sys.__stdout__ - - except Exception as e: - print(f"Error reading puzzle {i}: {str(e)}") - continue - - # print(f"Number of solvable puzzles: {num_sols}") - # print(f"Number of timeouts: {num_timeouts}") - - - - - # filepath = f"generated_QAs/birds_923.txt" - # ground_truth, raw_constraints = read_file(filepath) - # birds = deepcopy(ground_truth) - # random.shuffle(birds) - # solver = PuzzleSolver(birds=birds, puzzle_constraints= raw_constraints) - # solver.parse_constraints() - # final_arrangement = solver.solve() - # if not final_arrangement: - # print("No solution found.") - # else: - # print("Solution found.") - - -def solve_puzzle(ground_truth, raw_constraints): - timeout_seconds = 30 - birds = deepcopy(ground_truth) - random.shuffle(birds) - solver = PuzzleSolver(birds=birds, puzzle_constraints=raw_constraints) - solver.parse_constraints() - start_time = time.time() - final_arrangement = solve_with_timeout(solver, timeout_seconds) - solve_time = time.time() - start_time - if final_arrangement: - return solver.steps - else: - return None - - - -if __name__ == "__main__": - solve_all_puzzles(1000) \ No newline at end of file diff --git a/reasoners/data_preprocess/logic/graph_dataset_gen/core/LICENSE b/reasoners/data_preprocess/logic/graph_dataset_gen/core/LICENSE deleted file mode 100644 index 924bb45c8..000000000 --- a/reasoners/data_preprocess/logic/graph_dataset_gen/core/LICENSE +++ /dev/null @@ -1,201 +0,0 @@ - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - - 1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - - 7. Disclaimer of Warranty. Unless required by applicable law or - agreed to in writing, Licensor provides the Work (and each - Contributor provides its Contributions) on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or - implied, including, without limitation, any warranties or conditions - of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A - PARTICULAR PURPOSE. You are solely responsible for determining the - appropriateness of using or redistributing the Work and assume any - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - - END OF TERMS AND CONDITIONS - - APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "{}" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - - Copyright 2017 Abulhair Saparov - - 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. diff --git a/reasoners/data_preprocess/logic/graph_dataset_gen/core/README.md b/reasoners/data_preprocess/logic/graph_dataset_gen/core/README.md deleted file mode 100644 index 7a253f0f4..000000000 --- a/reasoners/data_preprocess/logic/graph_dataset_gen/core/README.md +++ /dev/null @@ -1,21 +0,0 @@ -Documentation and code examples are available at . The repository is located at . - -This repository is a small collection of general-purpose data structures and algorithms. Originally, I developed this library to explore and tinker with different programming paradigms and design patterns. Since then, the code has evolved into a collection of generic data structures, using [templates](http://en.cppreference.com/w/cpp/language/templates) for polymorphism. - -### Usage and Dependencies - -To use the code, simply download the files into a folder named "core". Add this folder to the include path for the compiler, for example by using the `-I` flag. - -The code makes use of the [xxhash](https://github.com/Cyan4973/xxHash) library as its default hash function implementation, but the required files are included in the repository. Otherwise, there are no dependencies on external libraries. The code makes use of `C++11` and is regularly tested with `gcc 11` but I have previously compiled it with `gcc 4.8`, `clang 4.0`, and `Microsoft Visual C++ 14.0 (2015)`. The code is intended to be platform-independent, so please create an issue if there are any compilation bugs. - -### Overview - -The library currently implements the following data structures: - - [array](https://asaparov.org/docs/core/array.h.html#struct%20array) is a self-expanding sequential array, - - [hash_set](https://asaparov.org/docs/core/map.h.html#struct%20hash_set) is an unordered associative set of unique elements, implemented using a hashtable, - - [hash_map](https://asaparov.org/docs/core/map.h.html#struct%20hash_map) is an unordered associative map of key-value pairs, where the keys are unique, implemented using a hashtable, - - [array_map](https://asaparov.org/docs/core/map.h.html#struct%20array_map) is a sequential map of key-value pairs, where the keys are unique. - -In addition, the library implements a handful of sorting algorithms in [array.h](https://asaparov.org/docs/core/array.h.html). Set operations, such as union, intersection, and subtraction, can also be found in [array.h](https://asaparov.org/docs/core/array.h.html), which operate on sets represented as sorted arrays. The library implements serialization/deserialization with the functions `read`, `write`, and `print` for all data structures using a very regular pattern involving ["scribes"](https://asaparov.org/docs/core/io.h.html#scribes) as described in [io.h](https://asaparov.org/docs/core/io.h.html). Functions and structures useful for lexical analysis are provided in [lex.h](https://asaparov.org/docs/core/lex.h.html). [random.h](https://asaparov.org/docs/core/random.h.html) provides pseudo-random number generation and sampling from a handful of built-in distributions. - -The implementations of every data structure are transparent, with `public` visibility for the underlying fields, which gives the user more control over data structure behavior. All procedures and structures do *not* use exceptions. Rather, errors are returned, typically as a `bool` type, where `true` indicates success and `false` indicates failure. I observed, in practice, that the overhead of this kind of error handling is negligible when no errors are thrown, and the overhead is much smaller than exception handling when errors do occur. The library does not provide automatic memory management, and expects the user to manage memory appropriately. The `free` function is implemented for all data structures, but it does not automatically deallocate child objects in the container classes. The library does not use any thread synchronization methods to guarantee thread safety, and the user must implement their own synchronization or avoid simultaneous writes to data structures. diff --git a/reasoners/data_preprocess/logic/graph_dataset_gen/core/array.h b/reasoners/data_preprocess/logic/graph_dataset_gen/core/array.h deleted file mode 100644 index d361f6faa..000000000 --- a/reasoners/data_preprocess/logic/graph_dataset_gen/core/array.h +++ /dev/null @@ -1,2429 +0,0 @@ -/** - * \file array.h - * - * This file contains the implementation of the `array` data structure, which - * is a resizeable ordered collection of elements with generic type `T`. In - * addition, the file contains global helper functions for working with arrays, - * such as resizing, linear and binary search, insertion sort, and quicksort. - * The `pair` structure is also defined here. - * - * This file also contains functions that perform set operations on sorted - * arrays and native arrays, such as union, intersection, subtraction, and - * subset. - * - * - */ - -#ifndef ARRAY_H_ -#define ARRAY_H_ - -#include -#include -#include -#include - -#include - -/** - * The multiplicative factor by which array capacity is changed. - */ -#define RESIZE_FACTOR 2 - - -namespace core { - - -namespace detail { - template static auto test_resizeable(int32_t) -> - decltype(void(std::declval().on_resize()), std::true_type{}); - template static auto test_resizeable(int64_t) -> std::false_type; -} - -/** - * This type trait is [true_type](http://en.cppreference.com/w/cpp/types/integral_constant) - * if and only if `T` is class with a public function `void on_resize()`. - */ -template struct is_resizeable : decltype(core::detail::test_resizeable(0)){}; - -/** - * Resizes the given native array `data` with the requested capacity `new_capacity`. - * \tparam SizeType a type that satisfies [is_integral](http://en.cppreference.com/w/cpp/types/is_integral). - * \tparam T the generic type of the elements in `data`. - * \param data the array to resize. - * \param new_capacity the requested size. - * \return `true` on success; `data` may point to a different memory address. - * \return `false` on failure; `data` is unchanged. - */ -template::value>::type* = nullptr> -inline bool resize(T*& data, const SizeType& new_capacity) { - T* new_data = (T*) realloc(static_cast(data), new_capacity * sizeof(T)); - if (new_data == NULL) { - fprintf(stderr, "resize ERROR: Out of memory.\n"); - return false; - } - data = new_data; - return true; -} - -/** - * This function multiplies `capacity` by #RESIZE_FACTOR. It then repeats this - * until `capacity >= new_length`. - */ -template::value>::type* = nullptr> -inline void expand_capacity(SizeType& capacity, size_t new_length) { - do { - /* increase the size of the underlying array */ - capacity *= RESIZE_FACTOR; - } while (new_length > capacity); -} - -/** - * For a given requested length `new_length`, this function calls - * expand_capacity() to determine the new `capacity` of the native array - * `data`. The function then attempts to resize the array with this capacity. - * Note this function does not check whether `new_length <= capacity`. - */ -template::value>::type* = nullptr> -inline bool expand(T*& data, SizeType& capacity, size_t new_length) { - expand_capacity(capacity, new_length); - return resize(data, capacity); -} - -/** - * For a given requested length `new_length`, this function expands the native - * array `data` by factors of #RESIZE_FACTOR until `capacity >= new_length`. - * If initially `new_length <= capacity`, this function does nothing. - */ -template::value>::type* = nullptr> -inline bool ensure_capacity(T*& data, SizeType& capacity, size_t new_length) -{ - if (new_length <= capacity) - return true; - SizeType new_capacity = capacity; - if (!expand(data, new_capacity, new_length)) - return false; - capacity = new_capacity; - return true; -} - -/** - * Performs a linear search through the array `data` to find the smallest index - * `i >= start` such that `element == data[i]`. - * \tparam Key a generic type for which operator `==` is defined for arguments of type `Key` and `T`. - * \tparam T a generic type for which operator `==` is defined for arguments of type `Key` and `T`. - * \return an index in `start, start + 1, ..., length - 1` if the element was found. - * \return `length` if the element was not found. - */ -template::value>::type* = nullptr> -inline SizeType index_of(const Key& element, const T* data, - const SizeType& length, const SizeType& start = 0) -{ - for (SizeType i = start; i < length; i++) - if (element == data[i]) - return i; - return length; -} - -/** - * Performs a linear search through the array `data` to find the largest index - * `i` such that `element == data[i]`. - * \tparam Key a generic type for which operator `==` is defined for arguments of type `Key` and `T`. - * \tparam T a generic type for which operator `==` is defined for arguments of type `Key` and `T`. - * \return an index in `0, 1, ..., length - 1` if the element was found. - * \return `static_cast(-1)` if the element was not found. - */ -template::value>::type* = nullptr> -inline unsigned int last_index_of(const Key& element, const T* data, const SizeType& length) -{ - unsigned int i = length; - while (i != 0) { - i--; - if (element == data[i]) - return i; - } - return static_cast(-1); -} - -/** - * A resizeable sequence of objects, stored contiguously, each with generic - * type `T`. This structure does not automatically initialize or free its - * elements, and so the user must appropriately free each element before the - * array is destroyed. - * - * In the following example, we demonstrate a simple use-case of array. Here, - * `a` is automatically freed by the destructor since it was initialized on the - * stack. The expected output is `-1 -1 0 3 `. - * - * ```{.cpp} - * #include - * #include - * using namespace core; - * - * int main() { - * array a = array(8); - * a.add(-1); a.add(-4); - * a.add(3); a.add(0); - * a.remove(1); - * - * printf("%d ", a[0]); - * for (int element : a) - * printf("%d ", element); - * } - * ``` - * - * - * However, if `a` is not allocated on the stack, the destructor will not be - * automatically called, and so it must be freed manually using `core::free` or - * `array::free`. - * - * ```{.cpp} - * #include - * #include - * using namespace core; - * - * int main() { - * array& a = *((array*) alloca(sizeof(array))); - * array_init(a, 8); - * a.add(-1); a.add(-4); - * a.add(3); a.add(0); - * a.remove(1); - * - * printf("%d ", a[0]); - * for (int element : a) - * printf("%d ", element); - * free(a); - * } - * ``` - * - * - * Also note that a number of member functions require that `T` be - * [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). - * In other cases, elements should be added manually to the underlying native - * array array::data. This structure also defines array::begin and array::end, - * similar to Standard Template Library iterators, which enables the use of the - * range-based for loop in the example below. In this example, the expected - * output is `first second `. - * - * ```{.cpp} - * #include - * #include - * #include - * using namespace core; - * - * struct custom_string { - * char* buffer; - * - * static void free(custom_string& s) { - * core::free(s.buffer); - * } - * }; - * - * bool init(custom_string& s, const char* src) { - * s.buffer = (char*) malloc(sizeof(char) * (strlen(src) + 1)); - * if (s.buffer == NULL) - * return false; - * memcpy(s.buffer, src, sizeof(char) * (strlen(src) + 1)); - * return true; - * } - * - * int main() { - * array a = array(8); - * init(a[0], "first"); - * init(a[1], "second"); - * a.length = 2; - * - * for (const custom_string& s : a) - * printf("%s ", s.buffer); - * for (custom_string& s : a) - * free(s); - * } - * ``` - * Note in the above example that since the array struct does not automatically - * free its elements, they must be freed manually. - */ -template -struct array { - /** - * The underlying native array of elements. - */ - T* data; - - /** - * The length of the array. - */ - size_t length; - - /** - * The capacity of array::data. - */ - size_t capacity; - - /** - * Constructs an array with zero size and the given `initial_capacity`. - */ - array(size_t initial_capacity) - { - if (!initialize(initial_capacity)) - exit(EXIT_FAILURE); - } - - ~array() { free(); } - - /** - * Returns a reference to the element at the given `index`. No bounds-checking is performed. - */ - inline T& operator[] (size_t index) { - return data[index]; - } - - /** - * Returns a const reference to the element at the given `index`. No bounds-checking is performed. - */ - inline const T& operator[] (size_t index) const { - return data[index]; - } - - /** - * Sets the length of the array to `0`. Any elements are not freed. - */ - inline void clear() - { - length = 0; - } - - /** - * Moves the last element in the array to the position given by `index` and - * decrements array::length by `1`. The element initially at `index` is not - * freed. - */ - void remove(size_t index) - { - core::move(data[length - 1], data[index]); - length--; - } - - /** - * For a given requested length `new_length`, this function expands the - * array by factors of #RESIZE_FACTOR until `array::capacity >= new_length`. - * If initially `new_length <= array::capacity`, this function does - * nothing. If the resize operation moved array::data to a new memory - * address, and `T` satisfies is_resizeable, then `x.on_resize()` is called - * for every element `x` in the array. - */ - template::value && is_resizeable::value>::type* = nullptr> - bool ensure_capacity(size_t new_length) { - const T* old_data = data; - if (!core::ensure_capacity(data, capacity, new_length)) return false; - if (data != old_data) { - for (unsigned int i = 0; i < length; i++) - data[i].on_resize(); - } - return true; - } - - template::value && !is_resizeable::value>::type* = nullptr> - bool ensure_capacity(size_t new_length) { - return core::ensure_capacity(data, capacity, new_length); - } - - /** - * Adds the given native array of elements to this structure. This function - * uses [memcpy](http://en.cppreference.com/w/cpp/string/byte/memcpy), - * and so it should not be used if the elements are not - * [TriviallyCopyable](https://en.cppreference.com/w/cpp/named_req/TriviallyCopyable). - * In such a case, addition should be performed manually using the public fields. - * \tparam T is [TriviallyCopyable](https://en.cppreference.com/w/cpp/named_req/TriviallyCopyable). - */ - bool append(const T* elements, size_t size) - { - if (!ensure_capacity(length + size)) - return false; - memcpy(&data[length], elements, sizeof(T) * size); - length += size; - return true; - } - - /** - * Calls array::index_of to determine whether `element` exists in this array. - * \tparam Key a generic type for which operator `==` is defined for arguments of type `Key` and `T`. - */ - template - inline bool contains(const Key& element) const { - return index_of(element) < length; - } - - /** - * Performs a linear search of the array to find the smallest index `i` - * such that `element == array::data[i]`. - * \tparam Key a generic type for which operator `==` is defined for arguments of type `Key` and `T`. - * \return an index in `0, 1, ..., array::length - 1` if the element was found. - * \return `array::length` if the element was not found. - */ - template - inline unsigned int index_of(const Key& element) const { - return core::index_of(element, data, (unsigned int) length); - } - - /** - * Performs a linear search through the array to find the smallest index - * `i >= start` such that `element == array::data[i]`. - * \tparam Key a generic type for which operator `==` is defined for arguments of type `Key` and `T`. - * \return an index in `start, start + 1, ..., length - 1` if the element was found. - * \return `length` if the element was not found. - */ - template - inline unsigned int index_of(const Key& element, SizeType start) const { - return core::index_of(element, data, (unsigned int) length, start); - } - - /** - * Returns a reference to `array::data[0]`, ignoring any bounds. - */ - T& first() - { - return data[0]; - } - - /** - * Returns a const reference to `array::data[0]`, ignoring any bounds. - */ - const T& first() const - { - return data[0]; - } - - /** - * Returns a reference to `array::data[array::length - 1]`, ignoring any bounds. - */ - T& last() - { - return data[length - 1]; - } - - /** - * Returns a const reference to `array::data[array::length - 1]`, ignoring any bounds. - */ - const T& last() const - { - return data[length - 1]; - } - - /** - * Adds the given element to this array, increasing its capacity if - * necessary. The assignment operator performs the addition, and so this - * function should not be used if `T` is not - * [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). - * In such a case, addition should be performed manually using the public fields. - * \tparam T is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). - */ - bool add(const T& element) - { - if (!ensure_capacity(length + 1)) - return false; - data[length] = element; - length++; - return true; - } - - /** - * Returns the element at `array::length - 1` and decrements array::length - * by `1`, ignoring any bounds. - */ - T pop() - { - length--; - return data[length]; - } - - /** - * Returns an iterator to the beginning of the array. - */ - inline T* begin() { - return data; - } - - /** - * Returns an iterator to the end of the array. - */ - inline T* end() { - return data + length; - } - - /** - * Returns a const iterator to the beginning of the array. - */ - inline const T* begin() const { - return data; - } - - /** - * Returns a const iterator to the end of the array. - */ - inline const T* end() const { - return data + length; - } - - /** - * Copies the underlying fields from `src` into `dst`, effectively moving - * the array from `src` into `dst`. - */ - static inline void move(const array& src, array& dst) { - dst.length = src.length; - dst.capacity = src.capacity; - dst.data = src.data; - } - - template - static inline long unsigned int size_of(const array& a, const Metric& metric) { - long unsigned int sum = core::size_of(a.capacity) + core::size_of(a.length); - for (unsigned int i = 0; i < a.length; i++) - sum += core::size_of(a.data[i], metric); - return sum + (a.capacity - a.length) * sizeof(T); - } - - /** - * Frees array::data. This should not be used if `a` was constructed on the - * stack, as the destructor will automatically free array::data. The - * elements of `a` are not freed. - */ - static inline void free(array& a) { a.free(); } - -private: - inline bool initialize(size_t initial_capacity) - { -#if !defined(NDEBUG) - if (initial_capacity == 0) - fprintf(stderr, "array.initialize WARNING: Initial capacity is zero.\n"); -#endif - - capacity = initial_capacity; - length = 0; - data = (T*) malloc(sizeof(T) * capacity); - if (data == NULL) { - fprintf(stderr, "array.initialize ERROR: Out of memory.\n"); - return false; - } - return true; - } - - inline void free() { - if (data != NULL) { - core::free(data); - data = NULL; - } - } - - template - friend bool array_init(array& m, size_t initial_capacity); -}; - -/** - * Initializes the array `m` with the given `initial_capacity`. - */ -template -bool array_init(array& m, size_t initial_capacity) { - return m.initialize(initial_capacity); -} - -/** - * Returns array::length of `m`. - */ -template -inline size_t size(const array& m) { - return m.length; -} - -/** - * Swaps the underlying buffers of the given arrays. - */ -template -inline void swap(array& a, array& b) -{ - T* temp = a.data; - a.data = b.data; - b.data = temp; - - size_t swap = a.length; - a.length = b.length; - b.length = swap; - - swap = a.capacity; - a.capacity = b.capacity; - b.capacity = swap; -} - -template -inline long unsigned int size(const array& a, const Metric& metric) { - long unsigned int sum = size(a.capacity) + size(a.length); - for (unsigned int i = 0; i < a.length; i++) - sum += size(a.data[i], metric); - return sum + (a.capacity - a.length) * sizeof(T); -} - -/** - * Compares the two given arrays and returns true if and only if `a.length == b.length` - * and there is no index `i` such that `a.data[i] != b.data[i]`. - */ -template -inline bool operator == (const array& a, const array& b) { - if (a.length != b.length) - return false; - for (unsigned int i = 0; i < a.length; i++) - if (a.data[i] != b.data[i]) - return false; - return true; -} - -/** - * Compares the two given arrays and returns true if and only if `a.length != b.length` - * or there is some index `i` such that `a.data[i] != b.data[i]`. - */ -template -inline bool operator != (const array& a, const array& b) { - if (a.length != b.length) - return true; - for (unsigned int i = 0; i < a.length; i++) - if (a.data[i] != b.data[i]) - return true; - return false; -} - -/** - * Performs insertion sort on the given native array `keys` with given `length`. - * \tparam T is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable), - * [CopyConstructible](https://en.cppreference.com/w/cpp/named_req/CopyConstructible), - * and [LessThanComparable](https://en.cppreference.com/w/cpp/named_req/LessThanComparable). - */ -template -void insertion_sort(T* keys, unsigned int length) -{ - for (unsigned int i = 1; i < length; i++) { - T item = keys[i]; - unsigned int hole = i; - - while (hole > 0 && item < keys[hole - 1]) { - keys[hole] = keys[hole - 1]; - hole--; - } - - keys[hole] = item; - } -} - -/** - * Performs insertion sort on the given native array `keys` with given `length` - * and `sorter`. - * \tparam T satisfies is_moveable, and for which a function - * `bool less_than(const T&, const T&, const Sorter&)` is defined. - */ -template::value>::type* = nullptr> -void insertion_sort(T* keys, unsigned int length, const Sorter& sorter) -{ - T& item = *((T*) malloc(sizeof(T))); - for (unsigned int i = 1; i < length; i++) { - move(keys[i], item); - unsigned int hole = i; - - while (hole > 0 && less_than(item, keys[hole - 1], sorter)) { - move(keys[hole - 1], keys[hole]); - hole--; - } - - move(item, keys[hole]); - } - free(&item); -} - -/** - * Performs insertion sort on the given native arrays `keys` and `values` with given `length`. - * \tparam K is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable), - * [CopyConstructible](https://en.cppreference.com/w/cpp/named_req/CopyConstructible), - * and [LessThanComparable](https://en.cppreference.com/w/cpp/named_req/LessThanComparable). - * \tparam V is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable), - * [CopyConstructible](https://en.cppreference.com/w/cpp/named_req/CopyConstructible). - */ -template -void insertion_sort(K* keys, V* values, unsigned int length) -{ - for (unsigned int i = 1; i < length; i++) { - K item = keys[i]; - V value = values[i]; - unsigned int hole = i; - - while (hole > 0 && item < keys[hole - 1]) { - keys[hole] = keys[hole - 1]; - values[hole] = values[hole - 1]; - hole--; - } - - keys[hole] = item; - values[hole] = value; - } -} - -/** - * Performs insertion sort on the given native arrays `keys` and `values` with - * given `length` and `sorter`. - * \tparam K satisfies is_moveable, and for which a function - * `bool less_than(const K&, const K&, const Sorter&)` is defined. - * \tparam V satisfies is_moveable. - */ -template::value>::type* = nullptr> -void insertion_sort(K* keys, V* values, unsigned int length, const Sorter& sorter) -{ - K& item = *((K*) malloc(sizeof(K))); - V& value = *((V*) malloc(sizeof(V))); - for (unsigned int i = 1; i < length; i++) { - move(keys[i], item); - move(values[i], value); - unsigned int hole = i; - - while (hole > 0 && less_than(item, keys[hole - 1], sorter)) { - move(keys[hole - 1], keys[hole]); - move(values[hole - 1], values[hole]); - hole--; - } - - move(item, keys[hole]); - move(value, values[hole]); - } - free(&item); free(&value); -} - -/** - * Performs insertion sort on the given array `keys`. - * \tparam T is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable), - * [CopyConstructible](https://en.cppreference.com/w/cpp/named_req/CopyConstructible), - * and [LessThanComparable](https://en.cppreference.com/w/cpp/named_req/LessThanComparable). - */ -template -inline void insertion_sort(array& keys) { - insertion_sort(keys.data, (unsigned int) keys.length); -} - -/** - * Performs insertion sort on the given array `keys` with the given `sorter`. - * \tparam T satisfies is_moveable, and for which a function - * `bool less_than(const T&, const T&, const Sorter&)` is defined. - */ -template::value>::type* = nullptr> -inline void insertion_sort(array& keys, const Sorter& sorter) -{ - insertion_sort(keys.data, (unsigned int) keys.length, sorter); -} - -/** - * Performs insertion sort on the given arrays `keys` and `values`. - * \tparam K is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable), - * [CopyConstructible](https://en.cppreference.com/w/cpp/named_req/CopyConstructible), - * and [LessThanComparable](https://en.cppreference.com/w/cpp/named_req/LessThanComparable). - * \tparam V is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable), - * [CopyConstructible](https://en.cppreference.com/w/cpp/named_req/CopyConstructible). - */ -template -inline void insertion_sort(array& keys, array& values) { - insertion_sort(keys.data, values.data, (unsigned int) keys.length); -} - -/** - * Performs insertion sort on the given arrays `keys` and `values` with the - * given `sorter`. - * \tparam K satisfies is_moveable, and for which a function - * `bool less_than(const K&, const K&, const Sorter&)` is defined. - * \tparam V satisfies is_moveable. - */ -template::value>::type* = nullptr> -inline void insertion_sort(array& keys, array& values, const Sorter& sorter) -{ - insertion_sort(keys.data, values.data, (unsigned int) keys.length, sorter); -} - -/** - * Reverses the order of the elements in the given native `array` with given `length`. - * \tparam T satisfies is_swappable. - */ -template -void reverse(T* array, unsigned int length) { - for (unsigned int i = 0; i < length / 2; i++) { - unsigned int other = length - i - 1; - swap(array[i], array[other]); - } -} - -/** - * Reverses the order of the elements in the given `array`. - * \tparam T satisfies is_swappable. - */ -template -inline void reverse(array& array) { - reverse(array.data, (unsigned int) array.length); -} - -template -inline const T& get_pivot(T* array, unsigned int start, unsigned int end) { - return array[(end + start) / 2]; -} - -template -inline void quick_sort_partition(T* array, - unsigned int start, unsigned int end, - unsigned int& l, unsigned int& r) -{ - const T p = get_pivot(array, start, end); - while (true) { - while (array[l] < p) - l++; - while (p < array[r]) - r--; - if (l == r) { - l++; - if (r > 0) r--; - return; - } else if (l > r) return; - swap(array[l++], array[r--]); - } -} - -template::value>::type* = nullptr> -inline void quick_sort_partition( - T* array, unsigned int start, unsigned int end, - unsigned int& l, unsigned int& r, const Sorter& sorter) -{ - T& p = *((T*) malloc(sizeof(T))); - move(get_pivot(array, start, end), p); - while (true) { - while (less_than(array[l], p, sorter)) - l++; - while (less_than(p, array[r], sorter)) - r--; - if (l == r) { - l++; - if (r > 0) r--; - break; - } else if (l > r) break; - swap(array[l++], array[r--]); - } - free(&p); -} - -template -inline void quick_sort_partition(K* keys, V* values, - unsigned int start, unsigned int end, unsigned int& l, unsigned int& r) -{ - const K p = get_pivot(keys, start, end); - while (true) { - while (keys[l] < p) - l++; - while (p < keys[r]) - r--; - if (l == r) { - l++; - if (r > 0) r--; - return; - } else if (l > r) return; - swap(values[l], values[r]); - swap(keys[l++], keys[r--]); - } -} - -template::value>::type* = nullptr> -inline void quick_sort_partition( - K* keys, V* values, unsigned int start, unsigned int end, - unsigned int& l, unsigned int& r, const Sorter& sorter) -{ - K& p = *((K*) malloc(sizeof(K))); - move(get_pivot(keys, start, end), p); - while(true) { - while (less_than(keys[l], p, sorter)) - l++; - while (less_than(p, keys[r], sorter)) - r--; - if (l == r) { - l++; - if (r > 0) r--; - break; - } else if (l > r) break; - swap(values[l], values[r]); - swap(keys[l++], keys[r--]); - } - free(&p); -} - -template -void quick_sort(T* array, unsigned int start, unsigned int end) -{ - if (start >= end) - return; - unsigned int l = start, r = end; - quick_sort_partition(array, start, end, l, r); - quick_sort(array, start, r); - quick_sort(array, l, end); -} - -template::value>::type* = nullptr> -void quick_sort(T* array, unsigned int start, unsigned int end, const Sorter& sorter) -{ - if (start >= end) - return; - unsigned int l = start, r = end; - quick_sort_partition(array, start, end, l, r, sorter); - quick_sort(array, start, r, sorter); - quick_sort(array, l, end, sorter); -} - -template -void quick_sort(K* keys, V* values, unsigned int start, unsigned int end) -{ - if (start >= end) - return; - unsigned int l = start, r = end; - quick_sort_partition(keys, values, start, end, l, r); - quick_sort(keys, values, start, r); - quick_sort(keys, values, l, end); -} - -template::value>::type* = nullptr> -void quick_sort(K* keys, V* values, - unsigned int start, unsigned int end, const Sorter& sorter) -{ - if (start >= end) - return; - unsigned int l = start, r = end; - quick_sort_partition(keys, values, start, end, l, r, sorter); - quick_sort(keys, values, start, r, sorter); - quick_sort(keys, values, l, end, sorter); -} - -/** - * Performs Quicksort on the given native array `keys` with given `length`. - * This function assumes `length > 0`. If the preprocessor `NDEBUG` is not - * defined, this function outputs a warning when `length == 0`. - * \tparam T is [CopyConstructible](https://en.cppreference.com/w/cpp/named_req/CopyConstructible), - * [LessThanComparable](https://en.cppreference.com/w/cpp/named_req/LessThanComparable), - * and is_swappable. - */ -template -inline void quick_sort(T* keys, unsigned int length) { -#if !defined(NDEBUG) - if (length == 0) { - fprintf(stderr, "quick_sort WARNING: Length is zero.\n"); - return; - } -#endif - quick_sort(keys, 0, length - 1); -} - -/** - * Performs Quicksort on the given native array `keys` with given `length` and - * `sorter`. This function assumes `length > 0`. If the preprocessor `NDEBUG` - * is not defined, this function outputs a warning when `length == 0`. - * \tparam T a generic type that satisfies is_swappable, and for which a - * function `bool less_than(const T&, const T&, const Sorter&)` is - * defined. - */ -template::value>::type* = nullptr> -inline void quick_sort(T* keys, unsigned int length, const Sorter& sorter) -{ -#if !defined(NDEBUG) - if (length == 0) { - fprintf(stderr, "quick_sort WARNING: Length is zero.\n"); - return; - } -#endif - quick_sort(keys, 0, length - 1, sorter); -} - -/** - * Performs Quicksort on the given native arrays `keys` and `values` with given - * `length`. This function assumes `length > 0`. If the preprocessor `NDEBUG` - * is not defined, this function outputs a warning when `length == 0`. - * \tparam K is [CopyConstructible](https://en.cppreference.com/w/cpp/named_req/CopyConstructible), - * [LessThanComparable](https://en.cppreference.com/w/cpp/named_req/LessThanComparable), - * and is_swappable. - * \tparam V satisfies is_swappable. - */ -template -inline void quick_sort(K* keys, V* values, unsigned int length) { -#if !defined(NDEBUG) - if (length == 0) { - fprintf(stderr, "quick_sort WARNING: Length is zero.\n"); - return; - } -#endif - quick_sort(keys, values, 0, length - 1); -} - -/** - * Performs Quicksort on the given native arrays `keys` and `values` with given - * `length` and `sorter`. This function assumes `length > 0`. If the preprocessor - * `NDEBUG` is not defined, this function outputs a warning when `length == 0`. - * \tparam K a generic type that satisfies is_swappable, and for which a - * function `bool less_than(const K&, const K&, const Sorter&)` is - * defined. - * \tparam V satisfies is_swappable. - */ -template::value>::type* = nullptr> -inline void quick_sort(K* keys, V* values, unsigned int length, const Sorter& sorter) -{ -#if !defined(NDEBUG) - if (length == 0) { - fprintf(stderr, "quick_sort WARNING: Length is zero.\n"); - return; - } -#endif - quick_sort(keys, values, 0, length - 1, sorter); -} - -/** - * Performs Quicksort on the given array `keys`. This function assumes - * `length > 0`. If the preprocessor `NDEBUG` is not defined, this function - * outputs a warning when `length == 0`. - * \tparam T is [CopyConstructible](https://en.cppreference.com/w/cpp/named_req/CopyConstructible), - * [LessThanComparable](https://en.cppreference.com/w/cpp/named_req/LessThanComparable), - * and is_swappable. - */ -template -inline void quick_sort(array& keys) { -#if !defined(NDEBUG) - if (keys.length == 0) { - fprintf(stderr, "quick_sort WARNING: Length is zero.\n"); - return; - } -#endif - quick_sort(keys.data, 0, (unsigned int) keys.length - 1); -} - -/** - * Performs Quicksort on the given native array `keys` with the given `sorter`. - * This function assumes `length > 0`. If the preprocessor `NDEBUG` is not - * defined, this function outputs a warning when `length == 0`. - * \tparam T a generic type that satisfies is_swappable, and for which a - * function `bool less_than(const T&, const T&, const Sorter&)` is - * defined. - */ -template::value>::type* = nullptr> -inline void quick_sort(array& keys, const Sorter& sorter) -{ -#if !defined(NDEBUG) - if (keys.length == 0) { - fprintf(stderr, "quick_sort WARNING: Length is zero.\n"); - return; - } -#endif - quick_sort(keys.data, 0, (unsigned int) keys.length - 1, sorter); -} - -/** - * Performs Quicksort on the given arrays `keys` and `values`. This function - * assumes `length > 0`. If the preprocessor `NDEBUG` is not defined, this - * function outputs a warning when `length == 0`. - * \tparam K is [CopyConstructible](https://en.cppreference.com/w/cpp/named_req/CopyConstructible), - * [LessThanComparable](https://en.cppreference.com/w/cpp/named_req/LessThanComparable), - * and is_swappable. - * \tparam V satisfies is_swappable. - */ -template -inline void quick_sort(array& keys, array& values) { -#if !defined(NDEBUG) - if (keys.length == 0) { - fprintf(stderr, "quick_sort WARNING: Length is zero.\n"); - return; - } -#endif - quick_sort(keys.data, values.data, 0, (unsigned int) keys.length - 1); -} - -/** - * Performs Quicksort on the given arrays `keys` and `values` with the given - * `sorter`. This function assumes `length > 0`. If the preprocessor `NDEBUG` - * is not defined, this function outputs a warning when `length == 0`. - * \tparam K a generic type that satisfies is_swappable, and for which a - * function `bool less_than(const K&, const K&, const Sorter&)` is - * defined. - * \tparam V satisfies is_swappable. - */ -template::value>::type* = nullptr> -inline void quick_sort(array& keys, array& values, const Sorter& sorter) -{ -#if !defined(NDEBUG) - if (keys.length == 0) { - fprintf(stderr, "quick_sort WARNING: Length is zero.\n"); - return; - } -#endif - quick_sort(keys.data, values.data, 0, (unsigned int) keys.length - 1, sorter); -} - - -/** - * - */ - -template -void sort(T* array, unsigned int start, unsigned int end) -{ - if (start >= end) - return; - else if (start + 16 >= end) { - insertion_sort(&array[start], end - start + 1); - return; - } - unsigned int l = start, r = end; - quick_sort_partition(array, start, end, l, r); - sort(array, start, r); - sort(array, l, end); -} - -template::value>::type* = nullptr> -void sort(T* array, unsigned int start, unsigned int end, const Sorter& sorter) -{ - if (start >= end) - return; - else if (start + 16 >= end) { - insertion_sort(&array[start], end - start + 1, sorter); - return; - } - unsigned int l = start, r = end; - quick_sort_partition(array, start, end, l, r, sorter); - sort(array, start, r, sorter); - sort(array, l, end, sorter); -} - -template -void sort(K* keys, V* values, unsigned int start, unsigned int end) -{ - if (start >= end) - return; - else if (start + 16 >= end) { - insertion_sort(&keys[start], &values[start], end - start + 1); - return; - } - unsigned int l = start, r = end; - quick_sort_partition(keys, values, start, end, l, r); - sort(keys, values, start, r); - sort(keys, values, l, end); -} - -template::value>::type* = nullptr> -void sort(K* keys, V* values, unsigned int start, unsigned int end, const Sorter& sorter) -{ - if (start >= end) - return; - else if (start + 16 >= end) { - insertion_sort(&keys[start], &values[start], end - start + 1, sorter); - return; - } - unsigned int l = start, r = end; - quick_sort_partition(keys, values, start, end, l, r, sorter); - sort(keys, values, start, r, sorter); - sort(keys, values, l, end, sorter); -} - -/** - * Performs hybrid Quicksort-insertion sort on the given native array `keys` - * with given `length`. To improve performance, the Quicksort switches to - * insertion sort for small partitions. This function assumes `length > 0`. If - * the preprocessor `NDEBUG` is not defined, this function outputs a warning - * when `length == 0`. - * \tparam T is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable), - * [CopyConstructible](https://en.cppreference.com/w/cpp/named_req/CopyConstructible) - * [LessThanComparable](https://en.cppreference.com/w/cpp/named_req/LessThanComparable), - * and is_swappable. - */ -template -inline void sort(T* keys, unsigned int length) { -#if !defined(NDEBUG) - if (length == 0) { - fprintf(stderr, "sort WARNING: Length is zero.\n"); - return; - } -#endif - sort(keys, 0, length - 1); -} - -/** - * Performs hybrid Quicksort-insertion sort on the given native array `keys` - * with given `length` and `sorter`. To improve performance, the Quicksort - * switches to insertion sort for small partitions. This function assumes - * `length > 0`. If the preprocessor `NDEBUG` is not defined, this function - * outputs a warning when `length == 0`. - * \tparam T a generic type that satisfies is_swappable and is_moveable, and - * for which a function `bool less_than(const T&, const T&, const Sorter&)` - * is defined. - */ -template::value>::type* = nullptr> -inline void sort(T* keys, unsigned int length, const Sorter& sorter) -{ -#if !defined(NDEBUG) - if (length == 0) { - fprintf(stderr, "sort WARNING: Length is zero.\n"); - return; - } -#endif - sort(keys, 0, length - 1, sorter); -} - -/** - * Performs hybrid Quicksort-insertion sort on the given native arrays `keys` - * and `values` with given `length`. To improve performance, the Quicksort - * switches to insertion sort for small partitions. This function assumes - * `length > 0`. If the preprocessor `NDEBUG` is not defined, this function - * outputs a warning when `length == 0`. - * \tparam K is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable), - * [CopyConstructible](https://en.cppreference.com/w/cpp/named_req/CopyConstructible) - * [LessThanComparable](https://en.cppreference.com/w/cpp/named_req/LessThanComparable), - * and is_swappable. - * \tparam V satisfies is_swappable and is_moveable. - */ -template -inline void sort(K* keys, V* values, unsigned int length) { -#if !defined(NDEBUG) - if (length == 0) { - fprintf(stderr, "sort WARNING: Length is zero.\n"); - return; - } -#endif - sort(keys, values, 0, length - 1); -} - -/** - * Performs hybrid Quicksort-insertion sort on the given native arrays `keys` - * and `values` with given `length` and `sorter`. To improve performance, the - * Quicksort switches to insertion sort for small partitions. This function - * assumes `length > 0`. If the preprocessor `NDEBUG` is not defined, this - * function outputs a warning when `length == 0`. - * \tparam K a generic type that satisfies is_swappable and is_moveable, and - * for which a function `bool less_than(const K&, const K&, const Sorter&)` - * is defined. - * \tparam V satisfies is_swappable and is_moveable. - */ -template::value>::type* = nullptr> -inline void sort(K* keys, V* values, unsigned int length, const Sorter& sorter) -{ -#if !defined(NDEBUG) - if (length == 0) { - fprintf(stderr, "sort WARNING: Length is zero.\n"); - return; - } -#endif - sort(keys, values, 0, length - 1, sorter); -} - -/** - * Performs hybrid Quicksort-insertion sort on the given array `keys`. To - * improve performance, the Quicksort switches to insertion sort for small - * partitions. This function assumes `length > 0`. If the preprocessor `NDEBUG` - * is not defined, this function outputs a warning when `length == 0`. - * \tparam T is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable), - * [CopyConstructible](https://en.cppreference.com/w/cpp/named_req/CopyConstructible) - * [LessThanComparable](https://en.cppreference.com/w/cpp/named_req/LessThanComparable), - * and is_swappable. - */ -template -inline void sort(array& keys) { -#if !defined(NDEBUG) - if (keys.length == 0) { - fprintf(stderr, "sort WARNING: Length is zero.\n"); - return; - } -#endif - sort(keys.data, 0, (unsigned int) keys.length - 1); -} - -/** - * Performs hybrid Quicksort-insertion sort on the given array `keys` with the - * given `sorter`. To improve performance, the Quicksort switches to insertion - * sort for small partitions. This function assumes `length > 0`. If the - * preprocessor `NDEBUG` is not defined, this function outputs a warning when - * `length == 0`. - * \tparam T a generic type that satisfies is_swappable and is_moveable, and - * for which a function `bool less_than(const T&, const T&, const Sorter&)` - * is defined. - */ -template::value>::type* = nullptr> -inline void sort(array& keys, const Sorter& sorter) -{ -#if !defined(NDEBUG) - if (keys.length == 0) { - fprintf(stderr, "sort WARNING: Length is zero.\n"); - return; - } -#endif - sort(keys.data, 0, (unsigned int) keys.length - 1, sorter); -} - -/** - * Performs hybrid Quicksort-insertion sort on the given arrays `keys` and - * `values`. To improve performance, the Quicksort switches to insertion sort - * for small partitions. This function assumes `length > 0`. If the - * preprocessor `NDEBUG` is not defined, this function outputs a warning when - * `length == 0`. - * \tparam K is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable), - * [CopyConstructible](https://en.cppreference.com/w/cpp/named_req/CopyConstructible) - * [LessThanComparable](https://en.cppreference.com/w/cpp/named_req/LessThanComparable), - * and is_swappable. - * \tparam V satisfies is_swappable and is_moveable. - */ -template -inline void sort(array& keys, array& values) { -#if !defined(NDEBUG) - if (keys.length == 0) { - fprintf(stderr, "sort WARNING: Length is zero.\n"); - return; - } -#endif - sort(keys.data, values.data, 0, (unsigned int) keys.length - 1); -} - -/** - * Performs hybrid Quicksort-insertion sort on the given arrays `keys` and - * `values` with the given `sorter`. To improve performance, the Quicksort - * switches to insertion sort for small partitions. This function assumes - * `length > 0`. If the preprocessor `NDEBUG` is not defined, this function - * outputs a warning when `length == 0`. - * \tparam K a generic type that satisfies is_swappable and is_moveable, and - * for which a function `bool less_than(const K&, const K&, const Sorter&)` - * is defined. - * \tparam V satisfies is_swappable and is_moveable. - */ -template::value>::type* = nullptr> -inline void sort(array& keys, array& values, const Sorter& sorter) -{ -#if !defined(NDEBUG) - if (keys.length == 0) { - fprintf(stderr, "sort WARNING: Length is zero.\n"); - return; - } -#endif - sort(keys.data, values.data, 0, (unsigned int) keys.length - 1, sorter); -} - -/** - * Returns `true` if and only if the given native array `items` with the given - * `length` is sorted in non-decreasing order. - * \tparam T a generic type for which a function `bool less_than(const T&, const T&, const Sorter&)` - * is defined. - */ -template -inline bool is_sorted(const T* items, size_t length, const Sorter& sorter) { - if (length == 0) - return true; - size_t first = 0, next = 0; - while (++next != length) { - if (less_than(items[next], items[first], sorter)) - return false; - first = next; - } - return true; -} - -/** - * Returns `true` if and only if the given array `items` is sorted in - * non-decreasing order. - * \tparam T a generic type for which a function `bool less_than(const T&, const T&, const Sorter&)` - * is defined. - */ -template -inline bool is_sorted(const array& items, const Sorter& sorter) { - return is_sorted(items.data, items.length, sorter); -} - -/** - * The default_sorter provides a default implementation of - * `bool less_than(const T&, const T&, const default_sorter&)` - * where the comparison is done using the operator `<`. - */ -struct default_sorter { }; - -template -inline bool less_than(const T& first, const T& second, const default_sorter& sorter) { - return (first < second); -} - -/** - * The pointer_sorter compares items using the `<` operator on the dereferenced - * arguments. - */ -struct pointer_sorter { }; - -template -inline bool less_than(const T* first, const T* second, const pointer_sorter& sorter) { - return (*first < *second); -} - -/** - * Deletes consecutive duplicates in the given native `array` with given - * `length` and returns the new length. Note the deleted elements are not freed. - * \tparam T is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). - */ -template -unsigned int unique(T* array, size_t length) -{ - unsigned int result = 0; - for (unsigned int i = 1; i < length; i++) { - if (array[result] != array[i]) - array[++result] = array[i]; - } - return result + 1; -} - -/** - * Deletes consecutive duplicates in the given `array` with given and returns - * the new length. Note the deleted elements are not freed. - * \tparam T is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). - */ -template -inline void unique(array& a) { - a.length = unique(a.data, a.length); -} - -/* forward declarations */ - -unsigned int sample_uniform(unsigned int); - -/** - * Performs a Knuth shuffle on the given native `array` with given `length`. - * \tparam T satisfies is_swappable. - */ -template -void shuffle(T* array, unsigned int length) { -#if !defined(NDEBUG) - if (length == 0) { - fprintf(stderr, "shuffle WARNING: Length is zero.\n"); - return; - } -#endif - for (unsigned int i = length - 1; i > 0; i--) { - unsigned int next = sample_uniform(i + 1); - if (next != i) - core::swap(array[next], array[i]); - } -} - -/** - * Performs a Knuth shuffle on the given native arrays `keys` and `values` with given `length`. - * \tparam T satisfies is_swappable. - */ -template -void shuffle(K* keys, V* values, unsigned int length) { -#if !defined(NDEBUG) - if (length == 0) { - fprintf(stderr, "shuffle WARNING: Length is zero.\n"); - return; - } -#endif - for (unsigned int i = length - 1; i > 0; i--) { - unsigned int next = sample_uniform(i + 1); - if (next != i) { - core::swap(keys[next], keys[i]); - core::swap(values[next], values[i]); - } - } -} - -/** - * Performs a Knuth shuffle on the given `array`. - * \tparam T satisfies is_swappable. - */ -template -inline void shuffle(array& items) { - shuffle(items.data, items.length); -} - -/** - * Given sorted array `a`, this function finds the smallest index `i` such that - * `a[i] >= b` and `i >= start` and `i < end`. - */ -template -unsigned int linear_search( - const T* a, const T& b, - unsigned int start, - unsigned int end) -{ - for (unsigned int i = start; i < end; i++) - if (a[i] >= b) return i; - return end; -} - -/** - * Given sorted array `a`, this function finds the smallest index `i` such that - * `a[i] > b` and `i >= start` and `i < end`. - */ -template -unsigned int strict_linear_search( - const T* a, const T& b, - unsigned int start, - unsigned int end) -{ - for (unsigned int i = start; i < end; i++) - if (a[i] > b) return i; - return end; -} - -/** - * Given sorted array `a`, this function finds the smallest index `i` such that - * `a[i] > b` and `i >= start` and `i < end`. - */ -template -unsigned int reverse_strict_linear_search( - const T* a, const T& b, - unsigned int start, - unsigned int end) -{ - for (unsigned int i = end; i > start; i--) - if (a[i - 1] <= b) return i; - return start; -} - -/** - * Given sorted array `a`, this function finds the smallest index `i` such that - * `a[i] >= b` and `i >= min` and `i <= max`. - */ -/* TODO: implement a strict variant */ -template -unsigned int binary_search( - const T* a, const T& b, - unsigned int min, - unsigned int max) -{ - if (a[max] < b) - return max + 1; - - while (min < max) { - unsigned int mid = (max + min) / 2; - if (a[mid] < b) - min = mid + 1; - else max = mid; - } - - return min; -} - - -/** - * A simple pair data structure. - */ - -template -struct pair { - /** - * The key object in the field. - */ - K key; - - /** - * The value object in the field. - */ - V value; - - pair(const K& key, const V& value) : key(key), value(value) { } - - inline bool operator == (const pair& other) const { - return key == other.key && value == other.value; - } - - inline bool operator != (const pair& other) const { - return key != other.key || value != other.value; - } - - static inline void move(const pair& src, pair& dst) { - core::move(src.key, dst.key); - core::move(src.value, dst.value); - } - - static inline void swap(pair& first, pair& second) { - core::swap(first.key, second.key); - core::swap(first.value, second.value); - } - - static inline unsigned int hash(const pair& pair) { - return hasher::hash(pair.key) + hasher::hash(pair.value); - } - - static inline bool is_empty(const pair& pair) { - return hasher::is_empty(pair.key); - } - - static inline void set_empty(pair& pair) { - hasher::set_empty(pair.key); - } - - static inline void free(pair& pair) { - core::free(pair.key); - core::free(pair.value); - } -}; - -template -inline bool operator < (const pair& first, const pair& second) { - return first.key < second.key; -} - -template -inline bool operator <= (const pair& first, const pair& second) { - return first.key <= second.key; -} - -template -inline bool operator > (const pair& first, const pair& second) { - return first.key > second.key; -} - -template -inline bool operator >= (const pair& first, const pair& second) { - return first.key >= second.key; -} - -template -constexpr pair make_pair(const K& key, const V& value) { - return pair(key, value); -} - -template -inline void swap(pair& first, pair& second) { - swap(first.key, second.key); - swap(first.value, second.value); -} - - -/** - * - */ - -/** - * For every index `i >= index`, this function moves each element at `i` to index `i + 1`. - * \tparam T satisfies is_moveable. - */ -template -void shift_right(T* list, unsigned int length, unsigned int index) -{ - for (unsigned int i = length; i > index; i--) - move(list[i - 1], list[i]); -} - -/** - * For every index `i < index`, this function moves each element at `i + 1` to index `i`. - * \tparam T satisfies is_moveable. - */ -template -void shift_left(T* list, unsigned int index) -{ - for (unsigned int i = 0; i < index; i++) - move(list[i + 1], list[i]); -} - -template -inline void set_union_helper(UnionFunc do_union, const T& item, - unsigned int i, unsigned int j, const T*& prev) -{ - if (RemoveDuplicates) { - if (prev == NULL || *prev != item) { - do_union(item, i, j); - prev = &item; - } - } else { - do_union(item, i, j); - } -} - -/** - * Given ordered native arrays `first` and `second`, this function visits their - * elements sequentially, and: - * 1. For every element `x` that appears in both arrays, `union_both(x, i, j)` - * is called where `first[i] == x` and `second[j] == x`. - * 2. For every element `x` that appears in `first` but not `second`, - * `union_first(x, i, j)` is called where `first[i] == x`, and `j` is the - * smallest such index where `second[j] > x`. - * 3. For every element `x` that appears in `second` but not `first`, - * `union_second(x, i, j)` is called where `second[j] == x`, and `i` is - * the smallest such index where `first[i] > x`. - * - * The visited elements are also ordered. - * - * \tparam T a generic type for which the operators `==`, `!=` and `<` are implemented. - * \tparam RemoveDuplicates if `true`, this function will avoid calling the - * union functions more than once for each element. - */ -template -void set_union(UnionBoth union_both, - UnionFirst union_first, UnionSecond union_second, - const T* first, unsigned int first_length, - const T* second, unsigned int second_length) -{ - unsigned int i = 0, j = 0; - const T* prev = NULL; - while (i < first_length && j < second_length) - { - if (first[i] == second[j]) { - set_union_helper(union_both, first[i], i, j, prev); - i++; j++; - } else if (first[i] < second[j]) { - set_union_helper(union_first, first[i], i, j, prev); - i++; - } else { - set_union_helper(union_second, second[j], i, j, prev); - j++; - } - } - - while (i < first_length) { - set_union_helper(union_first, first[i], i, j, prev); - i++; - } while (j < second_length) { - set_union_helper(union_second, second[j], i, j, prev); - j++; - } -} - -template::value>::type* = nullptr> -inline void set_union_helper(T* dst, SizeType& dst_length, const T& item) { - if (!RemoveDuplicates || dst_length == 0 || dst[dst_length - 1] != item) { - dst[dst_length] = item; - dst_length++; - } -} - -/** - * Given ordered native arrays `first` and `second`, compute their union and - * appends the result to `dst`. The union is also ordered. This function - * assumes `dst` has sufficient capacity to store the union. - * \tparam T satisfies [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable) - * and implements the operators `==` and `<`. - * \tparam RemoveDuplicates if `true`, this function ignores duplicate elements. - */ -template::value>::type* = nullptr> -void set_union(T* dst, SizeType& dst_length, - const T* first, unsigned int first_length, - const T* second, unsigned int second_length) -{ - unsigned int i = 0, j = 0; - while (i < first_length && j < second_length) - { - if (first[i] == second[j]) { - set_union_helper(dst, dst_length, first[i]); - i++; j++; - } else if (first[i] < second[j]) { - set_union_helper(dst, dst_length, first[i]); - i++; - } else { - set_union_helper(dst, dst_length, second[j]); - j++; - } - } - - while (i < first_length) { - set_union_helper(dst, dst_length, first[i]); - i++; - } while (j < second_length) { - set_union_helper(dst, dst_length, second[j]); - j++; - } -} - -/** - * Given ordered native arrays `first` and `second`, compute their union and - * appends the result to `dst`. The union is also ordered. - * \tparam T satisfies [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable) - * and implements the operators `==` and `<`. - * \tparam RemoveDuplicates if `true`, this function ignore duplicate elements. - */ -template -inline bool set_union(array& dst, - const T* first, unsigned int first_length, - const T* second, unsigned int second_length) -{ - if (!dst.ensure_capacity(dst.length + first_length + second_length)) - return false; - set_union(dst.data, dst.length, first, first_length, second, second_length); - return true; -} - -/** - * Given ordered arrays `first` and `second`, compute their union and appends - * the result to `dst`. The union is also ordered. - * \tparam T satisfies [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable) - * and implements the operators `==` and `<`. - * \tparam RemoveDuplicates if `true`, this function ignore duplicate elements. - */ -template -inline bool set_union(array& dst, const array& first, const array& second) { - return set_union(dst, - first.data, (unsigned int) first.length, - second.data, (unsigned int) second.length); -} - -template -struct array_position { - unsigned int array_id; - unsigned int position; - T* element; -}; - -/** - * Given a collection of ordered arrays `arrays`, compute their union and - * appends the result to `dst`. The union is also ordered. - * NOTE: this function assumes the given arrays are all non-empty. - * \tparam T satisfies [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable) - * and implements the operators `==` and `<`. - * \tparam ArraySetCollection a collection type where each element is accessed - * using `arrays[i]` and the size of each array can be obtained using - * `size(arrays[i])`. The element at index `j` for array `i` is accessed - * using `arrays[i][j]`. - */ -template -bool set_union(array& dst, const ArraySetCollection& arrays, unsigned int array_count) -{ - /* first ensure the destination array has enough space */ - unsigned int total_size = dst.length; - for (unsigned int i = 0; i < array_count; i++) { -#if !defined(NDEBUG) - if (size(arrays[i]) == 0) - fprintf(stderr, "set_union WARNING: Input array %u is empty.\n", i); -#endif - total_size += size(arrays[i]); - } - if (!dst.ensure_capacity(total_size)) - return false; - - /* TODO: we can probably use a faster heap structure */ - array_position* heap = (array_position*) malloc(array_count * sizeof(array_position)); - if (heap == NULL) { - fprintf(stderr, "set_union ERROR: Out of memory.\n"); - return false; - } - for (unsigned int i = 0; i < array_count; i++) - heap[i] = { i, 0, &arrays[i][0] }; - std::make_heap(heap, heap + array_count); - - /* add the first item to the destination set */ - unsigned int heap_size = array_count; - std::pop_heap(heap, heap + heap_size); - const array_position& first = heap[heap_size]; - dst.data[dst.length] = first.key; - dst.length++; - if (size(arrays[first.value]) > 1) { - heap[heap_size] = { first.array_id, 1, &arrays[first.array_id][1] }; - std::push_heap(heap, heap + heap_size); - } else { heap_size--; } - - while (heap_size > 0) - { - std::pop_heap(heap, heap + heap_size); - const array_position& next = heap[heap_size]; - if (next.key != dst.last()) { - dst.data[dst.length] = next.key; - dst.length++; - } - if (next.value + 1 < size(arrays[next.value])) { - heap[heap_size] = { next.array_id, next.position + 1, arrays[next.array_id][next.position + 1] }; - std::push_heap(heap, heap + heap_size); - } else { heap_size--; } - } - free(heap); - return true; -} - -/** - * Given ordered native arrays `first` and `second`, this function visits their - * elements sequentially, and for every element `x` that exists in both arrays, - * `intersect(i, j)` is called where `first[i] == x` and `second[j] == x`. The - * visited elements are also ordered. - * \tparam T a generic type for which the operators `==` and `<` are implemented. - * \tparam BinarySearch if `true`, binary search is used to find indices of - * identical elements rather than linear search. - */ -template -bool set_intersect(Intersect intersect, - const T* first, unsigned int first_length, - const T* second, unsigned int second_length) -{ - unsigned int i = 0, j = 0; - while (i < first_length && j < second_length) - { - if (first[i] == second[j]) { - intersect(i, j); - i++; j++; - } else if (first[i] < second[j]) { - if (BinarySearch) { - /* use binary search to find the value of i - such that first.data[i] >= second.data[j] */ - i = binary_search(first, second[j], i, first_length - 1); - } else { - i++; - } - } else { - if (BinarySearch) { - /* use binary search to find the value of j - such that second.data[j] >= first.data[i] */ - j = binary_search(second, first[i], j, second_length - 1); - } else { - j++; - } - } - } - return true; -} - -/** - * Given ordered native arrays `first` and `second`, compute the intersection - * and append it to the native array `intersection`. The computed intersection - * is also ordered. This function assumes `intersection` has sufficient capacity. - * \tparam T satisfies [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable) - * and implements the operators `==` and `<`. - * \tparam BinarySearch if `true`, binary search is used to find indices of - * identical elements rather than linear search. - */ -template::value>::type* = nullptr> -bool set_intersect( - T* intersection, SizeType& intersection_length, - const T* first, unsigned int first_length, - const T* second, unsigned int second_length) -{ - unsigned int i = 0, j = 0; - while (i < first_length && j < second_length) - { - if (first[i] == second[j]) { - intersection[intersection_length] = first[i]; - intersection_length++; - i++; j++; - } else if (first[i] < second[j]) { - if (BinarySearch) { - /* use binary search to find the value of i - such that first.data[i] >= second.data[j] */ - i = binary_search(first, second[j], i, first_length - 1); - } else { - i++; - } - } else { - if (BinarySearch) { - /* use binary search to find the value of j - such that second.data[j] >= first.data[i] */ - j = binary_search(second, first[i], j, second_length - 1); - } else { - j++; - } - } - } - return true; -} - -/** - * Given ordered native arrays `first` and `second`, compute the intersection - * and append it to the array `intersection`. The computed intersection is also - * ordered. - * \tparam T satisfies [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable) - * and implements the operators `==` and `<`. - * \tparam BinarySearch if `true`, binary search is used to find indices of - * identical elements rather than linear search. - */ -template -inline bool set_intersect( - array& intersection, - const T* first, unsigned int first_length, - const T* second, unsigned int second_length) -{ - if (!intersection.ensure_capacity(intersection.length + max(first_length, second_length))) - return false; - - return set_intersect( - intersection.data, intersection.length, - first, first_length, second, second_length); -} - -/** - * Given ordered arrays `first` and `second`, compute the intersection and - * append it to the array `intersection`. The computed intersection is also - * ordered. - * \tparam T satisfies [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable) - * and implements the operators `==` and `<`. - * \tparam BinarySearch if `true`, binary search is used to find indices of - * identical elements rather than linear search. - */ -template -inline bool set_intersect( - array& intersection, - const array& first, - const array& second) -{ - return set_intersect(intersection, - first.data, (unsigned int) first.length, - second.data, (unsigned int) second.length); -} - -/** - * Given ordered native arrays `first` and `second`, compute the intersection - * in-place and store it in `first`. The computed intersection is also ordered. - * \tparam T satisfies [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable) - * and implements the operators `==` and `<`. - * \tparam BinarySearch if `true`, binary search is used to find indices of - * identical elements rather than linear search. - */ -template::value>::type* = nullptr> -void set_intersect( - T* first, SizeType& first_length, - const T* second, unsigned int second_length) -{ - unsigned int index = 0; - unsigned int i = 0, j = 0; - while (i < first_length && j < second_length) - { - if (first[i] == second[j]) { - first[index] = first[i]; - index++; i++; j++; - } else if (first[i] < second[j]) { - if (BinarySearch) { - /* use binary search to find the value of i - such that first.data[i] >= second.data[j] */ - i = binary_search(first, second[j], i, first_length - 1); - } else { - i++; - } - } else { - if (BinarySearch) { - /* use binary search to find the value of j - such that second.data[j] >= first.data[i] */ - j = binary_search(second, first[i], j, second_length - 1); - } else { - j++; - } - } - } - first_length = index; -} - -/** - * Given ordered array `first` and ordered native array `second`, compute the - * intersection in-place and store it in `first`. The computed intersection is - * also ordered. - * \tparam T satisfies [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable) - * and implements the operators `==` and `<`. - * \tparam BinarySearch if `true`, binary search is used to find indices of - * identical elements rather than linear search. - */ -template -inline void set_intersect(array& first, - const T* second, unsigned int second_length) -{ - return set_intersect( - first.data, first.length, second, second_length); -} - -/** - * Given ordered arrays `first` and `second`, compute the intersection in-place - * and store it in `first`. The computed intersection is also ordered. - * \tparam T satisfies [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable) - * and implements the operators `==` and `<`. - * \tparam BinarySearch if `true`, binary search is used to find indices of - * identical elements rather than linear search. - */ -template -inline void set_intersect(array& first, const array& second) { - return set_intersect(first, second.data, second.length); -} - -/** - * Returns true if the intersection of `first` and `second` is non-empty, - * where `first` and `second` are ordered native arrays. - * \tparam T a generic type that implements the operators `==` and `<`. - * \tparam BinarySearch if `true`, binary search is used to find indices of - * identical elements rather than linear search. - */ -template -bool has_intersection( - const T* first, unsigned int first_length, - const T* second, unsigned int second_length) -{ - unsigned int i = 0, j = 0; - while (i < first_length && j < second_length) - { - if (first[i] == second[j]) { - return true; - i++; j++; - } else if (first[i] < second[j]) { - if (BinarySearch) { - /* use binary search to find the value of i - such that first.data[i] >= second.data[j] */ - i = binary_search(first, second[j], i, first_length - 1); - } else { - i++; - } - } else { - if (BinarySearch) { - /* use binary search to find the value of j - such that second.data[j] >= first.data[i] */ - j = binary_search(second, first[i], j, second_length - 1); - } else { - j++; - } - } - } - return false; -} - -/** - * Returns true if the intersection of `first` and `second` is non-empty, - * where `first` and `second` are ordered arrays. - * \tparam T a generic type that implements the operators `==` and `<`. - * \tparam BinarySearch if `true`, binary search is used to find indices of - * identical elements rather than linear search. - */ -template -inline bool has_intersection(const array& first, const array& second) { - return has_intersection( - first.data, (unsigned int) first.length, - second.data, (unsigned int) second.length); -} - -/** - * Returns true if `first` is a subset of `second`, where `first` and `second` - * are ordered native arrays. - * \tparam T a generic type that implements the operators `==` and `<`. - * \tparam BinarySearch if `true`, binary search is used to find indices of - * identical elements rather than linear search. - */ -template -bool is_subset( - const T* first, unsigned int first_length, - const T* second, unsigned int second_length) -{ - unsigned int i = 0, j = 0; - while (i < first_length && j < second_length) - { - if (first[i] == second[j]) { - i++; j++; - } else if (first[i] < second[j]) { - return false; - } else { - if (BinarySearch) { - /* use binary search to find the value of j - such that second.data[j] >= first.data[i] */ - j = binary_search(second, first[i], j, second_length - 1); - } else { - j++; - } - } - } - return (i == first_length); -} - -/** - * Given ordered native arrays `first` and `second`, this function visits their - * elements sequentially, and for every element `x` that exists in `first` but - * not in `second`, `emit(i)` is called where `first[i] == x`. The visited - * elements are also ordered. - * \tparam T a generic type that implements the operators `==` and `<`. - * \tparam BinarySearch if `true`, binary search is used to find indices of - * identical elements rather than linear search. - */ -template -void set_subtract(EmitFunction emit, - const T* first, unsigned int first_length, - const T* second, unsigned int second_length) -{ - unsigned int i = 0, j = 0; - while (i < first_length && j < second_length) - { - if (first[i] == second[j]) { - i++; j++; - } else if (first[i] < second[j]) { - if (BinarySearch) { - /* use binary search to find the value of i - such that first.data[i] >= second.data[j] */ - unsigned int next_i = binary_search(first, second[j], i, first_length - 1); - for (; i < next_i; i++) - emit(i); - i = next_i; - } else { - emit(i); - i++; - } - } else { - if (BinarySearch) { - /* use binary search to find the value of j - such that second.data[j] >= first.data[i] */ - j = binary_search(second, first[i], j, second_length - 1); - } else { - j++; - } - } - } - - while (i < first_length) { - emit(i); - i++; - } -} - -/** - * Given ordered native arrays `first` and `second`, this function computes the - * set difference between `first` and `second` and stores the result in `dst`. - * This function assumes `dst` has sufficient capacity. The set difference is - * also ordered. - * \tparam T satisfies [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable) - * and implements the operators `==` and `<`. - * \tparam BinarySearch if `true`, binary search is used to find indices of - * identical elements rather than linear search. - */ -template::value>::type* = nullptr> -void set_subtract(T* dst, SizeType& dst_length, - const T* first, unsigned int first_length, - const T* second, unsigned int second_length) -{ - unsigned int i = 0, j = 0; - while (i < first_length && j < second_length) - { - if (first[i] == second[j]) { - i++; j++; - } else if (first[i] < second[j]) { - if (BinarySearch) { - /* use binary search to find the value of i - such that first.data[i] >= second.data[j] */ - unsigned int next_i = binary_search(first, second[j], i, first_length - 1); - for (; i < next_i; i++) { - dst[dst_length] = first[i]; - dst_length++; - } - i = next_i; - } else { - dst[dst_length] = first[i]; - dst_length++; - i++; - } - } else { - if (BinarySearch) { - /* use binary search to find the value of j - such that second.data[j] >= first.data[i] */ - j = binary_search(second, first[i], j, second_length - 1); - } else { - j++; - } - } - } - - memcpy(dst + dst_length, first + i, (first_length - i) * sizeof(T)); - dst_length += first_length - i; -} - -/** - * Given ordered native arrays `first` and `second`, this function computes the - * set difference between `first` and `second` and stores the result in `dst`. - * The set difference is also ordered. - * \tparam T satisfies [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable) - * and implements the operators `==` and `<`. - * \tparam BinarySearch if `true`, binary search is used to find indices of - * identical elements rather than linear search. - */ -template -bool set_subtract(array& dst, - const T* first, unsigned int first_length, - const T* second, unsigned int second_length) -{ - if (!dst.ensure_capacity(dst.length + max(first_length, second_length))) - return false; - set_subtract(dst.data, dst.length, first, first_length, second, second_length); - return true; -} - -/** - * Given ordered arrays `first` and `second`, this function computes the set - * difference between `first` and `second` and stores the result in `dst`. The - * set difference is also ordered. - * \tparam T satisfies [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable) - * and implements the operators `==` and `<`. - * \tparam BinarySearch if `true`, binary search is used to find indices of - * identical elements rather than linear search. - */ -template -inline bool set_subtract(array& dst, - const array& first, - const array& second) -{ - return set_subtract(dst, - first.data, (unsigned int) first.length, - second.data, (unsigned int) second.length); -} - -/** - * Given ordered native arrays `first` and `second`, this function computes the - * set difference in-place between `first` and `second` and stores the result - * in `first`. The set difference is also ordered. - * \tparam T satisfies [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable) - * and implements the operators `==` and `<`. - * \tparam BinarySearch if `true`, binary search is used to find indices of - * identical elements rather than linear search. - */ -template::value>::type* = nullptr> -void set_subtract( - T* first, SizeType& first_length, - const T* second, unsigned int second_length) -{ - unsigned int index = 0; - unsigned int i = 0, j = 0; - while (i < first_length && j < second_length) - { - if (first[i] == second[j]) { - i++; j++; - } else if (first[i] < second[j]) { - if (BinarySearch) { - /* use binary search to find the value of i - such that first.data[i] >= second.data[j] */ - unsigned int next_i = binary_search(first, second[j], i, first_length - 1); - for (; i < next_i; i++) { - first[index] = first[i]; - index++; - } - i = next_i; - } else { - first[index] = first[i]; - index++; i++; - } - } else { - if (BinarySearch) { - /* use binary search to find the value of j - such that second.data[j] >= first.data[i] */ - j = binary_search(second, first[i], j, second_length - 1); - } else { - j++; - } - } - } - - while (i < first_length) { - first[index] = first[i]; - index++; i++; - } - first_length = index; -} - -/** - * Given ordered arrays `first` and `second`, this function computes the set - * difference in-place between `first` and `second` and stores the result in - * `first`. The set difference is also ordered. - * \tparam T satisfies [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable) - * and implements the operators `==` and `<`. - * \tparam BinarySearch if `true`, binary search is used to find indices of - * identical elements rather than linear search. - */ -template -inline void set_subtract(array& first, const array& second) -{ - return set_subtract(first.data, first.length, second.data, second.length); -} - -inline void array_test(void) -{ - array buf(1); - buf.append("0123456789 ", 11); - - buf[(unsigned int) buf.length] = '\0'; - if (buf.length != 11) - fprintf(stderr, "array test ERROR: First array length test failed.\n"); - if (buf.capacity != 16) - fprintf(stderr, "array test ERROR: First array capacity test failed.\n"); - if (strcmp(buf.data, "0123456789 ") != 0) - fprintf(stderr, "array test ERROR: First string comparison test failed.\n"); - - buf.append("abcdefghijklmnopqrstuvwxyz ", 27); - - buf[(unsigned int) buf.length] = '\0'; - if (buf.length != 11 + 27) - fprintf(stderr, "array test ERROR: Second array length test failed.\n"); - if (buf.capacity != 64) - fprintf(stderr, "array test ERROR: Second array capacity test failed.\n"); - if (strcmp(buf.data, "0123456789 abcdefghijklmnopqrstuvwxyz ") != 0) - fprintf(stderr, "array test ERROR: Second string comparison test failed.\n"); - - buf.append("9876543210 ", 11); - - buf[(unsigned int) buf.length] = '\0'; - if (buf.length != 11 + 27 + 11) - fprintf(stderr, "array test ERROR: Third array length test failed.\n"); - if (buf.capacity != 64) - fprintf(stderr, "array test ERROR: Third array capacity test failed.\n"); - if (strcmp(buf.data, "0123456789 abcdefghijklmnopqrstuvwxyz 9876543210 ") != 0) - fprintf(stderr, "array test ERROR: Third string comparison test failed.\n"); - - /* test some of the helper functions */ - array numbers = array(10); - numbers.add(4); - numbers.add(-6); - numbers.add(4); - numbers.add(2); - numbers.add(0); - numbers.add(-6); - numbers.add(1); - numbers.add(4); - numbers.add(2); - - array numbers_copy = array(10); - numbers_copy.append(numbers.data, numbers.length); - insertion_sort(numbers); - quick_sort(numbers_copy); - - int expected[] = {-6, -6, 0, 1, 2, 2, 4, 4, 4}; - for (unsigned int i = 0; i < 9; i++) { - if (numbers[i] != expected[i]) { - fprintf(stderr, "array test ERROR: insertion_sort failed.\n"); - return; - } - if (numbers_copy[i] != expected[i]) { - fprintf(stderr, "array test ERROR: quick_sort failed.\n"); - return; - } - } - - int expected_unique[] = {-6, 0, 1, 2, 4}; - unique(numbers); - if (numbers.length != 5) { - fprintf(stderr, "array test ERROR: unique failed.\n"); - return; - } - for (unsigned int i = 0; i < 5; i++) { - if (numbers[i] != expected_unique[i]) { - fprintf(stderr, "array test ERROR: unique failed.\n"); - return; - } - } - - printf("array test completed.\n"); -} - -} /* namespace core */ - -#endif /* ARRAY_H_ */ diff --git a/reasoners/data_preprocess/logic/graph_dataset_gen/core/core.h b/reasoners/data_preprocess/logic/graph_dataset_gen/core/core.h deleted file mode 100644 index 18986681f..000000000 --- a/reasoners/data_preprocess/logic/graph_dataset_gen/core/core.h +++ /dev/null @@ -1,479 +0,0 @@ -/** - * \file core.h - * - * This file contains functions widely used by the core library. It contains - * functions for object movement, swapping, copying, and hashing, as well as - * the `free` function, which deinitializes objects and releases its - * memory resources. - * - * Memory management - * ----------------- - * - * This library does not manage memory, and leaves it to the user to ensure - * memory is allocated and freed correctly. The core data structures, such as - * core::array, core::hash_set, core::hash_map, core::array_map do not - * automatically free their elements upon deallocation. The library uses the - * core::free() function to free memory resources. When developing a new type, - * the class/struct should implement the public static function `free`, as in - * the example below. - * - * ```{.cpp} - * #include - * #include - * #include - * using namespace core; - * - * struct custom_string { - * char* buffer; - * - * custom_string() { } - * - * custom_string(const char* src) { - * buffer = (char*) malloc(sizeof(char) * (strlen(src) + 1)); - * memcpy(buffer, src, sizeof(char) * (strlen(src) + 1)); - * } - * - * ~custom_string() { - * core::free(buffer); - * } - * - * static void free(custom_string& s) { - * core::free(s.buffer); - * } - * }; - * - * bool init(custom_string& s, const char* src) { - * s.buffer = (char*) malloc(sizeof(char) * (strlen(src) + 1)); - * if (s.buffer == NULL) - * return false; - * memcpy(s.buffer, src, sizeof(char) * (strlen(src) + 1)); - * return true; - * } - * - * int main() { - * custom_string first = custom_string("first"); - * custom_string& second = *((custom_string*) alloca(sizeof(custom_string))); - * init(second, "second"); - * - * printf("%s ", first.buffer); - * printf("%s", second.buffer); - * - * core::free(second); - * } - * ``` - * In this example, core::free() is not called on `first` since it was - * initialized on the stack, and so its destructor will automatically free - * `buffer`. However, this was not the case for `second`, and so we must call - * core::free() directly. The expected output of the program is `first second`. - * - * - */ - -#ifndef CORE_H_ -#define CORE_H_ - -#include -#include -#include -#include -#include - -#if defined(_WIN32) -#include -#endif - - -namespace core { - - -/** - * This function returns the compile-time length of the static array `array`. - */ -template -inline constexpr std::size_t array_length(const T (&array)[N]) { - return N; -} - -/** - * Returns `first` if `first < second`, and returns `second` otherwise. - */ -template -inline typename std::common_type::type min(const A& first, const B& second) { - typedef typename std::common_type::type C; - return (static_cast(first) < static_cast(second)) ? first : second; -} - -/** - * Returns `second` if `first < second`, and returns `first` otherwise. - */ -template -inline typename std::common_type::type max(const A& first, const B& second) { - typedef typename std::common_type::type C; - return (static_cast(first) < static_cast(second)) ? second : first; -} - -/** - * If `T` satisfies [is_fundamental](http://en.cppreference.com/w/cpp/types/is_fundamental), - * [is_enum](http://en.cppreference.com/w/cpp/types/is_enum), or - * [is_pointer](http://en.cppreference.com/w/cpp/types/is_pointer), this - * function assigns `a = b`. - */ -template::value || std::is_enum::value || std::is_pointer::value>::type* = nullptr> -inline bool init(T& a, const T& b) { - a = b; - return true; -} - -/** - * If `T` satisfies [is_fundamental](http://en.cppreference.com/w/cpp/types/is_fundamental), - * [is_enum](http://en.cppreference.com/w/cpp/types/is_enum), - * or [is_pointer](http://en.cppreference.com/w/cpp/types/is_pointer), - * `dst` is assigned to `src`. Otherwise, this function calls - * `T::move(src, dst)`. This function is intended to be used to effectively - * move the object stored in `src` into the location specified by `dst`. - */ -template::value && !std::is_enum::value && !std::is_pointer::value>::type* = nullptr> -inline void move(const T& src, T& dst) { - T::move(src, dst); -} - -template::value || std::is_enum::value || std::is_pointer::value>::type* = nullptr> -inline void move(const T& a, T& b) { - b = a; -} - -/** - * If `T` satisfies [is_fundamental](http://en.cppreference.com/w/cpp/types/is_fundamental), - * [is_enum](http://en.cppreference.com/w/cpp/types/is_enum), - * or [is_pointer](http://en.cppreference.com/w/cpp/types/is_pointer), - * `dst` is assigned to `src`. Otherwise, this function calls and returns the - * result of `T::copy(src, dst)`. This function is intended to be used to copy - * the object stored in `src` into the location specified by `dst`. - */ -template::value && !std::is_enum::value && !std::is_pointer::value>::type* = nullptr> -inline bool copy(const T& src, T& dst) { - return T::copy(src, dst); -} - -template::value || std::is_enum::value || std::is_pointer::value>::type* = nullptr> -inline bool copy(const T& a, T& b) { - b = a; - return true; -} - -/** - * If `T` satisfies [is_fundamental](http://en.cppreference.com/w/cpp/types/is_fundamental), - * [is_enum](http://en.cppreference.com/w/cpp/types/is_enum), - * or [is_pointer](http://en.cppreference.com/w/cpp/types/is_pointer), - * the values of `a` and `b` are swapped. Otherwise, this function calls - * `T::swap(a, b)`. This function is intended to be used to swap the object - * stored in `a` with the object stored in `b`. - */ -template::value && !std::is_enum::value && !std::is_pointer::value>::type* = nullptr> -inline void swap(T& a, T& b) { - T::swap(a, b); -} - -template::value || std::is_enum::value || std::is_pointer::value>::type* = nullptr> -inline void swap(T& a, T& b) { - T temp = a; - a = b; - b = temp; -} - -template::value || std::is_enum::value>::type* = nullptr> -inline constexpr T size_of(const T& a) { - return sizeof(T); -} - -template::value>::type* = nullptr> -inline constexpr long unsigned int size_of(const E& a) { - return sizeof(a); -} - -struct default_metric { }; - -template::value && !std::is_enum::value && !std::is_pointer::value>::type* = nullptr> -inline long unsigned int size_of(const T& a) { - return T::size_of(a, default_metric()); -} - -template -inline constexpr long unsigned int size_of(const T* const& a) { - return sizeof(T*); -} - -template -inline long unsigned int size_of(const T (&array)[N]) { - long unsigned int size = 0; - for (size_t i = 0; i < N; i++) - size += size_of(array[i]); - return size; -} - -template::value || std::is_enum::value>::type* = nullptr> -inline long unsigned int size_of(const T& a, const default_metric& metric) { - return size_of(a); -} - -template -inline long unsigned int size_of(const T& a, const Metric& metric) { - return T::size_of(a, metric); -} - -template::value || std::is_enum::value>::type* = nullptr> -inline void free(T& a) { } - -/** - * If `T` satisfies [is_fundamental](http://en.cppreference.com/w/cpp/types/is_fundamental) - * or [is_enum](http://en.cppreference.com/w/cpp/types/is_enum), - * this function does nothing. Otherwise, this function calls - * `T::free(a, std::forward(args)...)`. This function is intended to be - * used to free the object stored in `a`, passing along any additional - * arguments. Note that since `args` is variadic, it may also be empty. - */ -template::value && !std::is_enum::value>::type* = nullptr> -inline void free(T& a, Args&&... args) { - T::free(a, std::forward(args)...); -} - -/** - * This function calls [free](http://en.cppreference.com/w/c/memory/free) - * on `a`. - */ -template -inline void free(T* a) { - ::free(a); -} - - -/** - * - */ - -namespace detail { - template static auto test_swappable(int32_t) -> - decltype(void(core::swap(std::declval(), std::declval())), std::true_type{}); - template static auto test_swappable(int64_t) -> std::false_type; - - template static auto test_moveable(int32_t) -> - decltype(void(core::move(std::declval(), std::declval())), std::true_type{}); - template static auto test_moveable(int64_t) -> std::false_type; - - template static auto test_copyable(int32_t) -> - decltype(void(core::copy(std::declval(), std::declval())), std::true_type{}); - template static auto test_copyable(int64_t) -> std::false_type; -} - -/** - * This type trait is [true_type](http://en.cppreference.com/w/cpp/types/integral_constant) - * if and only if `T` satisfies any of the following: - * 1. [is_fundamental](http://en.cppreference.com/w/cpp/types/is_fundamental), - * 2. [is_enum](http://en.cppreference.com/w/cpp/types/is_enum), - * 3. [is_pointer](http://en.cppreference.com/w/cpp/types/is_pointer), - * 4. or implements the public static method `void T::move(const T&, T&)`. - */ -template struct is_moveable : decltype(core::detail::test_moveable(0)){}; - -/** - * This type trait is [true_type](http://en.cppreference.com/w/cpp/types/integral_constant) - * if and only if `T` satisfies any of the following: - * 1. [is_fundamental](http://en.cppreference.com/w/cpp/types/is_fundamental), - * 2. [is_enum](http://en.cppreference.com/w/cpp/types/is_enum), - * 3. [is_pointer](http://en.cppreference.com/w/cpp/types/is_pointer), - * 4. or implements the public static method `void T::swap(T&, T&)`. - */ -template struct is_swappable : decltype(core::detail::test_swappable(0)){}; - -/** - * This type trait is [true_type](http://en.cppreference.com/w/cpp/types/integral_constant) - * if and only if `T` satisfies any of the following: - * 1. [is_fundamental](http://en.cppreference.com/w/cpp/types/is_fundamental), - * 2. [is_enum](http://en.cppreference.com/w/cpp/types/is_enum), - * 3. [is_pointer](http://en.cppreference.com/w/cpp/types/is_pointer), - * 4. or implements the public static method `bool T::copy(const T&, T&)`. - */ -template struct is_copyable : decltype(core::detail::test_copyable(0)){}; - -/** - * This type trait is [true_type](http://en.cppreference.com/w/cpp/types/integral_constant) - * if and only if every template argument is also `true_type`. - */ -template struct and_type; -template<> struct and_type<> { - typedef std::true_type type; - static constexpr bool value = true; -}; -template struct and_type { - typedef typename and_type::type type; - static constexpr bool value = and_type::value; -}; -template struct and_type { - typedef std::false_type type; - static constexpr bool value = false; -}; - - -/** - * - */ - -#define DEFAULT_HASH_SEED 0 - -/* these are defined in map.h */ -#if defined(__LP64__) || defined(_WIN64) || defined(__x86_64__) || defined(__ppc64__) - -template -inline uint_fast32_t default_hash(const K& key); - -template -inline uint_fast32_t default_hash(const K* keys, unsigned int length); - -#else -template -inline unsigned int default_hash(const K& key); - -template -inline unsigned int default_hash(const K* keys, unsigned int length); -#endif - -template -inline void set_all_empty(K* keys, unsigned int length) { - memset(keys, 0, sizeof(K) * length); -} - -template -struct hasher { - static inline bool is_empty(const K& key) { - return K::is_empty(key); - } - - static inline void set_empty(K& key) { - K::set_empty(key); - } - - static inline void set_empty(K* keys, unsigned int length) { - K::set_empty(keys, length); - } - - static inline unsigned int hash(const K& key) { - return K::hash(key); - } -}; - -template -struct hasher::value || std::is_enum::value || std::is_pointer::value>::type> { - static inline bool is_empty(const K& key) { - return (key == static_cast(0)); - } - - static inline void set_empty(K& key) { - key = static_cast(0); - } - - static inline void set_empty(K* keys, unsigned int length) { - set_all_empty(keys, length); - } - - static inline unsigned int hash(const K& key) { - return default_hash(key); - } -}; - -/** - * Hashtables in this library require the type `K` to define a special - * "empty value" to indicate that a bucket is empty. For `K` that satisfies - * [is_fundamental](http://en.cppreference.com/w/cpp/types/is_fundamental), - * [is_enum](http://en.cppreference.com/w/cpp/types/is_enum), - * or [is_pointer](http://en.cppreference.com/w/cpp/types/is_pointer), - * this function returns whether `key == static_cast(0)`. Otherwise, this - * function returns `K::is_empty(key)`. Thus, to enable the use of a custom - * struct/class as a hashtable key, it must implement the public static - * function `is_empty`. - */ -template -inline bool is_empty(const K& key) { - return hasher::is_empty(key); -} - -/** - * Hashtables in this library require the type `K` to define a special - * "empty value" to indicate that a bucket is empty. For `K` that satisfies - * [is_fundamental](http://en.cppreference.com/w/cpp/types/is_fundamental), - * [is_enum](http://en.cppreference.com/w/cpp/types/is_enum), - * or [is_pointer](http://en.cppreference.com/w/cpp/types/is_pointer), - * this function sets `key` to `static_cast(0)`. Otherwise, this function - * calls `K::set_empty(key)`. Some hashtable operations use this operation, and - * therefore, if a custom struct/class is used as a hashtable key, it must - * implement the public static function `set_empty` in order to use the - * aforementioned hashtable operations. - */ -template -inline void set_empty(K& key) { - hasher::set_empty(key); -} - - -/** - * Returns the base-2 logarithm of the given 32-bit unsigned integer argument. - * This function assumes the argument is not `0`. - */ -template::value && std::is_integral::value && sizeof(T) == sizeof(uint32_t)>::type* = nullptr> -inline unsigned int log2(T x) { -#if defined(_WIN32) - unsigned long index; - _BitScanReverse(&index, x); - return index; -#else - return (unsigned int) sizeof(uint32_t) * CHAR_BIT - __builtin_clz(x) - 1; -#endif -} - -/** - * Returns the base-2 logarithm of the given 64-bit unsigned integer argument. - * This function assumes the argument is not `0`. - */ -template::value && std::is_integral::value && sizeof(T) == sizeof(uint64_t)>::type* = nullptr> -inline unsigned int log2(T x) { -#if defined(_WIN64) - unsigned long index; - _BitScanReverse64(&index, x); - return index; -#elif defined(_WIN32) - uint32_t left = (x >> 32); - if (left == 0) return log2((uint32_t) x); - else return 32 + log2(left); -#else - return (unsigned int) sizeof(uint64_t) * CHAR_BIT - __builtin_clzll(x) - 1; -#endif -} - -/** - * Returns the base-2 logarithm of the given unsigned integer argument. This - * function assumes the argument is not `0`. - */ -template::value && std::is_integral::value>::type* = nullptr> -constexpr unsigned int static_log2(T x) { - return (x < 2) ? 1 : (1 + static_log2(x)); -} - - -} /* namespace core */ - -#endif /* CORE_H_ */ diff --git a/reasoners/data_preprocess/logic/graph_dataset_gen/core/io.h b/reasoners/data_preprocess/logic/graph_dataset_gen/core/io.h deleted file mode 100644 index 0ddad8cac..000000000 --- a/reasoners/data_preprocess/logic/graph_dataset_gen/core/io.h +++ /dev/null @@ -1,1496 +0,0 @@ -/** - * \file io.h - * - * This file defines serialization and deserialization functions `read` and - * `write`, as well as the `print` function for all fundamental types and core - * data structures. - * - * The file also contains the definition and implementation of the - * `core::memory_stream` class, which may be used to read/write to an in-memory - * buffer. - * - * Scribes - * ------- - * - * The `read`, `write`, and `print` functions in this library follow a very - * regular argument structure: The first argument is the object to - * read/write/print. The second argument is the stream to read from/write - * to/print to. Most functions also require a third (optional) argument called - * the `scribe`, which controls how the subobjects are read/written/printed. - * The scribe is typically passed to the read/write/print functions when - * operating on subobjects of the given object. - * - * Calling read/write/print with core::default_scribe will call the same - * function without the third argument. This largely corresponds to the - * "default" behavior of those functions. - * - * For example, `write(const core::array& a, Stream& out, Writer&&... writer)` - * will call the function `write(a[i], out, writer)` for every element `a[i]` - * in `a`. This enables users to define their own scribes, and define new ways - * to read/write/print objects without having to re-implement the - * read/write/print functions for container structures such as core::array. The - * following example demonstrates how the behavior of the `print` function for - * an array of integers can be altered using a custom scribe. - * - * ```{.cpp} - * #include - * using namespace core; - * - * struct my_string_scribe { - * const char* strings[3]; - * }; - * - * template - * bool print(int i, Stream& out, const my_string_scribe& printer) { - * return print(printer.strings[i], out); - * } - * - * int main() { - * array a = array(8); - * a.add(1); a.add(2); a.add(0); - * - * print(a, stdout); print(' ', stdout); - * - * default_scribe def; - * print(a, stdout, def); print(' ', stdout); - * - * my_string_scribe printer; - * printer.strings[0] = "vici"; - * printer.strings[1] = "veni"; - * printer.strings[2] = "vidi"; - * - * print(a, stdout, printer); - * } - * ``` - * - * This example has expected output `[1, 2, 0] [1, 2, 0] [veni, vidi, vici]`. - * - * - */ - -#ifndef IO_H_ -#define IO_H_ - -#include "array.h" -#include "map.h" - -#include -#include -#include - -#if defined(__APPLE__) -#include - -inline size_t mbrtoc32(char32_t* wc, const char* s, size_t n, mbstate_t* ps) { - wchar_t c; - size_t result = mbrtowc(&c, s, n, ps); - if (result < (size_t) -2 && wc != nullptr) - *wc = c; - return result; -} - -#else /* __APPLE__ */ -#include -#endif /* __APPLE__ */ - -namespace core { - -/** - * Reads `sizeof(T)` bytes from `in` and writes them to the memory referenced - * by `value`. This function does not perform endianness transformations. - * \param in the stream given by a [FILE](https://en.cppreference.com/w/c/io) pointer. - * \tparam T satisfies [is_fundamental](https://en.cppreference.com/w/cpp/types/is_fundamental). - */ -template::value>::type* = nullptr> -inline bool read(T& value, FILE* in) { - return (fread(&value, sizeof(T), 1, in) == 1); -} - -/** - * Reads `length` elements from `in` and writes them to the native array - * `values`. This function does not perform endianness transformations. - * \param in the stream given by a [FILE](https://en.cppreference.com/w/c/io) pointer. - * \tparam T satisfies [is_fundamental](https://en.cppreference.com/w/cpp/types/is_fundamental). - */ -template::value>::type* = nullptr> -inline bool read(T* values, FILE* in, SizeType length) { - return (fread(values, sizeof(T), length, in) == length); -} - -/** - * Writes `sizeof(T)` bytes to `out` from the memory referenced by `value`. - * This function does not perform endianness transformations. - * \param out the stream given by a [FILE](https://en.cppreference.com/w/c/io) pointer. - * \tparam T satisfies [is_fundamental](https://en.cppreference.com/w/cpp/types/is_fundamental). - */ -template::value>::type* = nullptr> -inline bool write(const T& value, FILE* out) { - return (fwrite(&value, sizeof(T), 1, out) == 1); -} - -/** - * Writes `length` elements to `out` from the native array `values`. This - * function does not perform endianness transformations. - * \param out the stream given by a [FILE](https://en.cppreference.com/w/c/io) pointer. - * \tparam T satisfies [is_fundamental](https://en.cppreference.com/w/cpp/types/is_fundamental). - */ -template::value>::type* = nullptr> -inline bool write(const T* values, FILE* out, SizeType length) { - return (fwrite(values, sizeof(T), length, out) == length); -} - -/** - * Prints the character `value` to the stream given by the - * [FILE](https://en.cppreference.com/w/c/io) pointer `out`. - */ -inline bool print(const char& value, FILE* out) { - return (fputc(value, out) != EOF); -} - -/** - * Prints the int `value` to the stream given by the - * [FILE](https://en.cppreference.com/w/c/io) pointer `out`. - */ -inline bool print(const int& value, FILE* out) { - return (fprintf(out, "%d", value) > 0); -} - -/** - * Prints the long `value` to the stream given by the - * [FILE](https://en.cppreference.com/w/c/io) pointer `out`. - */ -inline bool print(const long& value, FILE* out) { - return (fprintf(out, "%ld", value) > 0); -} - -/** - * Prints the long long `value` to the stream given by the - * [FILE](https://en.cppreference.com/w/c/io) pointer `out`. - */ -inline bool print(const long long& value, FILE* out) { - return (fprintf(out, "%lld", value) > 0); -} - -/** - * Prints the unsigned int `value` to the stream given by the - * [FILE](https://en.cppreference.com/w/c/io) pointer `out`. - */ -inline bool print(const unsigned int& value, FILE* out) { - return (fprintf(out, "%u", value) > 0); -} - -/** - * Prints the unsigned long `value` to the stream given by the - * [FILE](https://en.cppreference.com/w/c/io) pointer `out`. - */ -inline bool print(const unsigned long& value, FILE* out) { - return (fprintf(out, "%lu", value) > 0); -} - -/** - * Prints the unsigned long long `value` to the stream given by the - * [FILE](https://en.cppreference.com/w/c/io) pointer `out`. - */ -inline bool print(const unsigned long long& value, FILE* out) { - return (fprintf(out, "%llu", value) > 0); -} - -/** - * Prints the float `value` to the stream given by the - * [FILE](https://en.cppreference.com/w/c/io) pointer `out`. - */ -inline bool print(const float& value, FILE* out) { - return (fprintf(out, "%f", (double) value) > 0); -} - -/** - * Prints the double `value` to the stream given by the - * [FILE](https://en.cppreference.com/w/c/io) pointer `out`. - */ -inline bool print(const double& value, FILE* out) { - return (fprintf(out, "%lf", value) > 0); -} - -/** - * Prints the float `value` to the stream given by the - * [FILE](https://en.cppreference.com/w/c/io) pointer `out` - * with the given `precision`. - */ -inline bool print(const float& value, FILE* out, unsigned int precision) { - return (fprintf(out, "%.*f", precision, (double) value) > 0); -} - -/** - * Prints the double `value` to the stream given by the - * [FILE](https://en.cppreference.com/w/c/io) pointer `out` - * with the given `precision`. - */ -inline bool print(const double& value, FILE* out, unsigned int precision) { - return (fprintf(out, "%.*lf", precision, value) > 0); -} - -/** - * Prints the null-terminated C string `value` to the stream given by the - * [FILE](https://en.cppreference.com/w/c/io) pointer `out`. - */ -inline bool print(const char* values, FILE* out) { - return (fprintf(out, "%s", values) >= 0); -} - -namespace detail { - template static auto test_readable(int32_t) -> - decltype(bool(read(std::declval(), std::declval())), std::true_type{}); - template static auto test_readable(int64_t) -> std::false_type; - - template static auto test_writeable(int32_t) -> - decltype(bool(write(std::declval(), std::declval())), std::true_type{}); - template static auto test_writeable(int64_t) -> std::false_type; - - template static auto test_printable(int32_t) -> - decltype(bool(print(std::declval(), std::declval())), std::true_type{}); - template static auto test_printable(int64_t) -> std::false_type; -} - -/** - * This type trait is [true_type](https://en.cppreference.com/w/cpp/types/integral_constant) - * if and only if the function `bool read(integral&, T&)` is defined where - * `integral` is any integral type. - */ -template struct is_readable : and_type< - decltype(core::detail::test_readable(0)), - decltype(core::detail::test_readable(0)), - decltype(core::detail::test_readable(0)), - decltype(core::detail::test_readable(0)), - decltype(core::detail::test_readable(0)), - decltype(core::detail::test_readable(0)), - decltype(core::detail::test_readable(0)), - decltype(core::detail::test_readable(0)), - decltype(core::detail::test_readable(0)), - decltype(core::detail::test_readable(0)), - decltype(core::detail::test_readable(0)), - decltype(core::detail::test_readable(0)), - decltype(core::detail::test_readable(0)), - decltype(core::detail::test_readable(0))>::type {}; - -/** - * This type trait is [true_type](https://en.cppreference.com/w/cpp/types/integral_constant) - * if and only if the function `bool fwrite(const integral&, T&)` is defined where - * `integral` is any integral type. - */ -template struct is_writeable : and_type< - decltype(core::detail::test_writeable(0)), - decltype(core::detail::test_writeable(0)), - decltype(core::detail::test_writeable(0)), - decltype(core::detail::test_writeable(0)), - decltype(core::detail::test_writeable(0)), - decltype(core::detail::test_writeable(0)), - decltype(core::detail::test_writeable(0)), - decltype(core::detail::test_writeable(0)), - decltype(core::detail::test_writeable(0)), - decltype(core::detail::test_writeable(0)), - decltype(core::detail::test_writeable(0)), - decltype(core::detail::test_writeable(0)), - decltype(core::detail::test_writeable(0)), - decltype(core::detail::test_writeable(0))>::type {}; - -/** - * This type trait is [true_type](https://en.cppreference.com/w/cpp/types/integral_constant) - * if and only if the function `bool print(value, T&)` is defined. - */ -template struct is_printable : and_type< - decltype(core::detail::test_printable(0)), - decltype(core::detail::test_printable(0)), - decltype(core::detail::test_printable(0)), - decltype(core::detail::test_printable(0)), - decltype(core::detail::test_printable(0)), - decltype(core::detail::test_printable(0)), - decltype(core::detail::test_printable(0))>::type {}; - -/** - * Represents a stream to read/write from an in-memory buffer. - */ -struct memory_stream { - /** - * The size of the stream. - */ - unsigned int length; - - /** - * The current position of the stream in the buffer. - */ - unsigned int position; - - /** - * The underlying buffer. - */ - char* buffer; - - /** - * The default constructor does not initialize any fields. - */ - memory_stream() { } - - /** - * Initializes the stream with memory_stream::length given by - * `initial_capacity` and memory_stream::position set to `0`. - * memory_stream::buffer is allocated but not initialized to any value. - */ - memory_stream(unsigned int initial_capacity) : length(initial_capacity), position(0) { - buffer = (char*) malloc(sizeof(char) * length); - if (buffer == NULL) { - fprintf(stderr, "memory_stream ERROR: Unable to initialize buffer.\n"); - exit(EXIT_FAILURE); - } - } - - /** - * Initializes the stream with the memory_stream::buffer given by `buf`, - * memory_stream::length given by `length`, and memory_stream::position set - * to `0`. - */ - memory_stream(const char* buf, unsigned int length) : length(length), position(0) { - buffer = (char*) malloc(sizeof(char) * length); - if (buffer == NULL) { - fprintf(stderr, "memory_stream ERROR: Unable to initialize buffer.\n"); - exit(EXIT_FAILURE); - } - memcpy(buffer, buf, sizeof(char) * length); - } - - ~memory_stream() { - free(buffer); - } - - /** - * Reads a number of bytes given by `bytes` from the memory_stream and - * writes them to `dst`. This function assumes `dst` has sufficient capacity. - */ - inline bool read(void* dst, unsigned int bytes) { - if (position + bytes >= length) - return false; - memcpy(dst, buffer + position, bytes); - position += bytes; - return true; - } - - /** - * Checks whether the stream has sufficient size for an additional number - * of bytes given by `bytes` at its current memory_stream::position. If - * not, this function attempts to expand the buffer to a new size computed - * as `memory_stream::position + bytes`. - */ - inline bool ensure_capacity(unsigned int bytes) { - if (position + bytes <= length) - return true; - - unsigned int new_length = length; - if (!core::expand(buffer, new_length, position + bytes)) - return false; - length = new_length; - return true; - } - - /** - * Writes a number of bytes given by `bytes` from the given native array - * `src` to the current position in this stream. memory_stream::ensure_capacity - * is called to ensure the underlying buffer has sufficient size. - */ - inline bool write(const void* src, unsigned int bytes) { - if (!ensure_capacity(bytes)) - return false; - memcpy(buffer + position, src, bytes); - position += bytes; - return true; - } -}; - -/** - * Reads `sizeof(T)` bytes from `in` and writes them to the memory referenced - * by `value`. This function does not perform endianness transformations. - * \param in a memory_stream. - * \tparam T satisfies [is_fundamental](https://en.cppreference.com/w/cpp/types/is_fundamental). - */ -template::value>::type* = nullptr> -inline bool read(T& value, memory_stream& in) { - return in.read(&value, sizeof(T)); -} - -/** - * Reads `length` elements from `in` and writes them to the native array - * `values`. This function does not perform endianness transformations. - * \param in a memory_stream. - * \tparam T satisfies [is_fundamental](https://en.cppreference.com/w/cpp/types/is_fundamental). - */ -template::value>::type* = nullptr> -inline bool read(T* values, memory_stream& in, unsigned int length) { - return in.read(values, (unsigned int) sizeof(T) * length); -} - -/** - * Writes `sizeof(T)` bytes to `out` from the memory referenced by `value`. - * This function does not perform endianness transformations. - * \param out a memory_stream. - * \tparam T satisfies [is_fundamental](https://en.cppreference.com/w/cpp/types/is_fundamental). - */ -template::value>::type* = nullptr> -inline bool write(const T& value, memory_stream& out) { - return out.write(&value, sizeof(T)); -} - -/** - * Writes `length` elements to `out` from the native array `values`. This - * function does not perform endianness transformations. - * \param out a memory_stream. - * \tparam T satisfies [is_fundamental](https://en.cppreference.com/w/cpp/types/is_fundamental). - */ -template::value>::type* = nullptr> -inline bool write(const T* values, memory_stream& out, unsigned int length) { - return out.write(values, (unsigned int) sizeof(T) * length); -} - -/** - * Reads an array of `n` elements, each with a size of `size` bytes, from the - * memory_stream `in`, to the memory address referenced by `dst`. - * \see This function mirrors the equivalent [fread](http://en.cppreference.com/w/cpp/io/c/fread) - * for [FILE](https://en.cppreference.com/w/c/io) pointer streams. - * \returns the number of elements read. - */ -inline size_t fread(void* dst, size_t size, size_t count, memory_stream& in) { - size_t num_bytes = size * count; - if (in.position + num_bytes > in.length) { - count = (in.length - in.position) / size; - num_bytes = size * count; - } - - memcpy(dst, in.buffer + in.position, num_bytes); - in.position += num_bytes; - return count; -} - -/** - * Writes the array of `n` elements, each with a size of `size` bytes, from the - * memory address referenced by `src` to the memory_stream `out`. - * \see This function mirrors the equivalent [fwrite](http://en.cppreference.com/w/cpp/io/c/fwrite) - * for [FILE](https://en.cppreference.com/w/c/io) pointer streams. - * \returns either `n` if the write is successful, or `0` upon failure. - */ -inline size_t fwrite(const void* src, size_t size, size_t n, memory_stream& out) { - if (out.write(src, (unsigned int) (size * n))) - return n; - else return 0; -} - -/** - * Retrieves the current position in the given memory_stream. - * \see This function mirrors the equivalent [fgetpos](https://en.cppreference.com/w/c/io/fgetpos) - * for [FILE](https://en.cppreference.com/w/c/io) pointer streams. - * \returns 0 on success; nonzero value otherwise. - */ -inline int fgetpos(const memory_stream& stream, fpos_t* pos) { -#if defined(_WIN32) || defined(__APPLE__) /* on Windows or Mac */ - *pos = (fpos_t) stream.position; -#else /* on Windows or Linux */ - pos->__pos = stream.position; -#endif - return 0; -} - -/** - * Sets the current position in the given memory_stream. - * \see This function mirrors the equivalent [fsetpos](https://en.cppreference.com/w/c/io/fsetpos) - * for [FILE](https://en.cppreference.com/w/c/io) pointer streams. - * \returns 0 on success; nonzero value otherwise. - */ -inline int fsetpos(memory_stream& stream, const fpos_t* pos) { -#if defined(_WIN32) || defined(__APPLE__) /* on Windows or Mac */ - stream.position = (unsigned int) *pos; -#else /* on Windows or Linux */ - stream.position = pos->__pos; -#endif - return 0; -} - -/** - * Writes the given character `c` to the memory_stream `out`. - * \see This function mirrors the equivalent [fputc](http://en.cppreference.com/w/cpp/io/c/fputc) - * for [FILE](https://en.cppreference.com/w/c/io) pointer streams. - */ -inline int fputc(int c, memory_stream& out) { - char ch = (char) c; - if (out.write(&ch, sizeof(char))) - return c; - else return EOF; -} - -/** - * Writes the given null-terminated C string `s` to the memory_stream `out`. - * \see This function mirrors the equivalent [fputs](http://en.cppreference.com/w/cpp/io/c/fputs) - * for [FILE](https://en.cppreference.com/w/c/io) pointer streams. - */ -inline int fputs(const char* s, memory_stream& out) { - if (out.write(s, (unsigned int) strlen(s))) - return 1; - else return EOF; -} - -/** - * Writes the given arguments according to the format string `format` to the - * memory_stream `out`. - * \see This function mirrors the equivalent [fprintf](http://en.cppreference.com/w/cpp/io/c/fprintf) - * for [FILE](https://en.cppreference.com/w/c/io) pointer streams. - * \returns the number of bytes written to the stream, or `-1` upon error. - */ -inline int fprintf(memory_stream& out, const char* format, ...) { - va_list argptr; - va_start(argptr, format); - -#if defined(_WIN32) - int required = _vscprintf(format, argptr); - if (!out.ensure_capacity((unsigned int) required + 1)) { - fprintf(stderr, "fprintf ERROR: Unable to expand memory stream buffer.\n"); - va_end(argptr); - return -1; - } - - int written = vsnprintf_s(out.buffer + out.position, out.length - out.position, (size_t) required, format, argptr); -#else - int written = vsnprintf(out.buffer + out.position, out.length - out.position, format, argptr); - if (written < 0) { - va_end(argptr); - return -1; - } else if ((unsigned) written < out.length - out.position) { - va_end(argptr); - out.position += written; - return written; - } - - if (!out.ensure_capacity(written + 1)) { - fprintf(stderr, "fprintf ERROR: Unable to expand memory stream buffer.\n"); - va_end(argptr); - return -1; - } - written = vsnprintf(out.buffer + out.position, out.length - out.position, format, argptr); -#endif - - va_end(argptr); - if (written < 0) return -1; - out.position += written; - return written; -} - -/** - * Prints the character `value` to the stream given by the memory_stream `out`. - */ -inline bool print(const char& value, memory_stream& out) { - return (fputc(value, out) != EOF); -} - -/** - * Prints the int `value` to the stream given by the memory_stream `out`. - */ -inline bool print(const int& value, memory_stream& out) { - return (fprintf(out, "%d", value) > 0); -} - -/** - * Prints the long `value` to the stream given by the memory_stream `out`. - */ -inline bool print(const long& value, memory_stream& out) { - return (fprintf(out, "%ld", value) > 0); -} - -/** - * Prints the long long `value` to the stream given by the memory_stream `out`. - */ -inline bool print(const long long& value, memory_stream& out) { - return (fprintf(out, "%lld", value) > 0); -} - -/** - * Prints the unsigned int `value` to the stream given by the memory_stream - * `out`. - */ -inline bool print(const unsigned int& value, memory_stream& out) { - return (fprintf(out, "%u", value) > 0); -} - -/** - * Prints the unsigned long `value` to the stream given by the memory_stream - * `out`. - */ -inline bool print(const unsigned long& value, memory_stream& out) { - return (fprintf(out, "%lu", value) > 0); -} - -/** - * Prints the unsigned long long `value` to the stream given by the - * memory_stream `out`. - */ -inline bool print(const unsigned long long& value, memory_stream& out) { - return (fprintf(out, "%llu", value) > 0); -} - -/** - * Prints the float `value` to the stream given by the memory_stream `out`. - */ -inline bool print(const float& value, memory_stream& out) { - return (fprintf(out, "%f", (double) value) > 0); -} - -/** - * Prints the double `value` to the stream given by the memory_stream `out`. - */ -inline bool print(const double& value, memory_stream& out) { - return (fprintf(out, "%lf", value) > 0); -} - -/** - * Prints the float `value` to the stream given by the memory_stream `out` with - * the given `precision`. - */ -inline bool print(const float& value, memory_stream& out, unsigned int precision) { - return (fprintf(out, "%.*f", precision, (double) value) > 0); -} - -/** - * Prints the double `value` to the stream given by the memory_stream `out` - * with the given `precision`. - */ -inline bool print(const double& value, memory_stream& out, unsigned int precision) { - return (fprintf(out, "%.*lf", precision, value) > 0); -} - -/** - * Prints the null-terminated C string `value` to the stream given by the - * memory_stream `out`. - */ -inline bool print(const char* values, memory_stream& out) { - return (fprintf(out, "%s", values) >= 0); -} - -/** - * A stream wrapper for reading UTF-32 characters from an underlying multibyte - * stream (such as UTF-8). - */ -template -struct buffered_stream { - Stream& underlying_stream; - char buffer[BufferSize]; - unsigned int position; - unsigned int length; - mbstate_t shift; - - buffered_stream(Stream& underlying_stream) : underlying_stream(underlying_stream), position(0) { - shift = {0}; - fill_buffer(); - } - - inline void fill_buffer() { - length = fread(buffer, sizeof(char), BufferSize, underlying_stream); - } - - /** - * Returns the next UTF-32 character (as a `char32_t`) from the stream. If - * there are no further bytes in the underlying stream or an error occurred, - * `static_cast(-1)` is returned. - */ - char32_t fgetc32() { - static_assert(BufferSize >= MB_LEN_MAX, "BufferSize must be at least MB_LEN_MAX"); - - while (true) - { - if (length == 0) - return static_cast(-1); - - char32_t c; - size_t status = mbrtoc32(&c, buffer + position, sizeof(char) * (length - position), &shift); - if (status == static_cast(-1)) { - /* encoding error occurred */ - return static_cast(-1); - } else if (status == static_cast(-2)) { - /* the character is valid but incomplete */ - position = 0; - fill_buffer(); - continue; - } else { - if (status == 0) - position += 1; - else position += status; - if (position == length) { - position = 0; - fill_buffer(); - } - return c; - } - } - } -}; - -template -struct buffered_stream { - memory_stream& underlying_stream; - mbstate_t shift; - - buffered_stream(memory_stream& underlying_stream) : underlying_stream(underlying_stream) { - shift = {0}; - } - - /** - * Returns the next UTF-32 character (as a `char32_t`) from the stream. If - * there are no further bytes in the underlying stream or an error occurred, - * `static_cast(-1)` is returned. - */ - char32_t fgetc32() { - if (underlying_stream.position == underlying_stream.length) - return static_cast(-1); - - char32_t c; - size_t status = mbrtoc32(&c, underlying_stream.buffer + underlying_stream.position, sizeof(char) * (underlying_stream.length - underlying_stream.position), &shift); - if (status == static_cast(-1) || status == static_cast(-2)) { - /* encoding error occurred or the character is incomplete */ - return static_cast(-1); - } - - if (status == 0) - underlying_stream.position += 1; - else underlying_stream.position += status; - return c; - } -}; - -/** - * Returns the next UTF-32 character (as a `char32_t`) from the buffered_stream - * `input`. If there are no further bytes in the underlying stream or an error - * occurred, `static_cast(-1)` is returned. - */ -template -inline char32_t fgetc32(buffered_stream& input) -{ - return input.fgetc32(); -} - -/** - * A stream wrapper for reading/writing integral types as fixed-width integral - * values. This is useful for cross-platform readability and writeability. - */ -template -struct fixed_width_stream { - Stream& stream; - - fixed_width_stream(Stream& stream) : stream(stream) { } - - template struct type { }; - template struct type { typedef BoolType value; }; - template struct type { typedef CharType value; }; - template struct type { typedef UCharType value; }; - template struct type { typedef ShortType value; }; - template struct type { typedef UShortType value; }; - template struct type { typedef IntType value; }; - template struct type { typedef UIntType value; }; - template struct type { typedef LongType value; }; - template struct type { typedef ULongType value; }; - template struct type { typedef LongLongType value; }; - template struct type { typedef ULongLongType value; }; - template struct type { typedef FloatType value; }; - template struct type { typedef DoubleType value; }; -}; - -/** - * Reads `size(K)` bytes from `in` where `K` is the appropriate template - * argument in the fixed_width_stream and writes them to the memory referenced - * by `value`. This function does not perform endianness transformations. - * \param in a fixed_width_stream. - * \tparam T satisfies [is_fundamental](https://en.cppreference.com/w/cpp/types/is_fundamental). - */ -template::value>::type* = nullptr> -inline bool read(T& value, fixed_width_stream& in) { - typedef typename fixed_width_stream::template type::value value_type; - value_type c; - if (!read(c, in.stream)) return false; - value = (T) c; - return true; -} - -/** - * Reads `length` elements from `in` and writes them to the native array - * `values`. This function does not perform endianness transformations. - * \param in a fixed_width_stream. - * \tparam T satisfies [is_fundamental](https://en.cppreference.com/w/cpp/types/is_fundamental). - */ -template::value>::type* = nullptr> -inline bool read(T* values, fixed_width_stream& in, SizeType length) { - for (SizeType i = 0; i < length; i++) - if (!read(values[i], in)) return false; - return true; -} - -/** - * Writes `sizeof(K)` bytes to `out` from the memory referenced by `value` - * where `K` is the appropriate template argument in the fixed_width_stream. - * This function does not perform endianness transformations. - * \param out a fixed_width_stream. - * \tparam T satisfies [is_fundamental](https://en.cppreference.com/w/cpp/types/is_fundamental). - */ -template::value>::type* = nullptr> -inline bool write(const T& value, fixed_width_stream& out) { - typedef typename fixed_width_stream::template type::value value_type; - return write((value_type) value, out.stream); -} - -/** - * Writes `length` elements to `out` from the native array `values`. This - * function does not perform endianness transformations. - * \param out a fixed_width_stream. - * \tparam T satisfies [is_fundamental](https://en.cppreference.com/w/cpp/types/is_fundamental). - */ -template::value>::type* = nullptr> -inline bool write(const T* values, fixed_width_stream& out, SizeType length) { - for (SizeType i = 0; i < length; i++) - if (!write(values[i], out)) return false; - return true; -} - -/** - * Writes the given null-terminated C string `values` to the stream `out`. - * \tparam Stream satisfies is_writeable. - */ -template::value>::type* = nullptr> -inline bool write(const char* values, Stream& out) { - return write(values, out, strlen(values)); -} - -/** - * The default scribe implementation that provides the default behavior for - * read/write/print functions. - * \see [Section on scribes](#scribes). - */ -struct default_scribe { }; - -/* a type trait for detecting whether the function 'print' is defined with a default_scribe argument */ -namespace detail { - template static auto test_default_read(int) -> - decltype(bool(read(std::declval(), std::declval(), std::declval())), std::true_type{}); - template static auto test_default_read(long) -> std::false_type; - template static auto test_default_write(int) -> - decltype(bool(write(std::declval(), std::declval(), std::declval())), std::true_type{}); - template static auto test_default_write(long) -> std::false_type; - template static auto test_default_print(int) -> - decltype(bool(print(std::declval(), std::declval(), std::declval())), std::true_type{}); - template static auto test_default_print(long) -> std::false_type; -} - -template struct has_default_read : decltype(core::detail::test_default_read(0)){}; -template struct has_default_write : decltype(core::detail::test_default_write(0)){}; -template struct has_default_print : decltype(core::detail::test_default_print(0)){}; - -/** - * Calls and returns `read(value, in)`, dropping the default_scribe argument. - * \tparam Stream satisfies is_readable. - */ -template::value && !has_default_read::value>::type** = nullptr> -inline bool read(T& value, Stream& in, default_scribe& scribe) { - return read(value, in); -} - -/** - * Calls and returns `write(value, out)`, dropping the default_scribe argument. - * \tparam Stream satisfies is_writeable. - */ -template::value && !has_default_write::value>::type* = nullptr> -inline bool write(const T& value, Stream& out, default_scribe& scribe) { - return write(value, out); -} - -/** - * Calls and returns `print(value, out)`, dropping the default_scribe argument. - * \tparam Stream satisfies is_printable. - */ -template::value && !has_default_print::value>::type* = nullptr> -inline auto print(const T& value, Stream& out, default_scribe& scribe) -> decltype(print(value, out)) { - return print(value, out); -} - -/** - * A scribe that prints pointers by dereferencing the pointer and calling the - * appropriate read/write/print function. - * \see [Section on scribes](#scribes). - */ -struct pointer_scribe { }; - -/** - * Allocates memory and stores the address in `value`, and then calls - * `read(*value, in, reader)`, dropping the pointer_scribe argument. Note that - * since `reader` is a variadic argument, it may be empty. - * \tparam Stream satisfies is_readable. - */ -template::value>::type* = nullptr> -inline bool read(T*& value, Stream& in, const pointer_scribe& scribe, Reader&&... reader) { - value = (T*) malloc(sizeof(T)); - if (value == NULL) { - fprintf(stderr, "read ERROR: Out of memory.\n"); - return false; - } else if (!read(*value, in, std::forward(reader)...)) { - free(value); - return false; - } - return true; -} - -/** - * Calls and returns `write(*value, out, writer)`, dropping the pointer_scribe - * argument. Note that since `writer` is a variadic argument, it may be empty. - * \tparam Stream satisfies is_writeable. - */ -template::value>::type* = nullptr> -inline bool write(const T* const value, Stream& out, const pointer_scribe& scribe, Writer&&... writer) { - return write(*value, out, std::forward(writer)...); -} - -/** - * Calls and returns `print(*value, out, printer)`, dropping the pointer_scribe - * argument. Note that since `printer` is a variadic argument, it may be empty. - * \tparam Stream satisfies is_printable. - */ -template::value>::type* = nullptr> -inline bool print(const T* const value, Stream& out, const pointer_scribe& scribe, Printer&&... printer) { - return print(*value, out, std::forward(printer)...); -} - -/** - * The default left delimitter "[" for the array print functions. - */ -char default_left_bracket[] = "["; - -/** - * The default right delimitter "]" for the array print functions. - */ -char default_right_bracket[] = "]"; - -/** - * The default separator between elements ", " for the array print functions. - */ -char default_array_separator[] = ", "; - -/** - * Prints the given native array of `values` each of type `T`, where `length` - * is the number of elements in the array. The output stream is `out`. - * \param printer a scribe for which the function `bool print(const T&, Stream&, Printer&)` - * is defined. Note that since this is a variadic argument, it may be empty. - * \tparam Stream satisfies is_printable. - */ -template::value>::type* = nullptr> -bool print(const T* values, SizeType length, Stream& out, Printer&&... printer) { - if (!print(LeftBracket, out)) return false; - if (length == 0) - return print(RightBracket, out); - if (!print(values[0], out, std::forward(printer)...)) return false; - for (SizeType i = 1; i < length; i++) { - if (!print(Separator, out) || !print(values[i], out, std::forward(printer)...)) - return false; - } - return print(RightBracket, out); -} - -/** - * Prints the given native static array of `values` each of type `T`, where `N` - * is the number of elements in the array. The output stream is `out`. - * \param printer a scribe for which the function `bool print(const T&, Stream&, Printer&)` - * is defined. Note that since this is a variadic argument, it may be empty. - * \tparam Stream satisfies is_printable. - */ -template::value>::type* = nullptr> -bool print(const T (&values)[N], Stream& out, Printer&&... printer) { - if (!print(LeftBracket, out)) return false; - if (N == 0) - return print(RightBracket, out); - if (!print(values[0], out, std::forward(printer)...)) return false; - for (size_t i = 1; i < N; i++) { - if (!print(Separator, out) || !print(values[i], out, std::forward(printer)...)) - return false; - } - return print(RightBracket, out); -} - -/** - * Reads an array of `length` elements from `in` and stores the result in the - * given native array `a`. This function assumes `a` has sufficient capacity. - * \param reader a scribe for which the function `bool read(T&, Stream&, Reader&&...)` - * is defined. Note that since this is a variadic argument, it may be empty. - * \tparam Stream satisfies is_readable. - */ -template::value>::type* = nullptr> -inline bool read(T* a, Stream& in, SizeType length, Reader&&... reader) { - for (SizeType i = 0; i < length; i++) - if (!read(a[i], in, std::forward(reader)...)) return false; - return true; -} - -/** - * Reads an array of `N` elements from `in` and stores the result in the given - * native array `a`. - * \param reader a scribe for which the function `bool read(T&, Stream&, Reader&&...)` - * is defined. Note that since this is a variadic argument, it may be empty. - * \tparam Stream satisfies is_readable. - */ -template::value>::type* = nullptr> -inline bool read(T (&a)[N], Stream& in, Reader&&... reader) { - for (size_t i = 0; i < N; i++) - if (!read(a[i], in, std::forward(reader)...)) return false; - return true; -} - -/** - * Reads a core::array structure from `in` and stores the result in `a`. - * \param a an uninitialized core::array structure. This function initializes - * `a`, and the caller is responsible for its memory and must call free - * to release its memory resources. - * \param reader a scribe for which the function `bool read(T&, Stream&, Reader&&...)` - * is defined. Note that since this is a variadic argument, it may be empty. - * \tparam Stream satisfies is_readable. - */ -template::value>::type* = nullptr> -bool read(array& a, Stream& in, Reader&&... reader) { - size_t length; - if (!read(length, in)) - return false; - size_t capacity = ((size_t) 1) << (core::log2(length == 0 ? 1 : length) + 1); - a.data = (T*) malloc(sizeof(T) * capacity); - if (a.data == NULL) return false; - if (!read(a.data, in, length, std::forward(reader)...)) { - free(a.data); - return false; - } - a.length = length; - a.capacity = capacity; - return true; -} - -/** - * Writes the given native array `a` of elements to `out`, each of type `T`, - * where the number of elements is given by `length`. - * \param writer a scribe for which the function `bool write(const T&, Stream&, Writer&&...)` - * is defined. Note that since this is a variadic argument, it may be empty. - * \tparam Stream satisfies is_writeable. - */ -template::value>::type* = nullptr> -inline bool write(const T* a, Stream& out, SizeType length, Writer&&... writer) { - for (SizeType i = 0; i < length; i++) - if (!write(a[i], out, std::forward(writer)...)) return false; - return true; -} - -/** - * Writes the given native array `a` of elements to `out`, each of type `T`, - * where the number of elements is given by `N`. - * \param writer a scribe for which the function `bool write(const T&, Stream&, Writer&&...)` - * is defined. Note that since this is a variadic argument, it may be empty. - * \tparam Stream satisfies is_writeable. - */ -template::value>::type* = nullptr> -inline bool write(const T (&a)[N], Stream& out, Writer&&... writer) { - for (size_t i = 0; i < N; i++) - if (!write(a[i], out, std::forward(writer)...)) return false; - return true; -} - -/** - * Writes the given core::array structure `a` of elements to `out`, each of type `T`. - * \param writer a scribe for which the function `bool write(const T&, Stream&, Writer&&...)` - * is defined. Note that since this is a variadic argument, it may be empty. - * \tparam Stream satisfies is_writeable. - */ -template::value>::type* = nullptr> -bool write(const array& a, Stream& out, Writer&&... writer) { - return write(a.length, out) - && write(a.data, out, a.length, std::forward(writer)...); -} - -/** - * Prints the given core::array structure `a` of elements to `out`, each of type `T`. - * \param printer a scribe for which the function `bool print(const T&, Stream&, Printer&&...)` - * is defined. Note that since this is a variadic argument, it may be empty. - * \tparam Stream satisfies is_printable. - */ -template::value>::type* = nullptr> -inline bool print(const array& a, Stream&& out, Printer&&... printer) { - return print(a.data, a.length, out, std::forward(printer)...); -} - -/** - * Reads a core::hash_set structure `set` from `in`. - * \param set an uninitialized core::hash_set structure. This function - * initializes `set`, and the caller is responsible for its memory and - * must call free to release its memory resources. - * \param alloc_keys a memory allocation function with prototype - * `void* alloc_keys(size_t count, size_t size)` that allocates space for - * `count` items, each with size `size`, and initializes them such that - * core::is_empty() returns `true` for each element. - * \param reader a scribe for which the function `bool read(T&, Stream&, Reader&&...)` - * is defined. Note that since this is a variadic argument, it may be empty. - * \tparam Stream satisfies is_readable. - */ -template::value>::type* = nullptr> -bool read(hash_set& set, Stream& in, alloc_keys_func alloc_keys, Reader&&... reader) { - unsigned int length; - if (!read(length, in)) return false; - - set.size = 0; - set.capacity = 1 << (core::log2(RESIZE_THRESHOLD_INVERSE * (length == 0 ? 1 : length)) + 1); - set.keys = (T*) alloc_keys(set.capacity, sizeof(T)); - if (set.keys == NULL) return false; - - for (unsigned int i = 0; i < length; i++) { - T& key = *((T*) alloca(sizeof(T))); - if (!read(key, in, std::forward(reader)...)) return false; - move(key, set.keys[set.index_to_insert(key)]); - set.size++; - } - return true; -} - -/** - * Reads a core::hash_set structure `set` from `in`. The keys in the hash_set - * are allocated using [calloc](http://en.cppreference.com/w/c/memory/calloc). - * \param set an uninitialized core::hash_set structure. This function - * initializes `set`, and the caller is responsible for its memory and - * must call free to release its memory resources. - * \param reader a scribe for which the function `bool read(T&, Stream&, Reader&&...)` - * is defined. Note that since this is a variadic argument, it may be empty. - * \tparam Stream satisfies is_readable. - */ -template::value>::type* = nullptr> -inline bool read(hash_set& set, Stream& in, Reader&&... reader) { - return read(set, in, calloc, std::forward(reader)...); -} - -/** - * Writes the given core::hash_set structure `set` to `out`. - * \param writer a scribe for which the function `bool write(const T&, Stream&, Writer&&...)` - * is defined. Note that since this is a variadic argument, it may be empty. - * \tparam Stream satisfies is_writeable. - */ -template::value>::type* = nullptr> -bool write(const hash_set& set, Stream& out, Writer&&... writer) { - if (!write(set.size, out)) return false; - for (unsigned int i = 0; i < set.capacity; i++) { - if (is_empty(set.keys[i])) continue; - if (!write(set.keys[i], out, std::forward(writer)...)) return false; - } - return true; -} - -/** - * Reads a core::hash_map structure `map` from `in`. - * \param map an uninitialized core::hash_map structure. This function - * initializes `map`, and the caller is responsible for its memory and - * must call free to release its memory resources. - * \param alloc_keys a memory allocation function with prototype - * `void* alloc_keys(size_t count, size_t size)` that allocates space for - * `count` items, each with size `size`, and initializes them such that - * core::is_empty() returns `true` for each element of type `K`. - * \param key_reader a scribe for which the function `bool read(K&, Stream&, KeyReader&)` is defined. - * \param value_reader a scribe for which the function `bool read(V&, Stream&, ValueReader&)` is defined. - * \tparam Stream satisfies is_readable. - */ -template::value>::type* = nullptr> -bool read(hash_map& map, - Stream& in, alloc_keys_func alloc_keys, - KeyReader& key_reader, ValueReader& value_reader) -{ - unsigned int length; - if (!read(length, in)) return false; - - map.table.size = 0; - map.table.capacity = 1 << (core::log2(RESIZE_THRESHOLD_INVERSE * (length == 0 ? 1 : length)) + 1); - map.table.keys = (K*) alloc_keys(map.table.capacity, sizeof(K)); - if (map.table.keys == NULL) return false; - map.values = (V*) malloc(sizeof(V) * map.table.capacity); - if (map.values == NULL) { - free(map.table.keys); - return false; - } - - for (unsigned int i = 0; i < length; i++) { - K& key = *((K*) alloca(sizeof(K))); - if (!read(key, in, key_reader)) return false; - - bool contains; - unsigned int bucket; - map.get(key, contains, bucket); - if (!read(map.values[bucket], in, value_reader)) - return false; - if (!contains) { - move(key, map.table.keys[bucket]); - map.table.size++; - } else { - fprintf(stderr, "read WARNING: Serialized hash_map contains duplicates.\n"); - } - } - return true; -} - -/** - * Reads a core::hash_map structure `map` from `in`. - * \param map an uninitialized core::hash_map structure. This function - * initializes `map`, and the caller is responsible for its memory and - * must call free to release its memory resources. - * \param alloc_keys a memory allocation function with prototype - * `void* alloc_keys(size_t count, size_t size)` that allocates space for - * `count` items, each with size `size`, and initializes them such that - * core::is_empty() returns `true` for each element of type `K`. - * \param key_reader a scribe for which the function `bool read(K&, Stream&, KeyReader&)` is defined. - * \tparam Stream satisfies is_readable. - */ -template::value>::type* = nullptr> -inline bool read(hash_map& map, Stream& in, - KeyReader& key_reader, - alloc_keys_func alloc_keys = calloc) -{ - default_scribe scribe; - return read(map, in, alloc_keys, key_reader, scribe); -} - -/** - * Reads a core::hash_map structure `map` from `in`. - * \param map an uninitialized core::hash_map structure. This function - * initializes `map`, and the caller is responsible for its memory and - * must call free to release its memory resources. - * \param alloc_keys a memory allocation function with prototype - * `void* alloc_keys(size_t count, size_t size)` that allocates space for - * `count` items, each with size `size`, and initializes them such that - * core::is_empty() returns `true` for each element of type `K`. - * \tparam Stream satisfies is_readable. - */ -template::value>::type* = nullptr> -inline bool read(hash_map& map, Stream& in, - alloc_keys_func alloc_keys = calloc) -{ - default_scribe scribe; - return read(map, in, alloc_keys, scribe, scribe); -} - -/** - * Writes the core::hash_map structure `map` to `out`. - * \param key_writer a scribe for which the function `bool write(const K&, Stream&, KeyWriter&)` is defined. - * \param value_writer a scribe for which the function `bool write(const V&, Stream&, ValueWriter&)` is defined. - * \tparam Stream satisfies is_writeable. - */ -template::value>::type* = nullptr> -bool write(const hash_map& map, Stream& out, - KeyWriter& key_writer, ValueWriter& value_writer) -{ - if (!write(map.table.size, out)) return false; - for (unsigned int i = 0; i < map.table.capacity; i++) { - if (is_empty(map.table.keys[i])) continue; - if (!write(map.table.keys[i], out, key_writer) - || !write(map.values[i], out, value_writer)) - return false; - } - return true; -} - -/** - * Writes the core::hash_map structure `map` to `out`. - * \param key_writer a scribe for which the function `bool write(const K&, Stream&, KeyWriter&)` is defined. - * \tparam Stream satisfies is_writeable. - */ -template::value>::type* = nullptr> -inline bool write(const hash_map& map, Stream& out, KeyWriter& key_writer) { - default_scribe scribe; - return write(map, out, key_writer, scribe); -} - -/** - * Writes the core::hash_map structure `map` to `out`. - * \tparam Stream satisfies is_writeable. - */ -template::value>::type* = nullptr> -inline bool write(const hash_map& map, Stream& out) { - default_scribe scribe; - return write(map, out, scribe, scribe); -} - -/** - * Reads a core::array_map structure `map` from `in`. - * \param map an uninitialized core::array_map structure. This function - * initializes `map`, and the caller is responsible for its memory and - * must call free to release its memory resources. - * \param key_reader a scribe for which the function `bool read(K&, Stream&, KeyReader&)` is defined. - * \param value_reader a scribe for which the function `bool read(V&, Stream&, ValueReader&)` is defined. - * \tparam Stream satisfies is_readable. - */ -template::value>::type* = nullptr> -bool read(array_map& map, Stream& in, - KeyReader& key_reader, ValueReader& value_reader) -{ - size_t length; - if (!read(length, in)) return false; - - map.size = 0; - map.capacity = 1 << (core::log2(length == 0 ? 1 : length) + 1); - map.keys = (K*) malloc(sizeof(K) * map.capacity); - if (map.keys == NULL) return false; - map.values = (V*) malloc(sizeof(V) * map.capacity); - if (map.values == NULL) { - free(map.keys); - return false; - } - - for (unsigned int i = 0; i < length; i++) { - if (!read(map.keys[i], in, key_reader)) return false; - if (!read(map.values[i], in, value_reader)) return false; - map.size++; - } - return true; -} - -/** - * Reads a core::array_map structure `map` from `in`. - * \param map an uninitialized core::array_map structure. This function - * initializes `map`, and the caller is responsible for its memory and - * must call free to release its memory resources. - * \param key_reader a scribe for which the function `bool read(K&, Stream&, KeyReader&)` is defined. - * \tparam Stream satisfies is_readable. - */ -template::value>::type* = nullptr> -inline bool read(array_map& map, Stream& in, KeyReader& key_reader) { - default_scribe scribe; - return read(map, in, key_reader, scribe); -} - -/** - * Reads a core::array_map structure `map` from `in`. - * \param map an uninitialized core::array_map structure. This function - * initializes `map`, and the caller is responsible for its memory and - * must call free to release its memory resources. - * \tparam Stream satisfies is_readable. - */ -template::value>::type* = nullptr> -inline bool read(array_map& map, Stream& in) { - default_scribe scribe; - return read(map, in, scribe, scribe); -} - -/** - * Writes the given core::array_map structure `map` to `out`. - * \param key_writer a scribe for which the function `bool write(const K&, Stream&, KeyWriter&)` is defined. - * \param value_writer a scribe for which the function `bool write(const V&, Stream&, ValueWriter&)` is defined. - * \tparam Stream satisfies is_writeable. - */ -template::value>::type* = nullptr> -bool write(const array_map& map, Stream& out, - KeyWriter& key_writer, ValueWriter& value_writer) -{ - if (!write(map.size, out)) return false; - for (unsigned int i = 0; i < map.size; i++) { - if (!write(map.keys[i], out, key_writer)) return false; - if (!write(map.values[i], out, value_writer)) return false; - } - return true; -} - -/** - * Writes the given core::array_map structure `map` to `out`. - * \param key_writer a scribe for which the function `bool write(const K&, Stream&, KeyWriter&)` is defined. - * \tparam Stream satisfies is_writeable. - */ -template::value>::type* = nullptr> -inline bool write(const array_map& map, Stream& out, KeyWriter& key_writer) { - default_scribe scribe; - return write(map, out, key_writer, scribe); -} - -/** - * Writes the given core::array_map structure `map` to `out`. - * \tparam Stream satisfies is_writeable. - */ -template::value>::type* = nullptr> -inline bool write(const array_map& map, Stream& out) { - default_scribe scribe; - return write(map, out, scribe, scribe); -} - -/** - * Reads a core::pair structure `p` from `stream` by calling `read(p.key, stream)` and `read(p.value, stream)`. - * \tparam Stream satisfies is_readable. - */ -template::value>::type* = nullptr> -inline bool read(pair& p, Stream& stream) -{ - return read(p.key, stream) && read(p.value, stream); -} - -/** - * Writes the given core::pair structure `p` to `stream` by calling - * `write(p.key, stream)` and `write(p.value, stream)`. - * \tparam Stream satisfies is_writeable. - */ -template::value>::type* = nullptr> -inline bool write(const pair& p, Stream& stream) -{ - return write(p.key, stream) && write(p.value, stream); -} - -} /* namespace core */ - -#endif /* IO_H_ */ diff --git a/reasoners/data_preprocess/logic/graph_dataset_gen/core/lex.h b/reasoners/data_preprocess/logic/graph_dataset_gen/core/lex.h deleted file mode 100644 index 1d38055b1..000000000 --- a/reasoners/data_preprocess/logic/graph_dataset_gen/core/lex.h +++ /dev/null @@ -1,504 +0,0 @@ -/** - * \file lex.h - * - * This file implements common functionality for lexical analysis, such as - * string comparison, tokenization (splitting strings by whitespace), and - * parsing arithmetic types from strings. - * - * - */ - -#ifndef LEX_H_ -#define LEX_H_ - -#include "map.h" -#include "utility.h" - -#include - -#if defined(__APPLE__) -#include - -inline size_t c32rtomb(char* s, char32_t wc, mbstate_t* ps) { - return wcrtomb(s, (wchar_t) wc, ps); -} - -#else /* __APPLE__ */ -#include -#endif /* __APPLE__ */ - - -namespace core { - -/** - * Compares the strings given by the array structure `first` and the - * null-terminated C string `second`. - * \returns `true` if the strings are equivalent, and `false` otherwise. - */ -/* TODO: test with benchmarks whether we should inline these functions */ -inline bool compare_strings(const array& first, const char* second) { - for (unsigned int i = 0; i < first.length; i++) { - if (first[i] != second[i]) - return false; - } - if (second[first.length] != '\0') - return false; - return true; -} - -/** - * Compares the strings given by the core::string `first` and the native char - * array `second` whose length is given by `second_length`. - * \returns `true` if the strings are equivalent, and `false` otherwise. - */ -inline bool compare_strings(const string& first, const char* second, unsigned int second_length) { - if (first.length != second_length) - return false; - for (unsigned int i = 0; i < first.length; i++) { - if (first[i] != second[i]) - return false; - } - return true; -} - -/** - * Tokenizes the given native char array `str` with length `length`, assigning - * to each unique token an `unsigned int` identifier. These identifiers are - * stored in the core::hash_map `names`. The tokenized identifiers are added to - * the core::array `tokens`. - */ -bool tokenize(const char* str, unsigned int length, - array& tokens, hash_map& names) -{ - bool whitespace = true; - unsigned int token_start = 0; - for (unsigned int i = 0; i < length; i++) { - if (whitespace) { - if (!isspace(str[i])) { - token_start = i; - whitespace = false; - } - } else { - if (isspace(str[i])) { - unsigned int id; - if (!get_token(string(str + token_start, i - token_start), id, names) - || !tokens.add(id)) - return false; - whitespace = true; - } - } - } - - if (!whitespace) { - unsigned int id; - if (!get_token(string(str + token_start, length - token_start), id, names) - || !tokens.add(id)) - return false; - } - return true; -} - -/** - * Attempts to parse the string given by `token` as a `double`. - * \tparam CharArray a string type that implements two fields: (1) `data` which - * returns the underlying `char*` array, and (2) `length` which returns - * the length of the string. - * \returns `true` if successful, or `false` if there is insufficient memory or - * `token` is not an appropriate string representation of a floating-point - * number. - */ -template -inline bool parse_float(const CharArray& token, double& value) { - char* buffer = (char*) malloc(sizeof(char) * (token.length + 1)); - if (buffer == NULL) { - fprintf(stderr, "parse_float ERROR: Unable to allocate temporary string buffer.\n"); - return false; - } - memcpy(buffer, token.data, sizeof(char) * token.length); - buffer[token.length] = '\0'; - - char* end_ptr; - value = strtod(buffer, &end_ptr); - if (*end_ptr != '\0') { - free(buffer); - return false; - } - free(buffer); - return true; -} - -/** - * Attempts to parse the string given by `token` as an `unsigned int`. - * \tparam CharArray a string type that implements two fields: (1) `data` which - * returns the underlying `char*` array, and (2) `length` which returns - * the length of the string. - * \returns `true` if successful, or `false` if there is insufficient memory or - * `token` is not an appropriate string representation of a unsigned - * integer. - */ -template -inline bool parse_uint(const CharArray& token, unsigned int& value, unsigned int base = 0) { - char* buffer = (char*) malloc(sizeof(char) * (token.length + 1)); - if (buffer == NULL) { - fprintf(stderr, "parse_uint ERROR: Unable to allocate temporary string buffer.\n"); - return false; - } - memcpy(buffer, token.data, sizeof(char) * token.length); - buffer[token.length] = '\0'; - - char* end_ptr; - value = strtoul(buffer, &end_ptr, base); - if (*end_ptr != '\0') { - free(buffer); - return false; - } - free(buffer); - return true; -} - -/** - * Attempts to parse the string given by `token` as an `unsigned int`. - * \tparam CharArray a string type that implements two fields: (1) `data` which - * returns the underlying `char*` array, and (2) `length` which returns - * the length of the string. - * \returns `true` if successful, or `false` if there is insufficient memory or - * `token` is not an appropriate string representation of a unsigned - * integer. - */ -template -inline bool parse_ulonglong(const CharArray& token, unsigned long long& value) { - char* buffer = (char*) malloc(sizeof(char) * (token.length + 1)); - if (buffer == NULL) { - fprintf(stderr, "parse_ulonglong ERROR: Unable to allocate temporary string buffer.\n"); - return false; - } - memcpy(buffer, token.data, sizeof(char) * token.length); - buffer[token.length] = '\0'; - - char* end_ptr; - value = strtoull(buffer, &end_ptr, 0); - if (*end_ptr != '\0') { - free(buffer); - return false; - } - free(buffer); - return true; -} - -/** - * Attempts to parse the string given by `token` as an `int`. - * \tparam CharArray a string type that implements two fields: (1) `data` which - * returns the underlying `char*` array, and (2) `length` which returns - * the length of the string. - * \returns `true` if successful, or `false` if there is insufficient memory or - * `token` is not an appropriate string representation of a integer. - */ -template -inline bool parse_int(const CharArray& token, int& value) { - char* buffer = (char*) malloc(sizeof(char) * (token.length + 1)); - if (buffer == NULL) { - fprintf(stderr, "parse_int ERROR: Unable to allocate temporary string buffer.\n"); - return false; - } - memcpy(buffer, token.data, sizeof(char) * token.length); - buffer[token.length] = '\0'; - - char* end_ptr; - value = strtol(buffer, &end_ptr, 0); - if (*end_ptr != '\0') { - free(buffer); - return false; - } - free(buffer); - return true; -} - -/** - * Attempts to parse the string given by `token` as a `long`. - * \tparam CharArray a string type that implements two fields: (1) `data` which - * returns the underlying `char*` array, and (2) `length` which returns - * the length of the string. - * \returns `true` if successful, or `false` if there is insufficient memory or - * `token` is not an appropriate string representation of a long. - */ -template -inline bool parse_long(const CharArray& token, long& value) { - char* buffer = (char*) malloc(sizeof(char) * (token.length + 1)); - if (buffer == NULL) { - fprintf(stderr, "parse_longlong ERROR: Unable to allocate temporary string buffer.\n"); - return false; - } - memcpy(buffer, token.data, sizeof(char) * token.length); - buffer[token.length] = '\0'; - - char* end_ptr; - value = strtol(buffer, &end_ptr, 0); - if (*end_ptr != '\0') { - free(buffer); - return false; - } - free(buffer); - return true; -} - -/** - * Attempts to parse the string given by `token` as a `long`. - * \tparam CharArray a string type that implements two fields: (1) `data` which - * returns the underlying `char*` array, and (2) `length` which returns - * the length of the string. - * \returns `true` if successful, or `false` if there is insufficient memory or - * `token` is not an appropriate string representation of a long. - */ -template -inline bool parse_long_long(const CharArray& token, long long& value) { - char* buffer = (char*)malloc(sizeof(char) * (token.length + 1)); - if (buffer == NULL) { - fprintf(stderr, "parse_longlong ERROR: Unable to allocate temporary string buffer.\n"); - return false; - } - memcpy(buffer, token.data, sizeof(char) * token.length); - buffer[token.length] = '\0'; - - char* end_ptr; - value = strtoll(buffer, &end_ptr, 0); - if (*end_ptr != '\0') { - free(buffer); - return false; - } - free(buffer); - return true; -} - -/** - * Attempts to parse the string given by `token` as an `unsigned int`. - * \param base if `0`, the numeric base of the integer is detected - * automatically in the same way as [strtoul](http://en.cppreference.com/w/cpp/string/byte/strtoul). - * Otherwise, the numeric base can be specified explicitly. - * \returns `true` if successful, or `false` if `token` is not an appropriate - * string representation of an unsigned integer. - */ -template -inline bool parse_uint(const char (&token)[N], unsigned int& value, unsigned int base = 0) { - char buffer[N + 1]; - for (unsigned int i = 0; i < N; i++) - buffer[i] = token[i]; - buffer[N] = '\0'; - - char* end_ptr; - value = strtoul(buffer, &end_ptr, base); - return (*end_ptr == '\0'); -} - -/** - * Represents a position in a file. This structure is typically used to provide - * informative errors during lexical analysis of data from a file. - */ -struct position { - /** - * The line number of the position in the file. - */ - unsigned int line; - - /** - * The column number of the position in the file. - */ - unsigned int column; - - /** - * Constructs the position structure with the given `line` and `column`. - */ - position(unsigned int line, unsigned int column) : - line(line), column(column) { } - - /** - * Constructs the position structure by copying from `p`. - */ - position(const position& p) : - line(p.line), column(p.column) { } - - /** - * Returns a position with the column number increased by `i`. - */ - position operator + (unsigned int i) const { - return position(line, column + i); - } - - /** - * Returns a position with the column number decreased by `i`. - */ - position operator - (unsigned int i) const { - return position(line, column - i); - } - - /** - * Copies the position structure from `src` into `dst`. - */ - static inline bool copy(const position& src, position& dst) { - dst.line = src.line; - dst.column = src.column; - return true; - } -}; - -/** - * A structure representing a single token during lexical analysis. This - * structure is generic, intended for use across multiple lexical analyzers. - */ -template -struct lexical_token { - /** - * The generic type of this token. - */ - TokenType type; - - /** - * The start position (inclusive) of the token in the source file. - */ - position start; - - /** - * The end position (exclusive) of the token in the source file. - */ - position end; - - /** - * An (optional) string representing the contents of the token. - */ - string text; -}; - -/** - * Prints the given lexical_token `token` to the output `stream`. - * \tparam Printer a scribe type for which the functions - * `print(const TokenType&, Stream&, Printer&)` and - * `print(const core::string& s, Stream&, Printer&)` are defined. - */ -template -bool print(const lexical_token& token, Stream& stream, Printer& printer) { - bool success = true; - success &= print(token.type, stream, printer); - if (!is_empty(token.text)) { - success &= print('(', stream); - success &= print(token.text, stream, printer); - success &= print(')', stream); - } - return success; -} - -/** - * Reports an error with the given message `error` as a null-terminated C - * string at the given source file position `pos` to [stderr](http://en.cppreference.com/w/cpp/io/c). - */ -inline void read_error(const char* error, const position& pos) { - fprintf(stderr, "ERROR at %d:%d: %s.\n", pos.line, pos.column, error); -} - -/** - * Constructs a lexical_token with the given `start` and `end` positions, and - * TokenType `type`, with an empty lexical_token::text message and appends it - * to the `tokens` array. - */ -template -bool emit_token(array>& tokens, - const position& start, const position& end, TokenType type) -{ - if (!tokens.ensure_capacity(tokens.length + 1)) { - fprintf(stderr, "emit_token ERROR: Unable to create token.\n"); - return false; - } - - lexical_token& new_token = tokens[(unsigned int) tokens.length]; - new_token.text.data = NULL; - new_token.type = type; - new_token.start = start; - new_token.end = end; - tokens.length++; - return true; -} - -/** - * Constructs a lexical_token with the given `start` and `end` positions, and - * TokenType `type`, with lexical_token::text copied from `token` and appends - * it to the `tokens` array. - */ -template -bool emit_token( - array>& tokens, array& token, - const position& start, const position& end, TokenType type) -{ - if (!tokens.ensure_capacity(tokens.length + 1)) { - fprintf(stderr, "emit_token ERROR: Unable to create token.\n"); - return false; - } - - lexical_token& new_token = tokens[(unsigned int) tokens.length]; - if (!init(new_token.text, token.data, (unsigned int) token.length)) { - fprintf(stderr, "emit_token ERROR: Unable to create string.\n"); - return false; - } - new_token.type = type; - new_token.start = start; - new_token.end = end; - tokens.length++; - token.clear(); - return true; -} - -/** - * Frees every element in the given `tokens` array. This function does not free - * the array itself. - */ -template -void free_tokens(array>& tokens) { - for (unsigned int i = 0; i < tokens.length; i++) - if (tokens[i].text.data != NULL) - core::free(tokens[i].text); -} - -/** - * Inspects the element at the given `index` in the `tokens` array. If `index` - * is not out of bounds, and the token at that index has type that matches the - * given `type`, the function returns `true`. Otherwise, an error message is - * printed to [stderr](http://en.cppreference.com/w/cpp/io/c) indicating that - * the expected token was missing, with its `name` as part of the error - * message, and `false` is returned. - */ -template -bool expect_token(const array>& tokens, - const unsigned int& index, TokenType type, const char* name) -{ - FILE* out = stderr; - if (index == tokens.length) { - /* unexpected end of input */ - fprintf(out, "ERROR: Unexpected end of input. Expected %s.\n", name); - return false; - } - else if (tokens[index].type != type) { - fprintf(out, "ERROR at %d:%d: Unexpected token ", tokens[index].start.line, tokens[index].start.column); - print(tokens[index].type, out); - fprintf(out, ". Expected %s.\n", name); - return false; - } - return true; -} - -/** - * Appends the given wide character `next` to the char array `token` which - * represents a multi-byte string. - */ -inline bool append_to_token( - array& token, char32_t next, mbstate_t& shift) -{ - if (!token.ensure_capacity(token.length + MB_CUR_MAX)) - return false; - size_t written = c32rtomb(token.data + token.length, next, &shift); - if (written == static_cast(-1)) - return false; - token.length += written; - return true; -} - -} /* namespace core */ - -#endif /* LEX_H_ */ diff --git a/reasoners/data_preprocess/logic/graph_dataset_gen/core/map.h b/reasoners/data_preprocess/logic/graph_dataset_gen/core/map.h deleted file mode 100644 index 9ddba2a02..000000000 --- a/reasoners/data_preprocess/logic/graph_dataset_gen/core/map.h +++ /dev/null @@ -1,2508 +0,0 @@ -/** - * \file map.h - * - * This file contains the hash_set, hash_map, and array_map data structures. It - * also defines the default_hash function, which provides a default hash - * algorithm (currently implemented using [xxhash](https://github.com/Cyan4973/xxHash)). - * - * - */ - -#ifndef MAP_H_ -#define MAP_H_ - -#include -#include - -#define XXH_INLINE_ALL -#include "array.h" -#include "xxhash.h" - - -namespace core { - -/** - * For hash_set and hash_map, `capacity` is the size of the underlying array - * (i.e. the number of buckets), whereas `size` is the number of elements (i.e. - * number of non-empty buckets). The functions hash_set::check_size and - * hash_map::check_size compute the load factor `size / capacity` and compare - * it to RESIZE_THRESHOLD. If the load factor is too large, the hashtable is - * resized and the capacity is increased. - */ -#define RESIZE_THRESHOLD 1 / 2 - -/** - * The multiplicative inverse of RESIZE_THRESHOLD. - */ -#define RESIZE_THRESHOLD_INVERSE 2 / 1 - -/** - * The multiplicative factor by which hash_set, hash_map, and array_map capacity is changed. - */ -#define RESIZE_FACTOR 2 - -/** - * A function pointer type describing a function that returns a pointer to - * allocated memory. The first argument is the number of elements to allocate, - * and the second argument is the number of bytes for each element. - * [calloc](http://en.cppreference.com/w/c/memory/calloc) is an example of a - * function with this type. - */ -typedef void*(alloc_keys_func)(size_t, size_t); - - -/* forward declarations */ -#if !defined(DOXYGEN_IGNORE) - -template -struct hash_set; - -template -struct hash_map; - -template -struct array_map; - -#endif /* DOXYGEN_IGNORE */ - - -#if defined(__LP64__) || defined(_WIN64) || defined(__x86_64__) || defined(__ppc64__) -template -inline uint_fast32_t default_hash(const K& key) { - return (uint_fast32_t) XXH64(&key, sizeof(K), Seed); -} - -template -inline uint_fast32_t default_hash(const K* keys, unsigned int length) { - return (uint_fast32_t) XXH64(keys, sizeof(K) * length, Seed); -} - -#else -/** - * Evaluates the hash function of the given value `key` with the given `Seed` using the default implementation. - */ -template -inline unsigned int default_hash(const K& key) { - return XXH32(&key, sizeof(K), Seed); -} - -/** - * Evaluates the hash function of the given native array of values `keys` with the given `Seed` using the default implementation. - */ -template -inline unsigned int default_hash(const K* keys, unsigned int length) { - return XXH32(keys, sizeof(K) * length, Seed); -} -#endif - -template -struct key_value_metric { - const KeyMetric& key_metric; - const ValueMetric& value_metric; - - constexpr key_value_metric(const KeyMetric& key_metric, const ValueMetric& value_metric) : - key_metric(key_metric), value_metric(value_metric) { } -}; - -template -inline constexpr key_value_metric make_key_value_metric( - const KeyMetric& key_metric, const ValueMetric& value_metric) { - return key_value_metric(key_metric, value_metric); -} - -inline constexpr key_value_metric make_key_value_metric() { - return make_key_value_metric(default_metric(), default_metric()); -} - - -/** - * - */ - -/** - * An iterator implementation, similar to those in the Standard Template - * Library, to enable iteration of elements in a hash_set. This iterator is - * typically initialized using hash_set::begin. - * - * This definition enables the use of range-based for loops. - */ -template -struct hash_set_iterator { - typedef typename std::conditional&, hash_set&>::type container_type; - - /** - * The type of the entries returned by this iterator. If this is a const - * iterator, `value_type` is `const T&`. Otherwise, `value_type` is `T&`. - */ - typedef typename std::conditional::type value_type; - - container_type set; - unsigned int position; - - hash_set_iterator(container_type& set, unsigned int position) : set(set), position(position) { } - - /** - * Returns whether this iterator is in the same position as `other`. This - * function assumes the two iterators were created from the same hash_set, - * and that it was not modified. - */ - inline bool operator != (const hash_set_iterator& other) const { - return position != other.position; - } - - /** - * Returns the element in the hash_set at the current iterator position. - * This function assumes the hash_set was not resized and no element was - * removed since the last call to either the operator `++` or the - * constructor of this iterator, whichever came later. - */ - inline value_type operator * () { - return set.keys[position]; - } - - /** - * Advances the position of the iterator to the next element in the hash_set. - */ - inline const hash_set_iterator& operator ++ () { - do { - ++position; - } while (position < set.capacity && hasher::is_empty(set.keys[position])); - return *this; - } -}; - -/** - * An iterator implementation, similar to those in the Standard Template - * Library, to enable iteration of elements in a hash_map. This iterator is - * typically initialized using hash_map::begin. - * - * This definition enables the use of range-based for loops. - */ -template -struct hash_map_iterator { - typedef typename std::conditional&, hash_map&>::type container_type; - - /** - * The type of the entries returned by this iterator. If this is a const - * iterator, `value_type` is `core::pair`. Otherwise, - * `value_type` is `core::pair`. - */ - typedef typename std::conditional, pair>::type value_type; - - container_type map; - unsigned int position; - - hash_map_iterator(container_type& map, unsigned int position) : map(map), position(position) { } - - /** - * Returns whether this iterator is in the same position as `other`. This - * function assumes the two iterators were created from the same hash_map, - * and that it was not modified. - */ - inline bool operator != (const hash_map_iterator& other) const { - return position != other.position; - } - - /** - * Returns the entry in the hash_map at the current iterator position. This - * function assumes the hash_map was not resized and no element was removed - * since the last call to either the operator `++` or the constructor of - * this iterator, whichever came later. - */ - inline value_type operator * () { - return { map.table.keys[position], map.values[position] }; - } - - /** - * Advances the position of the iterator to the next entry in the hash_map. - */ - inline const hash_map_iterator& operator ++ () { - do { - ++position; - } while (position < map.table.capacity && hasher::is_empty(map.table.keys[position])); - return *this; - } -}; - -/** - * An iterator implementation, similar to those in the Standard Template - * Library, to enable iteration of elements in an array_map. This iterator is - * typically initialized using array_map::begin. - * - * This definition enables the use of range-based for loops. - */ -template -struct array_map_iterator { - typedef typename std::conditional&, array_map&>::type container_type; - - /** - * The type of the entries returned by this iterator. If this is a const - * iterator, `value_type` is `core::pair`. Otherwise, - * `value_type` is `core::pair`. - */ - typedef typename std::conditional, pair>::type value_type; - - container_type map; - size_t position; - - array_map_iterator(container_type map, size_t position) : map(map), position(position) { } - - /** - * Returns whether this iterator is in the same position as `other`. This - * function assumes the two iterators were created from the same array_map, - * and that it was not modified. - */ - inline bool operator != (const array_map_iterator& other) const { - return position != other.position; - } - - /** - * Returns the entry in the array_map at the current iterator position. - * This function assumes the array_map was not resized and no element was - * removed since the last call to either the operator `++` or the - * constructor of this iterator, whichever came later. - */ - inline value_type operator * () { - return{ map.keys[position], map.values[position] }; - } - - /** - * Advances the position of the iterator to the next entry in the array_map. - */ - inline const array_map_iterator& operator ++ () { - ++position; - return *this; - } -}; - -/** - * Returns `true` only if `probe > start` and `probe <= end` where `probe`, `start`, - * and `end` are in the additive group of integers modulo the capacity of this set. - */ -inline bool index_between(unsigned int probe, unsigned int start, unsigned int end) { - if (end >= start) { - return (probe > start && probe <= end); - } else { - return (probe <= end || probe > start); - } -} - -/** - * An unordered associative container that contains a set of unique elements, - * each of type `T`. The elements are stored in the native array - * hash_set::keys, which has capacity hash_set::capacity. To compute the index - * of an element in hash_set::keys, we first compute its hash. To do so: - * 1. If `T` [is_pointer](http://en.cppreference.com/w/cpp/types/is_pointer), - * [is_fundamental](http://en.cppreference.com/w/cpp/types/is_fundamental), - * or [is_enum](http://en.cppreference.com/w/cpp/types/is_enum), - * core::default_hash() provides the hash. - * 2. For all other types, the `T::hash` function provides the hash. - * Once the hash is computed, the index into hash_set::keys is computed using - * `hash % hash_set::capacity` (modular division by the capacity). - * - * The above approach could produce the same index for distinct elements. This - * event is known as a *collision*. We use - * [linear probing](https://en.wikipedia.org/wiki/Linear_probing) - * to resolve collisions: When adding an element to the hash_set, we compute - * its `index` using the above procedure, then we inspect - * `hash_set::keys[index]` and use core::is_empty() to determine if another - * element already occupies that position. If so, we try the next position - * `(index + 1) % hash_set::capacity`. We continue until we find an empty - * index, and the element is inserted there. - * - * Thus, the function core::is_empty() is used to determine if a position in - * hash_set::keys is occupied or empty. The total number of occupied positions - * is given by hash_set::size. **WARNING:** If hash_set::keys becomes full, the - * above linear probing mechanism could lead to an infinite loop. The function - * hash_set::check_size should be used whenever adding new elements to avoid - * this scenario. - * - * **Performance:** This data structure provides constant-time access and - * modification, given that the load factor (`size / capacity`) is not too - * large. - * - * - * Below is an example of a simple use-case of hash_set, where the expected - * output is `a.contains(2): false, a.contains(3): true, -4 9 3`. - * - * ```{.cpp} - * #include - * #include - * using namespace core; - * - * int main() { - * hash_set a = hash_set(8); - * a.add(-1); a.add(-4); - * a.add(3); a.add(9); - * a.remove(-1); - * - * printf("a.contains(2): %s, ", a.contains(2) ? "true" : "false"); - * printf("a.contains(3): %s, ", a.contains(3) ? "true" : "false"); - * for (int element : a) - * printf("%d ", element); - * } - * ``` - * - * - * However, if `a` is not allocated on the stack, the destructor will not be - * automatically called, and so it must be freed manually using `core::free` or - * `hash_set::free`. In the example below, the expected output is the same as - * that of the program above: `a.contains(2): false, a.contains(3): true, -4 9 3`. - * - * ```{.cpp} - * #include - * #include - * using namespace core; - * - * int main() { - * hash_set& a = *((hash_set*) alloca(sizeof(hash_set))); - * hash_set_init(a, 8); - * a.add(-1); a.add(-4); - * a.add(3); a.add(9); - * a.remove(-1); - * - * printf("a.contains(2): %s, ", a.contains(2) ? "true" : "false"); - * printf("a.contains(3): %s, ", a.contains(3) ? "true" : "false"); - * for (int element : a) - * printf("%d ", element); - * free(a); - * } - * ``` - * - * - * A number of member functions require `T` to be - * [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). - * If this is not the case, those operations can be performed directly on the - * public fields, as in the following example. In addition, when using a custom - * struct/class with hash_set, it must implement public static functions, like - * `hash`, `is_empty`, and `move`, as well as the operator `==`. The expected - * output of the following example is `first second `. - * - * ```{.cpp} - * #include - * #include - * #include - * using namespace core; - * - * struct custom_string { - * char* buffer; - * - * static unsigned int hash(const custom_string& s) { - * return default_hash(s.buffer, strlen(s.buffer)); - * } - * - * static bool is_empty(const custom_string& s) { - * return s.buffer == NULL; - * } - * - * static void move(const custom_string& src, custom_string& dst) { - * dst.buffer = src.buffer; - * } - * - * static void free(custom_string& s) { - * core::free(s.buffer); - * } - * }; - * - * inline bool operator == (const custom_string& first, const custom_string& second) { - * if (first.buffer == NULL) - * return second.buffer == NULL; - * return strcmp(first.buffer, second.buffer) == 0; - * } - * - * bool init(custom_string& s, const char* src) { - * s.buffer = (char*) malloc(sizeof(char) * (strlen(src) + 1)); - * if (s.buffer == NULL) - * return false; - * memcpy(s.buffer, src, sizeof(char) * (strlen(src) + 1)); - * return true; - * } - * - * int main() { - * custom_string first, second; - * init(first, "first"); - * init(second, "second"); - * - * hash_set a = hash_set(8); - * a.check_size(a.size + 2); - * - * bool contains; unsigned int index; - * index = a.index_of(first, contains); - * if (!contains) { - * core::move(first, a.keys[index]); - * a.size++; - * } - * - * index = a.index_of(second, contains); - * if (!contains) { - * core::move(second, a.keys[index]); - * a.size++; - * } - * - * for (const custom_string& s : a) - * printf("%s ", s.buffer); - * for (custom_string& s : a) - * free(s); - * } - * ``` - * - * \tparam T the generic type of the elements in the set. `T` must satisfy either: - * 1. [is_fundamental](http://en.cppreference.com/w/cpp/types/is_fundamental), - * 2. [is_enum](http://en.cppreference.com/w/cpp/types/is_enum), - * 3. [is_pointer](http://en.cppreference.com/w/cpp/types/is_pointer), - * 4. implements the public static method `unsigned int hash(const T&)`, - * the public static method `void is_empty(const T&)`, implements the - * operators `==`, and satisfies is_moveable. Some operations also - * require the operator `!=`, and public static methods - * `void set_empty(T&)` and `void set_empty(T*, unsigned int)`. - * **NOTE:** The first argument to the `==` and `!=` operators may be - * empty. - */ -/* TODO: consider other collision resolution mechanisms */ -template -struct hash_set -{ - /** - * The native array of keys underlying the hashtable. - */ - T* keys; - - /** - * The capacity of hash_set::keys. - */ - unsigned int capacity; - - /** - * The number of elements in the hashtable (i.e. the number of non-empty buckets). - */ - unsigned int size; - - /** - * Constructs the hash_set with the given `initial_capacity`. - * \param alloc_keys a memory allocation function with prototype - * `void* alloc_keys(size_t count, size_t size)` that allocates space for - * `count` items, each with size `size`, and initializes them such that - * core::is_empty() returns `true` for each element. - */ - hash_set(unsigned int initial_capacity, alloc_keys_func alloc_keys = calloc) { - if (!initialize(initial_capacity, alloc_keys)) { - fprintf(stderr, "hash_set ERROR: Unable to allocate memory.\n"); - exit(EXIT_FAILURE); - } - } - - /** - * Constructs the hash_set and inserts the given native array of elements - * `set` with given `length`. - * \param alloc_keys a memory allocation function with prototype - * `void* alloc_keys(size_t count, size_t size)` that allocates space for - * `count` items, each with size `size`, and initializes them such that - * core::is_empty() returns `true` for each element. - */ - hash_set(const T* set, unsigned int length, - alloc_keys_func alloc_keys = calloc) : - hash_set(length * RESIZE_THRESHOLD_INVERSE + 1, alloc_keys) - { - for (unsigned int i = 0; i < length; i++) - insert(set[i]); - } - - /** - * Constructs the hash_set and inserts the given - * [initializer_list](http://en.cppreference.com/w/cpp/language/list_initialization) - * of elements `list`. - * \param alloc_keys a memory allocation function with prototype - * `void* alloc_keys(size_t count, size_t size)` that allocates space for - * `count` items, each with size `size`, and initializes them such that - * core::is_empty() returns `true` for each element. - */ - hash_set(const std::initializer_list& list, - alloc_keys_func alloc_keys = calloc) : - hash_set(list.size() * RESIZE_THRESHOLD_INVERSE + 1, alloc_keys) - { - typename std::initializer_list::iterator i; - for (i = list.begin(); i != list.end(); i++) - insert(*i); - } - - ~hash_set() { free(); } - - /** - * Forces the underlying hash_set::keys to be resized to the requested - * `capacity`. - * - * **WARNING:** If `new_capacity <= hash_set::size`, the hashtable could - * become full during the resize process, leading to an infinite loop due - * to the linear probing collision resolution mechanism. - * - * \param alloc_keys a memory allocation function with prototype - * `void* alloc_keys(size_t count, size_t size)` that allocates space for - * `count` items, each with size `size`, and initializes them such that - * core::is_empty() returns `true` for each element. - */ - bool resize(unsigned int new_capacity, - alloc_keys_func alloc_keys = calloc) - { - T* old_keys = keys; - - keys = (T*) alloc_keys(new_capacity, sizeof(T)); - if (keys == NULL) { - /* revert changes and return error */ - keys = old_keys; - return false; - } - - /* iterate through keys and re-hash the elements */ - unsigned int old_capacity = capacity; - capacity = new_capacity; - for (unsigned int i = 0; i < old_capacity; i++) { - if (!hasher::is_empty(old_keys[i])) - core::move(old_keys[i], keys[next_empty(old_keys[i])]); - } - core::free(old_keys); - return true; - } - - /** - * This function first determines whether `hash_set::size < hash_set::capacity * RESIZE_THRESHOLD`. - * If not, the capacity of the underlying hash_set::keys is increased by - * RESIZE_FACTOR until the condition is satisfied. This is useful to ensure - * the hashtable is sufficiently large when preparing to add new elements. - * \param alloc_keys a memory allocation function with prototype - * `void* alloc_keys(size_t count, size_t size)` that allocates space for - * `count` items, each with size `size`, and initializes them such that - * core::is_empty() returns `true` for each element. - * \returns `true` if the resize was successful, and `false` if there is insufficient memory. - */ - inline bool check_size(alloc_keys_func alloc_keys = calloc) { - return check_size(size, alloc_keys); - } - - /** - * For a requested number of elements `new_size`, this function first - * determines whether `new_size < hash_set::capacity * RESIZE_THRESHOLD`. - * If not, the capacity of the underlying hashtable is increased by - * RESIZE_FACTOR until the condition is satisfied. This is useful to ensure - * the hashtable is sufficiently large when preparing to add new elements. - * \param alloc_keys a memory allocation function with prototype - * `void* alloc_keys(size_t count, size_t size)` that allocates space for - * `count` items, each with size `size`, and initializes them such that - * core::is_empty() returns `true` for each element. - * \returns `true` if the resize was successful, and `false` if there is insufficient memory. - */ - inline bool check_size(unsigned int new_size, alloc_keys_func alloc_keys = calloc) - { - while (new_size >= capacity * RESIZE_THRESHOLD) { - if (!resize(RESIZE_FACTOR * capacity, alloc_keys)) { - fprintf(stderr, "hash_set.put ERROR: Unable to resize hashtable.\n"); - return false; - } - } - return true; - } - - /** - * Add the given `element` to this set. The assignment operator is used to - * insert each element, and so this function should not be used if `T` is not - * [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). - * In such a case, insertion should be performed manually using - * hash_set::index_of to find the appropriate index and directly modifying - * hash_set::keys and hash_set::size. See the example in the hash_set - * description. - * \param alloc_keys a memory allocation function with prototype - * `void* alloc_keys(size_t count, size_t size)` that allocates space for - * `count` items, each with size `size`, and initializes them such that - * core::is_empty() returns `true` for each element. - * \tparam T is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). - */ - bool add(const T& element, alloc_keys_func alloc_keys = calloc) - { - if (!check_size(size, alloc_keys)) return false; - insert(element); - return true; - } - - /** - * Adds all the elements in the hash_set `elements` to this set. The - * assignment operator is used to insert each element, and so this function - * should not be used if `T` is not - * [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). - * In such a case, insertion should be performed manually using - * hash_set::index_of or hash_set::index_to_insert to find the appropriate - * index and directly modifying hash_set::keys and hash_set::size. See the - * example in the hash_set description. - * \param alloc_keys a memory allocation function with prototype - * `void* alloc_keys(size_t count, size_t size)` that allocates space for - * `count` items, each with size `size`, and initializes them such that - * core::is_empty() returns `true` for each element. - * \tparam T is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). - */ - bool add_all(const hash_set& elements, - alloc_keys_func alloc_keys = calloc) - { - if (!check_size(size + elements.size, alloc_keys)) return false; - for (unsigned int i = 0; i < elements.capacity; i++) - if (!hasher::is_empty(elements.keys[i])) - insert(elements.keys[i]); - return true; - } - - /** - * Adds all the elements in the native array `elements` with length `count` - * to this set. The assignment operator is used to insert each element, and - * so this function should not be used if `T` is not - * [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). - * In such a case, insertion should be performed manually using - * hash_set::index_of or hash_set::index_to_insert and direct modification - * of the public fields. - * \param alloc_keys a memory allocation function with prototype - * `void* alloc_keys(size_t count, size_t size)` that allocates space for - * `count` items, each with size `size`, and initializes them such that - * core::is_empty() returns `true` for each element. - * \tparam T is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). - */ - bool add_all(const T* elements, unsigned int count, - alloc_keys_func alloc_keys = calloc) - { - if (!check_size(size + count, alloc_keys)) return false; - for (unsigned int i = 0; i < count; i++) - insert(elements[i]); - return true; - } - - /** - * This function removes the given `element` from the set. This function - * does not free the removed element. - * \returns `true` if the element is removed, and `false` if the set does not contain `element`. - */ - bool remove(const T& element) - { -#if !defined(NDEBUG) - if (is_empty(element)) - fprintf(stderr, "hash_set.remove WARNING: Specified key is empty.\n"); -#endif - - unsigned int index = hasher::hash(element) % capacity; - while (true) { - if (keys[index] == element) { - break; - } else if (hasher::is_empty(keys[index])) - return false; - index = (index + 1) % capacity; - } - - remove_at(index); - return true; - } - - template - bool remove(const T& element, V* values) - { -#if !defined(NDEBUG) - if (is_empty(element)) - fprintf(stderr, "hash_set.remove WARNING: Specified key is empty.\n"); -#endif - - unsigned int index = hasher::hash(element) % capacity; - while (true) { - if (keys[index] == element) { - break; - } else if (hasher::is_empty(keys[index])) - return false; - index = (index + 1) % capacity; - } - - remove_at(values, index); - return true; - } - - /** - * This function removes the element at the bucket given by `index`. This - * function assumes that an element is located at the given bucket with the - * correct provided hash value. This function does not free the removed - * element. - */ - void remove_at(unsigned int index) - { - unsigned int last = index; - unsigned int search = (index + 1) % capacity; - while (!hasher::is_empty(keys[search])) - { - unsigned int search_hash = hasher::hash(keys[search]) % capacity; - if (!index_between(search_hash, last, search)) { - core::move(keys[search], keys[last]); - last = search; - } - search = (search + 1) % capacity; - } - - hasher::set_empty(keys[last]); - size--; - } - - /** - * Returns `true` if `element` exists in the set, and `false` otherwise. - */ - bool contains(const T& element) const - { -#if !defined(NDEBUG) - if (is_empty(element)) - fprintf(stderr, "hashtable.contains WARNING: Specified key is empty.\n"); - if (size == capacity) - fprintf(stderr, "hashtable.contains WARNING: Hashtable is full!\n"); -#endif - - unsigned int index = hasher::hash(element) % capacity; - while (true) { - if (keys[index] == element) { - return true; - } else if (hasher::is_empty(keys[index])) { - return false; - } - index = (index + 1) % capacity; - } - } - - /** - * If the given `element` exists in this set, this function returns the - * index of the bucket that contains it. If not, this function returns the - * index where the key would be located, if it had existed in the set. - */ - unsigned int index_of(const T& element) const - { -#if !defined(NDEBUG) - if (is_empty(element)) - fprintf(stderr, "hashtable.index_of WARNING: Specified key is empty.\n"); - if (size == capacity) - fprintf(stderr, "hashtable.index_of WARNING: Hashtable is full!\n"); -#endif - - unsigned int index = hasher::hash(element) % capacity; - while (keys[index] != element && !hasher::is_empty(keys[index])) - index = (index + 1) % capacity; - return index; - } - - /** - * If the given `element` exists in this set, this function returns the - * index of the bucket that contains it, and sets `contains` to `true`. - * If `element` is not in the set, `contains` is set to false, and this - * function returns the index where the key would be located, if it had - * existed in the set. - */ - inline unsigned int index_of( - const T& element, bool& contains) const - { - unsigned int hash_value; - return index_of(element, contains, hash_value); - } - - /** - * If the given `element` exists in this set, this function returns the - * index of the bucket that contains it, and sets `contains` to `true`. - * If `element` is not in the set, `contains` is set to false, and this - * function returns the index where the key would be located, if it had - * existed in the set. In any case, the evaluated hash function of - * `element` is stored in `hash_value`. - */ - unsigned int index_of(const T& element, - bool& contains, unsigned int& hash_value) const - { -#if !defined(NDEBUG) - if (is_empty(element)) - fprintf(stderr, "hashtable.index_of WARNING: Specified key is empty.\n"); - if (size == capacity) - fprintf(stderr, "hashtable.index_of WARNING: Hashtable is full!\n"); -#endif - - hash_value = hasher::hash(element); - unsigned int index = hash_value % capacity; - while (true) { - if (keys[index] == element) { - contains = true; - return index; - } else if (hasher::is_empty(keys[index])) { - contains = false; - return index; - } - index = (index + 1) % capacity; - } - } - - /** - * For a given `element`, this function computes and returns the index of - * the bucket where the element would be inserted, for example by a call to - * hash_set::add. `contains` is set to `true` if and only if the given - * `element` already exists in the set. - */ - inline unsigned int index_to_insert(const T& element, bool& contains) - { -#if !defined(NDEBUG) - if (size == capacity) - fprintf(stderr, "hashtable.index_to_insert WARNING: Hashtable is full!\n"); -#endif - unsigned int index = hasher::hash(element) % capacity; - while (true) { - if (hasher::is_empty(keys[index])) { - contains = false; break; - } if (keys[index] == element) { - contains = true; break; - } - index = (index + 1) % capacity; - } - return index; - } - - /** - * For a given `element`, this function computes and returns the index of - * the bucket where the element would be inserted, for example by a call to - * hash_set::add, **assuming** the given element is not in the set. - */ - inline unsigned int index_to_insert(const T& element) - { -#if !defined(NDEBUG) - if (size == capacity) - fprintf(stderr, "hashtable.index_to_insert WARNING: Hashtable is full!\n"); -#endif - unsigned int index = hasher::hash(element) % capacity; - while (true) { - if (hasher::is_empty(keys[index])) break; - index = (index + 1) % capacity; - } - return index; - } - - /** - * Removes all elements from this hash_set. Note that this function does - * not free each element beforehand. - */ - void clear() { - hasher::set_empty(keys, capacity); - size = 0; - } - - /** - * Returns `true` if this hash_set is a subset of `other`. - */ - bool is_subset(const hash_set& other) const - { - for (unsigned int i = 0; i < capacity; i++) - if (!hasher::is_empty(keys[i]) && !other.contains(keys[i])) - return false; - return true; - } - - bool equals(const hash_set& other) const - { - if (size != other.size) return false; - return is_subset(other); - } - - /** - * Returns a hash_set_iterator pointing to the first element in this container. - * - * **NOTE:** Unlike the libstdc++ [unordered_set](http://en.cppreference.com/w/cpp/container/unordered_set) - * and [unordered_map](http://en.cppreference.com/w/cpp/container/unordered_map) - * iterators, we do not keep a linked list among the elements, so if rapid - * iteration over elements is more critical than rapid queries, consider - * using an array_map. - */ - inline hash_set_iterator begin() { - return hash_set_iterator(*this, first_empty()); - } - - /** - * Returns a hash_set_iterator pointing to the end of this container. - * - * **NOTE:** Unlike the libstdc++ [unordered_set](http://en.cppreference.com/w/cpp/container/unordered_set) - * and [unordered_map](http://en.cppreference.com/w/cpp/container/unordered_map) - * iterators, we do not keep a linked list among the elements, so if rapid - * iteration over elements is more critical than rapid queries, consider - * using an array_map. - */ - inline hash_set_iterator end() { - return hash_set_iterator(*this, capacity); - } - - /** - * Returns a const hash_set_iterator pointing to the first element in this container. - * - * **NOTE:** Unlike the libstdc++ [unordered_set](http://en.cppreference.com/w/cpp/container/unordered_set) - * and [unordered_map](http://en.cppreference.com/w/cpp/container/unordered_map) - * iterators, we do not keep a linked list among the elements, so if rapid - * iteration over elements is more critical than rapid queries, consider - * using an array_map. - */ - inline hash_set_iterator begin() const { - return hash_set_iterator(*this, first_empty()); - } - - /** - * Returns a const hash_set_iterator pointing to the end of this container. - * - * **NOTE:** Unlike the libstdc++ [unordered_set](http://en.cppreference.com/w/cpp/container/unordered_set) - * and [unordered_map](http://en.cppreference.com/w/cpp/container/unordered_map) - * iterators, we do not keep a linked list among the elements, so if rapid - * iteration over elements is more critical than rapid queries, consider - * using an array_map. - */ - inline hash_set_iterator end() const { - return hash_set_iterator(*this, capacity); - } - - /** - * Swaps the contents of the hash_set `first` with that of `second`. - */ - static void swap(hash_set& first, hash_set& second) { - core::swap(first.keys, second.keys); - core::swap(first.capacity, second.capacity); - core::swap(first.size, second.size); - } - - /** - * Moves the contents of the hash_set `src` into `dst`. Note this function - * does not copy the contents of the underlying hash_set::keys, it merely - * copies the pointer. - */ - static inline void move(const hash_set& src, hash_set& dst) { - dst.keys = src.keys; - dst.capacity = src.capacity; - dst.size = src.size; - } - - /** - * Copies the contents of the hash_set `src` into `dst`. - * \param alloc_keys a memory allocation function with prototype - * `void* alloc_keys(size_t count, size_t size)` that allocates space for - * `count` items, each with size `size`, and initializes them such that - * core::is_empty() returns `true` for each element. - */ - static inline bool copy(const hash_set& src, hash_set& dst, alloc_keys_func alloc_keys = calloc) { - dst.capacity = src.capacity; - dst.size = src.size; - dst.keys = (T*) alloc_keys(src.capacity, sizeof(T)); - if (dst.keys == NULL) return false; - - for (unsigned int i = 0; i < src.capacity; i++) { - if (is_empty(src.keys[i])) continue; - if (!copy(src.keys[i], dst.keys[i])) { - free(dst); return false; - } - } - return true; - } - - template - static inline long unsigned int size_of(const hash_set& set, const Metric& metric) - { - long unsigned int sum = core::size_of(set.capacity) + core::size_of(set.size); - for (unsigned int i = 0; i < set.capacity; i++) { - if (is_empty(set.keys[i])) - sum += sizeof(T); - else sum += core::size_of(set.keys[i], metric); - } - return sum; - } - - /** - * Frees hash_set::keys. This should not be used if `set` was constructed - * on the stack, as the destructor will automatically free hash_set::data. - * The elements of `set` are not freed. - */ - static inline void free(hash_set& set) { set.free(); } - -private: - bool initialize( - unsigned int initial_capacity, - alloc_keys_func alloc_keys) - { -#if !defined(NDEBUG) - if (initial_capacity == 0) - fprintf(stderr, "hashtable.initialize WARNING: Initial capacity is zero.\n"); -#endif - - size = 0; - capacity = initial_capacity; - - keys = (T*) alloc_keys(capacity, sizeof(T)); - return (keys != NULL); - } - - inline void free() { - core::free(keys); - } - - inline void place( - const T& element, unsigned int index) - { -#if !defined(NDEBUG) - if (is_empty(element)) - fprintf(stderr, "hashtable.place WARNING: Specified key is empty.\n"); -#endif - - keys[index] = element; - } - - inline void insert_unique(const T& element) - { - place(element, next_empty(element)); - } - - inline unsigned int first_empty() const { - unsigned int index = 0; - while (index < capacity && hasher::is_empty(keys[index])) - index++; - return index; - } - - inline unsigned int next_empty(const T& element) - { -#if !defined(NDEBUG) - if (size == capacity) - fprintf(stderr, "hashtable.next_empty WARNING: Hashtable is full!\n"); -#endif - unsigned int index = hasher::hash(element) % capacity; - while (!hasher::is_empty(keys[index])) - index = (index + 1) % capacity; - return index; - } - - inline void insert(const T& element) - { - bool contains; - unsigned int index = index_to_insert(element, contains); - if (!contains) { - place(element, index); - size++; - } - } - - template - void remove_at(V* values, unsigned int index) - { - unsigned int last = index; - unsigned int search = (index + 1) % capacity; - if (!hasher::is_empty(keys[search])) - { - do { - unsigned int search_hash = hasher::hash(keys[search]) % capacity; - if (!index_between(search_hash, last, search)) { - core::move(keys[search], keys[last]); - core::move(values[search], values[last]); - last = search; - } - search = (search + 1) % capacity; - } while (!hasher::is_empty(keys[search])); - } - - hasher::set_empty(keys[last]); - size--; - } - - template - friend struct hash_map; - - template - friend bool hash_set_init(hash_set& set, - unsigned int capacity, alloc_keys_func alloc_keys); - - template - friend bool hash_set_init(hash_set& set, - const std::initializer_list& list, - alloc_keys_func alloc_keys); -}; - -/** - * Initializes the hash_set `set` with the given initial `capacity`. - * \param alloc_keys a memory allocation function with prototype - * `void* alloc_keys(size_t count, size_t size)` that allocates space for - * `count` items, each with size `size`, and initializes them such that - * core::is_empty() returns `true` for each element. - */ -template -bool hash_set_init( - hash_set& set, unsigned int capacity, - alloc_keys_func alloc_keys = calloc) -{ - if (!set.initialize(capacity, alloc_keys)) { - fprintf(stderr, "hash_set_init ERROR: Unable to allocate memory.\n"); - return false; - } - return true; -} - -/** - * Swaps the underlying buffers of the given hash_sets `first` and `second`. - */ -template -void swap(hash_set& first, hash_set& second) { - T* keys_swap = first.keys; - first.keys = second.keys; - second.keys = keys_swap; - - unsigned int ui_swap = first.size; - first.size = second.size; - second.size = ui_swap; - - ui_swap = first.capacity; - first.capacity = second.capacity; - second.capacity = ui_swap; -} - -/** - * An unordered associative container that contains a set of key-value pairs, - * where the keys are unique and have type `K` and the values have type `V`. - * The keys are stored in a hash_set structure, and the values are stored in - * hash_map::values, which is a native array parallel to the key array in - * hash_set::keys. To compute the index of a key-value pair, we compute the - * index of the key using the algorithm described in hash_set. The value has - * the same index in hash_map::values. This structure uses the same linear - * probing approach to resolve collisions (i.e. when two distinct keys are - * computed to have the same index). - * - * **WARNING:** As with hash_set, if the map becomes full, the linear probing - * mechanism could lead to an infinite loop in many hash_map operations. The - * function hash_map::check_size should be used whenever adding new elements - * to avoid this scenario. - * - * **Performance:** This data structure provides constant-time access and - * modification, given that the load factor (`size / capacity`) is not too - * large. - * - * Below is a simple example using a hash_map, where the expected output is - * `a.get(3): c, a.get(4): d, 1:a 4:d 3:c`. - * - * ```{.cpp} - * #include - * #include - * using namespace core; - * - * int main() { - * hash_map a = hash_map(8); - * a.put(1, 'a'); a.put(2, 'b'); - * a.put(3, 'c'); a.put(4, 'd'); - * a.remove(2); - * - * printf("a.get(3): %c, ", a.get(3)); - * printf("a.get(4): %c, ", a.get(4)); - * for (const auto& entry : a) - * printf("%d:%c ", entry.key, entry.value); - * } - * ``` - * - * - * However, if `a` is not allocated on the stack, the destructor will not be - * automatically called, and so it must be freed manually using `core::free` - * or `hash_map::free`. In the example below, the expected output is the same - * as that of the program above: `a.get(3): c, a.get(4): d, 1:a 4:d 3:c`. - * - * ```{.cpp} - * #include - * #include - * using namespace core; - * - * int main() { - * hash_map& a = *((hash_map*) malloc(sizeof(hash_map))); - * hash_map_init(a, 8); - * a.put(1, 'a'); a.put(2, 'b'); - * a.put(3, 'c'); a.put(4, 'd'); - * a.remove(2); - * - * printf("a.get(3): %c, ", a.get(3)); - * printf("a.get(4): %c, ", a.get(4)); - * for (const auto& entry : a) - * printf("%d:%c ", entry.key, entry.value); - * free(a); - * } - * ``` - * - * - * A number of member functions require `K` and `V` to be - * [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). - * If this is not the case, those operations can be performed directly on the - * public fields, as in the following example. In addition, when using a custom - * struct/class as the key type in hash_map, it must implement public static - * functions, like `hash`, `is_empty`, and `move`, as well as the operator - * `==`. The expected output of the following example is `first:1 second:2 `. - * - * ```{.cpp} - * #include - * #include - * #include - * using namespace core; - * - * struct custom_string { - * char* buffer; - * - * static unsigned int hash(const custom_string& s) { - * return default_hash(s.buffer, strlen(s.buffer)); - * } - * - * static bool is_empty(const custom_string& s) { - * return s.buffer == NULL; - * } - * - * static void move(const custom_string& src, custom_string& dst) { - * dst.buffer = src.buffer; - * } - * - * static void free(custom_string& s) { - * core::free(s.buffer); - * } - * }; - * - * inline bool operator == (const custom_string& first, const custom_string& second) { - * if (first.buffer == NULL) - * return second.buffer == NULL; - * return strcmp(first.buffer, second.buffer) == 0; - * } - * - * bool init(custom_string& s, const char* src) { - * s.buffer = (char*) malloc(sizeof(char) * (strlen(src) + 1)); - * if (s.buffer == NULL) - * return false; - * memcpy(s.buffer, src, sizeof(char) * (strlen(src) + 1)); - * return true; - * } - * - * int main() { - * custom_string key_one, key_two; - * int value_one = 1, value_two = 2; - * init(key_one, "first"); - * init(key_two, "second"); - * - * hash_map a = hash_map(8); - * a.check_size(a.table.size + 2); - * - * bool contains; unsigned int index; - * a.get(key_one, contains, index); - * if (!contains) { - * core::move(key_one, a.table.keys[index]); - * a.values[index] = value_one; - * a.table.size++; - * } - * - * a.get(key_two, contains, index); - * if (!contains) { - * core::move(key_two, a.table.keys[index]); - * a.values[index] = value_two; - * a.table.size++; - * } - * - * for (const auto& entry : a) - * printf("%s:%d ", entry.key.buffer, entry.value); - * for (auto entry : a) - * free(entry.key); - * } - * ``` - * - * \tparam K the generic type of the keys in the map. `K` must satisfy either: - * 1. [is_fundamental](http://en.cppreference.com/w/cpp/types/is_fundamental), - * 2. [is_enum](http://en.cppreference.com/w/cpp/types/is_enum), - * 3. [is_pointer](http://en.cppreference.com/w/cpp/types/is_pointer), - * 4. implements the public static method `unsigned int hash(const T&)`, - * the public static method `void is_empty(const T&)`, implements the - * operators `==`, and satisfies is_moveable. Some operations also - * require the operator `!=`, and public static methods - * `void set_empty(T&)` and `void set_empty(T*, unsigned int)`. - * **NOTE:** The first argument to the `==` and `!=` operators may be - * empty. - * \tparam V the generic type of the values in the map. `V` must satisfy is_moveable. - */ -/* TODO: consider other collision resolution mechanisms */ -template -struct hash_map -{ - /** - * The type of the keys in this hash_map. - */ - typedef K key_type; - - /** - * The type of the values in this hash_map. - */ - typedef V value_type; - - /** - * The underlying hashtable containing the keys. - */ - hash_set table; - - /** - * An array parallel to hash_set::keys in hash_map::table, - * containing a value at every non-empty bucket index. - */ - V* values; - - /** - * Constructs the hash_map with the given initial `capacity`. - * \param alloc_keys a memory allocation function with prototype - * `void* alloc_keys(size_t count, size_t size)` that allocates space for - * `count` items, each with size `size`, and initializes them such that - * core::is_empty() returns `true` for each key. - */ - hash_map(unsigned int capacity, alloc_keys_func alloc_keys = calloc) : - table(capacity, alloc_keys) - { - if (!initialize_values()) { - fprintf(stderr, "hash_map ERROR: Unable to allocate memory.\n"); - exit(EXIT_FAILURE); - } - } - - /** - * Constructs the hash_map and inserts the array of keys `map`, where each - * element in `map` is interpreted as a key, and its corresponding value is - * the index of the element. - * \param alloc_keys a memory allocation function with prototype - * `void* alloc_keys(size_t count, size_t size)` that allocates space for - * `count` items, each with size `size`, and initializes them such that - * core::is_empty() returns `true` for each key. - * \tparam V satisfies [is_integral](http://en.cppreference.com/w/cpp/types/is_integral). - */ - hash_map(const K* map, unsigned int length, - alloc_keys_func alloc_keys = calloc) : - table(length * RESIZE_THRESHOLD_INVERSE + 1, alloc_keys) - { - if (!initialize_values()) { - fprintf(stderr, "hash_map ERROR: Unable to allocate memory.\n"); - exit(EXIT_FAILURE); - } - for (unsigned int i = 0; i < length; i++) - insert(map[i], i); - } - - /** - * Constructs the hash_map and inserts the list of key-value pairs given by - * the native arrays `keys` and `values`. - * \param alloc_keys a memory allocation function with prototype - * `void* alloc_keys(size_t count, size_t size)` that allocates space for - * `count` items, each with size `size`, and initializes them such that - * core::is_empty() returns `true` for each key. - */ - hash_map(const K* keys, const V* values, - unsigned int length, alloc_keys_func alloc_keys = calloc) : - table(length * RESIZE_THRESHOLD_INVERSE + 1, alloc_keys) - { - if (!initialize_values()) { - fprintf(stderr, "hash_map ERROR: Unable to allocate memory.\n"); - exit(EXIT_FAILURE); - } - for (unsigned int i = 0; i < length; i++) - insert(keys[i], values[i]); - } - - /** - * Constructs the hash_map and inserts the given - * [initializer_list](http://en.cppreference.com/w/cpp/language/list_initialization) - * of core::pair entries. - * \param alloc_keys a memory allocation function with prototype - * `void* alloc_keys(size_t count, size_t size)` that allocates space for - * `count` items, each with size `size`, and initializes them such that - * core::is_empty() returns `true` for each key. - */ - hash_map(const std::initializer_list>& list, - alloc_keys_func alloc_keys = calloc) : - table(list.size() * RESIZE_THRESHOLD_INVERSE + 1, alloc_keys) - { - if (!initialize_values()) { - fprintf(stderr, "hash_map ERROR: Unable to allocate memory.\n"); - exit(EXIT_FAILURE); - } - typename std::initializer_list>::iterator i; - for (i = list.begin(); i != list.end(); i++) - insert(i->key, i->value); - } - - ~hash_map() { - ::free(values); - } - - /** - * Forces the underlying hash_map::table.keys and hash_map::values to be - * resized to the requested `capacity`. - * - * **WARNING:** If `new_capacity <= hash_map::table.size`, the hashtable - * could become full during the resize process, leading to an infinite loop - * due to the linear probing collision resolution mechanism. - * - * \param alloc_keys a memory allocation function with prototype - * `void* alloc_keys(size_t count, size_t size)` that allocates space for - * `count` items, each with size `size`, and initializes them such that - * core::is_empty() returns `true` for each key. - */ - bool resize(unsigned int new_capacity, - alloc_keys_func alloc_keys = calloc) - { - K* old_keys = table.keys; - V* old_values = values; - - table.keys = (K*) alloc_keys(new_capacity, sizeof(K)); - if (table.keys == NULL) { - /* revert changes and return error */ - table.keys = old_keys; - return false; - } - - values = (V*) malloc(sizeof(V) * new_capacity); - if (values == NULL) { - /* revert changes and return error */ - ::free(table.keys); - table.keys = old_keys; - values = old_values; - return false; - } - - /* iterate through keys and re-hash the elements */ - unsigned int old_capacity = table.capacity; - table.capacity = new_capacity; - for (unsigned int i = 0; i < old_capacity; i++) { - if (!hasher::is_empty(old_keys[i])) { - unsigned int new_bucket = table.next_empty(old_keys[i]); - core::move(old_keys[i], table.keys[new_bucket]); - core::move(old_values[i], values[new_bucket]); - } - } - ::free(old_keys); - ::free(old_values); - return true; - } - - /** - * This function first determines whether `hash_map::table.size < hash_map::table.capacity * RESIZE_THRESHOLD`. - * If not, the capacities of the underlying hash_map::table and - * hash_map::values is increased by RESIZE_FACTOR until the condition is - * satisfied. This is useful to ensure the hashtable is sufficiently large - * when preparing to add new elements. - * \param alloc_keys a memory allocation function with prototype - * `void* alloc_keys(size_t count, size_t size)` that allocates space for - * `count` items, each with size `size`, and initializes them such that - * core::is_empty() returns `true` for each key. - * \returns `true` if the resize was successful, and `false` if there is insufficient memory. - */ - inline bool check_size(alloc_keys_func alloc_keys = calloc) { - return check_size(table.size, alloc_keys); - } - - /** - * For a requested number of elements `new_size`, this function first - * determines whether `new_size < hash_map::table.capacity * RESIZE_THRESHOLD`. - * If not, the capacities of the underlying hash_map::table and - * hash_map::values is increased by RESIZE_FACTOR until the condition is - * satisfied. This is useful to ensure the hashtable is sufficiently large - * when preparing to add new elements. - * \param alloc_keys a memory allocation function with prototype - * `void* alloc_keys(size_t count, size_t size)` that allocates space for - * `count` items, each with size `size`, and initializes them such that - * core::is_empty() returns `true` for each key. - * \returns `true` if the resize was successful, and `false` if there is insufficient memory. - */ - inline bool check_size(unsigned int new_size, alloc_keys_func alloc_keys = calloc) - { - while (new_size >= table.capacity * RESIZE_THRESHOLD) { - if (!resize(RESIZE_FACTOR * table.capacity, alloc_keys)) { - fprintf(stderr, "hash_map.put ERROR: Unable to resize hashtable.\n"); - return false; - } - } - return true; - } - - /** - * Adds the given key-value pair to this map. The assignment operator is - * used insert the entry, and so this function should not be used if `K` - * and `V` are not [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). - * In such a case, insertion should be performed manually by using - * hash_map::get or hash_set::index_to_insert to find the appropriate index - * and directly modifying hash_map::table.keys, hash_map::table.size, and - * hash_map::values. See the example in the hash_map description. - * \param alloc_keys a memory allocation function with prototype - * `void* alloc_keys(size_t count, size_t size)` that allocates space for - * `count` items, each with size `size`, and initializes them such that - * core::is_empty() returns `true` for each key. - * \tparam K is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). - * \tparam V is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). - */ - bool put(const K& key, const V& value, - alloc_keys_func alloc_keys = calloc) - { - if (!check_size(table.size, alloc_keys)) return false; - insert(key, value); - return true; - } - - /** - * Adds the given key-value pairs in `elements` to this map. The assignment - * operator is used insert each entry, and so this function should not be - * used if `K` and `V` are not [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). - * In such a case, insertion should be performed manually by using - * hash_map::get or hash_map::index_to_insert to find the appropriate index - * and directly modifying hash_map::table.keys, hash_map::table.size, and - * hash_map::values. See the example in the hash_map description. - * \param alloc_keys a memory allocation function with prototype - * `void* alloc_keys(size_t count, size_t size)` that allocates space for - * `count` items, each with size `size`, and initializes them such that - * core::is_empty() returns `true` for each key. - * \tparam K is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). - * \tparam V is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). - */ - bool put_all(const hash_map& elements, - alloc_keys_func alloc_keys = calloc) - { - if (!check_size(table.size + elements.table.size, alloc_keys)) - return false; - for (unsigned int i = 0; i < elements.table.capacity; i++) - if (!hasher::is_empty(elements.table.keys[i])) - insert(elements.table.keys[i], elements.values[i]); - return true; - } - - /** - * This function removes `key` and associated value from the map. This - * function does not free the removed element. - * \returns `true` if the element is removed, and `false` if the set does not contain `element`. - */ - inline bool remove(const K& key) - { - return table.remove(key, values); - } - - /** - * This function removes the entry at the bucket given by `index`. This - * function assumes that an entry is located at the given bucket with the - * correct provided hash value. This function does not free the removed key or - * value. - */ - inline void remove_at(unsigned int index) - { - table.remove_at(values, index); - } - - /** - * Retrieves the value associated with the given `key`. This function - * assumes the given key exists in the map. - */ - V& get(const K& key) const - { - return values[table.index_of(key)]; - } - - /** - * If the given `key` exists in this map, this function returns the value - * associated with the key, and sets `contains` to `true`. If `key` is not - * in the map, `contains` is set to `false`. - */ - inline V& get(const K& key, bool& contains) const - { - unsigned int index; - return get(key, contains, index); - } - - /** - * If the given `key` exists in this map, this function returns the value - * associated with the key, sets `contains` to `true`, and sets `index` to - * the bucket containing the key. If `key` is not in the map, `contains` is - * set to `false`, and `index` is set to the index where the key would be - * located, if it had existed in the map. - */ - V& get(const K& key, bool& contains, unsigned int& index) const - { - index = table.index_of(key, contains); - return values[index]; - } - - /** - * If the given `key` exists in this map, this function returns the value - * associated with the key, sets `contains` to `true`, and sets `index` to - * the bucket containing the key. If `key` is not in the map, `contains` is - * set to `false`, and `index` is set to the index where the key would be - * located, if it had existed in the map. In any case, the evaluated hash - * function of `key` is stored in `hash_value`. - */ - V& get(const K& key, bool& contains, - unsigned int& index, - unsigned int& hash_value) const - { - index = table.index_of(key, contains, hash_value); - return values[index]; - } - - /** - * Removes all entries from this hash_map. Note that this function does - * not free each entry beforehand. - */ - inline void clear() { - table.clear(); - } - - /** - * Returns a hash_map_iterator pointing to the first entry in this container. - * - * **NOTE:** Unlike the libstdc++ [unordered_set](http://en.cppreference.com/w/cpp/container/unordered_set) - * and [unordered_map](http://en.cppreference.com/w/cpp/container/unordered_map) - * iterators, we do not keep a linked list among the elements, so if rapid - * iteration over elements is more critical than rapid queries, consider - * using an array_map. - */ - inline hash_map_iterator begin() { - return hash_map_iterator(*this, table.first_empty()); - } - - /** - * Returns a hash_map_iterator pointing to the end of this container. - * - * **NOTE:** Unlike the libstdc++ [unordered_set](http://en.cppreference.com/w/cpp/container/unordered_set) - * and [unordered_map](http://en.cppreference.com/w/cpp/container/unordered_map) - * iterators, we do not keep a linked list among the elements, so if rapid - * iteration over elements is more critical than rapid queries, consider - * using an array_map. - */ - inline hash_map_iterator end() { - return hash_map_iterator(*this, table.capacity); - } - - /** - * Returns a const hash_map_iterator pointing to the first entry in this container. - * - * **NOTE:** Unlike the libstdc++ [unordered_set](http://en.cppreference.com/w/cpp/container/unordered_set) - * and [unordered_map](http://en.cppreference.com/w/cpp/container/unordered_map) - * iterators, we do not keep a linked list among the elements, so if rapid - * iteration over elements is more critical than rapid queries, consider - * using an array_map. - */ - inline hash_map_iterator begin() const { - return hash_map_iterator(*this, table.first_empty()); - } - - /** - * Returns a const hash_map_iterator pointing to the end of this container. - * - * **NOTE:** Unlike the libstdc++ [unordered_set](http://en.cppreference.com/w/cpp/container/unordered_set) - * and [unordered_map](http://en.cppreference.com/w/cpp/container/unordered_map) - * iterators, we do not keep a linked list among the elements, so if rapid - * iteration over elements is more critical than rapid queries, consider - * using an array_map. - */ - inline hash_map_iterator end() const { - return hash_map_iterator(*this, table.capacity); - } - - /** - * Swaps the contents of the hash_map `first` with that of `second`. - */ - static inline void swap(hash_map& first, hash_map& second) { - hash_set::swap(first.table, second.table); - core::swap(first.values, second.values); - } - - /** - * Moves the contents of the hash_map `src` into `dst`. Note this function - * does not copy the contents of the underlying hash_map::table or - * hash_map::values, it merely copies the pointers. - */ - static inline void move(const hash_map& src, hash_map& dst) { - hash_set::move(src.table, dst.table); - dst.values = src.values; - } - - /** - * Copies the contents of the hash_map `src` into `dst`. - * \param alloc_keys a memory allocation function with prototype - * `void* alloc_keys(size_t count, size_t size)` that allocates space for - * `count` items, each with size `size`, and initializes them such that - * core::is_empty() returns `true` for each key. - */ - static inline bool copy(const hash_map& src, hash_map& dst, alloc_keys_func alloc_keys = calloc) { - dst.table.capacity = src.table.capacity; - dst.table.size = src.table.size; - dst.table.keys = (K*) alloc_keys(src.table.capacity, sizeof(K)); - if (dst.table.keys == NULL) return false; - dst.values = (V*) malloc(sizeof(V) * src.table.capacity); - if (dst.values == NULL) { - core::free(dst.table.keys); - return false; - } - - for (unsigned int i = 0; i < src.table.capacity; i++) { - if (is_empty(src.table.keys[i])) continue; - if (!core::copy(src.table.keys[i], dst.table.keys[i])) { - free(dst); return false; - } else if (!core::copy(src.values[i], dst.values[i])) { - core::free(dst.table.keys[i]); - set_empty(dst.table.keys[i]); - free(dst); return false; - } - } - return true; - } - - template - static inline long unsigned int size_of(const hash_map& map, - const key_value_metric& metric) - { - long unsigned int sum = core::size_of(map.table.capacity) + core::size_of(map.table.size); - for (unsigned int i = 0; i < map.table.capacity; i++) { - if (is_empty(map.table.keys[i])) - sum += sizeof(K) + sizeof(V); - else sum += core::size_of(map.table.keys[i], metric.key_metric) + core::size_of(map.values[i], metric.value_metric); - } - return sum; - } - - /** - * Frees hash_map::table and hash_map::values. This should not be used if - * `map` was constructed on the stack, as the destructor will automatically - * free hash_map::table and hash_map::values. The existing entries of `map` - * are not freed. - */ - static inline void free(hash_map& map) { - core::free(map.table); - core::free(map.values); - } - -private: - /* NOTE: this function assumes table is initialized */ - bool initialize_values() { - values = (V*) malloc(sizeof(V) * table.capacity); - if (values == NULL) { - core::free(table); - return false; - } - return true; - } - - bool initialize( - unsigned int initial_capacity, - alloc_keys_func alloc_keys) - { - if (!table.initialize(initial_capacity, alloc_keys)) - return false; - return initialize_values(); - } - - inline void place(const K& key, const V& value, unsigned int index) - { - table.place(key, index); - values[index] = value; - } - - inline void insert_unique( - const K& key, const V& value) - { - place(key, value, table.next_empty(key)); - } - - inline void insert(const K& key, const V& value) - { - bool contains; - unsigned int index = table.index_to_insert(key, contains); - if (!contains) { - place(key, value, index); - table.size++; - } - } - - template - friend bool hash_map_init(hash_map& map, - unsigned int capacity, alloc_keys_func alloc_keys); -}; - -/** - * Initializes the hash_map `map` with the given initial `capacity`. - * \param alloc_keys a memory allocation function with prototype - * `void* alloc_keys(size_t count, size_t size)` that allocates space for - * `count` items, each with size `size`, and initializes them such that - * core::is_empty() returns `true` for each key. - */ -template -bool hash_map_init( - hash_map& map, - unsigned int capacity, - alloc_keys_func alloc_keys = calloc) -{ - if (!map.initialize(capacity, alloc_keys)) { - fprintf(stderr, "hash_map_init ERROR: Unable to allocate memory.\n"); - return false; - } - return true; -} - -/** - * Swaps the underlying buffers of the given hash_maps `first` and `second`. - */ -template -void swap(hash_map& first, hash_map& second) { - K* keys_swap = first.table.keys; - first.table.keys = second.table.keys; - second.table.keys = keys_swap; - - V* values_swap = first.values; - first.values = second.values; - second.values = values_swap; - - unsigned int ui_swap = first.table.size; - first.table.size = second.table.size; - second.table.size = ui_swap; - - ui_swap = first.table.capacity; - first.table.capacity = second.table.capacity; - second.table.capacity = ui_swap; -} - -/** - * A sequentially-ordered associative container that contains a set of - * key-value pairs, where the keys are unique and have type `K` and the values - * have type `V`. The keys are stored sequentially in a native array - * array_map::keys, and the values are stored sequentially in a parallel native - * array array_map::values. - * - * When inserting a key-value pair into an array_map, a linear search is - * performed to determine if the key already exists in the map. If so, its - * corresponding value is replaced by the new value. If not, the new key and - * value are inserted at the ends of their respective arrays, and - * array_map::size is incremented. - * - * **Performance:** This data structure provides linear-time access and - * modification. However, due to locality of reference, this data structure - * may perform operations more quickly than hash_map if the size of the map - * is small. - * - * A simple example of the use of array_map is given below, where the expected - * output is `a.get(3): c, a.get(4): d, 1:a 4:d 3:c `. - * - * ```{.cpp} - * #include - * #include - * using namespace core; - * - * int main() { - * array_map a = array_map(8); - * a.put(1, 'a'); a.put(2, 'b'); - * a.put(3, 'c'); a.put(4, 'd'); - * a.remove(2); - * - * printf("a.get(3): %c, ", a.get(3)); - * printf("a.get(4): %c, ", a.get(4)); - * for (const auto& entry : a) - * printf("%d:%c ", entry.key, entry.value); - * } - * ``` - * - * - * However, if `a` is not allocated on the stack, the destructor will not be - * automatically called, and so it must be freed using `core::free` or - * `hash_map::free`. In the example below, the expected output is the same as - * that of the program above: `a.get(3): c, a.get(4): d, 1:a 4:d 3:c `. - * - * ```{.cpp} - * #include - * #include - * using namespace core; - * - * int main() { - * array_map& a = *((array_map*) alloca(sizeof(array_map))); - * array_map_init(a, 8); - * a.put(1, 'a'); a.put(2, 'b'); - * a.put(3, 'c'); a.put(4, 'd'); - * a.remove(2); - * - * printf("a.get(3): %c, ", a.get(3)); - * printf("a.get(4): %c, ", a.get(4)); - * for (const auto& entry : a) - * printf("%d:%c ", entry.key, entry.value); - * free(a); - * } - * ``` - * - * - * A number of member functions require `K` and `V` to be - * [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). - * If this is not the case, those operations can be performed directly on the - * public fields, as in the following example. The expected output is - * `first:1 second:2 `. - * - * ```{.cpp} - * #include - * #include - * #include - * using namespace core; - * - * struct custom_string { - * char* buffer; - * - * static void move(const custom_string& src, custom_string& dst) { - * dst.buffer = src.buffer; - * } - * - * static void free(custom_string& s) { - * core::free(s.buffer); - * } - * }; - * - * inline bool operator == (const custom_string& first, const custom_string& second) { - * if (first.buffer == NULL) - * return second.buffer == NULL; - * return strcmp(first.buffer, second.buffer) == 0; - * } - * - * bool init(custom_string& s, const char* src) { - * s.buffer = (char*) malloc(sizeof(char) * (strlen(src) + 1)); - * if (s.buffer == NULL) - * return false; - * memcpy(s.buffer, src, sizeof(char) * (strlen(src) + 1)); - * return true; - * } - * - * int main() { - * custom_string key_one, key_two; - * int value_one = 1, value_two = 2; - * init(key_one, "first"); - * init(key_two, "second"); - * - * array_map a = array_map(8); - * a.ensure_capacity(2); - * - * unsigned int index; - * a.get(key_one, index); - * if (index == a.size) { - * core::move(key_one, a.keys[index]); - * a.values[index] = value_one; - * a.size++; - * } - * - * a.get(key_two, index); - * if (index == a.size) { - * core::move(key_two, a.keys[index]); - * a.values[index] = value_two; - * a.size++; - * } - * - * for (const auto& entry : a) - * printf("%s:%d ", entry.key.buffer, entry.value); - * for (auto entry : a) - * free(entry.key); - * } - * ``` - */ -template -struct array_map { - /** - * The type of the keys in this array_map. - */ - typedef K key_type; - - /** - * The type of the values in this array_map. - */ - typedef V value_type; - - /** - * The native array of keys. - */ - K* keys; - - /** - * The native array of values parallel to array_map::keys. - */ - V* values; - - /** - * The capacity of array_map::keys and array_map::values. - */ - size_t capacity; - - /** - * The number of entries in the array_map. - */ - size_t size; - - /** - * Constructs the hash_map with the given `initial_capacity`. - */ - array_map(size_t initial_capacity) : size(0) { - if (!initialize(initial_capacity)) { - fprintf(stderr, "array_map ERROR: Error during initialization.\n"); - exit(EXIT_FAILURE); - } - } - - ~array_map() { free(); } - - /** - * Given the requested number of elements `new_length`, this function - * determines whether array_map::capacity is sufficient. If not, it - * attempts to increase its capacity by factors of RESIZE_FACTOR. - * \returns `true` if the resize was successful, and `false` if there is insufficient memory. - */ - bool ensure_capacity(size_t new_length) { - if (new_length <= capacity) - return true; - - size_t new_capacity = capacity; - if (!expand(keys, new_capacity, new_length)) - return false; - if (!resize(values, new_capacity)) - return false; - capacity = new_capacity; - return true; - } - - /** - * Adds the given key-value pair to this map. The assignment operator is - * used insert the entry, and so this function should not be used if `K` - * and `V` are not [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). - * In such a case, insertion should be performed manually by using - * array_map::get to compute the appropriate index and directly modifying - * the the public fields. See the example in the description of array_map. - * \tparam K is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). - * \tparam V is [CopyAssignable](https://en.cppreference.com/w/cpp/named_req/CopyAssignable). - */ - bool put(const K& key, const V& value) { - size_t index = index_of(key); - if (index < size) { - values[index] = value; - return true; - } - - if (!ensure_capacity(size + 1)) - return false; - keys[size] = key; - values[size] = value; - size++; - return true; - } - - /** - * Performs a linear search to find the index of the given `key`. If the - * `key` is not in this map, array_map::size is returned. - */ - template - inline size_t index_of(const Key& key) const { - return core::index_of(key, keys, size); - } - - /** - * Performs a linear search to find the index of the given `key`, with the - * search beginning at the index `start`. If the `key` is not in this map, - * array_map::size is returned. - */ - template - inline unsigned int index_of(const Key& key, unsigned int start) const { - for (unsigned int i = start; i < size; i++) - if (keys[i] == key) - return i; - return (unsigned int) size; - } - - /** - * Performs a reverse linear search to find the index of the given `key`. - * If the `key` is not in this map, `static_cast(-1)` is - * returned. - */ - template - inline unsigned int last_index_of(const Key& key) const { - return core::last_index_of(key, keys, size); - } - - /** - * Returns `true` if the given `key` exists in the map, and `false` otherwise. - */ - inline bool contains(const K& key) const { - return index_of(key) < size; - } - - /** - * Retrieves the value associated with the given `key`. This function - * assumes the given key exists in the map. - */ - template - inline V& get(const Key& key) { - return values[index_of(key)]; - } - - /** - * Retrieves the const value associated with the given `key`. This function - * assumes the given key exists in the map. - */ - template - inline const V& get(const Key& key) const { - return values[index_of(key)]; - } - - /** - * If the given `key` exists in the map, this function returns the value - * associated with the key, and sets `index` to the index of - * array_map::keys where the key is located. If `key` does not exist in the - * map, `index` is set to array_map::size. - */ - template - inline V& get(const Key& key, unsigned int& index) { - index = index_of(key); - return values[index]; - } - - /** - * If the given `key` exists in the map, this function returns the value - * associated with the key, and sets `contains` to `true`. If `key` does - * not exist in the map, `contains` is set to `false`. - */ - template - inline V& get(const Key& key, bool& contains) { - size_t index = index_of(key); - contains = (index != size); - return values[index]; - } - - /** - * If the given `key` exists in the map, this function returns the const - * value associated with the key, and sets `contains` to `true`. If `key` - * does not exist in the map, `contains` is set to `false`. - */ - template - inline const V& get(const Key& key, bool& contains) const { - unsigned int index = index_of(key); - contains = (index != size); - return values[index]; - } - - /** - * This function removes `key` and associated value from the map. This - * function does not free the removed element. - * \returns `true` if the element is removed, and `false` if the set does not contain `element`. - */ - bool remove(const K& key) { - size_t index = index_of(key); - if (index == size) - return false; - remove_at(index); - return true; - } - - /** - * This function removes the key-value pair located at the given `index` in - * the map. This function does not free the removed element. This function - * assumes `0 <= index < array_map::size`. - * \tparam K satisfies is_moveable. - * \tparam V satisfies is_moveable. - */ - inline void remove_at(size_t index) { - size--; - if (index == size) - return; - - /* move the last item into the position of the removed item */ - core::move(keys[size], keys[index]); - core::move(values[size], values[index]); - } - - /** - * Removes all entries from this array_map. Note that this function does - * not free each entry beforehand. - */ - inline void clear() { - size = 0; - } - - /** - * Returns an array_map_iterator pointing to the first entry in this container. - */ - inline array_map_iterator begin() { - return array_map_iterator(*this, 0); - } - - /** - * Returns an array_map_iterator pointing to the end of this container. - */ - inline array_map_iterator end() { - return array_map_iterator(*this, size); - } - - /** - * Returns a const array_map_iterator pointing to the first entry in this container. - */ - inline array_map_iterator begin() const { - return array_map_iterator(*this, 0); - } - - /** - * Returns a const array_map_iterator pointing to the end of this container. - */ - inline array_map_iterator end() const { - return array_map_iterator(*this, size); - } - - /** - * Swaps the contents of the array_map `first` with that of `second`. - */ - static inline void swap(array_map& first, array_map& second) { - core::swap(first.keys, second.keys); - core::swap(first.values, second.values); - core::swap(first.capacity, second.capacity); - core::swap(first.size, second.size); - } - - /** - * Moves the contents of the array_map `src` into `dst`. Note this function - * does not copy the contents of the underlying array_map::keys or - * array_map::values, it merely copies the pointers. - */ - static inline void move(const array_map& src, array_map& dst) { - dst.keys = src.keys; - dst.values = src.values; - dst.capacity = src.capacity; - dst.size = src.size; - } - - template - static inline long unsigned int size_of(const array_map& map, - const key_value_metric& metric) - { - long unsigned int sum = core::size_of(map.capacity) + core::size_of(map.size); - for (unsigned int i = 0; i < map.size; i++) - sum += core::size_of(map.keys[i], metric.key_metric) + core::size_of(map.values[i], metric.value_metric); - return sum + (map.capacity - map.size) * (sizeof(K) + sizeof(V)); - } - - static inline long unsigned int size_of(const array_map& map, const default_metric& metric) { - return size_of(map, make_key_value_metric(default_metric(), default_metric())); - } - - /** - * Frees array_map::keys and array_map::values. This should not be used if - * `map` was constructed on the stack, as the destructor will automatically - * free array_map::keys and array_map::values. The existing entries of - * `map` are not freed. - */ - static inline void free(array_map& map) { map.free(); } - -private: - inline bool initialize(size_t initial_capacity) { - capacity = initial_capacity; - keys = (K*) malloc(sizeof(K) * capacity); - if (keys == NULL) { - fprintf(stderr, "array_map.initialize ERROR: Out of memory.\n"); - return false; - } - values = (V*) malloc(sizeof(V) * capacity); - if (values == NULL) { - core::free(keys); - fprintf(stderr, "array_map.initialize ERROR: Out of memory.\n"); - return false; - } - return true; - } - - inline void free() { - core::free(keys); - core::free(values); - } - - template - friend bool array_map_init(array_map&, size_t); -}; - -/** - * Initializes the array_map `map` with the given `initial_capacity`. - */ -template -bool array_map_init(array_map& map, size_t initial_capacity) { - map.size = 0; - return map.initialize(initial_capacity); -} - -/** - * Returns the number of elements in the given hash_set `set`. - */ -template -inline unsigned int size(const hash_set& set) { - return set.size; -} - -/** - * Returns the number of elements in the given hash_map `map`. - */ -template -inline unsigned int size(const hash_map& map) { - return map.table.size; -} - -/** - * Returns the number of elements in the given array_map `map`. - */ -template -inline unsigned int size(const array_map& map) { - return map.size; -} - -/** - * Assuming the given `map` has value type that satisfies - * [is_integral](http://en.cppreference.com/w/cpp/types/is_integral), - * and the values in the map are unique and no larger than `size(map)`, - * this function returns an array of pointers to the keys in `map`, where the - * pointer at index `i` references the key in `map` that corresponds to the - * value `i`. The caller of this function is responsible for the memory of the - * returned native array, and must call free to release it. - * \tparam MapType a map type that allows range-based for iteration, and for - * which the function `unsigned int size(const MapType&)` is defined. - */ -template -inline const typename MapType::key_type** invert(const MapType& map) { - const typename MapType::key_type** inverse = - (const typename MapType::key_type**) calloc(size(map) + 1, sizeof(typename MapType::key_type*)); - if (inverse == NULL) { - fprintf(stderr, "invert ERROR: Unable to invert map. Out of memory.\n"); - return NULL; - } - for (const auto& entry : map) - inverse[entry.value] = &entry.key; - return inverse; -} - - -inline bool hash_map_test(void) -{ - hash_map map = hash_map(4); - bool contains_four = true, contains_minus_seven = true; - map.get(4, contains_four); - map.get(-7, contains_minus_seven); - if (contains_four || contains_minus_seven) { - fprintf(stderr, "hash_map_test ERROR: Map with no inserted elements should be empty.\n"); - return false; - } - - /* test insertion and retrieval operations */ - map.put(-7, "negative seven"); - map.put(4, "four"); - - if (map.table.size != 2) { - fprintf(stderr, "hash_map_test ERROR: Map size is %d after adding two elements.\n", map.table.size); - return false; - } - if (strcmp(map.get(-7), "negative seven") - || strcmp(map.get(4), "four")) - { - fprintf(stderr, "hash_map_test ERROR: Simple hashtable insertion failed.\n"); - return false; - } - - map.put(4, "new four"); - map.put(5, "five"); - map.put(12, "twelve"); - map.put(7, "seven"); - map.put(13, "thirteen"); - - /* test automatic capacity resizing */ - if (map.table.size != 6) { - fprintf(stderr, "hash_map_test ERROR: Map size is %d after adding six elements.\n", map.table.size); - return false; - } - if (map.table.capacity != 4 * RESIZE_FACTOR * RESIZE_FACTOR) { - fprintf(stderr, "hash_map_test ERROR: Unexpected hashtable capacity.\n"); - return false; - } - if (strcmp(map.get(4), "new four") - || strcmp(map.get(12), "twelve") - || strcmp(map.get(7), "seven") - || strcmp(map.get(-7), "negative seven")) - { - fprintf(stderr, "hash_map_test ERROR: Additional hashtable insertion failed.\n"); - return false; - } - - /* test removal operation */ - if (!map.remove(4)) { - fprintf(stderr, "hash_map_test ERROR: Removal of key '4' failed.\n"); - return false; - } - if (strcmp(map.get(12), "twelve") - || strcmp(map.get(13), "thirteen") - || strcmp(map.get(5), "five") - || strcmp(map.get(-7), "negative seven") - || strcmp(map.get(7), "seven")) { - fprintf(stderr, "hash_map_test ERROR: Hashtable lookup failed after removal.\n"); - return false; - } - - bool contains = true; - map.get(4, contains); - if (contains) { - fprintf(stderr, "hash_map_test ERROR: Retrieval of removed key failed.\n"); - return false; - } - - return true; -} - -inline bool array_map_test(void) -{ - array_map map = array_map(4); - bool contains_four = true, contains_minus_seven = true; - map.get(4, contains_four); - map.get(-7, contains_minus_seven); - if (contains_four || contains_minus_seven) { - fprintf(stderr, "array_map_test ERROR: Map with no inserted elements should be empty.\n"); - return false; - } - - /* test insertion and retrieval operations */ - map.put(-7, "negative seven"); - map.put(4, "four"); - - if (map.size != 2) { - fprintf(stderr, "array_map_test ERROR: Map size is %zu after adding two elements.\n", map.size); - return false; - } - if (strcmp(map.get(-7), "negative seven") - || strcmp(map.get(4), "four")) - { - fprintf(stderr, "array_map_test ERROR: Simple hashtable insertion failed.\n"); - return false; - } - - map.put(4, "new four"); - map.put(5, "five"); - map.put(12, "twelve"); - map.put(7, "seven"); - map.put(13, "thirteen"); - - if (map.size != 6) { - fprintf(stderr, "array_map_test ERROR: Map size is %zu after adding six elements.\n", map.size); - return false; - } - if (strcmp(map.get(4), "new four") - || strcmp(map.get(12), "twelve") - || strcmp(map.get(7), "seven") - || strcmp(map.get(-7), "negative seven")) - { - fprintf(stderr, "array_map_test ERROR: Additional hashtable insertion failed.\n"); - return false; - } - - /* test removal operation */ - if (!map.remove(4)) { - fprintf(stderr, "array_map_test ERROR: Removal of key '4' failed.\n"); - return false; - } - if (strcmp(map.get(12), "twelve") - || strcmp(map.get(13), "thirteen") - || strcmp(map.get(5), "five") - || strcmp(map.get(-7), "negative seven") - || strcmp(map.get(7), "seven")) { - fprintf(stderr, "array_map_test ERROR: Hashtable lookup failed after removal.\n"); - return false; - } - - bool contains = true; - map.get(4, contains); - if (contains) { - fprintf(stderr, "array_map_test ERROR: Retrieval of removed key failed.\n"); - return false; - } - - return true; -} - -} /* namespace core */ - -#endif /* MAP_H_ */ diff --git a/reasoners/data_preprocess/logic/graph_dataset_gen/core/random.h b/reasoners/data_preprocess/logic/graph_dataset_gen/core/random.h deleted file mode 100644 index 6ba64e0ef..000000000 --- a/reasoners/data_preprocess/logic/graph_dataset_gen/core/random.h +++ /dev/null @@ -1,321 +0,0 @@ -/** - * \file random.h - * - * This file defines the basic pseudorandom number generation functionality - * used by the core library, as well as functions to sample from a number of - * built-in distributions such as categorical, discrete uniform, continuous - * uniform, Bernoulli, geometric, beta, gamma, and Dirichlet. - * - * - */ - -#ifndef RANDOM_H_ -#define RANDOM_H_ - -#include -#include - -#include "timer.h" -#include "array.h" -#include "io.h" - -namespace core { - - -#if defined(NDEBUG) -static unsigned int seed = (unsigned int) milliseconds(); -#else -static unsigned int seed = 0; -#endif - -static thread_local std::minstd_rand engine = std::minstd_rand(seed); - -/** - * Returns the initial random seed used by all core functions that require pseudorandom number generation. - */ -inline unsigned int get_seed() { - return seed; -} - -/** - * Sets the seed of the underlying pseudorandom number generator. - */ -inline void set_seed(unsigned int new_seed) { - engine.seed(new_seed); - seed = new_seed; -} - -/** - * Reads the state of the pseudorandom number generator from `in`. This is - * useful to persist the state of the pseudorandom number generator. - */ -template -inline bool read_random_state(Stream& in) -{ - size_t length; - if (!read(length, in)) return false; - char* state = (char*) alloca(sizeof(char) * length); - if (state == NULL || !read(state, in, (unsigned int) length)) - return false; - - std::stringstream buffer(std::string(state, length)); - buffer >> engine; - return true; -} - -/** - * Writes the state of the pseudorandom number generator to `out`. This is - * useful to persist the state of the pseudorandom number generator. - */ -template -inline bool write_random_state(Stream& out) { - std::stringstream buffer; - buffer << engine; - std::string data = buffer.str(); - return write(data.length(), out) && write(data.c_str(), out, (unsigned int) data.length()); -} - -/** - * Samples from a categorical distribution, where the unnormalized probability - * of returning the index `i` is given by `probability[i]`. This function - * normalizes and overwrites `probability` with its cumulative distribution - * function. - */ -template::value>::type* = nullptr> -unsigned int sample_categorical(V* probability, unsigned int length) -{ -#if !defined(NDEBUG) - if (length == 0) - fprintf(stderr, "sample_categorical WARNING: Specified length is zero.\n"); -#endif - - for (unsigned int j = 1; j < length; j++) - probability[j] += probability[j - 1]; - - /* select the new table assignment */ - V random = probability[length - 1] * ((V) engine() / engine.max()); - unsigned int selected_table = length - 1; - for (unsigned int j = 0; j < length; j++) { - if (random < probability[j]) { - selected_table = j; - break; - } - } - return selected_table; -} - -/** - * Returns the smallest index `i` such that `random < sum from j=0 to i-1 of probability[j]`. - * Thus, this function implements the inverse cumulative distribution function - * for the categorical distribution with unnormalized probabilities given by - * `probability`. - */ -template -inline unsigned int select_categorical( - const U* probability, V random, unsigned int length) -{ -#if !defined(NDEBUG) - if (length == 0) - fprintf(stderr, "select_categorical WARNING: Specified length is zero.\n"); -#endif - - V aggregator = (V) 0; - unsigned int selected_table = length - 1; - for (unsigned int j = 0; j < length; j++) { - aggregator += probability[j]; - if (random <= aggregator) { - selected_table = j; - break; - } - } - return selected_table; -} - -/** - * Samples from a categorical distribution, where the unnormalized probability - * of returning the index `i` is given by `probability[i]` and its sum is given - * by the floating-point `sum`. This function doesn't modify `probability`. - */ -template::value>::type* = nullptr> -unsigned int sample_categorical(const V* probability, V sum, unsigned int length) -{ -#if !defined(NDEBUG) - if (length == 0) - fprintf(stderr, "sample_categorical WARNING: Specified length is zero.\n"); -#endif - - /* select the new table assignment */ - V random = sum * ((V) engine() / engine.max()); - return select_categorical(probability, random, length); -} - -/** - * Samples from a categorical distribution, where the unnormalized probability - * of returning the index `i` is given by `probability[i]` and its sum is given - * by the unsigned integer `sum`. This function doesn't modify `probability`. - */ -unsigned int sample_categorical( - const unsigned int* probability, - unsigned int sum, unsigned int length) -{ -#if !defined(NDEBUG) - if (length == 0) - fprintf(stderr, "sample_categorical WARNING: Specified length is zero.\n"); -#endif - - /* select the new table assignment */ - unsigned int random = engine() % sum; - return select_categorical(probability, random, length); -} - - -/** - */ - -/* forward declarations */ - -template struct array; - -/** - * Returns a sample from the discrete uniform distribution over `{0, ..., n - 1}`. - */ -inline unsigned int sample_uniform(unsigned int n) { - return engine() % n; -} - -/** - * Returns a sample from the discrete uniform distribution over `elements`. - */ -template -inline const T& sample_uniform(const T* elements, unsigned int length) { - return elements[engine() % length]; -} - -/** - * Returns a sample from the discrete uniform distribution over `elements`. - */ -template -inline const T& sample_uniform(const T (&elements)[N]) { - return elements[engine() % N]; -} - -/** - * Returns a sample from the discrete uniform distribution over `elements`. - */ -template -inline const T& sample_uniform(const array& elements) { - return sample_uniform(elements.data, (unsigned int) elements.length); -} - -/** - * Returns a sample from the continuous uniform distribution over [0, 1]. - */ -template -inline V sample_uniform() { - return (V) engine() / engine.max(); -} - -/** - * Returns a sample from the Bernoulli distribution: with probability 0.5, - * `true` is returned, otherwise `false` is returned. - */ -template -inline bool sample_bernoulli(const V& p) { - return sample_uniform() < p; -} - -/** - * Returns a sample from the geometric distribution with success probability `p`. - */ -template -inline unsigned int sample_geometric(const V& p) { - auto geom = std::geometric_distribution(p); - return geom(engine); -} - -/** - * Returns a sample from the negative binomial distribution with number of - * failures `r` and success probability `p`. - */ -template -inline double log_probability_negative_binomial(unsigned int x, const V& r, const V& p) { - return lgamma(x + r) - lgamma(x + 1) - lgamma(r) + r*log(1 - p) + x*log(p); -} - -/** - * Returns a sample from the beta distribution with shape parameter `alpha` and - * scale parameter 1. This function assumes `alpha > 0`. - */ -template -inline V sample_beta(const V& alpha) { - static std::gamma_distribution first_gamma = std::gamma_distribution(1.0); - std::gamma_distribution second_gamma = std::gamma_distribution(alpha); - V first = first_gamma(engine); - V second = second_gamma(engine); - return first / (first + second); -} - -/** - * Returns a sample from the beta distribution with shape parameter `alpha` and - * scale parameter `beta`. This function assumes `alpha > 0` and `beta > 0`. - */ -template -inline V sample_beta(const V& alpha, const V& beta) { - std::gamma_distribution first_gamma = std::gamma_distribution(alpha, 1.0); - std::gamma_distribution second_gamma = std::gamma_distribution(beta, 1.0); - V first = first_gamma(engine); - V second = second_gamma(engine); - return first / (first + second); -} - -/** - * Returns a sample from the gamma distribution with shape parameter `alpha` - * and rate parameter `beta`. This function assumes `alpha > 0` and `beta > 0`. - */ -template -inline V sample_gamma(const V& alpha, const V& beta) { - std::gamma_distribution gamma = std::gamma_distribution(alpha, 1.0 / beta); - return gamma(engine); -} - -/** - * Returns the log probability of drawing the observation `x` from a gamma - * distribution with shape parameter `alpha` and rate parameter `beta`. This - * function assumes `x > 0`, `alpha > 0`, and `beta > 0`. - */ -template -inline V log_probability_gamma(const V& x, const V& alpha, const V& beta) { - return alpha * log(beta) - lgamma(alpha) + (alpha - 1) * log(x) - beta * x; -} - -/** - * Samples from the Dirichlet distribution with parameter `alpha` and dimension - * `length`. The sample is written to `dst`. This function assumes `dst` has - * size at least `length`, `alpha[i] > 0` for all `i = 0, ..., length - 1`. - */ -template -inline void sample_dirichlet(V* dst, const V* alpha, unsigned int length) { - V sum = 0.0; - for (unsigned int i = 0; i < length; i++) { - if (alpha[i] == 0.0) { - dst[i] = 0.0; - } else { - std::gamma_distribution gamma_dist(alpha[i], 1.0); - V value = gamma_dist(engine); - dst[i] = value; - sum += value; - } - } - - for (unsigned int i = 0; i < length; i++) - dst[i] /= sum; -} - - -} /* namespace core */ - -#endif /* RANDOM_H_ */ diff --git a/reasoners/data_preprocess/logic/graph_dataset_gen/core/timer.h b/reasoners/data_preprocess/logic/graph_dataset_gen/core/timer.h deleted file mode 100644 index 306f7fd34..000000000 --- a/reasoners/data_preprocess/logic/graph_dataset_gen/core/timer.h +++ /dev/null @@ -1,130 +0,0 @@ -/** - * \file timer.h - * - * - */ - -#ifndef TIMER_H_ -#define TIMER_H_ - -#if defined(_WIN32) /* on Windows */ -#define _WINSOCKAPI_ -#include -#undef max -#undef min -#elif defined(__APPLE__) /* on Mac */ -#include -#else /* on Linux */ -#include -#endif - - -namespace core { - - -#if defined(__APPLE__) -struct timebase_info { - mach_timebase_info_data_t info; - - timebase_info() { - info = {0, 0}; - mach_timebase_info(&info); - } -}; - -inline timebase_info& get_timebase() { - static timebase_info timebase; - return timebase; -} -#endif - - -/** - * A simple structure that measures time. - */ -struct timer -{ -#if defined(_WIN32) /* on Windows */ - ULONGLONG start_time; -#elif defined(__APPLE__) /* on Mac */ - uint64_t start_time; -#else /* on Linux */ - timespec start_time; -#endif - - /** - * Constructor that starts the timer. - */ - timer() { - start(); - } - - /** - * Starts the timer. - */ - inline void start() { -#if defined(_WIN32) /* on Windows */ - start_time = GetTickCount64(); -#elif defined(__APPLE__) /* on Mac */ - start_time = mach_absolute_time(); -#else /* on Linux */ - clock_gettime(CLOCK_MONOTONIC, &start_time); -#endif - } - - /** - * Returns the number of milliseconds elapsed since the timer was last started. - */ - inline unsigned long long milliseconds() { -#if defined(_WIN32) /* on Windows */ - ULONGLONG end_time = GetTickCount64(); - return (unsigned long long) (end_time - start_time); -#elif defined(__APPLE__) /* on Mac */ - uint64_t end_time = mach_absolute_time(); - return (end_time - start_time) * get_timebase().info.numer / get_timebase().info.denom / 1000000; -#else /* on Linux */ - timespec end_time; - clock_gettime(CLOCK_MONOTONIC, &end_time); - return (unsigned long long) (end_time.tv_sec - start_time.tv_sec) * 1000 + (end_time.tv_nsec - start_time.tv_nsec) / 1000000; -#endif - } - - /** - * Returns the number of nanoseconds elapsed since the timer was last started. - */ - inline double nanoseconds() { -#if defined(_WIN32) /* on Windows */ - ULONGLONG end_time = GetTickCount64(); - return (end_time - start_time) * 1.0e6; -#elif defined(__APPLE__) /* on Mac */ - uint64_t end_time = mach_absolute_time(); - return (double) (end_time - start_time) * get_timebase().info.numer / get_timebase().info.denom; -#else /* on Linux */ - timespec end_time; - clock_gettime(CLOCK_MONOTONIC, &end_time); - return (end_time.tv_sec - start_time.tv_sec) * 1.0e9 + (end_time.tv_nsec - start_time.tv_nsec); -#endif - } -}; - -/** - * On Windows and Mac, this returns the number of milliseconds elapsed since - * the system was started. On Linux, it returns the system's best estimate of - * the current time of day. - */ -inline unsigned long long milliseconds() { -#if defined(_WIN32) /* on Windows */ - return GetTickCount64(); -#elif defined(__APPLE__) /* on Mac */ - return mach_absolute_time() * get_timebase().info.numer / get_timebase().info.denom / 1000000; -#else /* on Linux */ - timespec time; - clock_gettime(CLOCK_REALTIME, &time); - return (unsigned long long) time.tv_sec * 1000 + time.tv_nsec / 1000000; -#endif -} - - -} /* namespace core */ - -#endif /* TIMER_H_ */ diff --git a/reasoners/data_preprocess/logic/graph_dataset_gen/core/utility.h b/reasoners/data_preprocess/logic/graph_dataset_gen/core/utility.h deleted file mode 100644 index 25bfd99d6..000000000 --- a/reasoners/data_preprocess/logic/graph_dataset_gen/core/utility.h +++ /dev/null @@ -1,799 +0,0 @@ -/** - * \file utility.h - * - * This file contains a number of useful miscellaneous definitions and - * implementations, such as a core::string structure, a binary logarithm - * function, and a function to access the directory structure in the - * filesystem. - * - * - */ - -#ifndef UTILITY_H_ -#define UTILITY_H_ - -#include -#include "io.h" - -#if defined(_WIN32) -#if !defined(UNICODE) -#error "Unicode support required. Please compile with UNICODE preprocessor definition." -#endif -#define _WINSOCKAPI_ -#include -#include -#undef max -#undef min -#else -#include -#include -#endif - - -namespace core { - -/* a useful type trait for detecting whether the function 'print_special_string' is defined */ -namespace detail { - template static auto test_print_special_string(int32_t) -> - decltype(bool(print_special_string(0u, std::declval(), std::declval()...)), std::true_type{}); - template static auto test_print_special_string(int64_t) -> std::false_type; -} - -template struct has_print_special_string : decltype(core::detail::test_print_special_string(0)){}; - - -/** - * A basic string structure, containing a native array of `char` elements and an `unsigned int length`. - */ -struct string { - /** - * The length of the string in characters. - */ - unsigned int length; - - /** - * The native `char` array containing the string data. - */ - char* data; - - /** - * A constructor that does not initialize any fields. **WARNING:** The - * destructor will free string::data, which this constructor does not - * initialize. The user must initialize string::data using `init` or - * manually before the string is destroyed. - */ - string() { } - - /** - * Constructs the string by copying from the given null-terminated C string `src`. - */ - string(const char* src) { - if (!initialize(src, (unsigned int) strlen(src))) - exit(EXIT_FAILURE); - } - - /** - * Constructs the string by copying from the given native char array `src` - * with given `length`. - */ - string(const char* src, unsigned int length) { - if (!initialize(src, length)) - exit(EXIT_FAILURE); - } - - /** - * Constructs the string by initializing string::data with size `length` - * but without settings its contents. - */ - explicit string(unsigned int length) { - if (!initialize(length)) - exit(EXIT_FAILURE); - } - - ~string() { - core::free(data); - } - - /** - * Accesses the character at the given `index`. - */ - inline char& operator [] (unsigned int index) { - return data[index]; - } - - /** - * Accesses the character at the given `index`. - */ - inline const char& operator [] (unsigned int index) const { - return data[index]; - } - - /** - * Initializes this string by copying from `s`. Note that if this string - * was previously initialized, it is not freed. - */ - inline void operator = (const string& s) { - initialize(s.data, s.length); - } - - /** - * Appends the given null-terminated C string to this string. - */ - inline void operator += (const char* src) { - unsigned int src_length = (unsigned int) strlen(src); - char* new_data = (char*) realloc(data, sizeof(char) * (length + src_length)); - if (new_data == NULL) { - fprintf(stderr, "string.operator += ERROR: Unable to expand string.\n"); - exit(EXIT_FAILURE); - } - - data = new_data; - memcpy(data + length, src, sizeof(char) * src_length); - length += src_length; - } - - /** - * Appends the given string to this string. - */ - inline void operator += (const string& src) { - unsigned int src_length = src.length; - char* new_data = (char*) realloc(data, sizeof(char) * (length + src_length)); - if (new_data == NULL) { - fprintf(stderr, "string.operator += ERROR: Unable to expand string.\n"); - exit(EXIT_FAILURE); - } - - data = new_data; - memcpy(data + length, src.data, sizeof(char) * src_length); - length += src_length; - } - - /** - * Returns whether the current string precedes `other` in lexicographical order. - */ - inline bool operator < (const string& other) const { - for (unsigned int i = 0; ; i++) { - if (i == length) { - if (i == other.length) return false; - else return true; - } else if (i == other.length) { - return false; - } - - if (data[i] > other.data[i]) - return false; - else if (data[i] < other.data[i]) - return true; - } - } - - /** - * Returns the smallest index `i` such that `string::data[i] == c`. - */ - inline unsigned int index_of(char c) const { - for (unsigned int i = 0; i < length; i++) { - if (data[i] == c) - return i; - } - return length; - } - - /** - * Returns whether string::data is NULL. This enables strings to be used as keys in hashtables. - */ - static inline bool is_empty(const string& key) { - return key.data == NULL; - } - - /** - * Sets string::data to NULL. This enables strings to be used as keys in hashtables. - */ - static inline void set_empty(string& key) { - key.data = NULL; - } - - /** - * Sets string::data to NULL for every element in `keys`. This enables - * strings to be used as keys in hashtables. - */ - static inline void set_empty(string* keys, unsigned int length) { - memset(static_cast(keys), 0, sizeof(string) * length); - } - - /** - * Returns the hash of the given `key`. - */ - static inline unsigned int hash(const string& key) { - return default_hash(key.data, key.length); - } - - /** - * Copies the string::length and string::data pointer from `src` to `dst`. - * Note this function does not create a new pointer and copy the character - * contents, it simply copies the `char*` pointer. - */ - static inline void move(const string& src, string& dst) { - dst.length = src.length; - dst.data = src.data; - } - - /** - * Copies the string in `src` to `dst`. This function initializes a new - * string::data array in `dst` and copies the contents from `src.data`. - */ - static inline bool copy(const string& src, string& dst) { - return dst.initialize(src.data, src.length); - } - - /** - * Swaps the contents and lengths of `first` and `second`. - */ - static inline void swap(string& first, string& second) { - core::swap(first.length, second.length); - core::swap(first.data, second.data); - } - - template - static inline long unsigned int size_of(const string& str, const Metric& metric) { - return core::size_of(str.length) + sizeof(char) * str.length; - } - - /** - * Frees the underlying `char` array in `str`. - */ - static inline void free(string& str) { - core::free(str.data); - } - -private: - bool initialize(unsigned int src_length) { - length = src_length; - data = (char*) malloc(sizeof(char) * (length == 0 ? 1 : length)); - if (data == NULL) { - fprintf(stderr, "string.initialize ERROR: Unable to initialize string.\n"); - return false; - } - return true; - } - - bool initialize(const char* src, unsigned int src_length) { - if (!initialize(src_length)) - return false; - memcpy(data, src, sizeof(char) * length); - return true; - } - - friend bool init(string&, const char*, unsigned int); - friend bool init(string&, unsigned int); -}; - -/** - * Initializes the string `dst` with the given native `char` array `src` and the given `length`. - */ -inline bool init(string& dst, const char* src, unsigned int length) { - return dst.initialize(src, length); -} - -/** - * Initializes the string `dst` with the given string `src`. - */ -inline bool init(string& dst, const string& src) { - return init(dst, src.data, src.length); -} - -/** - * Initializes the string `dst` by allocating string::data with size `length`, - * but this function does not set its contents. - */ -inline bool init(string& dst, unsigned int length) { - return dst.initialize(length); -} - -/** - * Reads a string `s` from `in`. - * \param s an uninitialized string structure. This function initializes `s`, - * and the caller is responsible for its memory and must call free to - * release its memory resources. - */ -template -inline bool read(string& s, Stream& in) { - if (!read(s.length, in)) return false; - s.data = (char*) malloc(sizeof(char) * s.length); - if (s.data == NULL) - return false; - return read(s.data, in, s.length); -} - -/** - * Writes the string `s` to `out`. - */ -template -inline bool write(const string& s, Stream& out) { - if (!write(s.length, out)) return false; - return write(s.data, out, s.length); -} - -/** - * Prints the string `s` to `stream`. - */ -template -inline bool print(const string& s, Stream&& stream) { - return fwrite(s.data, sizeof(char), s.length, stream) == s.length; -} - -/** - * Compares the string `first` to the null-terminated C string `second` and - * returns `true` if they are equivalent, and `false` otherwise. - */ -inline bool operator == (const string& first, const char* second) { - for (unsigned int i = 0; i < first.length; i++) { - if (first[i] != second[i]) - return false; - } - if (second[first.length] != '\0') - return false; - return true; -} - -/** - * Compares the null-terminated C string `first` to the string `second` and - * returns `true` if they are equivalent, and `false` otherwise. - */ -inline bool operator == (const char* first, const string& second) { - return (second == first); -} - -/** - * Compares the string `first` to the string `second` and returns `true` if - * they are equivalent, and `false` otherwise. - */ -inline bool operator == (const string& first, const string& second) { - if (first.length != second.length) return false; - /* we are guaranteed that only the first may be uninitialized */ - if (first.data == NULL) - return false; - return memcmp(first.data, second.data, first.length * sizeof(char)) == 0; -} - -/** - * Compares the string `first` to the null-terminated C string `second` and - * returns `false` if they are equivalent, and `true` otherwise. - */ -inline bool operator != (const string& first, const char* second) { - return !(first == second); -} - -/** - * Compares the null-terminated C string `first` to the string `second` and - * returns `false` if they are equivalent, and `true` otherwise. - */ -inline bool operator != (const char* first, const string& second) { - return !(second == first); -} - -/** - * Compares the string `first` to the string `second` and returns `false` if - * they are equivalent, and `true` otherwise. - */ -inline bool operator != (const string& first, const string& second) { - if (first.length != second.length) return true; - if (first.data == NULL) { - if (second.data == NULL) return true; - else return false; - } else if (second.data == NULL) - return false; - return memcmp(first.data, second.data, first.length * sizeof(char)) != 0; -} - -struct sequence { - unsigned int* tokens; - unsigned int length; - - sequence(unsigned int* src, unsigned int length) : tokens(src), length(length) { } - - inline bool operator = (const sequence& src) { - return initialize(src); - } - - inline unsigned int& operator [] (unsigned int index) { - return tokens[index]; - } - - inline unsigned int operator [] (unsigned int index) const { - return tokens[index]; - } - - static inline unsigned int hash(const sequence& key) { - return default_hash(key.tokens, key.length); - } - - static inline bool is_empty(const sequence& key) { - return key.tokens == NULL; - } - - static inline void set_empty(sequence& key) { - key.tokens = NULL; - } - - static inline void move(const sequence& src, sequence& dst) { - dst.tokens = src.tokens; - dst.length = src.length; - } - - static inline bool copy(const sequence& src, sequence& dst) { - return dst.initialize(src); - } - - static inline void swap(sequence& first, sequence& second) { - core::swap(first.tokens, second.tokens); - core::swap(first.length, second.length); - } - - static inline void free(sequence& seq) { - core::free(seq.tokens); - } - -private: - inline bool initialize(const sequence& src) { - length = src.length; - tokens = (unsigned int*) malloc(sizeof(unsigned int) * length); - if (tokens == NULL) { - fprintf(stderr, "sequence.initialize ERROR: Out of memory.\n"); - return false; - } - memcpy(tokens, src.tokens, sizeof(unsigned int) * length); - return true; - } -}; - -inline bool init(sequence& seq, unsigned int length) { - seq.length = length; - seq.tokens = (unsigned int*) malloc(sizeof(unsigned int) * length); - if (seq.tokens == NULL) { - fprintf(stderr, "init ERROR: Insufficient memory for token array in sequence.\n"); - return false; - } - return true; -} - -inline bool init(sequence& seq, const sequence& src) { - return sequence::copy(src, seq); -} - -inline bool operator == (const sequence& first, const sequence& second) { - /* only the first argument can be uninitialized */ - if (first.tokens == NULL) return false; - if (first.length != second.length) return false; - for (unsigned int i = 0; i < first.length; i++) - if (first.tokens[i] != second.tokens[i]) return false; - return true; -} - -inline bool operator != (const sequence& first, const sequence& second) { - return !(first == second); -} - -inline bool operator < (const sequence& first, const sequence& second) { - if (first.length < second.length) return true; - else if (first.length > second.length) return false; - for (unsigned int i = 0; i < first.length; i++) { - if (first[i] < second[i]) return true; - else if (first[i] > second[i]) return false; - } - return false; -} - -template -inline bool read(sequence& item, Stream& in) { - if (!read(item.length, in)) return false; - item.tokens = (unsigned int*) malloc(sizeof(unsigned int) * item.length); - if (item.tokens == NULL) return false; - if (!read(item.tokens, in, item.length)) { - free(item.tokens); - return false; - } - return true; -} - -template -inline bool write(const sequence& item, Stream& out) { - return write(item.length, out) && write(item.tokens, out, item.length); -} - -template -inline bool print(const sequence& item, - Stream&& out, Printer&&... printer) -{ - if (item.length == 0) return true; - bool success = print(item[0], out, std::forward(printer)...); - for (unsigned int i = 1; i < item.length; i++) - success &= print(' ', out) && print(item[i], out, std::forward(printer)...); - return success; -} - -/** - * A scribe that maps unsigned integer indices to core::string pointers. - * - * ```{.cpp} - * #include - * using namespace core; - * - * int main() { - * string first = "first"; - * string second = "second"; - * - * string_map_scribe scribe; - * scribe.map = (const string**) malloc(sizeof(const string*) * 2); - * scribe.map[0] = &first; scribe.map[1] = &second; - * scribe.length = 2; - * - * print(0, stdout, scribe); print(' ', stdout); - * print(1, stdout, scribe); - * - * free(scribe.map); - * } - * ``` - * The expected output of this program is `first second`. - * - * Another way to construct this structure is to convert a - * hash_map into a core::string** array using the - * core::invert() function. - * - * ```{.cpp} - * #include - * using namespace core; - * - * int main() { - * hash_map token_map(16); - * token_map.put("first", 0); - * token_map.put("second", 1); - * - * string_map_scribe scribe; - * scribe.map = invert(token_map); - * scribe.length = 2; - * - * print(0, stdout, scribe); print(' ', stdout); - * print(1, stdout, scribe); - * - * free(scribe.map); - * for (auto entry : token_map) - * free(entry.key); - * } - * ``` - * The expected output of this program is `first second`. Notice that since - * hash_map does not automatically free its elements, we do it manually using a - * range-based for loop. - */ -struct string_map_scribe { - /** - * The native array of `const string*` elements that represents the - * effective map from unsigned integers to strings. - */ - const string** map; - - /** - * The length of the native array string_map_scribe::map. - */ - unsigned int length; -}; - -template::value>::type* = nullptr> -inline bool print_special_string_helper(unsigned int item, Stream& out, Printer&&... printer) { - return print_special_string(item, out, std::forward(printer)...); -} - -template::value>::type* = nullptr> -inline bool print_special_string_helper(unsigned int item, Stream& out, Printer&&... printer) { - fprintf(stderr, "print ERROR: The unsigned int %u exceeds the bounds of the " - "string_map_scribe. Did you forget to implement print_special_string?\n", item); - return true; -} - -/** - * Prints `item` to `out` using the string_map_scribe `printer`. If - * `item < printer.length`, the string at index `item` is accessed from - * string_map_scribe::map and printed to `out`. Otherwise, there are two cases: - * 1. If the function `bool print_special_string(unsigned int, Stream&)` is - * defined, this function calls it with arguments `item` and `out`. - * 2. If such a function is not defined, an error message is printed to - * [stderr](http://en.cppreference.com/w/cpp/io/c) and the function - * returns `true`. - * - * \param string_printer a scribe for which the function - * `bool print(const string&, Stream&, Printer&&...)` is defined, which is - * used to print the string itself. Note that since this is a variadic - * argument, it may be empty. - */ -template -inline bool print(unsigned int item, Stream&& out, const string_map_scribe& printer, Printer&&... string_printer) -{ - if (item < printer.length) - return print(*printer.map[item], out, std::forward(string_printer)...); - else return print_special_string_helper(item, out, std::forward(string_printer)...); -} - -/** - * Looks for the token `identifier` in the given hash_map `map`. If such a key - * exists in the map, `id` is set to its corresponding value and `true` is - * returned. If not, a new entry is added to the map with `identifier` as the - * key and `map.table.size + 1` as its value, and `id` is set to this new value. - * \returns `true` upon success, or `false` if the hash_map `map` could not be - * resized to accommodate the new entry. - */ -bool get_token(const string& identifier, unsigned int& id, hash_map& map) { - if (!map.check_size()) { - fprintf(stderr, "get_token ERROR: Unable to expand token map.\n"); - return false; - } - - bool contains; - unsigned int bucket; - unsigned int& value = map.get(identifier, contains, bucket); - if (!contains) { - map.table.keys[bucket] = identifier; - map.table.size++; - value = map.table.size; - } - id = value; - return true; -} - -/** - * Opens the file with the given `filename` with the given `mode` and returns - * either a [FILE](http://en.cppreference.com/w/c/io) pointer on success, or - * `NULL` on failure. [fclose](http://en.cppreference.com/w/c/io/fclose) should - * be used to close the stream once reading is complete. - * - * On Windows, this function uses - * [fopen_s](https://msdn.microsoft.com/en-us/library/z5hh6ee9.aspx) to open - * the stream. - */ -inline FILE* open_file(const char* filename, const char* mode) { -#if defined(_WIN32) - FILE* file; - if (fopen_s(&file, filename, mode) != 0) - return NULL; - return file; -#else - return fopen(filename, mode); -#endif -} - -/** - * Reads the contents of the file whose path is given by `filename`. If the - * file cannot be opened for reading, or if there is insufficient memory to - * allocate a buffer, `NULL` is returned. `bytes_read` is set to the number of - * bytes read, and the file contents are returned. The caller is responsible - * for the memory of the returned buffer and must call free to release its - * memory resources. - * \tparam AppendNull if `true`, a null terminating character is appended. - */ -template -inline char* read_file(const char* filename, size_t& bytes_read) -{ - FILE* fin = open_file(filename, "rb"); - if (fin == NULL || fseek(fin, 0, SEEK_END) != 0) { - if (fin != NULL) { -#if defined(_WIN32) - errno = (int) GetLastError(); -#endif - perror("read_file ERROR"); - fclose(fin); - } - return NULL; - } - - long int filesize = ftell(fin); - if (filesize == -1L) { - fprintf(stderr, "read_file ERROR: `ftell` returned error.\n"); - fclose(fin); - return NULL; - } - - if (fseek(fin, 0, SEEK_SET) != 0) { - fprintf(stderr, "read_file ERROR: `seek` returned error.\n"); - fclose(fin); - return NULL; - } - - char* data = (char*) malloc(sizeof(char) * (AppendNull ? (filesize + 1) : filesize)); - if (data == NULL) { - fprintf(stderr, "read_file ERROR: Out of memory.\n"); - fclose(fin); - return NULL; - } - bytes_read = fread(data, sizeof(char), filesize, fin); - fclose(fin); - if (AppendNull) - data[filesize] = '\0'; - return data; -} - -/** - * This function inspects the directory given by the path `directory` and adds - * the list of filenames in that directory to the string array `out`. - * \see http://stackoverflow.com/questions/306533/how-do-i-get-a-list-of-files-in-a-directory-in-c - */ -inline bool get_files_in_directory(array& out, const char* directory) -{ -#if defined(_WIN32) - HANDLE dir; - WIN32_FIND_DATA file_data; - - size_t required; - if (mbstowcs_s(&required, NULL, 0, directory, 0) != 0) - return false; - required--; /* ignore the null terminator */ - wchar_t* dir_prefix = (wchar_t*) malloc(sizeof(wchar_t) * (required + 3)); - if (dir_prefix == NULL) { - fprintf(stderr, "get_files_in_directory ERROR: Out of memory.\n"); - return false; - } - if (mbstowcs_s(&required, dir_prefix, required + 3, directory, required) != 0) { - free(dir_prefix); - return false; - } - required--; /* ignore the null terminator */ - dir_prefix[required] = '/'; - dir_prefix[required + 1] = '*'; - dir_prefix[required + 2] = '\0'; - if ((dir = FindFirstFile(dir_prefix, &file_data)) == INVALID_HANDLE_VALUE) { - free(dir_prefix); - return false; - } - free(dir_prefix); - - do { - const bool is_directory = (file_data.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY) != 0; - if (file_data.cFileName[0] == '.' || is_directory) - continue; - - if (!out.ensure_capacity(out.length + 1)) - return false; - - if (wcstombs_s(&required, NULL, 0, file_data.cFileName, 0) != 0 - || !init(out[(unsigned int) out.length], (unsigned int) required)) - return false; - out.length++; - if (wcstombs_s(&required, out.last().data, required, file_data.cFileName, wcslen(file_data.cFileName)) != 0) - return false; - out.last().length--; /* remove the null terminator */ - } while (FindNextFile(dir, &file_data)); - - FindClose(dir); - return true; -#else - DIR* dir; - dirent* ent; - struct stat st; - - dir = opendir(directory); - if (dir == NULL) - return false; - while ((ent = readdir(dir)) != NULL) { - string full_file_name = string(directory); - full_file_name += "/"; - full_file_name += ent->d_name; - full_file_name += " "; - full_file_name[full_file_name.length - 1] = '\0'; - if (ent->d_name[0] == '.' || stat(full_file_name.data, &st) == -1) - continue; - - const bool is_directory = (st.st_mode & S_IFDIR) != 0; - if (is_directory) - continue; - - if (!out.ensure_capacity(out.length + 1)) - return false; - init(out[out.length], ent->d_name); - out.length++; - } - closedir(dir); - return true; -#endif -} - -} /* namespace core */ - -#endif /* UTILITY_H_ */ diff --git a/reasoners/data_preprocess/logic/graph_dataset_gen/core/xxhash.h b/reasoners/data_preprocess/logic/graph_dataset_gen/core/xxhash.h deleted file mode 100644 index 08ab79457..000000000 --- a/reasoners/data_preprocess/logic/graph_dataset_gen/core/xxhash.h +++ /dev/null @@ -1,5580 +0,0 @@ -/* - * xxHash - Extremely Fast Hash algorithm - * Header File - * Copyright (C) 2012-2020 Yann Collet - * - * BSD 2-Clause License (https://www.opensource.org/licenses/bsd-license.php) - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are - * met: - * - * * Redistributions of source code must retain the above copyright - * notice, this list of conditions and the following disclaimer. - * * Redistributions in binary form must reproduce the above - * copyright notice, this list of conditions and the following disclaimer - * in the documentation and/or other materials provided with the - * distribution. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS - * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT - * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR - * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT - * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, - * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT - * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, - * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY - * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT - * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE - * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - * - * You can contact the author at: - * - xxHash homepage: https://www.xxhash.com - * - xxHash source repository: https://github.com/Cyan4973/xxHash - */ -/*! - * @mainpage xxHash - * - * @file xxhash.h - * xxHash prototypes and implementation - */ -/* TODO: update */ -/* Notice extracted from xxHash homepage: - -xxHash is an extremely fast hash algorithm, running at RAM speed limits. -It also successfully passes all tests from the SMHasher suite. - -Comparison (single thread, Windows Seven 32 bits, using SMHasher on a Core 2 Duo @3GHz) - -Name Speed Q.Score Author -xxHash 5.4 GB/s 10 -CrapWow 3.2 GB/s 2 Andrew -MurmurHash 3a 2.7 GB/s 10 Austin Appleby -SpookyHash 2.0 GB/s 10 Bob Jenkins -SBox 1.4 GB/s 9 Bret Mulvey -Lookup3 1.2 GB/s 9 Bob Jenkins -SuperFastHash 1.2 GB/s 1 Paul Hsieh -CityHash64 1.05 GB/s 10 Pike & Alakuijala -FNV 0.55 GB/s 5 Fowler, Noll, Vo -CRC32 0.43 GB/s 9 -MD5-32 0.33 GB/s 10 Ronald L. Rivest -SHA1-32 0.28 GB/s 10 - -Q.Score is a measure of quality of the hash function. -It depends on successfully passing SMHasher test set. -10 is a perfect score. - -Note: SMHasher's CRC32 implementation is not the fastest one. -Other speed-oriented implementations can be faster, -especially in combination with PCLMUL instruction: -https://fastcompression.blogspot.com/2019/03/presenting-xxh3.html?showComment=1552696407071#c3490092340461170735 - -A 64-bit version, named XXH64, is available since r35. -It offers much better speed, but for 64-bit applications only. -Name Speed on 64 bits Speed on 32 bits -XXH64 13.8 GB/s 1.9 GB/s -XXH32 6.8 GB/s 6.0 GB/s -*/ - -#if defined (__cplusplus) -extern "C" { -#endif - -/* **************************** - * INLINE mode - ******************************/ -/*! - * XXH_INLINE_ALL (and XXH_PRIVATE_API) - * Use these build macros to inline xxhash into the target unit. - * Inlining improves performance on small inputs, especially when the length is - * expressed as a compile-time constant: - * - * https://fastcompression.blogspot.com/2018/03/xxhash-for-small-keys-impressive-power.html - * - * It also keeps xxHash symbols private to the unit, so they are not exported. - * - * Usage: - * #define XXH_INLINE_ALL - * #include "xxhash.h" - * - * Do not compile and link xxhash.o as a separate object, as it is not useful. - */ -#if (defined(XXH_INLINE_ALL) || defined(XXH_PRIVATE_API)) \ - && !defined(XXH_INLINE_ALL_31684351384) - /* this section should be traversed only once */ -# define XXH_INLINE_ALL_31684351384 - /* give access to the advanced API, required to compile implementations */ -# undef XXH_STATIC_LINKING_ONLY /* avoid macro redef */ -# define XXH_STATIC_LINKING_ONLY - /* make all functions private */ -# undef XXH_PUBLIC_API -# if defined(__GNUC__) -# define XXH_PUBLIC_API static __inline __attribute__((unused)) -# elif defined (__cplusplus) || (defined (__STDC_VERSION__) && (__STDC_VERSION__ >= 199901L) /* C99 */) -# define XXH_PUBLIC_API static inline -# elif defined(_MSC_VER) -# define XXH_PUBLIC_API static __inline -# else - /* note: this version may generate warnings for unused static functions */ -# define XXH_PUBLIC_API static -# endif - - /* - * This part deals with the special case where a unit wants to inline xxHash, - * but "xxhash.h" has previously been included without XXH_INLINE_ALL, - * such as part of some previously included *.h header file. - * Without further action, the new include would just be ignored, - * and functions would effectively _not_ be inlined (silent failure). - * The following macros solve this situation by prefixing all inlined names, - * avoiding naming collision with previous inclusions. - */ - /* Before that, we unconditionally #undef all symbols, - * in case they were already defined with XXH_NAMESPACE. - * They will then be redefined for XXH_INLINE_ALL - */ -# undef XXH_versionNumber - /* XXH32 */ -# undef XXH32 -# undef XXH32_createState -# undef XXH32_freeState -# undef XXH32_reset -# undef XXH32_update -# undef XXH32_digest -# undef XXH32_copyState -# undef XXH32_canonicalFromHash -# undef XXH32_hashFromCanonical - /* XXH64 */ -# undef XXH64 -# undef XXH64_createState -# undef XXH64_freeState -# undef XXH64_reset -# undef XXH64_update -# undef XXH64_digest -# undef XXH64_copyState -# undef XXH64_canonicalFromHash -# undef XXH64_hashFromCanonical - /* XXH3_64bits */ -# undef XXH3_64bits -# undef XXH3_64bits_withSecret -# undef XXH3_64bits_withSeed -# undef XXH3_64bits_withSecretandSeed -# undef XXH3_createState -# undef XXH3_freeState -# undef XXH3_copyState -# undef XXH3_64bits_reset -# undef XXH3_64bits_reset_withSeed -# undef XXH3_64bits_reset_withSecret -# undef XXH3_64bits_update -# undef XXH3_64bits_digest -# undef XXH3_generateSecret - /* XXH3_128bits */ -# undef XXH128 -# undef XXH3_128bits -# undef XXH3_128bits_withSeed -# undef XXH3_128bits_withSecret -# undef XXH3_128bits_reset -# undef XXH3_128bits_reset_withSeed -# undef XXH3_128bits_reset_withSecret -# undef XXH3_128bits_reset_withSecretandSeed -# undef XXH3_128bits_update -# undef XXH3_128bits_digest -# undef XXH128_isEqual -# undef XXH128_cmp -# undef XXH128_canonicalFromHash -# undef XXH128_hashFromCanonical - /* Finally, free the namespace itself */ -# undef XXH_NAMESPACE - - /* employ the namespace for XXH_INLINE_ALL */ -# define XXH_NAMESPACE XXH_INLINE_ - /* - * Some identifiers (enums, type names) are not symbols, - * but they must nonetheless be renamed to avoid redeclaration. - * Alternative solution: do not redeclare them. - * However, this requires some #ifdefs, and has a more dispersed impact. - * Meanwhile, renaming can be achieved in a single place. - */ -# define XXH_IPREF(Id) XXH_NAMESPACE ## Id -# define XXH_OK XXH_IPREF(XXH_OK) -# define XXH_ERROR XXH_IPREF(XXH_ERROR) -# define XXH_errorcode XXH_IPREF(XXH_errorcode) -# define XXH32_canonical_t XXH_IPREF(XXH32_canonical_t) -# define XXH64_canonical_t XXH_IPREF(XXH64_canonical_t) -# define XXH128_canonical_t XXH_IPREF(XXH128_canonical_t) -# define XXH32_state_s XXH_IPREF(XXH32_state_s) -# define XXH32_state_t XXH_IPREF(XXH32_state_t) -# define XXH64_state_s XXH_IPREF(XXH64_state_s) -# define XXH64_state_t XXH_IPREF(XXH64_state_t) -# define XXH3_state_s XXH_IPREF(XXH3_state_s) -# define XXH3_state_t XXH_IPREF(XXH3_state_t) -# define XXH128_hash_t XXH_IPREF(XXH128_hash_t) - /* Ensure the header is parsed again, even if it was previously included */ -# undef XXHASH_H_5627135585666179 -# undef XXHASH_H_STATIC_13879238742 -#endif /* XXH_INLINE_ALL || XXH_PRIVATE_API */ - - - -/* **************************************************************** - * Stable API - *****************************************************************/ -#ifndef XXHASH_H_5627135585666179 -#define XXHASH_H_5627135585666179 1 - - -/*! - * @defgroup public Public API - * Contains details on the public xxHash functions. - * @{ - */ -/* specific declaration modes for Windows */ -#if !defined(XXH_INLINE_ALL) && !defined(XXH_PRIVATE_API) -# if defined(WIN32) && defined(_MSC_VER) && (defined(XXH_IMPORT) || defined(XXH_EXPORT)) -# ifdef XXH_EXPORT -# define XXH_PUBLIC_API __declspec(dllexport) -# elif XXH_IMPORT -# define XXH_PUBLIC_API __declspec(dllimport) -# endif -# else -# define XXH_PUBLIC_API /* do nothing */ -# endif -#endif - -#ifdef XXH_DOXYGEN -/*! - * @brief Emulate a namespace by transparently prefixing all symbols. - * - * If you want to include _and expose_ xxHash functions from within your own - * library, but also want to avoid symbol collisions with other libraries which - * may also include xxHash, you can use XXH_NAMESPACE to automatically prefix - * any public symbol from xxhash library with the value of XXH_NAMESPACE - * (therefore, avoid empty or numeric values). - * - * Note that no change is required within the calling program as long as it - * includes `xxhash.h`: Regular symbol names will be automatically translated - * by this header. - */ -# define XXH_NAMESPACE /* YOUR NAME HERE */ -# undef XXH_NAMESPACE -#endif - -#ifdef XXH_NAMESPACE -# define XXH_CAT(A,B) A##B -# define XXH_NAME2(A,B) XXH_CAT(A,B) -# define XXH_versionNumber XXH_NAME2(XXH_NAMESPACE, XXH_versionNumber) -/* XXH32 */ -# define XXH32 XXH_NAME2(XXH_NAMESPACE, XXH32) -# define XXH32_createState XXH_NAME2(XXH_NAMESPACE, XXH32_createState) -# define XXH32_freeState XXH_NAME2(XXH_NAMESPACE, XXH32_freeState) -# define XXH32_reset XXH_NAME2(XXH_NAMESPACE, XXH32_reset) -# define XXH32_update XXH_NAME2(XXH_NAMESPACE, XXH32_update) -# define XXH32_digest XXH_NAME2(XXH_NAMESPACE, XXH32_digest) -# define XXH32_copyState XXH_NAME2(XXH_NAMESPACE, XXH32_copyState) -# define XXH32_canonicalFromHash XXH_NAME2(XXH_NAMESPACE, XXH32_canonicalFromHash) -# define XXH32_hashFromCanonical XXH_NAME2(XXH_NAMESPACE, XXH32_hashFromCanonical) -/* XXH64 */ -# define XXH64 XXH_NAME2(XXH_NAMESPACE, XXH64) -# define XXH64_createState XXH_NAME2(XXH_NAMESPACE, XXH64_createState) -# define XXH64_freeState XXH_NAME2(XXH_NAMESPACE, XXH64_freeState) -# define XXH64_reset XXH_NAME2(XXH_NAMESPACE, XXH64_reset) -# define XXH64_update XXH_NAME2(XXH_NAMESPACE, XXH64_update) -# define XXH64_digest XXH_NAME2(XXH_NAMESPACE, XXH64_digest) -# define XXH64_copyState XXH_NAME2(XXH_NAMESPACE, XXH64_copyState) -# define XXH64_canonicalFromHash XXH_NAME2(XXH_NAMESPACE, XXH64_canonicalFromHash) -# define XXH64_hashFromCanonical XXH_NAME2(XXH_NAMESPACE, XXH64_hashFromCanonical) -/* XXH3_64bits */ -# define XXH3_64bits XXH_NAME2(XXH_NAMESPACE, XXH3_64bits) -# define XXH3_64bits_withSecret XXH_NAME2(XXH_NAMESPACE, XXH3_64bits_withSecret) -# define XXH3_64bits_withSeed XXH_NAME2(XXH_NAMESPACE, XXH3_64bits_withSeed) -# define XXH3_64bits_withSecretandSeed XXH_NAME2(XXH_NAMESPACE, XXH3_64bits_withSecretandSeed) -# define XXH3_createState XXH_NAME2(XXH_NAMESPACE, XXH3_createState) -# define XXH3_freeState XXH_NAME2(XXH_NAMESPACE, XXH3_freeState) -# define XXH3_copyState XXH_NAME2(XXH_NAMESPACE, XXH3_copyState) -# define XXH3_64bits_reset XXH_NAME2(XXH_NAMESPACE, XXH3_64bits_reset) -# define XXH3_64bits_reset_withSeed XXH_NAME2(XXH_NAMESPACE, XXH3_64bits_reset_withSeed) -# define XXH3_64bits_reset_withSecret XXH_NAME2(XXH_NAMESPACE, XXH3_64bits_reset_withSecret) -# define XXH3_64bits_reset_withSecretandSeed XXH_NAME2(XXH_NAMESPACE, XXH3_64bits_reset_withSecretandSeed) -# define XXH3_64bits_update XXH_NAME2(XXH_NAMESPACE, XXH3_64bits_update) -# define XXH3_64bits_digest XXH_NAME2(XXH_NAMESPACE, XXH3_64bits_digest) -# define XXH3_generateSecret XXH_NAME2(XXH_NAMESPACE, XXH3_generateSecret) -# define XXH3_generateSecret_fromSeed XXH_NAME2(XXH_NAMESPACE, XXH3_generateSecret_fromSeed) -/* XXH3_128bits */ -# define XXH128 XXH_NAME2(XXH_NAMESPACE, XXH128) -# define XXH3_128bits XXH_NAME2(XXH_NAMESPACE, XXH3_128bits) -# define XXH3_128bits_withSeed XXH_NAME2(XXH_NAMESPACE, XXH3_128bits_withSeed) -# define XXH3_128bits_withSecret XXH_NAME2(XXH_NAMESPACE, XXH3_128bits_withSecret) -# define XXH3_128bits_withSecretandSeed XXH_NAME2(XXH_NAMESPACE, XXH3_128bits_withSecretandSeed) -# define XXH3_128bits_reset XXH_NAME2(XXH_NAMESPACE, XXH3_128bits_reset) -# define XXH3_128bits_reset_withSeed XXH_NAME2(XXH_NAMESPACE, XXH3_128bits_reset_withSeed) -# define XXH3_128bits_reset_withSecret XXH_NAME2(XXH_NAMESPACE, XXH3_128bits_reset_withSecret) -# define XXH3_128bits_reset_withSecretandSeed XXH_NAME2(XXH_NAMESPACE, XXH3_128bits_reset_withSecretandSeed) -# define XXH3_128bits_update XXH_NAME2(XXH_NAMESPACE, XXH3_128bits_update) -# define XXH3_128bits_digest XXH_NAME2(XXH_NAMESPACE, XXH3_128bits_digest) -# define XXH128_isEqual XXH_NAME2(XXH_NAMESPACE, XXH128_isEqual) -# define XXH128_cmp XXH_NAME2(XXH_NAMESPACE, XXH128_cmp) -# define XXH128_canonicalFromHash XXH_NAME2(XXH_NAMESPACE, XXH128_canonicalFromHash) -# define XXH128_hashFromCanonical XXH_NAME2(XXH_NAMESPACE, XXH128_hashFromCanonical) -#endif - - -/* ************************************* -* Version -***************************************/ -#define XXH_VERSION_MAJOR 0 -#define XXH_VERSION_MINOR 8 -#define XXH_VERSION_RELEASE 1 -#define XXH_VERSION_NUMBER (XXH_VERSION_MAJOR *100*100 + XXH_VERSION_MINOR *100 + XXH_VERSION_RELEASE) - -/*! - * @brief Obtains the xxHash version. - * - * This is mostly useful when xxHash is compiled as a shared library, - * since the returned value comes from the library, as opposed to header file. - * - * @return `XXH_VERSION_NUMBER` of the invoked library. - */ -XXH_PUBLIC_API unsigned XXH_versionNumber (void); - - -/* **************************** -* Common basic types -******************************/ -#include /* size_t */ -typedef enum { XXH_OK=0, XXH_ERROR } XXH_errorcode; - - -/*-********************************************************************** -* 32-bit hash -************************************************************************/ -#if defined(XXH_DOXYGEN) /* Don't show include */ -/*! - * @brief An unsigned 32-bit integer. - * - * Not necessarily defined to `uint32_t` but functionally equivalent. - */ -typedef uint32_t XXH32_hash_t; - -#elif !defined (__VMS) \ - && (defined (__cplusplus) \ - || (defined (__STDC_VERSION__) && (__STDC_VERSION__ >= 199901L) /* C99 */) ) -# include - typedef uint32_t XXH32_hash_t; - -#else -# include -# if UINT_MAX == 0xFFFFFFFFUL - typedef unsigned int XXH32_hash_t; -# else -# if ULONG_MAX == 0xFFFFFFFFUL - typedef unsigned long XXH32_hash_t; -# else -# error "unsupported platform: need a 32-bit type" -# endif -# endif -#endif - -/*! - * @} - * - * @defgroup xxh32_family XXH32 family - * @ingroup public - * Contains functions used in the classic 32-bit xxHash algorithm. - * - * @note - * XXH32 is useful for older platforms, with no or poor 64-bit performance. - * Note that @ref xxh3_family provides competitive speed - * for both 32-bit and 64-bit systems, and offers true 64/128 bit hash results. - * - * @see @ref xxh64_family, @ref xxh3_family : Other xxHash families - * @see @ref xxh32_impl for implementation details - * @{ - */ - -/*! - * @brief Calculates the 32-bit hash of @p input using xxHash32. - * - * Speed on Core 2 Duo @ 3 GHz (single thread, SMHasher benchmark): 5.4 GB/s - * - * @param input The block of data to be hashed, at least @p length bytes in size. - * @param length The length of @p input, in bytes. - * @param seed The 32-bit seed to alter the hash's output predictably. - * - * @pre - * The memory between @p input and @p input + @p length must be valid, - * readable, contiguous memory. However, if @p length is `0`, @p input may be - * `NULL`. In C++, this also must be *TriviallyCopyable*. - * - * @return The calculated 32-bit hash value. - * - * @see - * XXH64(), XXH3_64bits_withSeed(), XXH3_128bits_withSeed(), XXH128(): - * Direct equivalents for the other variants of xxHash. - * @see - * XXH32_createState(), XXH32_update(), XXH32_digest(): Streaming version. - */ -XXH_PUBLIC_API XXH32_hash_t XXH32 (const void* input, size_t length, XXH32_hash_t seed); - -/*! - * Streaming functions generate the xxHash value from an incremental input. - * This method is slower than single-call functions, due to state management. - * For small inputs, prefer `XXH32()` and `XXH64()`, which are better optimized. - * - * An XXH state must first be allocated using `XXH*_createState()`. - * - * Start a new hash by initializing the state with a seed using `XXH*_reset()`. - * - * Then, feed the hash state by calling `XXH*_update()` as many times as necessary. - * - * The function returns an error code, with 0 meaning OK, and any other value - * meaning there is an error. - * - * Finally, a hash value can be produced anytime, by using `XXH*_digest()`. - * This function returns the nn-bits hash as an int or long long. - * - * It's still possible to continue inserting input into the hash state after a - * digest, and generate new hash values later on by invoking `XXH*_digest()`. - * - * When done, release the state using `XXH*_freeState()`. - * - * Example code for incrementally hashing a file: - * @code{.c} - * #include - * #include - * #define BUFFER_SIZE 256 - * - * // Note: XXH64 and XXH3 use the same interface. - * XXH32_hash_t - * hashFile(FILE* stream) - * { - * XXH32_state_t* state; - * unsigned char buf[BUFFER_SIZE]; - * size_t amt; - * XXH32_hash_t hash; - * - * state = XXH32_createState(); // Create a state - * assert(state != NULL); // Error check here - * XXH32_reset(state, 0xbaad5eed); // Reset state with our seed - * while ((amt = fread(buf, 1, sizeof(buf), stream)) != 0) { - * XXH32_update(state, buf, amt); // Hash the file in chunks - * } - * hash = XXH32_digest(state); // Finalize the hash - * XXH32_freeState(state); // Clean up - * return hash; - * } - * @endcode - */ - -/*! - * @typedef struct XXH32_state_s XXH32_state_t - * @brief The opaque state struct for the XXH32 streaming API. - * - * @see XXH32_state_s for details. - */ -typedef struct XXH32_state_s XXH32_state_t; - -/*! - * @brief Allocates an @ref XXH32_state_t. - * - * Must be freed with XXH32_freeState(). - * @return An allocated XXH32_state_t on success, `NULL` on failure. - */ -XXH_PUBLIC_API XXH32_state_t* XXH32_createState(void); -/*! - * @brief Frees an @ref XXH32_state_t. - * - * Must be allocated with XXH32_createState(). - * @param statePtr A pointer to an @ref XXH32_state_t allocated with @ref XXH32_createState(). - * @return XXH_OK. - */ -XXH_PUBLIC_API XXH_errorcode XXH32_freeState(XXH32_state_t* statePtr); -/*! - * @brief Copies one @ref XXH32_state_t to another. - * - * @param dst_state The state to copy to. - * @param src_state The state to copy from. - * @pre - * @p dst_state and @p src_state must not be `NULL` and must not overlap. - */ -XXH_PUBLIC_API void XXH32_copyState(XXH32_state_t* dst_state, const XXH32_state_t* src_state); - -/*! - * @brief Resets an @ref XXH32_state_t to begin a new hash. - * - * This function resets and seeds a state. Call it before @ref XXH32_update(). - * - * @param statePtr The state struct to reset. - * @param seed The 32-bit seed to alter the hash result predictably. - * - * @pre - * @p statePtr must not be `NULL`. - * - * @return @ref XXH_OK on success, @ref XXH_ERROR on failure. - */ -XXH_PUBLIC_API XXH_errorcode XXH32_reset (XXH32_state_t* statePtr, XXH32_hash_t seed); - -/*! - * @brief Consumes a block of @p input to an @ref XXH32_state_t. - * - * Call this to incrementally consume blocks of data. - * - * @param statePtr The state struct to update. - * @param input The block of data to be hashed, at least @p length bytes in size. - * @param length The length of @p input, in bytes. - * - * @pre - * @p statePtr must not be `NULL`. - * @pre - * The memory between @p input and @p input + @p length must be valid, - * readable, contiguous memory. However, if @p length is `0`, @p input may be - * `NULL`. In C++, this also must be *TriviallyCopyable*. - * - * @return @ref XXH_OK on success, @ref XXH_ERROR on failure. - */ -XXH_PUBLIC_API XXH_errorcode XXH32_update (XXH32_state_t* statePtr, const void* input, size_t length); - -/*! - * @brief Returns the calculated hash value from an @ref XXH32_state_t. - * - * @note - * Calling XXH32_digest() will not affect @p statePtr, so you can update, - * digest, and update again. - * - * @param statePtr The state struct to calculate the hash from. - * - * @pre - * @p statePtr must not be `NULL`. - * - * @return The calculated xxHash32 value from that state. - */ -XXH_PUBLIC_API XXH32_hash_t XXH32_digest (const XXH32_state_t* statePtr); - -/******* Canonical representation *******/ - -/* - * The default return values from XXH functions are unsigned 32 and 64 bit - * integers. - * This the simplest and fastest format for further post-processing. - * - * However, this leaves open the question of what is the order on the byte level, - * since little and big endian conventions will store the same number differently. - * - * The canonical representation settles this issue by mandating big-endian - * convention, the same convention as human-readable numbers (large digits first). - * - * When writing hash values to storage, sending them over a network, or printing - * them, it's highly recommended to use the canonical representation to ensure - * portability across a wider range of systems, present and future. - * - * The following functions allow transformation of hash values to and from - * canonical format. - */ - -/*! - * @brief Canonical (big endian) representation of @ref XXH32_hash_t. - */ -typedef struct { - unsigned char digest[4]; /*!< Hash bytes, big endian */ -} XXH32_canonical_t; - -/*! - * @brief Converts an @ref XXH32_hash_t to a big endian @ref XXH32_canonical_t. - * - * @param dst The @ref XXH32_canonical_t pointer to be stored to. - * @param hash The @ref XXH32_hash_t to be converted. - * - * @pre - * @p dst must not be `NULL`. - */ -XXH_PUBLIC_API void XXH32_canonicalFromHash(XXH32_canonical_t* dst, XXH32_hash_t hash); - -/*! - * @brief Converts an @ref XXH32_canonical_t to a native @ref XXH32_hash_t. - * - * @param src The @ref XXH32_canonical_t to convert. - * - * @pre - * @p src must not be `NULL`. - * - * @return The converted hash. - */ -XXH_PUBLIC_API XXH32_hash_t XXH32_hashFromCanonical(const XXH32_canonical_t* src); - - -#ifdef __has_attribute -# define XXH_HAS_ATTRIBUTE(x) __has_attribute(x) -#else -# define XXH_HAS_ATTRIBUTE(x) 0 -#endif - -/* C-language Attributes are added in C23. */ -#if defined(__STDC_VERSION__) && (__STDC_VERSION__ > 201710L) && defined(__has_c_attribute) -# define XXH_HAS_C_ATTRIBUTE(x) __has_c_attribute(x) -#else -# define XXH_HAS_C_ATTRIBUTE(x) 0 -#endif - -#if defined(__cplusplus) && defined(__has_cpp_attribute) -# define XXH_HAS_CPP_ATTRIBUTE(x) __has_cpp_attribute(x) -#else -# define XXH_HAS_CPP_ATTRIBUTE(x) 0 -#endif - -/* -Define XXH_FALLTHROUGH macro for annotating switch case with the 'fallthrough' attribute -introduced in CPP17 and C23. -CPP17 : https://en.cppreference.com/w/cpp/language/attributes/fallthrough -C23 : https://en.cppreference.com/w/c/language/attributes/fallthrough -*/ -#if XXH_HAS_C_ATTRIBUTE(x) -# define XXH_FALLTHROUGH [[fallthrough]] -#elif XXH_HAS_CPP_ATTRIBUTE(x) -# define XXH_FALLTHROUGH [[fallthrough]] -#elif XXH_HAS_ATTRIBUTE(__fallthrough__) -# define XXH_FALLTHROUGH __attribute__ ((fallthrough)) -#else -# define XXH_FALLTHROUGH -#endif - -/*! - * @} - * @ingroup public - * @{ - */ - -#ifndef XXH_NO_LONG_LONG -/*-********************************************************************** -* 64-bit hash -************************************************************************/ -#if defined(XXH_DOXYGEN) /* don't include */ -/*! - * @brief An unsigned 64-bit integer. - * - * Not necessarily defined to `uint64_t` but functionally equivalent. - */ -typedef uint64_t XXH64_hash_t; -#elif !defined (__VMS) \ - && (defined (__cplusplus) \ - || (defined (__STDC_VERSION__) && (__STDC_VERSION__ >= 199901L) /* C99 */) ) -# include - typedef uint64_t XXH64_hash_t; -#else -# include -# if defined(__LP64__) && ULONG_MAX == 0xFFFFFFFFFFFFFFFFULL - /* LP64 ABI says uint64_t is unsigned long */ - typedef unsigned long XXH64_hash_t; -# else - /* the following type must have a width of 64-bit */ - typedef unsigned long long XXH64_hash_t; -# endif -#endif - -/*! - * @} - * - * @defgroup xxh64_family XXH64 family - * @ingroup public - * @{ - * Contains functions used in the classic 64-bit xxHash algorithm. - * - * @note - * XXH3 provides competitive speed for both 32-bit and 64-bit systems, - * and offers true 64/128 bit hash results. - * It provides better speed for systems with vector processing capabilities. - */ - - -/*! - * @brief Calculates the 64-bit hash of @p input using xxHash64. - * - * This function usually runs faster on 64-bit systems, but slower on 32-bit - * systems (see benchmark). - * - * @param input The block of data to be hashed, at least @p length bytes in size. - * @param length The length of @p input, in bytes. - * @param seed The 64-bit seed to alter the hash's output predictably. - * - * @pre - * The memory between @p input and @p input + @p length must be valid, - * readable, contiguous memory. However, if @p length is `0`, @p input may be - * `NULL`. In C++, this also must be *TriviallyCopyable*. - * - * @return The calculated 64-bit hash. - * - * @see - * XXH32(), XXH3_64bits_withSeed(), XXH3_128bits_withSeed(), XXH128(): - * Direct equivalents for the other variants of xxHash. - * @see - * XXH64_createState(), XXH64_update(), XXH64_digest(): Streaming version. - */ -XXH_PUBLIC_API XXH64_hash_t XXH64(const void* input, size_t length, XXH64_hash_t seed); - -/******* Streaming *******/ -/*! - * @brief The opaque state struct for the XXH64 streaming API. - * - * @see XXH64_state_s for details. - */ -typedef struct XXH64_state_s XXH64_state_t; /* incomplete type */ -XXH_PUBLIC_API XXH64_state_t* XXH64_createState(void); -XXH_PUBLIC_API XXH_errorcode XXH64_freeState(XXH64_state_t* statePtr); -XXH_PUBLIC_API void XXH64_copyState(XXH64_state_t* dst_state, const XXH64_state_t* src_state); - -XXH_PUBLIC_API XXH_errorcode XXH64_reset (XXH64_state_t* statePtr, XXH64_hash_t seed); -XXH_PUBLIC_API XXH_errorcode XXH64_update (XXH64_state_t* statePtr, const void* input, size_t length); -XXH_PUBLIC_API XXH64_hash_t XXH64_digest (const XXH64_state_t* statePtr); - -/******* Canonical representation *******/ -typedef struct { unsigned char digest[sizeof(XXH64_hash_t)]; } XXH64_canonical_t; -XXH_PUBLIC_API void XXH64_canonicalFromHash(XXH64_canonical_t* dst, XXH64_hash_t hash); -XXH_PUBLIC_API XXH64_hash_t XXH64_hashFromCanonical(const XXH64_canonical_t* src); - -/*! - * @} - * ************************************************************************ - * @defgroup xxh3_family XXH3 family - * @ingroup public - * @{ - * - * XXH3 is a more recent hash algorithm featuring: - * - Improved speed for both small and large inputs - * - True 64-bit and 128-bit outputs - * - SIMD acceleration - * - Improved 32-bit viability - * - * Speed analysis methodology is explained here: - * - * https://fastcompression.blogspot.com/2019/03/presenting-xxh3.html - * - * Compared to XXH64, expect XXH3 to run approximately - * ~2x faster on large inputs and >3x faster on small ones, - * exact differences vary depending on platform. - * - * XXH3's speed benefits greatly from SIMD and 64-bit arithmetic, - * but does not require it. - * Any 32-bit and 64-bit targets that can run XXH32 smoothly - * can run XXH3 at competitive speeds, even without vector support. - * Further details are explained in the implementation. - * - * Optimized implementations are provided for AVX512, AVX2, SSE2, NEON, POWER8, - * ZVector and scalar targets. This can be controlled via the XXH_VECTOR macro. - * - * XXH3 implementation is portable: - * it has a generic C90 formulation that can be compiled on any platform, - * all implementations generage exactly the same hash value on all platforms. - * Starting from v0.8.0, it's also labelled "stable", meaning that - * any future version will also generate the same hash value. - * - * XXH3 offers 2 variants, _64bits and _128bits. - * - * When only 64 bits are needed, prefer invoking the _64bits variant, as it - * reduces the amount of mixing, resulting in faster speed on small inputs. - * It's also generally simpler to manipulate a scalar return type than a struct. - * - * The API supports one-shot hashing, streaming mode, and custom secrets. - */ - -/*-********************************************************************** -* XXH3 64-bit variant -************************************************************************/ - -/* XXH3_64bits(): - * default 64-bit variant, using default secret and default seed of 0. - * It's the fastest variant. */ -XXH_PUBLIC_API XXH64_hash_t XXH3_64bits(const void* data, size_t len); - -/* - * XXH3_64bits_withSeed(): - * This variant generates a custom secret on the fly - * based on default secret altered using the `seed` value. - * While this operation is decently fast, note that it's not completely free. - * Note: seed==0 produces the same results as XXH3_64bits(). - */ -XXH_PUBLIC_API XXH64_hash_t XXH3_64bits_withSeed(const void* data, size_t len, XXH64_hash_t seed); - -/*! - * The bare minimum size for a custom secret. - * - * @see - * XXH3_64bits_withSecret(), XXH3_64bits_reset_withSecret(), - * XXH3_128bits_withSecret(), XXH3_128bits_reset_withSecret(). - */ -#define XXH3_SECRET_SIZE_MIN 136 - -/* - * XXH3_64bits_withSecret(): - * It's possible to provide any blob of bytes as a "secret" to generate the hash. - * This makes it more difficult for an external actor to prepare an intentional collision. - * The main condition is that secretSize *must* be large enough (>= XXH3_SECRET_SIZE_MIN). - * However, the quality of the secret impacts the dispersion of the hash algorithm. - * Therefore, the secret _must_ look like a bunch of random bytes. - * Avoid "trivial" or structured data such as repeated sequences or a text document. - * Whenever in doubt about the "randomness" of the blob of bytes, - * consider employing "XXH3_generateSecret()" instead (see below). - * It will generate a proper high entropy secret derived from the blob of bytes. - * Another advantage of using XXH3_generateSecret() is that - * it guarantees that all bits within the initial blob of bytes - * will impact every bit of the output. - * This is not necessarily the case when using the blob of bytes directly - * because, when hashing _small_ inputs, only a portion of the secret is employed. - */ -XXH_PUBLIC_API XXH64_hash_t XXH3_64bits_withSecret(const void* data, size_t len, const void* secret, size_t secretSize); - - -/******* Streaming *******/ -/* - * Streaming requires state maintenance. - * This operation costs memory and CPU. - * As a consequence, streaming is slower than one-shot hashing. - * For better performance, prefer one-shot functions whenever applicable. - */ - -/*! - * @brief The state struct for the XXH3 streaming API. - * - * @see XXH3_state_s for details. - */ -typedef struct XXH3_state_s XXH3_state_t; -XXH_PUBLIC_API XXH3_state_t* XXH3_createState(void); -XXH_PUBLIC_API XXH_errorcode XXH3_freeState(XXH3_state_t* statePtr); -XXH_PUBLIC_API void XXH3_copyState(XXH3_state_t* dst_state, const XXH3_state_t* src_state); - -/* - * XXH3_64bits_reset(): - * Initialize with default parameters. - * digest will be equivalent to `XXH3_64bits()`. - */ -XXH_PUBLIC_API XXH_errorcode XXH3_64bits_reset(XXH3_state_t* statePtr); -/* - * XXH3_64bits_reset_withSeed(): - * Generate a custom secret from `seed`, and store it into `statePtr`. - * digest will be equivalent to `XXH3_64bits_withSeed()`. - */ -XXH_PUBLIC_API XXH_errorcode XXH3_64bits_reset_withSeed(XXH3_state_t* statePtr, XXH64_hash_t seed); -/* - * XXH3_64bits_reset_withSecret(): - * `secret` is referenced, it _must outlive_ the hash streaming session. - * Similar to one-shot API, `secretSize` must be >= `XXH3_SECRET_SIZE_MIN`, - * and the quality of produced hash values depends on secret's entropy - * (secret's content should look like a bunch of random bytes). - * When in doubt about the randomness of a candidate `secret`, - * consider employing `XXH3_generateSecret()` instead (see below). - */ -XXH_PUBLIC_API XXH_errorcode XXH3_64bits_reset_withSecret(XXH3_state_t* statePtr, const void* secret, size_t secretSize); - -XXH_PUBLIC_API XXH_errorcode XXH3_64bits_update (XXH3_state_t* statePtr, const void* input, size_t length); -XXH_PUBLIC_API XXH64_hash_t XXH3_64bits_digest (const XXH3_state_t* statePtr); - -/* note : canonical representation of XXH3 is the same as XXH64 - * since they both produce XXH64_hash_t values */ - - -/*-********************************************************************** -* XXH3 128-bit variant -************************************************************************/ - -/*! - * @brief The return value from 128-bit hashes. - * - * Stored in little endian order, although the fields themselves are in native - * endianness. - */ -typedef struct { - XXH64_hash_t low64; /*!< `value & 0xFFFFFFFFFFFFFFFF` */ - XXH64_hash_t high64; /*!< `value >> 64` */ -} XXH128_hash_t; - -XXH_PUBLIC_API XXH128_hash_t XXH3_128bits(const void* data, size_t len); -XXH_PUBLIC_API XXH128_hash_t XXH3_128bits_withSeed(const void* data, size_t len, XXH64_hash_t seed); -XXH_PUBLIC_API XXH128_hash_t XXH3_128bits_withSecret(const void* data, size_t len, const void* secret, size_t secretSize); - -/******* Streaming *******/ -/* - * Streaming requires state maintenance. - * This operation costs memory and CPU. - * As a consequence, streaming is slower than one-shot hashing. - * For better performance, prefer one-shot functions whenever applicable. - * - * XXH3_128bits uses the same XXH3_state_t as XXH3_64bits(). - * Use already declared XXH3_createState() and XXH3_freeState(). - * - * All reset and streaming functions have same meaning as their 64-bit counterpart. - */ - -XXH_PUBLIC_API XXH_errorcode XXH3_128bits_reset(XXH3_state_t* statePtr); -XXH_PUBLIC_API XXH_errorcode XXH3_128bits_reset_withSeed(XXH3_state_t* statePtr, XXH64_hash_t seed); -XXH_PUBLIC_API XXH_errorcode XXH3_128bits_reset_withSecret(XXH3_state_t* statePtr, const void* secret, size_t secretSize); - -XXH_PUBLIC_API XXH_errorcode XXH3_128bits_update (XXH3_state_t* statePtr, const void* input, size_t length); -XXH_PUBLIC_API XXH128_hash_t XXH3_128bits_digest (const XXH3_state_t* statePtr); - -/* Following helper functions make it possible to compare XXH128_hast_t values. - * Since XXH128_hash_t is a structure, this capability is not offered by the language. - * Note: For better performance, these functions can be inlined using XXH_INLINE_ALL */ - -/*! - * XXH128_isEqual(): - * Return: 1 if `h1` and `h2` are equal, 0 if they are not. - */ -XXH_PUBLIC_API int XXH128_isEqual(XXH128_hash_t h1, XXH128_hash_t h2); - -/*! - * XXH128_cmp(): - * - * This comparator is compatible with stdlib's `qsort()`/`bsearch()`. - * - * return: >0 if *h128_1 > *h128_2 - * =0 if *h128_1 == *h128_2 - * <0 if *h128_1 < *h128_2 - */ -XXH_PUBLIC_API int XXH128_cmp(const void* h128_1, const void* h128_2); - - -/******* Canonical representation *******/ -typedef struct { unsigned char digest[sizeof(XXH128_hash_t)]; } XXH128_canonical_t; -XXH_PUBLIC_API void XXH128_canonicalFromHash(XXH128_canonical_t* dst, XXH128_hash_t hash); -XXH_PUBLIC_API XXH128_hash_t XXH128_hashFromCanonical(const XXH128_canonical_t* src); - - -#endif /* XXH_NO_LONG_LONG */ - -/*! - * @} - */ -#endif /* XXHASH_H_5627135585666179 */ - - - -#if defined(XXH_STATIC_LINKING_ONLY) && !defined(XXHASH_H_STATIC_13879238742) -#define XXHASH_H_STATIC_13879238742 -/* **************************************************************************** - * This section contains declarations which are not guaranteed to remain stable. - * They may change in future versions, becoming incompatible with a different - * version of the library. - * These declarations should only be used with static linking. - * Never use them in association with dynamic linking! - ***************************************************************************** */ - -/* - * These definitions are only present to allow static allocation - * of XXH states, on stack or in a struct, for example. - * Never **ever** access their members directly. - */ - -/*! - * @internal - * @brief Structure for XXH32 streaming API. - * - * @note This is only defined when @ref XXH_STATIC_LINKING_ONLY, - * @ref XXH_INLINE_ALL, or @ref XXH_IMPLEMENTATION is defined. Otherwise it is - * an opaque type. This allows fields to safely be changed. - * - * Typedef'd to @ref XXH32_state_t. - * Do not access the members of this struct directly. - * @see XXH64_state_s, XXH3_state_s - */ -struct XXH32_state_s { - XXH32_hash_t total_len_32; /*!< Total length hashed, modulo 2^32 */ - XXH32_hash_t large_len; /*!< Whether the hash is >= 16 (handles @ref total_len_32 overflow) */ - XXH32_hash_t v[4]; /*!< Accumulator lanes */ - XXH32_hash_t mem32[4]; /*!< Internal buffer for partial reads. Treated as unsigned char[16]. */ - XXH32_hash_t memsize; /*!< Amount of data in @ref mem32 */ - XXH32_hash_t reserved; /*!< Reserved field. Do not read or write to it, it may be removed. */ -}; /* typedef'd to XXH32_state_t */ - - -#ifndef XXH_NO_LONG_LONG /* defined when there is no 64-bit support */ - -/*! - * @internal - * @brief Structure for XXH64 streaming API. - * - * @note This is only defined when @ref XXH_STATIC_LINKING_ONLY, - * @ref XXH_INLINE_ALL, or @ref XXH_IMPLEMENTATION is defined. Otherwise it is - * an opaque type. This allows fields to safely be changed. - * - * Typedef'd to @ref XXH64_state_t. - * Do not access the members of this struct directly. - * @see XXH32_state_s, XXH3_state_s - */ -struct XXH64_state_s { - XXH64_hash_t total_len; /*!< Total length hashed. This is always 64-bit. */ - XXH64_hash_t v[4]; /*!< Accumulator lanes */ - XXH64_hash_t mem64[4]; /*!< Internal buffer for partial reads. Treated as unsigned char[32]. */ - XXH32_hash_t memsize; /*!< Amount of data in @ref mem64 */ - XXH32_hash_t reserved32; /*!< Reserved field, needed for padding anyways*/ - XXH64_hash_t reserved64; /*!< Reserved field. Do not read or write to it, it may be removed. */ -}; /* typedef'd to XXH64_state_t */ - -#if defined(__STDC_VERSION__) && (__STDC_VERSION__ >= 201112L) /* >= C11 */ -# include -# define XXH_ALIGN(n) alignas(n) -#elif defined(__cplusplus) && (__cplusplus >= 201103L) /* >= C++11 */ -/* In C++ alignas() is a keyword */ -# define XXH_ALIGN(n) alignas(n) -#elif defined(__GNUC__) -# define XXH_ALIGN(n) __attribute__ ((aligned(n))) -#elif defined(_MSC_VER) -# define XXH_ALIGN(n) __declspec(align(n)) -#else -# define XXH_ALIGN(n) /* disabled */ -#endif - -/* Old GCC versions only accept the attribute after the type in structures. */ -#if !(defined(__STDC_VERSION__) && (__STDC_VERSION__ >= 201112L)) /* C11+ */ \ - && ! (defined(__cplusplus) && (__cplusplus >= 201103L)) /* >= C++11 */ \ - && defined(__GNUC__) -# define XXH_ALIGN_MEMBER(align, type) type XXH_ALIGN(align) -#else -# define XXH_ALIGN_MEMBER(align, type) XXH_ALIGN(align) type -#endif - -/*! - * @brief The size of the internal XXH3 buffer. - * - * This is the optimal update size for incremental hashing. - * - * @see XXH3_64b_update(), XXH3_128b_update(). - */ -#define XXH3_INTERNALBUFFER_SIZE 256 - -/*! - * @brief Default size of the secret buffer (and @ref XXH3_kSecret). - * - * This is the size used in @ref XXH3_kSecret and the seeded functions. - * - * Not to be confused with @ref XXH3_SECRET_SIZE_MIN. - */ -#define XXH3_SECRET_DEFAULT_SIZE 192 - -/*! - * @internal - * @brief Structure for XXH3 streaming API. - * - * @note This is only defined when @ref XXH_STATIC_LINKING_ONLY, - * @ref XXH_INLINE_ALL, or @ref XXH_IMPLEMENTATION is defined. - * Otherwise it is an opaque type. - * Never use this definition in combination with dynamic library. - * This allows fields to safely be changed in the future. - * - * @note ** This structure has a strict alignment requirement of 64 bytes!! ** - * Do not allocate this with `malloc()` or `new`, - * it will not be sufficiently aligned. - * Use @ref XXH3_createState() and @ref XXH3_freeState(), or stack allocation. - * - * Typedef'd to @ref XXH3_state_t. - * Do never access the members of this struct directly. - * - * @see XXH3_INITSTATE() for stack initialization. - * @see XXH3_createState(), XXH3_freeState(). - * @see XXH32_state_s, XXH64_state_s - */ -struct XXH3_state_s { - XXH_ALIGN_MEMBER(64, XXH64_hash_t acc[8]); - /*!< The 8 accumulators. Similar to `vN` in @ref XXH32_state_s::v1 and @ref XXH64_state_s */ - XXH_ALIGN_MEMBER(64, unsigned char customSecret[XXH3_SECRET_DEFAULT_SIZE]); - /*!< Used to store a custom secret generated from a seed. */ - XXH_ALIGN_MEMBER(64, unsigned char buffer[XXH3_INTERNALBUFFER_SIZE]); - /*!< The internal buffer. @see XXH32_state_s::mem32 */ - XXH32_hash_t bufferedSize; - /*!< The amount of memory in @ref buffer, @see XXH32_state_s::memsize */ - XXH32_hash_t useSeed; - /*!< Reserved field. Needed for padding on 64-bit. */ - size_t nbStripesSoFar; - /*!< Number or stripes processed. */ - XXH64_hash_t totalLen; - /*!< Total length hashed. 64-bit even on 32-bit targets. */ - size_t nbStripesPerBlock; - /*!< Number of stripes per block. */ - size_t secretLimit; - /*!< Size of @ref customSecret or @ref extSecret */ - XXH64_hash_t seed; - /*!< Seed for _withSeed variants. Must be zero otherwise, @see XXH3_INITSTATE() */ - XXH64_hash_t reserved64; - /*!< Reserved field. */ - const unsigned char* extSecret; - /*!< Reference to an external secret for the _withSecret variants, NULL - * for other variants. */ - /* note: there may be some padding at the end due to alignment on 64 bytes */ -}; /* typedef'd to XXH3_state_t */ - -#undef XXH_ALIGN_MEMBER - -/*! - * @brief Initializes a stack-allocated `XXH3_state_s`. - * - * When the @ref XXH3_state_t structure is merely emplaced on stack, - * it should be initialized with XXH3_INITSTATE() or a memset() - * in case its first reset uses XXH3_NNbits_reset_withSeed(). - * This init can be omitted if the first reset uses default or _withSecret mode. - * This operation isn't necessary when the state is created with XXH3_createState(). - * Note that this doesn't prepare the state for a streaming operation, - * it's still necessary to use XXH3_NNbits_reset*() afterwards. - */ -#define XXH3_INITSTATE(XXH3_state_ptr) { (XXH3_state_ptr)->seed = 0; } - - -/* XXH128() : - * simple alias to pre-selected XXH3_128bits variant - */ -XXH_PUBLIC_API XXH128_hash_t XXH128(const void* data, size_t len, XXH64_hash_t seed); - - -/* === Experimental API === */ -/* Symbols defined below must be considered tied to a specific library version. */ - -/* - * XXH3_generateSecret(): - * - * Derive a high-entropy secret from any user-defined content, named customSeed. - * The generated secret can be used in combination with `*_withSecret()` functions. - * The `_withSecret()` variants are useful to provide a higher level of protection than 64-bit seed, - * as it becomes much more difficult for an external actor to guess how to impact the calculation logic. - * - * The function accepts as input a custom seed of any length and any content, - * and derives from it a high-entropy secret of length @secretSize - * into an already allocated buffer @secretBuffer. - * @secretSize must be >= XXH3_SECRET_SIZE_MIN - * - * The generated secret can then be used with any `*_withSecret()` variant. - * Functions `XXH3_128bits_withSecret()`, `XXH3_64bits_withSecret()`, - * `XXH3_128bits_reset_withSecret()` and `XXH3_64bits_reset_withSecret()` - * are part of this list. They all accept a `secret` parameter - * which must be large enough for implementation reasons (>= XXH3_SECRET_SIZE_MIN) - * _and_ feature very high entropy (consist of random-looking bytes). - * These conditions can be a high bar to meet, so - * XXH3_generateSecret() can be employed to ensure proper quality. - * - * customSeed can be anything. It can have any size, even small ones, - * and its content can be anything, even "poor entropy" sources such as a bunch of zeroes. - * The resulting `secret` will nonetheless provide all required qualities. - * - * When customSeedSize > 0, supplying NULL as customSeed is undefined behavior. - */ -XXH_PUBLIC_API XXH_errorcode XXH3_generateSecret(void* secretBuffer, size_t secretSize, const void* customSeed, size_t customSeedSize); - - -/* - * XXH3_generateSecret_fromSeed(): - * - * Generate the same secret as the _withSeed() variants. - * - * The resulting secret has a length of XXH3_SECRET_DEFAULT_SIZE (necessarily). - * @secretBuffer must be already allocated, of size at least XXH3_SECRET_DEFAULT_SIZE bytes. - * - * The generated secret can be used in combination with - *`*_withSecret()` and `_withSecretandSeed()` variants. - * This generator is notably useful in combination with `_withSecretandSeed()`, - * as a way to emulate a faster `_withSeed()` variant. - */ -XXH_PUBLIC_API void XXH3_generateSecret_fromSeed(void* secretBuffer, XXH64_hash_t seed); - -/* - * *_withSecretandSeed() : - * These variants generate hash values using either - * @seed for "short" keys (< XXH3_MIDSIZE_MAX = 240 bytes) - * or @secret for "large" keys (>= XXH3_MIDSIZE_MAX). - * - * This generally benefits speed, compared to `_withSeed()` or `_withSecret()`. - * `_withSeed()` has to generate the secret on the fly for "large" keys. - * It's fast, but can be perceptible for "not so large" keys (< 1 KB). - * `_withSecret()` has to generate the masks on the fly for "small" keys, - * which requires more instructions than _withSeed() variants. - * Therefore, _withSecretandSeed variant combines the best of both worlds. - * - * When @secret has been generated by XXH3_generateSecret_fromSeed(), - * this variant produces *exactly* the same results as `_withSeed()` variant, - * hence offering only a pure speed benefit on "large" input, - * by skipping the need to regenerate the secret for every large input. - * - * Another usage scenario is to hash the secret to a 64-bit hash value, - * for example with XXH3_64bits(), which then becomes the seed, - * and then employ both the seed and the secret in _withSecretandSeed(). - * On top of speed, an added benefit is that each bit in the secret - * has a 50% chance to swap each bit in the output, - * via its impact to the seed. - * This is not guaranteed when using the secret directly in "small data" scenarios, - * because only portions of the secret are employed for small data. - */ -XXH_PUBLIC_API XXH64_hash_t -XXH3_64bits_withSecretandSeed(const void* data, size_t len, - const void* secret, size_t secretSize, - XXH64_hash_t seed); - -XXH_PUBLIC_API XXH128_hash_t -XXH3_128bits_withSecretandSeed(const void* data, size_t len, - const void* secret, size_t secretSize, - XXH64_hash_t seed64); - -XXH_PUBLIC_API XXH_errorcode -XXH3_64bits_reset_withSecretandSeed(XXH3_state_t* statePtr, - const void* secret, size_t secretSize, - XXH64_hash_t seed64); - -XXH_PUBLIC_API XXH_errorcode -XXH3_128bits_reset_withSecretandSeed(XXH3_state_t* statePtr, - const void* secret, size_t secretSize, - XXH64_hash_t seed64); - - -#endif /* XXH_NO_LONG_LONG */ -#if defined(XXH_INLINE_ALL) || defined(XXH_PRIVATE_API) -# define XXH_IMPLEMENTATION -#endif - -#endif /* defined(XXH_STATIC_LINKING_ONLY) && !defined(XXHASH_H_STATIC_13879238742) */ - - -/* ======================================================================== */ -/* ======================================================================== */ -/* ======================================================================== */ - - -/*-********************************************************************** - * xxHash implementation - *-********************************************************************** - * xxHash's implementation used to be hosted inside xxhash.c. - * - * However, inlining requires implementation to be visible to the compiler, - * hence be included alongside the header. - * Previously, implementation was hosted inside xxhash.c, - * which was then #included when inlining was activated. - * This construction created issues with a few build and install systems, - * as it required xxhash.c to be stored in /include directory. - * - * xxHash implementation is now directly integrated within xxhash.h. - * As a consequence, xxhash.c is no longer needed in /include. - * - * xxhash.c is still available and is still useful. - * In a "normal" setup, when xxhash is not inlined, - * xxhash.h only exposes the prototypes and public symbols, - * while xxhash.c can be built into an object file xxhash.o - * which can then be linked into the final binary. - ************************************************************************/ - -#if ( defined(XXH_INLINE_ALL) || defined(XXH_PRIVATE_API) \ - || defined(XXH_IMPLEMENTATION) ) && !defined(XXH_IMPLEM_13a8737387) -# define XXH_IMPLEM_13a8737387 - -/* ************************************* -* Tuning parameters -***************************************/ - -/*! - * @defgroup tuning Tuning parameters - * @{ - * - * Various macros to control xxHash's behavior. - */ -#ifdef XXH_DOXYGEN -/*! - * @brief Define this to disable 64-bit code. - * - * Useful if only using the @ref xxh32_family and you have a strict C90 compiler. - */ -# define XXH_NO_LONG_LONG -# undef XXH_NO_LONG_LONG /* don't actually */ -/*! - * @brief Controls how unaligned memory is accessed. - * - * By default, access to unaligned memory is controlled by `memcpy()`, which is - * safe and portable. - * - * Unfortunately, on some target/compiler combinations, the generated assembly - * is sub-optimal. - * - * The below switch allow selection of a different access method - * in the search for improved performance. - * - * @par Possible options: - * - * - `XXH_FORCE_MEMORY_ACCESS=0` (default): `memcpy` - * @par - * Use `memcpy()`. Safe and portable. Note that most modern compilers will - * eliminate the function call and treat it as an unaligned access. - * - * - `XXH_FORCE_MEMORY_ACCESS=1`: `__attribute__((packed))` - * @par - * Depends on compiler extensions and is therefore not portable. - * This method is safe _if_ your compiler supports it, - * and *generally* as fast or faster than `memcpy`. - * - * - `XXH_FORCE_MEMORY_ACCESS=2`: Direct cast - * @par - * Casts directly and dereferences. This method doesn't depend on the - * compiler, but it violates the C standard as it directly dereferences an - * unaligned pointer. It can generate buggy code on targets which do not - * support unaligned memory accesses, but in some circumstances, it's the - * only known way to get the most performance. - * - * - `XXH_FORCE_MEMORY_ACCESS=3`: Byteshift - * @par - * Also portable. This can generate the best code on old compilers which don't - * inline small `memcpy()` calls, and it might also be faster on big-endian - * systems which lack a native byteswap instruction. However, some compilers - * will emit literal byteshifts even if the target supports unaligned access. - * . - * - * @warning - * Methods 1 and 2 rely on implementation-defined behavior. Use these with - * care, as what works on one compiler/platform/optimization level may cause - * another to read garbage data or even crash. - * - * See http://fastcompression.blogspot.com/2015/08/accessing-unaligned-memory.html for details. - * - * Prefer these methods in priority order (0 > 3 > 1 > 2) - */ -# define XXH_FORCE_MEMORY_ACCESS 0 - -/*! - * @def XXH_FORCE_ALIGN_CHECK - * @brief If defined to non-zero, adds a special path for aligned inputs (XXH32() - * and XXH64() only). - * - * This is an important performance trick for architectures without decent - * unaligned memory access performance. - * - * It checks for input alignment, and when conditions are met, uses a "fast - * path" employing direct 32-bit/64-bit reads, resulting in _dramatically - * faster_ read speed. - * - * The check costs one initial branch per hash, which is generally negligible, - * but not zero. - * - * Moreover, it's not useful to generate an additional code path if memory - * access uses the same instruction for both aligned and unaligned - * addresses (e.g. x86 and aarch64). - * - * In these cases, the alignment check can be removed by setting this macro to 0. - * Then the code will always use unaligned memory access. - * Align check is automatically disabled on x86, x64 & arm64, - * which are platforms known to offer good unaligned memory accesses performance. - * - * This option does not affect XXH3 (only XXH32 and XXH64). - */ -# define XXH_FORCE_ALIGN_CHECK 0 - -/*! - * @def XXH_NO_INLINE_HINTS - * @brief When non-zero, sets all functions to `static`. - * - * By default, xxHash tries to force the compiler to inline almost all internal - * functions. - * - * This can usually improve performance due to reduced jumping and improved - * constant folding, but significantly increases the size of the binary which - * might not be favorable. - * - * Additionally, sometimes the forced inlining can be detrimental to performance, - * depending on the architecture. - * - * XXH_NO_INLINE_HINTS marks all internal functions as static, giving the - * compiler full control on whether to inline or not. - * - * When not optimizing (-O0), optimizing for size (-Os, -Oz), or using - * -fno-inline with GCC or Clang, this will automatically be defined. - */ -# define XXH_NO_INLINE_HINTS 0 - -/*! - * @def XXH32_ENDJMP - * @brief Whether to use a jump for `XXH32_finalize`. - * - * For performance, `XXH32_finalize` uses multiple branches in the finalizer. - * This is generally preferable for performance, - * but depending on exact architecture, a jmp may be preferable. - * - * This setting is only possibly making a difference for very small inputs. - */ -# define XXH32_ENDJMP 0 - -/*! - * @internal - * @brief Redefines old internal names. - * - * For compatibility with code that uses xxHash's internals before the names - * were changed to improve namespacing. There is no other reason to use this. - */ -# define XXH_OLD_NAMES -# undef XXH_OLD_NAMES /* don't actually use, it is ugly. */ -#endif /* XXH_DOXYGEN */ -/*! - * @} - */ - -#ifndef XXH_FORCE_MEMORY_ACCESS /* can be defined externally, on command line for example */ - /* prefer __packed__ structures (method 1) for gcc on armv7+ and mips */ -# if !defined(__clang__) && \ -( \ - (defined(__INTEL_COMPILER) && !defined(_WIN32)) || \ - ( \ - defined(__GNUC__) && ( \ - (defined(__ARM_ARCH) && __ARM_ARCH >= 7) || \ - ( \ - defined(__mips__) && \ - (__mips <= 5 || __mips_isa_rev < 6) && \ - (!defined(__mips16) || defined(__mips_mips16e2)) \ - ) \ - ) \ - ) \ -) -# define XXH_FORCE_MEMORY_ACCESS 1 -# endif -#endif - -#ifndef XXH_FORCE_ALIGN_CHECK /* can be defined externally */ -# if defined(__i386) || defined(__x86_64__) || defined(__aarch64__) \ - || defined(_M_IX86) || defined(_M_X64) || defined(_M_ARM64) /* visual */ -# define XXH_FORCE_ALIGN_CHECK 0 -# else -# define XXH_FORCE_ALIGN_CHECK 1 -# endif -#endif - -#ifndef XXH_NO_INLINE_HINTS -# if defined(__OPTIMIZE_SIZE__) /* -Os, -Oz */ \ - || defined(__NO_INLINE__) /* -O0, -fno-inline */ -# define XXH_NO_INLINE_HINTS 1 -# else -# define XXH_NO_INLINE_HINTS 0 -# endif -#endif - -#ifndef XXH32_ENDJMP -/* generally preferable for performance */ -# define XXH32_ENDJMP 0 -#endif - -/*! - * @defgroup impl Implementation - * @{ - */ - - -/* ************************************* -* Includes & Memory related functions -***************************************/ -/* - * Modify the local functions below should you wish to use - * different memory routines for malloc() and free() - */ -#include - -/*! - * @internal - * @brief Modify this function to use a different routine than malloc(). - */ -static void* XXH_malloc(size_t s) { return malloc(s); } - -/*! - * @internal - * @brief Modify this function to use a different routine than free(). - */ -static void XXH_free(void* p) { free(p); } - -#include - -/*! - * @internal - * @brief Modify this function to use a different routine than memcpy(). - */ -static void* XXH_memcpy(void* dest, const void* src, size_t size) -{ - return memcpy(dest,src,size); -} - -#include /* ULLONG_MAX */ - - -/* ************************************* -* Compiler Specific Options -***************************************/ -#ifdef _MSC_VER /* Visual Studio warning fix */ -# pragma warning(disable : 4127) /* disable: C4127: conditional expression is constant */ -#endif - -#if XXH_NO_INLINE_HINTS /* disable inlining hints */ -# if defined(__GNUC__) || defined(__clang__) -# define XXH_FORCE_INLINE static __attribute__((unused)) -# else -# define XXH_FORCE_INLINE static -# endif -# define XXH_NO_INLINE static -/* enable inlining hints */ -#elif defined(__GNUC__) || defined(__clang__) -# define XXH_FORCE_INLINE static __inline__ __attribute__((always_inline, unused)) -# define XXH_NO_INLINE static __attribute__((noinline)) -#elif defined(_MSC_VER) /* Visual Studio */ -# define XXH_FORCE_INLINE static __forceinline -# define XXH_NO_INLINE static __declspec(noinline) -#elif defined (__cplusplus) \ - || (defined (__STDC_VERSION__) && (__STDC_VERSION__ >= 199901L)) /* C99 */ -# define XXH_FORCE_INLINE static inline -# define XXH_NO_INLINE static -#else -# define XXH_FORCE_INLINE static -# define XXH_NO_INLINE static -#endif - - - -/* ************************************* -* Debug -***************************************/ -/*! - * @ingroup tuning - * @def XXH_DEBUGLEVEL - * @brief Sets the debugging level. - * - * XXH_DEBUGLEVEL is expected to be defined externally, typically via the - * compiler's command line options. The value must be a number. - */ -#ifndef XXH_DEBUGLEVEL -# ifdef DEBUGLEVEL /* backwards compat */ -# define XXH_DEBUGLEVEL DEBUGLEVEL -# else -# define XXH_DEBUGLEVEL 0 -# endif -#endif - -#if (XXH_DEBUGLEVEL>=1) -# include /* note: can still be disabled with NDEBUG */ -# define XXH_ASSERT(c) assert(c) -#else -# define XXH_ASSERT(c) ((void)0) -#endif - -/* note: use after variable declarations */ -#ifndef XXH_STATIC_ASSERT -# if defined(__STDC_VERSION__) && (__STDC_VERSION__ >= 201112L) /* C11 */ -# include -# define XXH_STATIC_ASSERT_WITH_MESSAGE(c,m) do { static_assert((c),m); } while(0) -# elif defined(__cplusplus) && (__cplusplus >= 201103L) /* C++11 */ -# define XXH_STATIC_ASSERT_WITH_MESSAGE(c,m) do { static_assert((c),m); } while(0) -# else -# define XXH_STATIC_ASSERT_WITH_MESSAGE(c,m) do { struct xxh_sa { char x[(c) ? 1 : -1]; }; } while(0) -# endif -# define XXH_STATIC_ASSERT(c) XXH_STATIC_ASSERT_WITH_MESSAGE((c),#c) -#endif - -/*! - * @internal - * @def XXH_COMPILER_GUARD(var) - * @brief Used to prevent unwanted optimizations for @p var. - * - * It uses an empty GCC inline assembly statement with a register constraint - * which forces @p var into a general purpose register (eg eax, ebx, ecx - * on x86) and marks it as modified. - * - * This is used in a few places to avoid unwanted autovectorization (e.g. - * XXH32_round()). All vectorization we want is explicit via intrinsics, - * and _usually_ isn't wanted elsewhere. - * - * We also use it to prevent unwanted constant folding for AArch64 in - * XXH3_initCustomSecret_scalar(). - */ -#if defined(__GNUC__) || defined(__clang__) -# define XXH_COMPILER_GUARD(var) __asm__ __volatile__("" : "+r" (var)) -#else -# define XXH_COMPILER_GUARD(var) ((void)0) -#endif - -/* ************************************* -* Basic Types -***************************************/ -#if !defined (__VMS) \ - && (defined (__cplusplus) \ - || (defined (__STDC_VERSION__) && (__STDC_VERSION__ >= 199901L) /* C99 */) ) -# include - typedef uint8_t xxh_u8; -#else - typedef unsigned char xxh_u8; -#endif -typedef XXH32_hash_t xxh_u32; - -#ifdef XXH_OLD_NAMES -# define BYTE xxh_u8 -# define U8 xxh_u8 -# define U32 xxh_u32 -#endif - -/* *** Memory access *** */ - -/*! - * @internal - * @fn xxh_u32 XXH_read32(const void* ptr) - * @brief Reads an unaligned 32-bit integer from @p ptr in native endianness. - * - * Affected by @ref XXH_FORCE_MEMORY_ACCESS. - * - * @param ptr The pointer to read from. - * @return The 32-bit native endian integer from the bytes at @p ptr. - */ - -/*! - * @internal - * @fn xxh_u32 XXH_readLE32(const void* ptr) - * @brief Reads an unaligned 32-bit little endian integer from @p ptr. - * - * Affected by @ref XXH_FORCE_MEMORY_ACCESS. - * - * @param ptr The pointer to read from. - * @return The 32-bit little endian integer from the bytes at @p ptr. - */ - -/*! - * @internal - * @fn xxh_u32 XXH_readBE32(const void* ptr) - * @brief Reads an unaligned 32-bit big endian integer from @p ptr. - * - * Affected by @ref XXH_FORCE_MEMORY_ACCESS. - * - * @param ptr The pointer to read from. - * @return The 32-bit big endian integer from the bytes at @p ptr. - */ - -/*! - * @internal - * @fn xxh_u32 XXH_readLE32_align(const void* ptr, XXH_alignment align) - * @brief Like @ref XXH_readLE32(), but has an option for aligned reads. - * - * Affected by @ref XXH_FORCE_MEMORY_ACCESS. - * Note that when @ref XXH_FORCE_ALIGN_CHECK == 0, the @p align parameter is - * always @ref XXH_alignment::XXH_unaligned. - * - * @param ptr The pointer to read from. - * @param align Whether @p ptr is aligned. - * @pre - * If @p align == @ref XXH_alignment::XXH_aligned, @p ptr must be 4 byte - * aligned. - * @return The 32-bit little endian integer from the bytes at @p ptr. - */ - -#if (defined(XXH_FORCE_MEMORY_ACCESS) && (XXH_FORCE_MEMORY_ACCESS==3)) -/* - * Manual byteshift. Best for old compilers which don't inline memcpy. - * We actually directly use XXH_readLE32 and XXH_readBE32. - */ -#elif (defined(XXH_FORCE_MEMORY_ACCESS) && (XXH_FORCE_MEMORY_ACCESS==2)) - -/* - * Force direct memory access. Only works on CPU which support unaligned memory - * access in hardware. - */ -static xxh_u32 XXH_read32(const void* memPtr) { return *(const xxh_u32*) memPtr; } - -#elif (defined(XXH_FORCE_MEMORY_ACCESS) && (XXH_FORCE_MEMORY_ACCESS==1)) - -/* - * __pack instructions are safer but compiler specific, hence potentially - * problematic for some compilers. - * - * Currently only defined for GCC and ICC. - */ -#ifdef XXH_OLD_NAMES -typedef union { xxh_u32 u32; } __attribute__((packed)) unalign; -#endif -static xxh_u32 XXH_read32(const void* ptr) -{ - typedef union { xxh_u32 u32; } __attribute__((packed)) xxh_unalign; - return ((const xxh_unalign*)ptr)->u32; -} - -#else - -/* - * Portable and safe solution. Generally efficient. - * see: http://fastcompression.blogspot.com/2015/08/accessing-unaligned-memory.html - */ -static xxh_u32 XXH_read32(const void* memPtr) -{ - xxh_u32 val; - XXH_memcpy(&val, memPtr, sizeof(val)); - return val; -} - -#endif /* XXH_FORCE_DIRECT_MEMORY_ACCESS */ - - -/* *** Endianness *** */ - -/*! - * @ingroup tuning - * @def XXH_CPU_LITTLE_ENDIAN - * @brief Whether the target is little endian. - * - * Defined to 1 if the target is little endian, or 0 if it is big endian. - * It can be defined externally, for example on the compiler command line. - * - * If it is not defined, - * a runtime check (which is usually constant folded) is used instead. - * - * @note - * This is not necessarily defined to an integer constant. - * - * @see XXH_isLittleEndian() for the runtime check. - */ -#ifndef XXH_CPU_LITTLE_ENDIAN -/* - * Try to detect endianness automatically, to avoid the nonstandard behavior - * in `XXH_isLittleEndian()` - */ -# if defined(_WIN32) /* Windows is always little endian */ \ - || defined(__LITTLE_ENDIAN__) \ - || (defined(__BYTE_ORDER__) && __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__) -# define XXH_CPU_LITTLE_ENDIAN 1 -# elif defined(__BIG_ENDIAN__) \ - || (defined(__BYTE_ORDER__) && __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__) -# define XXH_CPU_LITTLE_ENDIAN 0 -# else -/*! - * @internal - * @brief Runtime check for @ref XXH_CPU_LITTLE_ENDIAN. - * - * Most compilers will constant fold this. - */ -static int XXH_isLittleEndian(void) -{ - /* - * Portable and well-defined behavior. - * Don't use static: it is detrimental to performance. - */ - const union { xxh_u32 u; xxh_u8 c[4]; } one = { 1 }; - return one.c[0]; -} -# define XXH_CPU_LITTLE_ENDIAN XXH_isLittleEndian() -# endif -#endif - - - - -/* **************************************** -* Compiler-specific Functions and Macros -******************************************/ -#define XXH_GCC_VERSION (__GNUC__ * 100 + __GNUC_MINOR__) - -#ifdef __has_builtin -# define XXH_HAS_BUILTIN(x) __has_builtin(x) -#else -# define XXH_HAS_BUILTIN(x) 0 -#endif - -/*! - * @internal - * @def XXH_rotl32(x,r) - * @brief 32-bit rotate left. - * - * @param x The 32-bit integer to be rotated. - * @param r The number of bits to rotate. - * @pre - * @p r > 0 && @p r < 32 - * @note - * @p x and @p r may be evaluated multiple times. - * @return The rotated result. - */ -#if !defined(NO_CLANG_BUILTIN) && XXH_HAS_BUILTIN(__builtin_rotateleft32) \ - && XXH_HAS_BUILTIN(__builtin_rotateleft64) -# define XXH_rotl32 __builtin_rotateleft32 -# define XXH_rotl64 __builtin_rotateleft64 -/* Note: although _rotl exists for minGW (GCC under windows), performance seems poor */ -#elif defined(_MSC_VER) -# define XXH_rotl32(x,r) _rotl(x,r) -# define XXH_rotl64(x,r) _rotl64(x,r) -#else -# define XXH_rotl32(x,r) (((x) << (r)) | ((x) >> (32 - (r)))) -# define XXH_rotl64(x,r) (((x) << (r)) | ((x) >> (64 - (r)))) -#endif - -/*! - * @internal - * @fn xxh_u32 XXH_swap32(xxh_u32 x) - * @brief A 32-bit byteswap. - * - * @param x The 32-bit integer to byteswap. - * @return @p x, byteswapped. - */ -#if defined(_MSC_VER) /* Visual Studio */ -# define XXH_swap32 _byteswap_ulong -#elif XXH_GCC_VERSION >= 403 -# define XXH_swap32 __builtin_bswap32 -#else -static xxh_u32 XXH_swap32 (xxh_u32 x) -{ - return ((x << 24) & 0xff000000 ) | - ((x << 8) & 0x00ff0000 ) | - ((x >> 8) & 0x0000ff00 ) | - ((x >> 24) & 0x000000ff ); -} -#endif - - -/* *************************** -* Memory reads -*****************************/ - -/*! - * @internal - * @brief Enum to indicate whether a pointer is aligned. - */ -typedef enum { - XXH_aligned, /*!< Aligned */ - XXH_unaligned /*!< Possibly unaligned */ -} XXH_alignment; - -/* - * XXH_FORCE_MEMORY_ACCESS==3 is an endian-independent byteshift load. - * - * This is ideal for older compilers which don't inline memcpy. - */ -#if (defined(XXH_FORCE_MEMORY_ACCESS) && (XXH_FORCE_MEMORY_ACCESS==3)) - -XXH_FORCE_INLINE xxh_u32 XXH_readLE32(const void* memPtr) -{ - const xxh_u8* bytePtr = (const xxh_u8 *)memPtr; - return bytePtr[0] - | ((xxh_u32)bytePtr[1] << 8) - | ((xxh_u32)bytePtr[2] << 16) - | ((xxh_u32)bytePtr[3] << 24); -} - -XXH_FORCE_INLINE xxh_u32 XXH_readBE32(const void* memPtr) -{ - const xxh_u8* bytePtr = (const xxh_u8 *)memPtr; - return bytePtr[3] - | ((xxh_u32)bytePtr[2] << 8) - | ((xxh_u32)bytePtr[1] << 16) - | ((xxh_u32)bytePtr[0] << 24); -} - -#else -XXH_FORCE_INLINE xxh_u32 XXH_readLE32(const void* ptr) -{ - return XXH_CPU_LITTLE_ENDIAN ? XXH_read32(ptr) : XXH_swap32(XXH_read32(ptr)); -} - -static xxh_u32 XXH_readBE32(const void* ptr) -{ - return XXH_CPU_LITTLE_ENDIAN ? XXH_swap32(XXH_read32(ptr)) : XXH_read32(ptr); -} -#endif - -XXH_FORCE_INLINE xxh_u32 -XXH_readLE32_align(const void* ptr, XXH_alignment align) -{ - if (align==XXH_unaligned) { - return XXH_readLE32(ptr); - } else { - return XXH_CPU_LITTLE_ENDIAN ? *(const xxh_u32*)ptr : XXH_swap32(*(const xxh_u32*)ptr); - } -} - - -/* ************************************* -* Misc -***************************************/ -/*! @ingroup public */ -XXH_PUBLIC_API unsigned XXH_versionNumber (void) { return XXH_VERSION_NUMBER; } - - -/* ******************************************************************* -* 32-bit hash functions -*********************************************************************/ -/*! - * @} - * @defgroup xxh32_impl XXH32 implementation - * @ingroup impl - * @{ - */ - /* #define instead of static const, to be used as initializers */ -#define XXH_PRIME32_1 0x9E3779B1U /*!< 0b10011110001101110111100110110001 */ -#define XXH_PRIME32_2 0x85EBCA77U /*!< 0b10000101111010111100101001110111 */ -#define XXH_PRIME32_3 0xC2B2AE3DU /*!< 0b11000010101100101010111000111101 */ -#define XXH_PRIME32_4 0x27D4EB2FU /*!< 0b00100111110101001110101100101111 */ -#define XXH_PRIME32_5 0x165667B1U /*!< 0b00010110010101100110011110110001 */ - -#ifdef XXH_OLD_NAMES -# define PRIME32_1 XXH_PRIME32_1 -# define PRIME32_2 XXH_PRIME32_2 -# define PRIME32_3 XXH_PRIME32_3 -# define PRIME32_4 XXH_PRIME32_4 -# define PRIME32_5 XXH_PRIME32_5 -#endif - -/*! - * @internal - * @brief Normal stripe processing routine. - * - * This shuffles the bits so that any bit from @p input impacts several bits in - * @p acc. - * - * @param acc The accumulator lane. - * @param input The stripe of input to mix. - * @return The mixed accumulator lane. - */ -static xxh_u32 XXH32_round(xxh_u32 acc, xxh_u32 input) -{ - acc += input * XXH_PRIME32_2; - acc = XXH_rotl32(acc, 13); - acc *= XXH_PRIME32_1; -#if (defined(__SSE4_1__) || defined(__aarch64__)) && !defined(XXH_ENABLE_AUTOVECTORIZE) - /* - * UGLY HACK: - * A compiler fence is the only thing that prevents GCC and Clang from - * autovectorizing the XXH32 loop (pragmas and attributes don't work for some - * reason) without globally disabling SSE4.1. - * - * The reason we want to avoid vectorization is because despite working on - * 4 integers at a time, there are multiple factors slowing XXH32 down on - * SSE4: - * - There's a ridiculous amount of lag from pmulld (10 cycles of latency on - * newer chips!) making it slightly slower to multiply four integers at - * once compared to four integers independently. Even when pmulld was - * fastest, Sandy/Ivy Bridge, it is still not worth it to go into SSE - * just to multiply unless doing a long operation. - * - * - Four instructions are required to rotate, - * movqda tmp, v // not required with VEX encoding - * pslld tmp, 13 // tmp <<= 13 - * psrld v, 19 // x >>= 19 - * por v, tmp // x |= tmp - * compared to one for scalar: - * roll v, 13 // reliably fast across the board - * shldl v, v, 13 // Sandy Bridge and later prefer this for some reason - * - * - Instruction level parallelism is actually more beneficial here because - * the SIMD actually serializes this operation: While v1 is rotating, v2 - * can load data, while v3 can multiply. SSE forces them to operate - * together. - * - * This is also enabled on AArch64, as Clang autovectorizes it incorrectly - * and it is pointless writing a NEON implementation that is basically the - * same speed as scalar for XXH32. - */ - XXH_COMPILER_GUARD(acc); -#endif - return acc; -} - -/*! - * @internal - * @brief Mixes all bits to finalize the hash. - * - * The final mix ensures that all input bits have a chance to impact any bit in - * the output digest, resulting in an unbiased distribution. - * - * @param h32 The hash to avalanche. - * @return The avalanched hash. - */ -static xxh_u32 XXH32_avalanche(xxh_u32 h32) -{ - h32 ^= h32 >> 15; - h32 *= XXH_PRIME32_2; - h32 ^= h32 >> 13; - h32 *= XXH_PRIME32_3; - h32 ^= h32 >> 16; - return(h32); -} - -#define XXH_get32bits(p) XXH_readLE32_align(p, align) - -/*! - * @internal - * @brief Processes the last 0-15 bytes of @p ptr. - * - * There may be up to 15 bytes remaining to consume from the input. - * This final stage will digest them to ensure that all input bytes are present - * in the final mix. - * - * @param h32 The hash to finalize. - * @param ptr The pointer to the remaining input. - * @param len The remaining length, modulo 16. - * @param align Whether @p ptr is aligned. - * @return The finalized hash. - */ -static xxh_u32 -XXH32_finalize(xxh_u32 h32, const xxh_u8* ptr, size_t len, XXH_alignment align) -{ -#define XXH_PROCESS1 do { \ - h32 += (*ptr++) * XXH_PRIME32_5; \ - h32 = XXH_rotl32(h32, 11) * XXH_PRIME32_1; \ -} while (0) - -#define XXH_PROCESS4 do { \ - h32 += XXH_get32bits(ptr) * XXH_PRIME32_3; \ - ptr += 4; \ - h32 = XXH_rotl32(h32, 17) * XXH_PRIME32_4; \ -} while (0) - - if (ptr==NULL) XXH_ASSERT(len == 0); - - /* Compact rerolled version; generally faster */ - if (!XXH32_ENDJMP) { - len &= 15; - while (len >= 4) { - XXH_PROCESS4; - len -= 4; - } - while (len > 0) { - XXH_PROCESS1; - --len; - } - return XXH32_avalanche(h32); - } else { - switch(len&15) /* or switch(bEnd - p) */ { - case 12: XXH_PROCESS4; - XXH_FALLTHROUGH; - case 8: XXH_PROCESS4; - XXH_FALLTHROUGH; - case 4: XXH_PROCESS4; - return XXH32_avalanche(h32); - - case 13: XXH_PROCESS4; - XXH_FALLTHROUGH; - case 9: XXH_PROCESS4; - XXH_FALLTHROUGH; - case 5: XXH_PROCESS4; - XXH_PROCESS1; - return XXH32_avalanche(h32); - - case 14: XXH_PROCESS4; - XXH_FALLTHROUGH; - case 10: XXH_PROCESS4; - XXH_FALLTHROUGH; - case 6: XXH_PROCESS4; - XXH_PROCESS1; - XXH_PROCESS1; - return XXH32_avalanche(h32); - - case 15: XXH_PROCESS4; - XXH_FALLTHROUGH; - case 11: XXH_PROCESS4; - XXH_FALLTHROUGH; - case 7: XXH_PROCESS4; - XXH_FALLTHROUGH; - case 3: XXH_PROCESS1; - XXH_FALLTHROUGH; - case 2: XXH_PROCESS1; - XXH_FALLTHROUGH; - case 1: XXH_PROCESS1; - XXH_FALLTHROUGH; - case 0: return XXH32_avalanche(h32); - } - XXH_ASSERT(0); - return h32; /* reaching this point is deemed impossible */ - } -} - -#ifdef XXH_OLD_NAMES -# define PROCESS1 XXH_PROCESS1 -# define PROCESS4 XXH_PROCESS4 -#else -# undef XXH_PROCESS1 -# undef XXH_PROCESS4 -#endif - -/*! - * @internal - * @brief The implementation for @ref XXH32(). - * - * @param input , len , seed Directly passed from @ref XXH32(). - * @param align Whether @p input is aligned. - * @return The calculated hash. - */ -XXH_FORCE_INLINE xxh_u32 -XXH32_endian_align(const xxh_u8* input, size_t len, xxh_u32 seed, XXH_alignment align) -{ - xxh_u32 h32; - - if (input==NULL) XXH_ASSERT(len == 0); - - if (len>=16) { - const xxh_u8* const bEnd = input + len; - const xxh_u8* const limit = bEnd - 15; - xxh_u32 v1 = seed + XXH_PRIME32_1 + XXH_PRIME32_2; - xxh_u32 v2 = seed + XXH_PRIME32_2; - xxh_u32 v3 = seed + 0; - xxh_u32 v4 = seed - XXH_PRIME32_1; - - do { - v1 = XXH32_round(v1, XXH_get32bits(input)); input += 4; - v2 = XXH32_round(v2, XXH_get32bits(input)); input += 4; - v3 = XXH32_round(v3, XXH_get32bits(input)); input += 4; - v4 = XXH32_round(v4, XXH_get32bits(input)); input += 4; - } while (input < limit); - - h32 = XXH_rotl32(v1, 1) + XXH_rotl32(v2, 7) - + XXH_rotl32(v3, 12) + XXH_rotl32(v4, 18); - } else { - h32 = seed + XXH_PRIME32_5; - } - - h32 += (xxh_u32)len; - - return XXH32_finalize(h32, input, len&15, align); -} - -/*! @ingroup xxh32_family */ -XXH_PUBLIC_API XXH32_hash_t XXH32 (const void* input, size_t len, XXH32_hash_t seed) -{ -#if 0 - /* Simple version, good for code maintenance, but unfortunately slow for small inputs */ - XXH32_state_t state; - XXH32_reset(&state, seed); - XXH32_update(&state, (const xxh_u8*)input, len); - return XXH32_digest(&state); -#else - if (XXH_FORCE_ALIGN_CHECK) { - if ((((size_t)input) & 3) == 0) { /* Input is 4-bytes aligned, leverage the speed benefit */ - return XXH32_endian_align((const xxh_u8*)input, len, seed, XXH_aligned); - } } - - return XXH32_endian_align((const xxh_u8*)input, len, seed, XXH_unaligned); -#endif -} - - - -/******* Hash streaming *******/ -/*! - * @ingroup xxh32_family - */ -XXH_PUBLIC_API XXH32_state_t* XXH32_createState(void) -{ - return (XXH32_state_t*)XXH_malloc(sizeof(XXH32_state_t)); -} -/*! @ingroup xxh32_family */ -XXH_PUBLIC_API XXH_errorcode XXH32_freeState(XXH32_state_t* statePtr) -{ - XXH_free(statePtr); - return XXH_OK; -} - -/*! @ingroup xxh32_family */ -XXH_PUBLIC_API void XXH32_copyState(XXH32_state_t* dstState, const XXH32_state_t* srcState) -{ - XXH_memcpy(dstState, srcState, sizeof(*dstState)); -} - -/*! @ingroup xxh32_family */ -XXH_PUBLIC_API XXH_errorcode XXH32_reset(XXH32_state_t* statePtr, XXH32_hash_t seed) -{ - XXH32_state_t state; /* using a local state to memcpy() in order to avoid strict-aliasing warnings */ - memset(&state, 0, sizeof(state)); - state.v[0] = seed + XXH_PRIME32_1 + XXH_PRIME32_2; - state.v[1] = seed + XXH_PRIME32_2; - state.v[2] = seed + 0; - state.v[3] = seed - XXH_PRIME32_1; - /* do not write into reserved, planned to be removed in a future version */ - XXH_memcpy(statePtr, &state, sizeof(state) - sizeof(state.reserved)); - return XXH_OK; -} - - -/*! @ingroup xxh32_family */ -XXH_PUBLIC_API XXH_errorcode -XXH32_update(XXH32_state_t* state, const void* input, size_t len) -{ - if (input==NULL) { - XXH_ASSERT(len == 0); - return XXH_OK; - } - - { const xxh_u8* p = (const xxh_u8*)input; - const xxh_u8* const bEnd = p + len; - - state->total_len_32 += (XXH32_hash_t)len; - state->large_len |= (XXH32_hash_t)((len>=16) | (state->total_len_32>=16)); - - if (state->memsize + len < 16) { /* fill in tmp buffer */ - XXH_memcpy((xxh_u8*)(state->mem32) + state->memsize, input, len); - state->memsize += (XXH32_hash_t)len; - return XXH_OK; - } - - if (state->memsize) { /* some data left from previous update */ - XXH_memcpy((xxh_u8*)(state->mem32) + state->memsize, input, 16-state->memsize); - { const xxh_u32* p32 = state->mem32; - state->v[0] = XXH32_round(state->v[0], XXH_readLE32(p32)); p32++; - state->v[1] = XXH32_round(state->v[1], XXH_readLE32(p32)); p32++; - state->v[2] = XXH32_round(state->v[2], XXH_readLE32(p32)); p32++; - state->v[3] = XXH32_round(state->v[3], XXH_readLE32(p32)); - } - p += 16-state->memsize; - state->memsize = 0; - } - - if (p <= bEnd-16) { - const xxh_u8* const limit = bEnd - 16; - - do { - state->v[0] = XXH32_round(state->v[0], XXH_readLE32(p)); p+=4; - state->v[1] = XXH32_round(state->v[1], XXH_readLE32(p)); p+=4; - state->v[2] = XXH32_round(state->v[2], XXH_readLE32(p)); p+=4; - state->v[3] = XXH32_round(state->v[3], XXH_readLE32(p)); p+=4; - } while (p<=limit); - - } - - if (p < bEnd) { - XXH_memcpy(state->mem32, p, (size_t)(bEnd-p)); - state->memsize = (unsigned)(bEnd-p); - } - } - - return XXH_OK; -} - - -/*! @ingroup xxh32_family */ -XXH_PUBLIC_API XXH32_hash_t XXH32_digest(const XXH32_state_t* state) -{ - xxh_u32 h32; - - if (state->large_len) { - h32 = XXH_rotl32(state->v[0], 1) - + XXH_rotl32(state->v[1], 7) - + XXH_rotl32(state->v[2], 12) - + XXH_rotl32(state->v[3], 18); - } else { - h32 = state->v[2] /* == seed */ + XXH_PRIME32_5; - } - - h32 += state->total_len_32; - - return XXH32_finalize(h32, (const xxh_u8*)state->mem32, state->memsize, XXH_aligned); -} - - -/******* Canonical representation *******/ - -/*! - * @ingroup xxh32_family - * The default return values from XXH functions are unsigned 32 and 64 bit - * integers. - * - * The canonical representation uses big endian convention, the same convention - * as human-readable numbers (large digits first). - * - * This way, hash values can be written into a file or buffer, remaining - * comparable across different systems. - * - * The following functions allow transformation of hash values to and from their - * canonical format. - */ -XXH_PUBLIC_API void XXH32_canonicalFromHash(XXH32_canonical_t* dst, XXH32_hash_t hash) -{ - XXH_STATIC_ASSERT(sizeof(XXH32_canonical_t) == sizeof(XXH32_hash_t)); - if (XXH_CPU_LITTLE_ENDIAN) hash = XXH_swap32(hash); - XXH_memcpy(dst, &hash, sizeof(*dst)); -} -/*! @ingroup xxh32_family */ -XXH_PUBLIC_API XXH32_hash_t XXH32_hashFromCanonical(const XXH32_canonical_t* src) -{ - return XXH_readBE32(src); -} - - -#ifndef XXH_NO_LONG_LONG - -/* ******************************************************************* -* 64-bit hash functions -*********************************************************************/ -/*! - * @} - * @ingroup impl - * @{ - */ -/******* Memory access *******/ - -typedef XXH64_hash_t xxh_u64; - -#ifdef XXH_OLD_NAMES -# define U64 xxh_u64 -#endif - -#if (defined(XXH_FORCE_MEMORY_ACCESS) && (XXH_FORCE_MEMORY_ACCESS==3)) -/* - * Manual byteshift. Best for old compilers which don't inline memcpy. - * We actually directly use XXH_readLE64 and XXH_readBE64. - */ -#elif (defined(XXH_FORCE_MEMORY_ACCESS) && (XXH_FORCE_MEMORY_ACCESS==2)) - -/* Force direct memory access. Only works on CPU which support unaligned memory access in hardware */ -static xxh_u64 XXH_read64(const void* memPtr) -{ - return *(const xxh_u64*) memPtr; -} - -#elif (defined(XXH_FORCE_MEMORY_ACCESS) && (XXH_FORCE_MEMORY_ACCESS==1)) - -/* - * __pack instructions are safer, but compiler specific, hence potentially - * problematic for some compilers. - * - * Currently only defined for GCC and ICC. - */ -#ifdef XXH_OLD_NAMES -typedef union { xxh_u32 u32; xxh_u64 u64; } __attribute__((packed)) unalign64; -#endif -static xxh_u64 XXH_read64(const void* ptr) -{ - typedef union { xxh_u32 u32; xxh_u64 u64; } __attribute__((packed)) xxh_unalign64; - return ((const xxh_unalign64*)ptr)->u64; -} - -#else - -/* - * Portable and safe solution. Generally efficient. - * see: http://fastcompression.blogspot.com/2015/08/accessing-unaligned-memory.html - */ -static xxh_u64 XXH_read64(const void* memPtr) -{ - xxh_u64 val; - XXH_memcpy(&val, memPtr, sizeof(val)); - return val; -} - -#endif /* XXH_FORCE_DIRECT_MEMORY_ACCESS */ - -#if defined(_MSC_VER) /* Visual Studio */ -# define XXH_swap64 _byteswap_uint64 -#elif XXH_GCC_VERSION >= 403 -# define XXH_swap64 __builtin_bswap64 -#else -static xxh_u64 XXH_swap64(xxh_u64 x) -{ - return ((x << 56) & 0xff00000000000000ULL) | - ((x << 40) & 0x00ff000000000000ULL) | - ((x << 24) & 0x0000ff0000000000ULL) | - ((x << 8) & 0x000000ff00000000ULL) | - ((x >> 8) & 0x00000000ff000000ULL) | - ((x >> 24) & 0x0000000000ff0000ULL) | - ((x >> 40) & 0x000000000000ff00ULL) | - ((x >> 56) & 0x00000000000000ffULL); -} -#endif - - -/* XXH_FORCE_MEMORY_ACCESS==3 is an endian-independent byteshift load. */ -#if (defined(XXH_FORCE_MEMORY_ACCESS) && (XXH_FORCE_MEMORY_ACCESS==3)) - -XXH_FORCE_INLINE xxh_u64 XXH_readLE64(const void* memPtr) -{ - const xxh_u8* bytePtr = (const xxh_u8 *)memPtr; - return bytePtr[0] - | ((xxh_u64)bytePtr[1] << 8) - | ((xxh_u64)bytePtr[2] << 16) - | ((xxh_u64)bytePtr[3] << 24) - | ((xxh_u64)bytePtr[4] << 32) - | ((xxh_u64)bytePtr[5] << 40) - | ((xxh_u64)bytePtr[6] << 48) - | ((xxh_u64)bytePtr[7] << 56); -} - -XXH_FORCE_INLINE xxh_u64 XXH_readBE64(const void* memPtr) -{ - const xxh_u8* bytePtr = (const xxh_u8 *)memPtr; - return bytePtr[7] - | ((xxh_u64)bytePtr[6] << 8) - | ((xxh_u64)bytePtr[5] << 16) - | ((xxh_u64)bytePtr[4] << 24) - | ((xxh_u64)bytePtr[3] << 32) - | ((xxh_u64)bytePtr[2] << 40) - | ((xxh_u64)bytePtr[1] << 48) - | ((xxh_u64)bytePtr[0] << 56); -} - -#else -XXH_FORCE_INLINE xxh_u64 XXH_readLE64(const void* ptr) -{ - return XXH_CPU_LITTLE_ENDIAN ? XXH_read64(ptr) : XXH_swap64(XXH_read64(ptr)); -} - -static xxh_u64 XXH_readBE64(const void* ptr) -{ - return XXH_CPU_LITTLE_ENDIAN ? XXH_swap64(XXH_read64(ptr)) : XXH_read64(ptr); -} -#endif - -XXH_FORCE_INLINE xxh_u64 -XXH_readLE64_align(const void* ptr, XXH_alignment align) -{ - if (align==XXH_unaligned) - return XXH_readLE64(ptr); - else - return XXH_CPU_LITTLE_ENDIAN ? *(const xxh_u64*)ptr : XXH_swap64(*(const xxh_u64*)ptr); -} - - -/******* xxh64 *******/ -/*! - * @} - * @defgroup xxh64_impl XXH64 implementation - * @ingroup impl - * @{ - */ -/* #define rather that static const, to be used as initializers */ -#define XXH_PRIME64_1 0x9E3779B185EBCA87ULL /*!< 0b1001111000110111011110011011000110000101111010111100101010000111 */ -#define XXH_PRIME64_2 0xC2B2AE3D27D4EB4FULL /*!< 0b1100001010110010101011100011110100100111110101001110101101001111 */ -#define XXH_PRIME64_3 0x165667B19E3779F9ULL /*!< 0b0001011001010110011001111011000110011110001101110111100111111001 */ -#define XXH_PRIME64_4 0x85EBCA77C2B2AE63ULL /*!< 0b1000010111101011110010100111011111000010101100101010111001100011 */ -#define XXH_PRIME64_5 0x27D4EB2F165667C5ULL /*!< 0b0010011111010100111010110010111100010110010101100110011111000101 */ - -#ifdef XXH_OLD_NAMES -# define PRIME64_1 XXH_PRIME64_1 -# define PRIME64_2 XXH_PRIME64_2 -# define PRIME64_3 XXH_PRIME64_3 -# define PRIME64_4 XXH_PRIME64_4 -# define PRIME64_5 XXH_PRIME64_5 -#endif - -static xxh_u64 XXH64_round(xxh_u64 acc, xxh_u64 input) -{ - acc += input * XXH_PRIME64_2; - acc = XXH_rotl64(acc, 31); - acc *= XXH_PRIME64_1; - return acc; -} - -static xxh_u64 XXH64_mergeRound(xxh_u64 acc, xxh_u64 val) -{ - val = XXH64_round(0, val); - acc ^= val; - acc = acc * XXH_PRIME64_1 + XXH_PRIME64_4; - return acc; -} - -static xxh_u64 XXH64_avalanche(xxh_u64 h64) -{ - h64 ^= h64 >> 33; - h64 *= XXH_PRIME64_2; - h64 ^= h64 >> 29; - h64 *= XXH_PRIME64_3; - h64 ^= h64 >> 32; - return h64; -} - - -#define XXH_get64bits(p) XXH_readLE64_align(p, align) - -static xxh_u64 -XXH64_finalize(xxh_u64 h64, const xxh_u8* ptr, size_t len, XXH_alignment align) -{ - if (ptr==NULL) XXH_ASSERT(len == 0); - len &= 31; - while (len >= 8) { - xxh_u64 const k1 = XXH64_round(0, XXH_get64bits(ptr)); - ptr += 8; - h64 ^= k1; - h64 = XXH_rotl64(h64,27) * XXH_PRIME64_1 + XXH_PRIME64_4; - len -= 8; - } - if (len >= 4) { - h64 ^= (xxh_u64)(XXH_get32bits(ptr)) * XXH_PRIME64_1; - ptr += 4; - h64 = XXH_rotl64(h64, 23) * XXH_PRIME64_2 + XXH_PRIME64_3; - len -= 4; - } - while (len > 0) { - h64 ^= (*ptr++) * XXH_PRIME64_5; - h64 = XXH_rotl64(h64, 11) * XXH_PRIME64_1; - --len; - } - return XXH64_avalanche(h64); -} - -#ifdef XXH_OLD_NAMES -# define PROCESS1_64 XXH_PROCESS1_64 -# define PROCESS4_64 XXH_PROCESS4_64 -# define PROCESS8_64 XXH_PROCESS8_64 -#else -# undef XXH_PROCESS1_64 -# undef XXH_PROCESS4_64 -# undef XXH_PROCESS8_64 -#endif - -XXH_FORCE_INLINE xxh_u64 -XXH64_endian_align(const xxh_u8* input, size_t len, xxh_u64 seed, XXH_alignment align) -{ - xxh_u64 h64; - if (input==NULL) XXH_ASSERT(len == 0); - - if (len>=32) { - const xxh_u8* const bEnd = input + len; - const xxh_u8* const limit = bEnd - 31; - xxh_u64 v1 = seed + XXH_PRIME64_1 + XXH_PRIME64_2; - xxh_u64 v2 = seed + XXH_PRIME64_2; - xxh_u64 v3 = seed + 0; - xxh_u64 v4 = seed - XXH_PRIME64_1; - - do { - v1 = XXH64_round(v1, XXH_get64bits(input)); input+=8; - v2 = XXH64_round(v2, XXH_get64bits(input)); input+=8; - v3 = XXH64_round(v3, XXH_get64bits(input)); input+=8; - v4 = XXH64_round(v4, XXH_get64bits(input)); input+=8; - } while (inputtotal_len += len; - - if (state->memsize + len < 32) { /* fill in tmp buffer */ - XXH_memcpy(((xxh_u8*)state->mem64) + state->memsize, input, len); - state->memsize += (xxh_u32)len; - return XXH_OK; - } - - if (state->memsize) { /* tmp buffer is full */ - XXH_memcpy(((xxh_u8*)state->mem64) + state->memsize, input, 32-state->memsize); - state->v[0] = XXH64_round(state->v[0], XXH_readLE64(state->mem64+0)); - state->v[1] = XXH64_round(state->v[1], XXH_readLE64(state->mem64+1)); - state->v[2] = XXH64_round(state->v[2], XXH_readLE64(state->mem64+2)); - state->v[3] = XXH64_round(state->v[3], XXH_readLE64(state->mem64+3)); - p += 32 - state->memsize; - state->memsize = 0; - } - - if (p+32 <= bEnd) { - const xxh_u8* const limit = bEnd - 32; - - do { - state->v[0] = XXH64_round(state->v[0], XXH_readLE64(p)); p+=8; - state->v[1] = XXH64_round(state->v[1], XXH_readLE64(p)); p+=8; - state->v[2] = XXH64_round(state->v[2], XXH_readLE64(p)); p+=8; - state->v[3] = XXH64_round(state->v[3], XXH_readLE64(p)); p+=8; - } while (p<=limit); - - } - - if (p < bEnd) { - XXH_memcpy(state->mem64, p, (size_t)(bEnd-p)); - state->memsize = (unsigned)(bEnd-p); - } - } - - return XXH_OK; -} - - -/*! @ingroup xxh64_family */ -XXH_PUBLIC_API XXH64_hash_t XXH64_digest(const XXH64_state_t* state) -{ - xxh_u64 h64; - - if (state->total_len >= 32) { - h64 = XXH_rotl64(state->v[0], 1) + XXH_rotl64(state->v[1], 7) + XXH_rotl64(state->v[2], 12) + XXH_rotl64(state->v[3], 18); - h64 = XXH64_mergeRound(h64, state->v[0]); - h64 = XXH64_mergeRound(h64, state->v[1]); - h64 = XXH64_mergeRound(h64, state->v[2]); - h64 = XXH64_mergeRound(h64, state->v[3]); - } else { - h64 = state->v[2] /*seed*/ + XXH_PRIME64_5; - } - - h64 += (xxh_u64) state->total_len; - - return XXH64_finalize(h64, (const xxh_u8*)state->mem64, (size_t)state->total_len, XXH_aligned); -} - - -/******* Canonical representation *******/ - -/*! @ingroup xxh64_family */ -XXH_PUBLIC_API void XXH64_canonicalFromHash(XXH64_canonical_t* dst, XXH64_hash_t hash) -{ - XXH_STATIC_ASSERT(sizeof(XXH64_canonical_t) == sizeof(XXH64_hash_t)); - if (XXH_CPU_LITTLE_ENDIAN) hash = XXH_swap64(hash); - XXH_memcpy(dst, &hash, sizeof(*dst)); -} - -/*! @ingroup xxh64_family */ -XXH_PUBLIC_API XXH64_hash_t XXH64_hashFromCanonical(const XXH64_canonical_t* src) -{ - return XXH_readBE64(src); -} - -#ifndef XXH_NO_XXH3 - -/* ********************************************************************* -* XXH3 -* New generation hash designed for speed on small keys and vectorization -************************************************************************ */ -/*! - * @} - * @defgroup xxh3_impl XXH3 implementation - * @ingroup impl - * @{ - */ - -/* === Compiler specifics === */ - -#if ((defined(sun) || defined(__sun)) && __cplusplus) /* Solaris includes __STDC_VERSION__ with C++. Tested with GCC 5.5 */ -# define XXH_RESTRICT /* disable */ -#elif defined (__STDC_VERSION__) && __STDC_VERSION__ >= 199901L /* >= C99 */ -# define XXH_RESTRICT restrict -#else -/* Note: it might be useful to define __restrict or __restrict__ for some C++ compilers */ -# define XXH_RESTRICT /* disable */ -#endif - -#if (defined(__GNUC__) && (__GNUC__ >= 3)) \ - || (defined(__INTEL_COMPILER) && (__INTEL_COMPILER >= 800)) \ - || defined(__clang__) -# define XXH_likely(x) __builtin_expect(x, 1) -# define XXH_unlikely(x) __builtin_expect(x, 0) -#else -# define XXH_likely(x) (x) -# define XXH_unlikely(x) (x) -#endif - -#if defined(__GNUC__) -# if defined(__AVX2__) -# include -# elif defined(__SSE2__) -# include -# elif defined(__ARM_NEON__) || defined(__ARM_NEON) -# define inline __inline__ /* circumvent a clang bug */ -# include -# undef inline -# endif -#elif defined(_MSC_VER) -# include -#endif - -/* - * One goal of XXH3 is to make it fast on both 32-bit and 64-bit, while - * remaining a true 64-bit/128-bit hash function. - * - * This is done by prioritizing a subset of 64-bit operations that can be - * emulated without too many steps on the average 32-bit machine. - * - * For example, these two lines seem similar, and run equally fast on 64-bit: - * - * xxh_u64 x; - * x ^= (x >> 47); // good - * x ^= (x >> 13); // bad - * - * However, to a 32-bit machine, there is a major difference. - * - * x ^= (x >> 47) looks like this: - * - * x.lo ^= (x.hi >> (47 - 32)); - * - * while x ^= (x >> 13) looks like this: - * - * // note: funnel shifts are not usually cheap. - * x.lo ^= (x.lo >> 13) | (x.hi << (32 - 13)); - * x.hi ^= (x.hi >> 13); - * - * The first one is significantly faster than the second, simply because the - * shift is larger than 32. This means: - * - All the bits we need are in the upper 32 bits, so we can ignore the lower - * 32 bits in the shift. - * - The shift result will always fit in the lower 32 bits, and therefore, - * we can ignore the upper 32 bits in the xor. - * - * Thanks to this optimization, XXH3 only requires these features to be efficient: - * - * - Usable unaligned access - * - A 32-bit or 64-bit ALU - * - If 32-bit, a decent ADC instruction - * - A 32 or 64-bit multiply with a 64-bit result - * - For the 128-bit variant, a decent byteswap helps short inputs. - * - * The first two are already required by XXH32, and almost all 32-bit and 64-bit - * platforms which can run XXH32 can run XXH3 efficiently. - * - * Thumb-1, the classic 16-bit only subset of ARM's instruction set, is one - * notable exception. - * - * First of all, Thumb-1 lacks support for the UMULL instruction which - * performs the important long multiply. This means numerous __aeabi_lmul - * calls. - * - * Second of all, the 8 functional registers are just not enough. - * Setup for __aeabi_lmul, byteshift loads, pointers, and all arithmetic need - * Lo registers, and this shuffling results in thousands more MOVs than A32. - * - * A32 and T32 don't have this limitation. They can access all 14 registers, - * do a 32->64 multiply with UMULL, and the flexible operand allowing free - * shifts is helpful, too. - * - * Therefore, we do a quick sanity check. - * - * If compiling Thumb-1 for a target which supports ARM instructions, we will - * emit a warning, as it is not a "sane" platform to compile for. - * - * Usually, if this happens, it is because of an accident and you probably need - * to specify -march, as you likely meant to compile for a newer architecture. - * - * Credit: large sections of the vectorial and asm source code paths - * have been contributed by @easyaspi314 - */ -#if defined(__thumb__) && !defined(__thumb2__) && defined(__ARM_ARCH_ISA_ARM) -# warning "XXH3 is highly inefficient without ARM or Thumb-2." -#endif - -/* ========================================== - * Vectorization detection - * ========================================== */ - -#ifdef XXH_DOXYGEN -/*! - * @ingroup tuning - * @brief Overrides the vectorization implementation chosen for XXH3. - * - * Can be defined to 0 to disable SIMD or any of the values mentioned in - * @ref XXH_VECTOR_TYPE. - * - * If this is not defined, it uses predefined macros to determine the best - * implementation. - */ -# define XXH_VECTOR XXH_SCALAR -/*! - * @ingroup tuning - * @brief Possible values for @ref XXH_VECTOR. - * - * Note that these are actually implemented as macros. - * - * If this is not defined, it is detected automatically. - * @ref XXH_X86DISPATCH overrides this. - */ -enum XXH_VECTOR_TYPE /* fake enum */ { - XXH_SCALAR = 0, /*!< Portable scalar version */ - XXH_SSE2 = 1, /*!< - * SSE2 for Pentium 4, Opteron, all x86_64. - * - * @note SSE2 is also guaranteed on Windows 10, macOS, and - * Android x86. - */ - XXH_AVX2 = 2, /*!< AVX2 for Haswell and Bulldozer */ - XXH_AVX512 = 3, /*!< AVX512 for Skylake and Icelake */ - XXH_NEON = 4, /*!< NEON for most ARMv7-A and all AArch64 */ - XXH_VSX = 5, /*!< VSX and ZVector for POWER8/z13 (64-bit) */ -}; -/*! - * @ingroup tuning - * @brief Selects the minimum alignment for XXH3's accumulators. - * - * When using SIMD, this should match the alignment reqired for said vector - * type, so, for example, 32 for AVX2. - * - * Default: Auto detected. - */ -# define XXH_ACC_ALIGN 8 -#endif - -/* Actual definition */ -#ifndef XXH_DOXYGEN -# define XXH_SCALAR 0 -# define XXH_SSE2 1 -# define XXH_AVX2 2 -# define XXH_AVX512 3 -# define XXH_NEON 4 -# define XXH_VSX 5 -#endif - -#ifndef XXH_VECTOR /* can be defined on command line */ -# if defined(__AVX512F__) -# define XXH_VECTOR XXH_AVX512 -# elif defined(__AVX2__) -# define XXH_VECTOR XXH_AVX2 -# elif defined(__SSE2__) || defined(_M_AMD64) || defined(_M_X64) || (defined(_M_IX86_FP) && (_M_IX86_FP == 2)) -# define XXH_VECTOR XXH_SSE2 -# elif ( \ - defined(__ARM_NEON__) || defined(__ARM_NEON) /* gcc */ \ - || defined(_M_ARM64) || defined(_M_ARM_ARMV7VE) /* msvc */ \ - ) && ( \ - defined(_WIN32) || defined(__LITTLE_ENDIAN__) /* little endian only */ \ - || (defined(__BYTE_ORDER__) && __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__) \ - ) -# define XXH_VECTOR XXH_NEON -# elif (defined(__PPC64__) && defined(__POWER8_VECTOR__)) \ - || (defined(__s390x__) && defined(__VEC__)) \ - && defined(__GNUC__) /* TODO: IBM XL */ -# define XXH_VECTOR XXH_VSX -# else -# define XXH_VECTOR XXH_SCALAR -# endif -#endif - -/* - * Controls the alignment of the accumulator, - * for compatibility with aligned vector loads, which are usually faster. - */ -#ifndef XXH_ACC_ALIGN -# if defined(XXH_X86DISPATCH) -# define XXH_ACC_ALIGN 64 /* for compatibility with avx512 */ -# elif XXH_VECTOR == XXH_SCALAR /* scalar */ -# define XXH_ACC_ALIGN 8 -# elif XXH_VECTOR == XXH_SSE2 /* sse2 */ -# define XXH_ACC_ALIGN 16 -# elif XXH_VECTOR == XXH_AVX2 /* avx2 */ -# define XXH_ACC_ALIGN 32 -# elif XXH_VECTOR == XXH_NEON /* neon */ -# define XXH_ACC_ALIGN 16 -# elif XXH_VECTOR == XXH_VSX /* vsx */ -# define XXH_ACC_ALIGN 16 -# elif XXH_VECTOR == XXH_AVX512 /* avx512 */ -# define XXH_ACC_ALIGN 64 -# endif -#endif - -#if defined(XXH_X86DISPATCH) || XXH_VECTOR == XXH_SSE2 \ - || XXH_VECTOR == XXH_AVX2 || XXH_VECTOR == XXH_AVX512 -# define XXH_SEC_ALIGN XXH_ACC_ALIGN -#else -# define XXH_SEC_ALIGN 8 -#endif - -/* - * UGLY HACK: - * GCC usually generates the best code with -O3 for xxHash. - * - * However, when targeting AVX2, it is overzealous in its unrolling resulting - * in code roughly 3/4 the speed of Clang. - * - * There are other issues, such as GCC splitting _mm256_loadu_si256 into - * _mm_loadu_si128 + _mm256_inserti128_si256. This is an optimization which - * only applies to Sandy and Ivy Bridge... which don't even support AVX2. - * - * That is why when compiling the AVX2 version, it is recommended to use either - * -O2 -mavx2 -march=haswell - * or - * -O2 -mavx2 -mno-avx256-split-unaligned-load - * for decent performance, or to use Clang instead. - * - * Fortunately, we can control the first one with a pragma that forces GCC into - * -O2, but the other one we can't control without "failed to inline always - * inline function due to target mismatch" warnings. - */ -#if XXH_VECTOR == XXH_AVX2 /* AVX2 */ \ - && defined(__GNUC__) && !defined(__clang__) /* GCC, not Clang */ \ - && defined(__OPTIMIZE__) && !defined(__OPTIMIZE_SIZE__) /* respect -O0 and -Os */ -# pragma GCC push_options -# pragma GCC optimize("-O2") -#endif - - -#if XXH_VECTOR == XXH_NEON -/* - * NEON's setup for vmlal_u32 is a little more complicated than it is on - * SSE2, AVX2, and VSX. - * - * While PMULUDQ and VMULEUW both perform a mask, VMLAL.U32 performs an upcast. - * - * To do the same operation, the 128-bit 'Q' register needs to be split into - * two 64-bit 'D' registers, performing this operation:: - * - * [ a | b ] - * | '---------. .--------' | - * | x | - * | .---------' '--------. | - * [ a & 0xFFFFFFFF | b & 0xFFFFFFFF ],[ a >> 32 | b >> 32 ] - * - * Due to significant changes in aarch64, the fastest method for aarch64 is - * completely different than the fastest method for ARMv7-A. - * - * ARMv7-A treats D registers as unions overlaying Q registers, so modifying - * D11 will modify the high half of Q5. This is similar to how modifying AH - * will only affect bits 8-15 of AX on x86. - * - * VZIP takes two registers, and puts even lanes in one register and odd lanes - * in the other. - * - * On ARMv7-A, this strangely modifies both parameters in place instead of - * taking the usual 3-operand form. - * - * Therefore, if we want to do this, we can simply use a D-form VZIP.32 on the - * lower and upper halves of the Q register to end up with the high and low - * halves where we want - all in one instruction. - * - * vzip.32 d10, d11 @ d10 = { d10[0], d11[0] }; d11 = { d10[1], d11[1] } - * - * Unfortunately we need inline assembly for this: Instructions modifying two - * registers at once is not possible in GCC or Clang's IR, and they have to - * create a copy. - * - * aarch64 requires a different approach. - * - * In order to make it easier to write a decent compiler for aarch64, many - * quirks were removed, such as conditional execution. - * - * NEON was also affected by this. - * - * aarch64 cannot access the high bits of a Q-form register, and writes to a - * D-form register zero the high bits, similar to how writes to W-form scalar - * registers (or DWORD registers on x86_64) work. - * - * The formerly free vget_high intrinsics now require a vext (with a few - * exceptions) - * - * Additionally, VZIP was replaced by ZIP1 and ZIP2, which are the equivalent - * of PUNPCKL* and PUNPCKH* in SSE, respectively, in order to only modify one - * operand. - * - * The equivalent of the VZIP.32 on the lower and upper halves would be this - * mess: - * - * ext v2.4s, v0.4s, v0.4s, #2 // v2 = { v0[2], v0[3], v0[0], v0[1] } - * zip1 v1.2s, v0.2s, v2.2s // v1 = { v0[0], v2[0] } - * zip2 v0.2s, v0.2s, v1.2s // v0 = { v0[1], v2[1] } - * - * Instead, we use a literal downcast, vmovn_u64 (XTN), and vshrn_n_u64 (SHRN): - * - * shrn v1.2s, v0.2d, #32 // v1 = (uint32x2_t)(v0 >> 32); - * xtn v0.2s, v0.2d // v0 = (uint32x2_t)(v0 & 0xFFFFFFFF); - * - * This is available on ARMv7-A, but is less efficient than a single VZIP.32. - */ - -/*! - * Function-like macro: - * void XXH_SPLIT_IN_PLACE(uint64x2_t &in, uint32x2_t &outLo, uint32x2_t &outHi) - * { - * outLo = (uint32x2_t)(in & 0xFFFFFFFF); - * outHi = (uint32x2_t)(in >> 32); - * in = UNDEFINED; - * } - */ -# if !defined(XXH_NO_VZIP_HACK) /* define to disable */ \ - && defined(__GNUC__) \ - && !defined(__aarch64__) && !defined(__arm64__) && !defined(_M_ARM64) -# define XXH_SPLIT_IN_PLACE(in, outLo, outHi) \ - do { \ - /* Undocumented GCC/Clang operand modifier: %e0 = lower D half, %f0 = upper D half */ \ - /* https://github.com/gcc-mirror/gcc/blob/38cf91e5/gcc/config/arm/arm.c#L22486 */ \ - /* https://github.com/llvm-mirror/llvm/blob/2c4ca683/lib/Target/ARM/ARMAsmPrinter.cpp#L399 */ \ - __asm__("vzip.32 %e0, %f0" : "+w" (in)); \ - (outLo) = vget_low_u32 (vreinterpretq_u32_u64(in)); \ - (outHi) = vget_high_u32(vreinterpretq_u32_u64(in)); \ - } while (0) -# else -# define XXH_SPLIT_IN_PLACE(in, outLo, outHi) \ - do { \ - (outLo) = vmovn_u64 (in); \ - (outHi) = vshrn_n_u64 ((in), 32); \ - } while (0) -# endif -#endif /* XXH_VECTOR == XXH_NEON */ - -/* - * VSX and Z Vector helpers. - * - * This is very messy, and any pull requests to clean this up are welcome. - * - * There are a lot of problems with supporting VSX and s390x, due to - * inconsistent intrinsics, spotty coverage, and multiple endiannesses. - */ -#if XXH_VECTOR == XXH_VSX -# if defined(__s390x__) -# include -# else -/* gcc's altivec.h can have the unwanted consequence to unconditionally - * #define bool, vector, and pixel keywords, - * with bad consequences for programs already using these keywords for other purposes. - * The paragraph defining these macros is skipped when __APPLE_ALTIVEC__ is defined. - * __APPLE_ALTIVEC__ is _generally_ defined automatically by the compiler, - * but it seems that, in some cases, it isn't. - * Force the build macro to be defined, so that keywords are not altered. - */ -# if defined(__GNUC__) && !defined(__APPLE_ALTIVEC__) -# define __APPLE_ALTIVEC__ -# endif -# include -# endif - -typedef __vector unsigned long long xxh_u64x2; -typedef __vector unsigned char xxh_u8x16; -typedef __vector unsigned xxh_u32x4; - -# ifndef XXH_VSX_BE -# if defined(__BIG_ENDIAN__) \ - || (defined(__BYTE_ORDER__) && __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__) -# define XXH_VSX_BE 1 -# elif defined(__VEC_ELEMENT_REG_ORDER__) && __VEC_ELEMENT_REG_ORDER__ == __ORDER_BIG_ENDIAN__ -# warning "-maltivec=be is not recommended. Please use native endianness." -# define XXH_VSX_BE 1 -# else -# define XXH_VSX_BE 0 -# endif -# endif /* !defined(XXH_VSX_BE) */ - -# if XXH_VSX_BE -# if defined(__POWER9_VECTOR__) || (defined(__clang__) && defined(__s390x__)) -# define XXH_vec_revb vec_revb -# else -/*! - * A polyfill for POWER9's vec_revb(). - */ -XXH_FORCE_INLINE xxh_u64x2 XXH_vec_revb(xxh_u64x2 val) -{ - xxh_u8x16 const vByteSwap = { 0x07, 0x06, 0x05, 0x04, 0x03, 0x02, 0x01, 0x00, - 0x0F, 0x0E, 0x0D, 0x0C, 0x0B, 0x0A, 0x09, 0x08 }; - return vec_perm(val, val, vByteSwap); -} -# endif -# endif /* XXH_VSX_BE */ - -/*! - * Performs an unaligned vector load and byte swaps it on big endian. - */ -XXH_FORCE_INLINE xxh_u64x2 XXH_vec_loadu(const void *ptr) -{ - xxh_u64x2 ret; - XXH_memcpy(&ret, ptr, sizeof(xxh_u64x2)); -# if XXH_VSX_BE - ret = XXH_vec_revb(ret); -# endif - return ret; -} - -/* - * vec_mulo and vec_mule are very problematic intrinsics on PowerPC - * - * These intrinsics weren't added until GCC 8, despite existing for a while, - * and they are endian dependent. Also, their meaning swap depending on version. - * */ -# if defined(__s390x__) - /* s390x is always big endian, no issue on this platform */ -# define XXH_vec_mulo vec_mulo -# define XXH_vec_mule vec_mule -# elif defined(__clang__) && XXH_HAS_BUILTIN(__builtin_altivec_vmuleuw) -/* Clang has a better way to control this, we can just use the builtin which doesn't swap. */ -# define XXH_vec_mulo __builtin_altivec_vmulouw -# define XXH_vec_mule __builtin_altivec_vmuleuw -# else -/* gcc needs inline assembly */ -/* Adapted from https://github.com/google/highwayhash/blob/master/highwayhash/hh_vsx.h. */ -XXH_FORCE_INLINE xxh_u64x2 XXH_vec_mulo(xxh_u32x4 a, xxh_u32x4 b) -{ - xxh_u64x2 result; - __asm__("vmulouw %0, %1, %2" : "=v" (result) : "v" (a), "v" (b)); - return result; -} -XXH_FORCE_INLINE xxh_u64x2 XXH_vec_mule(xxh_u32x4 a, xxh_u32x4 b) -{ - xxh_u64x2 result; - __asm__("vmuleuw %0, %1, %2" : "=v" (result) : "v" (a), "v" (b)); - return result; -} -# endif /* XXH_vec_mulo, XXH_vec_mule */ -#endif /* XXH_VECTOR == XXH_VSX */ - - -/* prefetch - * can be disabled, by declaring XXH_NO_PREFETCH build macro */ -#if defined(XXH_NO_PREFETCH) -# define XXH_PREFETCH(ptr) (void)(ptr) /* disabled */ -#else -# if defined(_MSC_VER) && (defined(_M_X64) || defined(_M_IX86)) /* _mm_prefetch() not defined outside of x86/x64 */ -# include /* https://msdn.microsoft.com/fr-fr/library/84szxsww(v=vs.90).aspx */ -# define XXH_PREFETCH(ptr) _mm_prefetch((const char*)(ptr), _MM_HINT_T0) -# elif defined(__GNUC__) && ( (__GNUC__ >= 4) || ( (__GNUC__ == 3) && (__GNUC_MINOR__ >= 1) ) ) -# define XXH_PREFETCH(ptr) __builtin_prefetch((ptr), 0 /* rw==read */, 3 /* locality */) -# else -# define XXH_PREFETCH(ptr) (void)(ptr) /* disabled */ -# endif -#endif /* XXH_NO_PREFETCH */ - - -/* ========================================== - * XXH3 default settings - * ========================================== */ - -#define XXH_SECRET_DEFAULT_SIZE 192 /* minimum XXH3_SECRET_SIZE_MIN */ - -#if (XXH_SECRET_DEFAULT_SIZE < XXH3_SECRET_SIZE_MIN) -# error "default keyset is not large enough" -#endif - -/*! Pseudorandom secret taken directly from FARSH. */ -XXH_ALIGN(64) static const xxh_u8 XXH3_kSecret[XXH_SECRET_DEFAULT_SIZE] = { - 0xb8, 0xfe, 0x6c, 0x39, 0x23, 0xa4, 0x4b, 0xbe, 0x7c, 0x01, 0x81, 0x2c, 0xf7, 0x21, 0xad, 0x1c, - 0xde, 0xd4, 0x6d, 0xe9, 0x83, 0x90, 0x97, 0xdb, 0x72, 0x40, 0xa4, 0xa4, 0xb7, 0xb3, 0x67, 0x1f, - 0xcb, 0x79, 0xe6, 0x4e, 0xcc, 0xc0, 0xe5, 0x78, 0x82, 0x5a, 0xd0, 0x7d, 0xcc, 0xff, 0x72, 0x21, - 0xb8, 0x08, 0x46, 0x74, 0xf7, 0x43, 0x24, 0x8e, 0xe0, 0x35, 0x90, 0xe6, 0x81, 0x3a, 0x26, 0x4c, - 0x3c, 0x28, 0x52, 0xbb, 0x91, 0xc3, 0x00, 0xcb, 0x88, 0xd0, 0x65, 0x8b, 0x1b, 0x53, 0x2e, 0xa3, - 0x71, 0x64, 0x48, 0x97, 0xa2, 0x0d, 0xf9, 0x4e, 0x38, 0x19, 0xef, 0x46, 0xa9, 0xde, 0xac, 0xd8, - 0xa8, 0xfa, 0x76, 0x3f, 0xe3, 0x9c, 0x34, 0x3f, 0xf9, 0xdc, 0xbb, 0xc7, 0xc7, 0x0b, 0x4f, 0x1d, - 0x8a, 0x51, 0xe0, 0x4b, 0xcd, 0xb4, 0x59, 0x31, 0xc8, 0x9f, 0x7e, 0xc9, 0xd9, 0x78, 0x73, 0x64, - 0xea, 0xc5, 0xac, 0x83, 0x34, 0xd3, 0xeb, 0xc3, 0xc5, 0x81, 0xa0, 0xff, 0xfa, 0x13, 0x63, 0xeb, - 0x17, 0x0d, 0xdd, 0x51, 0xb7, 0xf0, 0xda, 0x49, 0xd3, 0x16, 0x55, 0x26, 0x29, 0xd4, 0x68, 0x9e, - 0x2b, 0x16, 0xbe, 0x58, 0x7d, 0x47, 0xa1, 0xfc, 0x8f, 0xf8, 0xb8, 0xd1, 0x7a, 0xd0, 0x31, 0xce, - 0x45, 0xcb, 0x3a, 0x8f, 0x95, 0x16, 0x04, 0x28, 0xaf, 0xd7, 0xfb, 0xca, 0xbb, 0x4b, 0x40, 0x7e, -}; - - -#ifdef XXH_OLD_NAMES -# define kSecret XXH3_kSecret -#endif - -#ifdef XXH_DOXYGEN -/*! - * @brief Calculates a 32-bit to 64-bit long multiply. - * - * Implemented as a macro. - * - * Wraps `__emulu` on MSVC x86 because it tends to call `__allmul` when it doesn't - * need to (but it shouldn't need to anyways, it is about 7 instructions to do - * a 64x64 multiply...). Since we know that this will _always_ emit `MULL`, we - * use that instead of the normal method. - * - * If you are compiling for platforms like Thumb-1 and don't have a better option, - * you may also want to write your own long multiply routine here. - * - * @param x, y Numbers to be multiplied - * @return 64-bit product of the low 32 bits of @p x and @p y. - */ -XXH_FORCE_INLINE xxh_u64 -XXH_mult32to64(xxh_u64 x, xxh_u64 y) -{ - return (x & 0xFFFFFFFF) * (y & 0xFFFFFFFF); -} -#elif defined(_MSC_VER) && defined(_M_IX86) -# include -# define XXH_mult32to64(x, y) __emulu((unsigned)(x), (unsigned)(y)) -#else -/* - * Downcast + upcast is usually better than masking on older compilers like - * GCC 4.2 (especially 32-bit ones), all without affecting newer compilers. - * - * The other method, (x & 0xFFFFFFFF) * (y & 0xFFFFFFFF), will AND both operands - * and perform a full 64x64 multiply -- entirely redundant on 32-bit. - */ -# define XXH_mult32to64(x, y) ((xxh_u64)(xxh_u32)(x) * (xxh_u64)(xxh_u32)(y)) -#endif - -/*! - * @brief Calculates a 64->128-bit long multiply. - * - * Uses `__uint128_t` and `_umul128` if available, otherwise uses a scalar - * version. - * - * @param lhs , rhs The 64-bit integers to be multiplied - * @return The 128-bit result represented in an @ref XXH128_hash_t. - */ -static XXH128_hash_t -XXH_mult64to128(xxh_u64 lhs, xxh_u64 rhs) -{ - /* - * GCC/Clang __uint128_t method. - * - * On most 64-bit targets, GCC and Clang define a __uint128_t type. - * This is usually the best way as it usually uses a native long 64-bit - * multiply, such as MULQ on x86_64 or MUL + UMULH on aarch64. - * - * Usually. - * - * Despite being a 32-bit platform, Clang (and emscripten) define this type - * despite not having the arithmetic for it. This results in a laggy - * compiler builtin call which calculates a full 128-bit multiply. - * In that case it is best to use the portable one. - * https://github.com/Cyan4973/xxHash/issues/211#issuecomment-515575677 - */ -#if defined(__GNUC__) && !defined(__wasm__) \ - && defined(__SIZEOF_INT128__) \ - || (defined(_INTEGRAL_MAX_BITS) && _INTEGRAL_MAX_BITS >= 128) - - __uint128_t const product = (__uint128_t)lhs * (__uint128_t)rhs; - XXH128_hash_t r128; - r128.low64 = (xxh_u64)(product); - r128.high64 = (xxh_u64)(product >> 64); - return r128; - - /* - * MSVC for x64's _umul128 method. - * - * xxh_u64 _umul128(xxh_u64 Multiplier, xxh_u64 Multiplicand, xxh_u64 *HighProduct); - * - * This compiles to single operand MUL on x64. - */ -#elif defined(_M_X64) || defined(_M_IA64) - -#ifndef _MSC_VER -# pragma intrinsic(_umul128) -#endif - xxh_u64 product_high; - xxh_u64 const product_low = _umul128(lhs, rhs, &product_high); - XXH128_hash_t r128; - r128.low64 = product_low; - r128.high64 = product_high; - return r128; - - /* - * MSVC for ARM64's __umulh method. - * - * This compiles to the same MUL + UMULH as GCC/Clang's __uint128_t method. - */ -#elif defined(_M_ARM64) - -#ifndef _MSC_VER -# pragma intrinsic(__umulh) -#endif - XXH128_hash_t r128; - r128.low64 = lhs * rhs; - r128.high64 = __umulh(lhs, rhs); - return r128; - -#else - /* - * Portable scalar method. Optimized for 32-bit and 64-bit ALUs. - * - * This is a fast and simple grade school multiply, which is shown below - * with base 10 arithmetic instead of base 0x100000000. - * - * 9 3 // D2 lhs = 93 - * x 7 5 // D2 rhs = 75 - * ---------- - * 1 5 // D2 lo_lo = (93 % 10) * (75 % 10) = 15 - * 4 5 | // D2 hi_lo = (93 / 10) * (75 % 10) = 45 - * 2 1 | // D2 lo_hi = (93 % 10) * (75 / 10) = 21 - * + 6 3 | | // D2 hi_hi = (93 / 10) * (75 / 10) = 63 - * --------- - * 2 7 | // D2 cross = (15 / 10) + (45 % 10) + 21 = 27 - * + 6 7 | | // D2 upper = (27 / 10) + (45 / 10) + 63 = 67 - * --------- - * 6 9 7 5 // D4 res = (27 * 10) + (15 % 10) + (67 * 100) = 6975 - * - * The reasons for adding the products like this are: - * 1. It avoids manual carry tracking. Just like how - * (9 * 9) + 9 + 9 = 99, the same applies with this for UINT64_MAX. - * This avoids a lot of complexity. - * - * 2. It hints for, and on Clang, compiles to, the powerful UMAAL - * instruction available in ARM's Digital Signal Processing extension - * in 32-bit ARMv6 and later, which is shown below: - * - * void UMAAL(xxh_u32 *RdLo, xxh_u32 *RdHi, xxh_u32 Rn, xxh_u32 Rm) - * { - * xxh_u64 product = (xxh_u64)*RdLo * (xxh_u64)*RdHi + Rn + Rm; - * *RdLo = (xxh_u32)(product & 0xFFFFFFFF); - * *RdHi = (xxh_u32)(product >> 32); - * } - * - * This instruction was designed for efficient long multiplication, and - * allows this to be calculated in only 4 instructions at speeds - * comparable to some 64-bit ALUs. - * - * 3. It isn't terrible on other platforms. Usually this will be a couple - * of 32-bit ADD/ADCs. - */ - - /* First calculate all of the cross products. */ - xxh_u64 const lo_lo = XXH_mult32to64(lhs & 0xFFFFFFFF, rhs & 0xFFFFFFFF); - xxh_u64 const hi_lo = XXH_mult32to64(lhs >> 32, rhs & 0xFFFFFFFF); - xxh_u64 const lo_hi = XXH_mult32to64(lhs & 0xFFFFFFFF, rhs >> 32); - xxh_u64 const hi_hi = XXH_mult32to64(lhs >> 32, rhs >> 32); - - /* Now add the products together. These will never overflow. */ - xxh_u64 const cross = (lo_lo >> 32) + (hi_lo & 0xFFFFFFFF) + lo_hi; - xxh_u64 const upper = (hi_lo >> 32) + (cross >> 32) + hi_hi; - xxh_u64 const lower = (cross << 32) | (lo_lo & 0xFFFFFFFF); - - XXH128_hash_t r128; - r128.low64 = lower; - r128.high64 = upper; - return r128; -#endif -} - -/*! - * @brief Calculates a 64-bit to 128-bit multiply, then XOR folds it. - * - * The reason for the separate function is to prevent passing too many structs - * around by value. This will hopefully inline the multiply, but we don't force it. - * - * @param lhs , rhs The 64-bit integers to multiply - * @return The low 64 bits of the product XOR'd by the high 64 bits. - * @see XXH_mult64to128() - */ -static xxh_u64 -XXH3_mul128_fold64(xxh_u64 lhs, xxh_u64 rhs) -{ - XXH128_hash_t product = XXH_mult64to128(lhs, rhs); - return product.low64 ^ product.high64; -} - -/*! Seems to produce slightly better code on GCC for some reason. */ -XXH_FORCE_INLINE xxh_u64 XXH_xorshift64(xxh_u64 v64, int shift) -{ - XXH_ASSERT(0 <= shift && shift < 64); - return v64 ^ (v64 >> shift); -} - -/* - * This is a fast avalanche stage, - * suitable when input bits are already partially mixed - */ -static XXH64_hash_t XXH3_avalanche(xxh_u64 h64) -{ - h64 = XXH_xorshift64(h64, 37); - h64 *= 0x165667919E3779F9ULL; - h64 = XXH_xorshift64(h64, 32); - return h64; -} - -/* - * This is a stronger avalanche, - * inspired by Pelle Evensen's rrmxmx - * preferable when input has not been previously mixed - */ -static XXH64_hash_t XXH3_rrmxmx(xxh_u64 h64, xxh_u64 len) -{ - /* this mix is inspired by Pelle Evensen's rrmxmx */ - h64 ^= XXH_rotl64(h64, 49) ^ XXH_rotl64(h64, 24); - h64 *= 0x9FB21C651E98DF25ULL; - h64 ^= (h64 >> 35) + len ; - h64 *= 0x9FB21C651E98DF25ULL; - return XXH_xorshift64(h64, 28); -} - - -/* ========================================== - * Short keys - * ========================================== - * One of the shortcomings of XXH32 and XXH64 was that their performance was - * sub-optimal on short lengths. It used an iterative algorithm which strongly - * favored lengths that were a multiple of 4 or 8. - * - * Instead of iterating over individual inputs, we use a set of single shot - * functions which piece together a range of lengths and operate in constant time. - * - * Additionally, the number of multiplies has been significantly reduced. This - * reduces latency, especially when emulating 64-bit multiplies on 32-bit. - * - * Depending on the platform, this may or may not be faster than XXH32, but it - * is almost guaranteed to be faster than XXH64. - */ - -/* - * At very short lengths, there isn't enough input to fully hide secrets, or use - * the entire secret. - * - * There is also only a limited amount of mixing we can do before significantly - * impacting performance. - * - * Therefore, we use different sections of the secret and always mix two secret - * samples with an XOR. This should have no effect on performance on the - * seedless or withSeed variants because everything _should_ be constant folded - * by modern compilers. - * - * The XOR mixing hides individual parts of the secret and increases entropy. - * - * This adds an extra layer of strength for custom secrets. - */ -XXH_FORCE_INLINE XXH64_hash_t -XXH3_len_1to3_64b(const xxh_u8* input, size_t len, const xxh_u8* secret, XXH64_hash_t seed) -{ - XXH_ASSERT(input != NULL); - XXH_ASSERT(1 <= len && len <= 3); - XXH_ASSERT(secret != NULL); - /* - * len = 1: combined = { input[0], 0x01, input[0], input[0] } - * len = 2: combined = { input[1], 0x02, input[0], input[1] } - * len = 3: combined = { input[2], 0x03, input[0], input[1] } - */ - { xxh_u8 const c1 = input[0]; - xxh_u8 const c2 = input[len >> 1]; - xxh_u8 const c3 = input[len - 1]; - xxh_u32 const combined = ((xxh_u32)c1 << 16) | ((xxh_u32)c2 << 24) - | ((xxh_u32)c3 << 0) | ((xxh_u32)len << 8); - xxh_u64 const bitflip = (XXH_readLE32(secret) ^ XXH_readLE32(secret+4)) + seed; - xxh_u64 const keyed = (xxh_u64)combined ^ bitflip; - return XXH64_avalanche(keyed); - } -} - -XXH_FORCE_INLINE XXH64_hash_t -XXH3_len_4to8_64b(const xxh_u8* input, size_t len, const xxh_u8* secret, XXH64_hash_t seed) -{ - XXH_ASSERT(input != NULL); - XXH_ASSERT(secret != NULL); - XXH_ASSERT(4 <= len && len <= 8); - seed ^= (xxh_u64)XXH_swap32((xxh_u32)seed) << 32; - { xxh_u32 const input1 = XXH_readLE32(input); - xxh_u32 const input2 = XXH_readLE32(input + len - 4); - xxh_u64 const bitflip = (XXH_readLE64(secret+8) ^ XXH_readLE64(secret+16)) - seed; - xxh_u64 const input64 = input2 + (((xxh_u64)input1) << 32); - xxh_u64 const keyed = input64 ^ bitflip; - return XXH3_rrmxmx(keyed, len); - } -} - -XXH_FORCE_INLINE XXH64_hash_t -XXH3_len_9to16_64b(const xxh_u8* input, size_t len, const xxh_u8* secret, XXH64_hash_t seed) -{ - XXH_ASSERT(input != NULL); - XXH_ASSERT(secret != NULL); - XXH_ASSERT(9 <= len && len <= 16); - { xxh_u64 const bitflip1 = (XXH_readLE64(secret+24) ^ XXH_readLE64(secret+32)) + seed; - xxh_u64 const bitflip2 = (XXH_readLE64(secret+40) ^ XXH_readLE64(secret+48)) - seed; - xxh_u64 const input_lo = XXH_readLE64(input) ^ bitflip1; - xxh_u64 const input_hi = XXH_readLE64(input + len - 8) ^ bitflip2; - xxh_u64 const acc = len - + XXH_swap64(input_lo) + input_hi - + XXH3_mul128_fold64(input_lo, input_hi); - return XXH3_avalanche(acc); - } -} - -XXH_FORCE_INLINE XXH64_hash_t -XXH3_len_0to16_64b(const xxh_u8* input, size_t len, const xxh_u8* secret, XXH64_hash_t seed) -{ - XXH_ASSERT(len <= 16); - { if (XXH_likely(len > 8)) return XXH3_len_9to16_64b(input, len, secret, seed); - if (XXH_likely(len >= 4)) return XXH3_len_4to8_64b(input, len, secret, seed); - if (len) return XXH3_len_1to3_64b(input, len, secret, seed); - return XXH64_avalanche(seed ^ (XXH_readLE64(secret+56) ^ XXH_readLE64(secret+64))); - } -} - -/* - * DISCLAIMER: There are known *seed-dependent* multicollisions here due to - * multiplication by zero, affecting hashes of lengths 17 to 240. - * - * However, they are very unlikely. - * - * Keep this in mind when using the unseeded XXH3_64bits() variant: As with all - * unseeded non-cryptographic hashes, it does not attempt to defend itself - * against specially crafted inputs, only random inputs. - * - * Compared to classic UMAC where a 1 in 2^31 chance of 4 consecutive bytes - * cancelling out the secret is taken an arbitrary number of times (addressed - * in XXH3_accumulate_512), this collision is very unlikely with random inputs - * and/or proper seeding: - * - * This only has a 1 in 2^63 chance of 8 consecutive bytes cancelling out, in a - * function that is only called up to 16 times per hash with up to 240 bytes of - * input. - * - * This is not too bad for a non-cryptographic hash function, especially with - * only 64 bit outputs. - * - * The 128-bit variant (which trades some speed for strength) is NOT affected - * by this, although it is always a good idea to use a proper seed if you care - * about strength. - */ -XXH_FORCE_INLINE xxh_u64 XXH3_mix16B(const xxh_u8* XXH_RESTRICT input, - const xxh_u8* XXH_RESTRICT secret, xxh_u64 seed64) -{ -#if defined(__GNUC__) && !defined(__clang__) /* GCC, not Clang */ \ - && defined(__i386__) && defined(__SSE2__) /* x86 + SSE2 */ \ - && !defined(XXH_ENABLE_AUTOVECTORIZE) /* Define to disable like XXH32 hack */ - /* - * UGLY HACK: - * GCC for x86 tends to autovectorize the 128-bit multiply, resulting in - * slower code. - * - * By forcing seed64 into a register, we disrupt the cost model and - * cause it to scalarize. See `XXH32_round()` - * - * FIXME: Clang's output is still _much_ faster -- On an AMD Ryzen 3600, - * XXH3_64bits @ len=240 runs at 4.6 GB/s with Clang 9, but 3.3 GB/s on - * GCC 9.2, despite both emitting scalar code. - * - * GCC generates much better scalar code than Clang for the rest of XXH3, - * which is why finding a more optimal codepath is an interest. - */ - XXH_COMPILER_GUARD(seed64); -#endif - { xxh_u64 const input_lo = XXH_readLE64(input); - xxh_u64 const input_hi = XXH_readLE64(input+8); - return XXH3_mul128_fold64( - input_lo ^ (XXH_readLE64(secret) + seed64), - input_hi ^ (XXH_readLE64(secret+8) - seed64) - ); - } -} - -/* For mid range keys, XXH3 uses a Mum-hash variant. */ -XXH_FORCE_INLINE XXH64_hash_t -XXH3_len_17to128_64b(const xxh_u8* XXH_RESTRICT input, size_t len, - const xxh_u8* XXH_RESTRICT secret, size_t secretSize, - XXH64_hash_t seed) -{ - XXH_ASSERT(secretSize >= XXH3_SECRET_SIZE_MIN); (void)secretSize; - XXH_ASSERT(16 < len && len <= 128); - - { xxh_u64 acc = len * XXH_PRIME64_1; - if (len > 32) { - if (len > 64) { - if (len > 96) { - acc += XXH3_mix16B(input+48, secret+96, seed); - acc += XXH3_mix16B(input+len-64, secret+112, seed); - } - acc += XXH3_mix16B(input+32, secret+64, seed); - acc += XXH3_mix16B(input+len-48, secret+80, seed); - } - acc += XXH3_mix16B(input+16, secret+32, seed); - acc += XXH3_mix16B(input+len-32, secret+48, seed); - } - acc += XXH3_mix16B(input+0, secret+0, seed); - acc += XXH3_mix16B(input+len-16, secret+16, seed); - - return XXH3_avalanche(acc); - } -} - -#define XXH3_MIDSIZE_MAX 240 - -XXH_NO_INLINE XXH64_hash_t -XXH3_len_129to240_64b(const xxh_u8* XXH_RESTRICT input, size_t len, - const xxh_u8* XXH_RESTRICT secret, size_t secretSize, - XXH64_hash_t seed) -{ - XXH_ASSERT(secretSize >= XXH3_SECRET_SIZE_MIN); (void)secretSize; - XXH_ASSERT(128 < len && len <= XXH3_MIDSIZE_MAX); - - #define XXH3_MIDSIZE_STARTOFFSET 3 - #define XXH3_MIDSIZE_LASTOFFSET 17 - - { xxh_u64 acc = len * XXH_PRIME64_1; - int const nbRounds = (int)len / 16; - int i; - for (i=0; i<8; i++) { - acc += XXH3_mix16B(input+(16*i), secret+(16*i), seed); - } - acc = XXH3_avalanche(acc); - XXH_ASSERT(nbRounds >= 8); -#if defined(__clang__) /* Clang */ \ - && (defined(__ARM_NEON) || defined(__ARM_NEON__)) /* NEON */ \ - && !defined(XXH_ENABLE_AUTOVECTORIZE) /* Define to disable */ - /* - * UGLY HACK: - * Clang for ARMv7-A tries to vectorize this loop, similar to GCC x86. - * In everywhere else, it uses scalar code. - * - * For 64->128-bit multiplies, even if the NEON was 100% optimal, it - * would still be slower than UMAAL (see XXH_mult64to128). - * - * Unfortunately, Clang doesn't handle the long multiplies properly and - * converts them to the nonexistent "vmulq_u64" intrinsic, which is then - * scalarized into an ugly mess of VMOV.32 instructions. - * - * This mess is difficult to avoid without turning autovectorization - * off completely, but they are usually relatively minor and/or not - * worth it to fix. - * - * This loop is the easiest to fix, as unlike XXH32, this pragma - * _actually works_ because it is a loop vectorization instead of an - * SLP vectorization. - */ - #pragma clang loop vectorize(disable) -#endif - for (i=8 ; i < nbRounds; i++) { - acc += XXH3_mix16B(input+(16*i), secret+(16*(i-8)) + XXH3_MIDSIZE_STARTOFFSET, seed); - } - /* last bytes */ - acc += XXH3_mix16B(input + len - 16, secret + XXH3_SECRET_SIZE_MIN - XXH3_MIDSIZE_LASTOFFSET, seed); - return XXH3_avalanche(acc); - } -} - - -/* ======= Long Keys ======= */ - -#define XXH_STRIPE_LEN 64 -#define XXH_SECRET_CONSUME_RATE 8 /* nb of secret bytes consumed at each accumulation */ -#define XXH_ACC_NB (XXH_STRIPE_LEN / sizeof(xxh_u64)) - -#ifdef XXH_OLD_NAMES -# define STRIPE_LEN XXH_STRIPE_LEN -# define ACC_NB XXH_ACC_NB -#endif - -XXH_FORCE_INLINE void XXH_writeLE64(void* dst, xxh_u64 v64) -{ - if (!XXH_CPU_LITTLE_ENDIAN) v64 = XXH_swap64(v64); - XXH_memcpy(dst, &v64, sizeof(v64)); -} - -/* Several intrinsic functions below are supposed to accept __int64 as argument, - * as documented in https://software.intel.com/sites/landingpage/IntrinsicsGuide/ . - * However, several environments do not define __int64 type, - * requiring a workaround. - */ -#if !defined (__VMS) \ - && (defined (__cplusplus) \ - || (defined (__STDC_VERSION__) && (__STDC_VERSION__ >= 199901L) /* C99 */) ) - typedef int64_t xxh_i64; -#else - /* the following type must have a width of 64-bit */ - typedef long long xxh_i64; -#endif - -/* - * XXH3_accumulate_512 is the tightest loop for long inputs, and it is the most optimized. - * - * It is a hardened version of UMAC, based off of FARSH's implementation. - * - * This was chosen because it adapts quite well to 32-bit, 64-bit, and SIMD - * implementations, and it is ridiculously fast. - * - * We harden it by mixing the original input to the accumulators as well as the product. - * - * This means that in the (relatively likely) case of a multiply by zero, the - * original input is preserved. - * - * On 128-bit inputs, we swap 64-bit pairs when we add the input to improve - * cross-pollination, as otherwise the upper and lower halves would be - * essentially independent. - * - * This doesn't matter on 64-bit hashes since they all get merged together in - * the end, so we skip the extra step. - * - * Both XXH3_64bits and XXH3_128bits use this subroutine. - */ - -#if (XXH_VECTOR == XXH_AVX512) \ - || (defined(XXH_DISPATCH_AVX512) && XXH_DISPATCH_AVX512 != 0) - -#ifndef XXH_TARGET_AVX512 -# define XXH_TARGET_AVX512 /* disable attribute target */ -#endif - -XXH_FORCE_INLINE XXH_TARGET_AVX512 void -XXH3_accumulate_512_avx512(void* XXH_RESTRICT acc, - const void* XXH_RESTRICT input, - const void* XXH_RESTRICT secret) -{ - __m512i* const xacc = (__m512i *) acc; - XXH_ASSERT((((size_t)acc) & 63) == 0); - XXH_STATIC_ASSERT(XXH_STRIPE_LEN == sizeof(__m512i)); - - { - /* data_vec = input[0]; */ - __m512i const data_vec = _mm512_loadu_si512 (input); - /* key_vec = secret[0]; */ - __m512i const key_vec = _mm512_loadu_si512 (secret); - /* data_key = data_vec ^ key_vec; */ - __m512i const data_key = _mm512_xor_si512 (data_vec, key_vec); - /* data_key_lo = data_key >> 32; */ - __m512i const data_key_lo = _mm512_shuffle_epi32 (data_key, (_MM_PERM_ENUM)_MM_SHUFFLE(0, 3, 0, 1)); - /* product = (data_key & 0xffffffff) * (data_key_lo & 0xffffffff); */ - __m512i const product = _mm512_mul_epu32 (data_key, data_key_lo); - /* xacc[0] += swap(data_vec); */ - __m512i const data_swap = _mm512_shuffle_epi32(data_vec, (_MM_PERM_ENUM)_MM_SHUFFLE(1, 0, 3, 2)); - __m512i const sum = _mm512_add_epi64(*xacc, data_swap); - /* xacc[0] += product; */ - *xacc = _mm512_add_epi64(product, sum); - } -} - -/* - * XXH3_scrambleAcc: Scrambles the accumulators to improve mixing. - * - * Multiplication isn't perfect, as explained by Google in HighwayHash: - * - * // Multiplication mixes/scrambles bytes 0-7 of the 64-bit result to - * // varying degrees. In descending order of goodness, bytes - * // 3 4 2 5 1 6 0 7 have quality 228 224 164 160 100 96 36 32. - * // As expected, the upper and lower bytes are much worse. - * - * Source: https://github.com/google/highwayhash/blob/0aaf66b/highwayhash/hh_avx2.h#L291 - * - * Since our algorithm uses a pseudorandom secret to add some variance into the - * mix, we don't need to (or want to) mix as often or as much as HighwayHash does. - * - * This isn't as tight as XXH3_accumulate, but still written in SIMD to avoid - * extraction. - * - * Both XXH3_64bits and XXH3_128bits use this subroutine. - */ - -XXH_FORCE_INLINE XXH_TARGET_AVX512 void -XXH3_scrambleAcc_avx512(void* XXH_RESTRICT acc, const void* XXH_RESTRICT secret) -{ - XXH_ASSERT((((size_t)acc) & 63) == 0); - XXH_STATIC_ASSERT(XXH_STRIPE_LEN == sizeof(__m512i)); - { __m512i* const xacc = (__m512i*) acc; - const __m512i prime32 = _mm512_set1_epi32((int)XXH_PRIME32_1); - - /* xacc[0] ^= (xacc[0] >> 47) */ - __m512i const acc_vec = *xacc; - __m512i const shifted = _mm512_srli_epi64 (acc_vec, 47); - __m512i const data_vec = _mm512_xor_si512 (acc_vec, shifted); - /* xacc[0] ^= secret; */ - __m512i const key_vec = _mm512_loadu_si512 (secret); - __m512i const data_key = _mm512_xor_si512 (data_vec, key_vec); - - /* xacc[0] *= XXH_PRIME32_1; */ - __m512i const data_key_hi = _mm512_shuffle_epi32 (data_key, (_MM_PERM_ENUM)_MM_SHUFFLE(0, 3, 0, 1)); - __m512i const prod_lo = _mm512_mul_epu32 (data_key, prime32); - __m512i const prod_hi = _mm512_mul_epu32 (data_key_hi, prime32); - *xacc = _mm512_add_epi64(prod_lo, _mm512_slli_epi64(prod_hi, 32)); - } -} - -XXH_FORCE_INLINE XXH_TARGET_AVX512 void -XXH3_initCustomSecret_avx512(void* XXH_RESTRICT customSecret, xxh_u64 seed64) -{ - XXH_STATIC_ASSERT((XXH_SECRET_DEFAULT_SIZE & 63) == 0); - XXH_STATIC_ASSERT(XXH_SEC_ALIGN == 64); - XXH_ASSERT(((size_t)customSecret & 63) == 0); - (void)(&XXH_writeLE64); - { int const nbRounds = XXH_SECRET_DEFAULT_SIZE / sizeof(__m512i); - __m512i const seed = _mm512_mask_set1_epi64(_mm512_set1_epi64((xxh_i64)seed64), 0xAA, (xxh_i64)(0U - seed64)); - - const __m512i* const src = (const __m512i*) ((const void*) XXH3_kSecret); - __m512i* const dest = ( __m512i*) customSecret; - int i; - XXH_ASSERT(((size_t)src & 63) == 0); /* control alignment */ - XXH_ASSERT(((size_t)dest & 63) == 0); - for (i=0; i < nbRounds; ++i) { - /* GCC has a bug, _mm512_stream_load_si512 accepts 'void*', not 'void const*', - * this will warn "discards 'const' qualifier". */ - union { - const __m512i* cp; - void* p; - } remote_const_void; - remote_const_void.cp = src + i; - dest[i] = _mm512_add_epi64(_mm512_stream_load_si512(remote_const_void.p), seed); - } } -} - -#endif - -#if (XXH_VECTOR == XXH_AVX2) \ - || (defined(XXH_DISPATCH_AVX2) && XXH_DISPATCH_AVX2 != 0) - -#ifndef XXH_TARGET_AVX2 -# define XXH_TARGET_AVX2 /* disable attribute target */ -#endif - -XXH_FORCE_INLINE XXH_TARGET_AVX2 void -XXH3_accumulate_512_avx2( void* XXH_RESTRICT acc, - const void* XXH_RESTRICT input, - const void* XXH_RESTRICT secret) -{ - XXH_ASSERT((((size_t)acc) & 31) == 0); - { __m256i* const xacc = (__m256i *) acc; - /* Unaligned. This is mainly for pointer arithmetic, and because - * _mm256_loadu_si256 requires a const __m256i * pointer for some reason. */ - const __m256i* const xinput = (const __m256i *) input; - /* Unaligned. This is mainly for pointer arithmetic, and because - * _mm256_loadu_si256 requires a const __m256i * pointer for some reason. */ - const __m256i* const xsecret = (const __m256i *) secret; - - size_t i; - for (i=0; i < XXH_STRIPE_LEN/sizeof(__m256i); i++) { - /* data_vec = xinput[i]; */ - __m256i const data_vec = _mm256_loadu_si256 (xinput+i); - /* key_vec = xsecret[i]; */ - __m256i const key_vec = _mm256_loadu_si256 (xsecret+i); - /* data_key = data_vec ^ key_vec; */ - __m256i const data_key = _mm256_xor_si256 (data_vec, key_vec); - /* data_key_lo = data_key >> 32; */ - __m256i const data_key_lo = _mm256_shuffle_epi32 (data_key, _MM_SHUFFLE(0, 3, 0, 1)); - /* product = (data_key & 0xffffffff) * (data_key_lo & 0xffffffff); */ - __m256i const product = _mm256_mul_epu32 (data_key, data_key_lo); - /* xacc[i] += swap(data_vec); */ - __m256i const data_swap = _mm256_shuffle_epi32(data_vec, _MM_SHUFFLE(1, 0, 3, 2)); - __m256i const sum = _mm256_add_epi64(xacc[i], data_swap); - /* xacc[i] += product; */ - xacc[i] = _mm256_add_epi64(product, sum); - } } -} - -XXH_FORCE_INLINE XXH_TARGET_AVX2 void -XXH3_scrambleAcc_avx2(void* XXH_RESTRICT acc, const void* XXH_RESTRICT secret) -{ - XXH_ASSERT((((size_t)acc) & 31) == 0); - { __m256i* const xacc = (__m256i*) acc; - /* Unaligned. This is mainly for pointer arithmetic, and because - * _mm256_loadu_si256 requires a const __m256i * pointer for some reason. */ - const __m256i* const xsecret = (const __m256i *) secret; - const __m256i prime32 = _mm256_set1_epi32((int)XXH_PRIME32_1); - - size_t i; - for (i=0; i < XXH_STRIPE_LEN/sizeof(__m256i); i++) { - /* xacc[i] ^= (xacc[i] >> 47) */ - __m256i const acc_vec = xacc[i]; - __m256i const shifted = _mm256_srli_epi64 (acc_vec, 47); - __m256i const data_vec = _mm256_xor_si256 (acc_vec, shifted); - /* xacc[i] ^= xsecret; */ - __m256i const key_vec = _mm256_loadu_si256 (xsecret+i); - __m256i const data_key = _mm256_xor_si256 (data_vec, key_vec); - - /* xacc[i] *= XXH_PRIME32_1; */ - __m256i const data_key_hi = _mm256_shuffle_epi32 (data_key, _MM_SHUFFLE(0, 3, 0, 1)); - __m256i const prod_lo = _mm256_mul_epu32 (data_key, prime32); - __m256i const prod_hi = _mm256_mul_epu32 (data_key_hi, prime32); - xacc[i] = _mm256_add_epi64(prod_lo, _mm256_slli_epi64(prod_hi, 32)); - } - } -} - -XXH_FORCE_INLINE XXH_TARGET_AVX2 void XXH3_initCustomSecret_avx2(void* XXH_RESTRICT customSecret, xxh_u64 seed64) -{ - XXH_STATIC_ASSERT((XXH_SECRET_DEFAULT_SIZE & 31) == 0); - XXH_STATIC_ASSERT((XXH_SECRET_DEFAULT_SIZE / sizeof(__m256i)) == 6); - XXH_STATIC_ASSERT(XXH_SEC_ALIGN <= 64); - (void)(&XXH_writeLE64); - XXH_PREFETCH(customSecret); - { __m256i const seed = _mm256_set_epi64x((xxh_i64)(0U - seed64), (xxh_i64)seed64, (xxh_i64)(0U - seed64), (xxh_i64)seed64); - - const __m256i* const src = (const __m256i*) ((const void*) XXH3_kSecret); - __m256i* dest = ( __m256i*) customSecret; - -# if defined(__GNUC__) || defined(__clang__) - /* - * On GCC & Clang, marking 'dest' as modified will cause the compiler: - * - do not extract the secret from sse registers in the internal loop - * - use less common registers, and avoid pushing these reg into stack - */ - XXH_COMPILER_GUARD(dest); -# endif - XXH_ASSERT(((size_t)src & 31) == 0); /* control alignment */ - XXH_ASSERT(((size_t)dest & 31) == 0); - - /* GCC -O2 need unroll loop manually */ - dest[0] = _mm256_add_epi64(_mm256_stream_load_si256(src+0), seed); - dest[1] = _mm256_add_epi64(_mm256_stream_load_si256(src+1), seed); - dest[2] = _mm256_add_epi64(_mm256_stream_load_si256(src+2), seed); - dest[3] = _mm256_add_epi64(_mm256_stream_load_si256(src+3), seed); - dest[4] = _mm256_add_epi64(_mm256_stream_load_si256(src+4), seed); - dest[5] = _mm256_add_epi64(_mm256_stream_load_si256(src+5), seed); - } -} - -#endif - -/* x86dispatch always generates SSE2 */ -#if (XXH_VECTOR == XXH_SSE2) || defined(XXH_X86DISPATCH) - -#ifndef XXH_TARGET_SSE2 -# define XXH_TARGET_SSE2 /* disable attribute target */ -#endif - -XXH_FORCE_INLINE XXH_TARGET_SSE2 void -XXH3_accumulate_512_sse2( void* XXH_RESTRICT acc, - const void* XXH_RESTRICT input, - const void* XXH_RESTRICT secret) -{ - /* SSE2 is just a half-scale version of the AVX2 version. */ - XXH_ASSERT((((size_t)acc) & 15) == 0); - { __m128i* const xacc = (__m128i *) acc; - /* Unaligned. This is mainly for pointer arithmetic, and because - * _mm_loadu_si128 requires a const __m128i * pointer for some reason. */ - const __m128i* const xinput = (const __m128i *) input; - /* Unaligned. This is mainly for pointer arithmetic, and because - * _mm_loadu_si128 requires a const __m128i * pointer for some reason. */ - const __m128i* const xsecret = (const __m128i *) secret; - - size_t i; - for (i=0; i < XXH_STRIPE_LEN/sizeof(__m128i); i++) { - /* data_vec = xinput[i]; */ - __m128i const data_vec = _mm_loadu_si128 (xinput+i); - /* key_vec = xsecret[i]; */ - __m128i const key_vec = _mm_loadu_si128 (xsecret+i); - /* data_key = data_vec ^ key_vec; */ - __m128i const data_key = _mm_xor_si128 (data_vec, key_vec); - /* data_key_lo = data_key >> 32; */ - __m128i const data_key_lo = _mm_shuffle_epi32 (data_key, _MM_SHUFFLE(0, 3, 0, 1)); - /* product = (data_key & 0xffffffff) * (data_key_lo & 0xffffffff); */ - __m128i const product = _mm_mul_epu32 (data_key, data_key_lo); - /* xacc[i] += swap(data_vec); */ - __m128i const data_swap = _mm_shuffle_epi32(data_vec, _MM_SHUFFLE(1,0,3,2)); - __m128i const sum = _mm_add_epi64(xacc[i], data_swap); - /* xacc[i] += product; */ - xacc[i] = _mm_add_epi64(product, sum); - } } -} - -XXH_FORCE_INLINE XXH_TARGET_SSE2 void -XXH3_scrambleAcc_sse2(void* XXH_RESTRICT acc, const void* XXH_RESTRICT secret) -{ - XXH_ASSERT((((size_t)acc) & 15) == 0); - { __m128i* const xacc = (__m128i*) acc; - /* Unaligned. This is mainly for pointer arithmetic, and because - * _mm_loadu_si128 requires a const __m128i * pointer for some reason. */ - const __m128i* const xsecret = (const __m128i *) secret; - const __m128i prime32 = _mm_set1_epi32((int)XXH_PRIME32_1); - - size_t i; - for (i=0; i < XXH_STRIPE_LEN/sizeof(__m128i); i++) { - /* xacc[i] ^= (xacc[i] >> 47) */ - __m128i const acc_vec = xacc[i]; - __m128i const shifted = _mm_srli_epi64 (acc_vec, 47); - __m128i const data_vec = _mm_xor_si128 (acc_vec, shifted); - /* xacc[i] ^= xsecret[i]; */ - __m128i const key_vec = _mm_loadu_si128 (xsecret+i); - __m128i const data_key = _mm_xor_si128 (data_vec, key_vec); - - /* xacc[i] *= XXH_PRIME32_1; */ - __m128i const data_key_hi = _mm_shuffle_epi32 (data_key, _MM_SHUFFLE(0, 3, 0, 1)); - __m128i const prod_lo = _mm_mul_epu32 (data_key, prime32); - __m128i const prod_hi = _mm_mul_epu32 (data_key_hi, prime32); - xacc[i] = _mm_add_epi64(prod_lo, _mm_slli_epi64(prod_hi, 32)); - } - } -} - -XXH_FORCE_INLINE XXH_TARGET_SSE2 void XXH3_initCustomSecret_sse2(void* XXH_RESTRICT customSecret, xxh_u64 seed64) -{ - XXH_STATIC_ASSERT((XXH_SECRET_DEFAULT_SIZE & 15) == 0); - (void)(&XXH_writeLE64); - { int const nbRounds = XXH_SECRET_DEFAULT_SIZE / sizeof(__m128i); - -# if defined(_MSC_VER) && defined(_M_IX86) && _MSC_VER < 1900 - /* MSVC 32bit mode does not support _mm_set_epi64x before 2015 */ - XXH_ALIGN(16) const xxh_i64 seed64x2[2] = { (xxh_i64)seed64, (xxh_i64)(0U - seed64) }; - __m128i const seed = _mm_load_si128((__m128i const*)seed64x2); -# else - __m128i const seed = _mm_set_epi64x((xxh_i64)(0U - seed64), (xxh_i64)seed64); -# endif - int i; - - const void* const src16 = XXH3_kSecret; - __m128i* dst16 = (__m128i*) customSecret; -# if defined(__GNUC__) || defined(__clang__) - /* - * On GCC & Clang, marking 'dest' as modified will cause the compiler: - * - do not extract the secret from sse registers in the internal loop - * - use less common registers, and avoid pushing these reg into stack - */ - XXH_COMPILER_GUARD(dst16); -# endif - XXH_ASSERT(((size_t)src16 & 15) == 0); /* control alignment */ - XXH_ASSERT(((size_t)dst16 & 15) == 0); - - for (i=0; i < nbRounds; ++i) { - dst16[i] = _mm_add_epi64(_mm_load_si128((const __m128i *)src16+i), seed); - } } -} - -#endif - -#if (XXH_VECTOR == XXH_NEON) - -XXH_FORCE_INLINE void -XXH3_accumulate_512_neon( void* XXH_RESTRICT acc, - const void* XXH_RESTRICT input, - const void* XXH_RESTRICT secret) -{ - XXH_ASSERT((((size_t)acc) & 15) == 0); - { - uint64x2_t* const xacc = (uint64x2_t *) acc; - /* We don't use a uint32x4_t pointer because it causes bus errors on ARMv7. */ - uint8_t const* const xinput = (const uint8_t *) input; - uint8_t const* const xsecret = (const uint8_t *) secret; - - size_t i; - for (i=0; i < XXH_STRIPE_LEN / sizeof(uint64x2_t); i++) { - /* data_vec = xinput[i]; */ - uint8x16_t data_vec = vld1q_u8(xinput + (i * 16)); - /* key_vec = xsecret[i]; */ - uint8x16_t key_vec = vld1q_u8(xsecret + (i * 16)); - uint64x2_t data_key; - uint32x2_t data_key_lo, data_key_hi; - /* xacc[i] += swap(data_vec); */ - uint64x2_t const data64 = vreinterpretq_u64_u8(data_vec); - uint64x2_t const swapped = vextq_u64(data64, data64, 1); - xacc[i] = vaddq_u64 (xacc[i], swapped); - /* data_key = data_vec ^ key_vec; */ - data_key = vreinterpretq_u64_u8(veorq_u8(data_vec, key_vec)); - /* data_key_lo = (uint32x2_t) (data_key & 0xFFFFFFFF); - * data_key_hi = (uint32x2_t) (data_key >> 32); - * data_key = UNDEFINED; */ - XXH_SPLIT_IN_PLACE(data_key, data_key_lo, data_key_hi); - /* xacc[i] += (uint64x2_t) data_key_lo * (uint64x2_t) data_key_hi; */ - xacc[i] = vmlal_u32 (xacc[i], data_key_lo, data_key_hi); - - } - } -} - -XXH_FORCE_INLINE void -XXH3_scrambleAcc_neon(void* XXH_RESTRICT acc, const void* XXH_RESTRICT secret) -{ - XXH_ASSERT((((size_t)acc) & 15) == 0); - - { uint64x2_t* xacc = (uint64x2_t*) acc; - uint8_t const* xsecret = (uint8_t const*) secret; - uint32x2_t prime = vdup_n_u32 (XXH_PRIME32_1); - - size_t i; - for (i=0; i < XXH_STRIPE_LEN/sizeof(uint64x2_t); i++) { - /* xacc[i] ^= (xacc[i] >> 47); */ - uint64x2_t acc_vec = xacc[i]; - uint64x2_t shifted = vshrq_n_u64 (acc_vec, 47); - uint64x2_t data_vec = veorq_u64 (acc_vec, shifted); - - /* xacc[i] ^= xsecret[i]; */ - uint8x16_t key_vec = vld1q_u8 (xsecret + (i * 16)); - uint64x2_t data_key = veorq_u64 (data_vec, vreinterpretq_u64_u8(key_vec)); - - /* xacc[i] *= XXH_PRIME32_1 */ - uint32x2_t data_key_lo, data_key_hi; - /* data_key_lo = (uint32x2_t) (xacc[i] & 0xFFFFFFFF); - * data_key_hi = (uint32x2_t) (xacc[i] >> 32); - * xacc[i] = UNDEFINED; */ - XXH_SPLIT_IN_PLACE(data_key, data_key_lo, data_key_hi); - { /* - * prod_hi = (data_key >> 32) * XXH_PRIME32_1; - * - * Avoid vmul_u32 + vshll_n_u32 since Clang 6 and 7 will - * incorrectly "optimize" this: - * tmp = vmul_u32(vmovn_u64(a), vmovn_u64(b)); - * shifted = vshll_n_u32(tmp, 32); - * to this: - * tmp = "vmulq_u64"(a, b); // no such thing! - * shifted = vshlq_n_u64(tmp, 32); - * - * However, unlike SSE, Clang lacks a 64-bit multiply routine - * for NEON, and it scalarizes two 64-bit multiplies instead. - * - * vmull_u32 has the same timing as vmul_u32, and it avoids - * this bug completely. - * See https://bugs.llvm.org/show_bug.cgi?id=39967 - */ - uint64x2_t prod_hi = vmull_u32 (data_key_hi, prime); - /* xacc[i] = prod_hi << 32; */ - xacc[i] = vshlq_n_u64(prod_hi, 32); - /* xacc[i] += (prod_hi & 0xFFFFFFFF) * XXH_PRIME32_1; */ - xacc[i] = vmlal_u32(xacc[i], data_key_lo, prime); - } - } } -} - -#endif - -#if (XXH_VECTOR == XXH_VSX) - -XXH_FORCE_INLINE void -XXH3_accumulate_512_vsx( void* XXH_RESTRICT acc, - const void* XXH_RESTRICT input, - const void* XXH_RESTRICT secret) -{ - /* presumed aligned */ - unsigned long long* const xacc = (unsigned long long*) acc; - xxh_u64x2 const* const xinput = (xxh_u64x2 const*) input; /* no alignment restriction */ - xxh_u64x2 const* const xsecret = (xxh_u64x2 const*) secret; /* no alignment restriction */ - xxh_u64x2 const v32 = { 32, 32 }; - size_t i; - for (i = 0; i < XXH_STRIPE_LEN / sizeof(xxh_u64x2); i++) { - /* data_vec = xinput[i]; */ - xxh_u64x2 const data_vec = XXH_vec_loadu(xinput + i); - /* key_vec = xsecret[i]; */ - xxh_u64x2 const key_vec = XXH_vec_loadu(xsecret + i); - xxh_u64x2 const data_key = data_vec ^ key_vec; - /* shuffled = (data_key << 32) | (data_key >> 32); */ - xxh_u32x4 const shuffled = (xxh_u32x4)vec_rl(data_key, v32); - /* product = ((xxh_u64x2)data_key & 0xFFFFFFFF) * ((xxh_u64x2)shuffled & 0xFFFFFFFF); */ - xxh_u64x2 const product = XXH_vec_mulo((xxh_u32x4)data_key, shuffled); - /* acc_vec = xacc[i]; */ - xxh_u64x2 acc_vec = vec_xl(0, xacc + 2 * i); - acc_vec += product; - - /* swap high and low halves */ -#ifdef __s390x__ - acc_vec += vec_permi(data_vec, data_vec, 2); -#else - acc_vec += vec_xxpermdi(data_vec, data_vec, 2); -#endif - /* xacc[i] = acc_vec; */ - vec_xst(acc_vec, 0, xacc + 2 * i); - } -} - -XXH_FORCE_INLINE void -XXH3_scrambleAcc_vsx(void* XXH_RESTRICT acc, const void* XXH_RESTRICT secret) -{ - XXH_ASSERT((((size_t)acc) & 15) == 0); - - { xxh_u64x2* const xacc = (xxh_u64x2*) acc; - const xxh_u64x2* const xsecret = (const xxh_u64x2*) secret; - /* constants */ - xxh_u64x2 const v32 = { 32, 32 }; - xxh_u64x2 const v47 = { 47, 47 }; - xxh_u32x4 const prime = { XXH_PRIME32_1, XXH_PRIME32_1, XXH_PRIME32_1, XXH_PRIME32_1 }; - size_t i; - for (i = 0; i < XXH_STRIPE_LEN / sizeof(xxh_u64x2); i++) { - /* xacc[i] ^= (xacc[i] >> 47); */ - xxh_u64x2 const acc_vec = xacc[i]; - xxh_u64x2 const data_vec = acc_vec ^ (acc_vec >> v47); - - /* xacc[i] ^= xsecret[i]; */ - xxh_u64x2 const key_vec = XXH_vec_loadu(xsecret + i); - xxh_u64x2 const data_key = data_vec ^ key_vec; - - /* xacc[i] *= XXH_PRIME32_1 */ - /* prod_lo = ((xxh_u64x2)data_key & 0xFFFFFFFF) * ((xxh_u64x2)prime & 0xFFFFFFFF); */ - xxh_u64x2 const prod_even = XXH_vec_mule((xxh_u32x4)data_key, prime); - /* prod_hi = ((xxh_u64x2)data_key >> 32) * ((xxh_u64x2)prime >> 32); */ - xxh_u64x2 const prod_odd = XXH_vec_mulo((xxh_u32x4)data_key, prime); - xacc[i] = prod_odd + (prod_even << v32); - } } -} - -#endif - -/* scalar variants - universal */ - -XXH_FORCE_INLINE void -XXH3_accumulate_512_scalar(void* XXH_RESTRICT acc, - const void* XXH_RESTRICT input, - const void* XXH_RESTRICT secret) -{ - xxh_u64* const xacc = (xxh_u64*) acc; /* presumed aligned */ - const xxh_u8* const xinput = (const xxh_u8*) input; /* no alignment restriction */ - const xxh_u8* const xsecret = (const xxh_u8*) secret; /* no alignment restriction */ - size_t i; - XXH_ASSERT(((size_t)acc & (XXH_ACC_ALIGN-1)) == 0); - for (i=0; i < XXH_ACC_NB; i++) { - xxh_u64 const data_val = XXH_readLE64(xinput + 8*i); - xxh_u64 const data_key = data_val ^ XXH_readLE64(xsecret + i*8); - xacc[i ^ 1] += data_val; /* swap adjacent lanes */ - xacc[i] += XXH_mult32to64(data_key & 0xFFFFFFFF, data_key >> 32); - } -} - -XXH_FORCE_INLINE void -XXH3_scrambleAcc_scalar(void* XXH_RESTRICT acc, const void* XXH_RESTRICT secret) -{ - xxh_u64* const xacc = (xxh_u64*) acc; /* presumed aligned */ - const xxh_u8* const xsecret = (const xxh_u8*) secret; /* no alignment restriction */ - size_t i; - XXH_ASSERT((((size_t)acc) & (XXH_ACC_ALIGN-1)) == 0); - for (i=0; i < XXH_ACC_NB; i++) { - xxh_u64 const key64 = XXH_readLE64(xsecret + 8*i); - xxh_u64 acc64 = xacc[i]; - acc64 = XXH_xorshift64(acc64, 47); - acc64 ^= key64; - acc64 *= XXH_PRIME32_1; - xacc[i] = acc64; - } -} - -XXH_FORCE_INLINE void -XXH3_initCustomSecret_scalar(void* XXH_RESTRICT customSecret, xxh_u64 seed64) -{ - /* - * We need a separate pointer for the hack below, - * which requires a non-const pointer. - * Any decent compiler will optimize this out otherwise. - */ - const xxh_u8* kSecretPtr = XXH3_kSecret; - XXH_STATIC_ASSERT((XXH_SECRET_DEFAULT_SIZE & 15) == 0); - -#if defined(__clang__) && defined(__aarch64__) - /* - * UGLY HACK: - * Clang generates a bunch of MOV/MOVK pairs for aarch64, and they are - * placed sequentially, in order, at the top of the unrolled loop. - * - * While MOVK is great for generating constants (2 cycles for a 64-bit - * constant compared to 4 cycles for LDR), long MOVK chains stall the - * integer pipelines: - * I L S - * MOVK - * MOVK - * MOVK - * MOVK - * ADD - * SUB STR - * STR - * By forcing loads from memory (as the asm line causes Clang to assume - * that XXH3_kSecretPtr has been changed), the pipelines are used more - * efficiently: - * I L S - * LDR - * ADD LDR - * SUB STR - * STR - * XXH3_64bits_withSeed, len == 256, Snapdragon 835 - * without hack: 2654.4 MB/s - * with hack: 3202.9 MB/s - */ - XXH_COMPILER_GUARD(kSecretPtr); -#endif - /* - * Note: in debug mode, this overrides the asm optimization - * and Clang will emit MOVK chains again. - */ - XXH_ASSERT(kSecretPtr == XXH3_kSecret); - - { int const nbRounds = XXH_SECRET_DEFAULT_SIZE / 16; - int i; - for (i=0; i < nbRounds; i++) { - /* - * The asm hack causes Clang to assume that kSecretPtr aliases with - * customSecret, and on aarch64, this prevented LDP from merging two - * loads together for free. Putting the loads together before the stores - * properly generates LDP. - */ - xxh_u64 lo = XXH_readLE64(kSecretPtr + 16*i) + seed64; - xxh_u64 hi = XXH_readLE64(kSecretPtr + 16*i + 8) - seed64; - XXH_writeLE64((xxh_u8*)customSecret + 16*i, lo); - XXH_writeLE64((xxh_u8*)customSecret + 16*i + 8, hi); - } } -} - - -typedef void (*XXH3_f_accumulate_512)(void* XXH_RESTRICT, const void*, const void*); -typedef void (*XXH3_f_scrambleAcc)(void* XXH_RESTRICT, const void*); -typedef void (*XXH3_f_initCustomSecret)(void* XXH_RESTRICT, xxh_u64); - - -#if (XXH_VECTOR == XXH_AVX512) - -#define XXH3_accumulate_512 XXH3_accumulate_512_avx512 -#define XXH3_scrambleAcc XXH3_scrambleAcc_avx512 -#define XXH3_initCustomSecret XXH3_initCustomSecret_avx512 - -#elif (XXH_VECTOR == XXH_AVX2) - -#define XXH3_accumulate_512 XXH3_accumulate_512_avx2 -#define XXH3_scrambleAcc XXH3_scrambleAcc_avx2 -#define XXH3_initCustomSecret XXH3_initCustomSecret_avx2 - -#elif (XXH_VECTOR == XXH_SSE2) - -#define XXH3_accumulate_512 XXH3_accumulate_512_sse2 -#define XXH3_scrambleAcc XXH3_scrambleAcc_sse2 -#define XXH3_initCustomSecret XXH3_initCustomSecret_sse2 - -#elif (XXH_VECTOR == XXH_NEON) - -#define XXH3_accumulate_512 XXH3_accumulate_512_neon -#define XXH3_scrambleAcc XXH3_scrambleAcc_neon -#define XXH3_initCustomSecret XXH3_initCustomSecret_scalar - -#elif (XXH_VECTOR == XXH_VSX) - -#define XXH3_accumulate_512 XXH3_accumulate_512_vsx -#define XXH3_scrambleAcc XXH3_scrambleAcc_vsx -#define XXH3_initCustomSecret XXH3_initCustomSecret_scalar - -#else /* scalar */ - -#define XXH3_accumulate_512 XXH3_accumulate_512_scalar -#define XXH3_scrambleAcc XXH3_scrambleAcc_scalar -#define XXH3_initCustomSecret XXH3_initCustomSecret_scalar - -#endif - - - -#ifndef XXH_PREFETCH_DIST -# ifdef __clang__ -# define XXH_PREFETCH_DIST 320 -# else -# if (XXH_VECTOR == XXH_AVX512) -# define XXH_PREFETCH_DIST 512 -# else -# define XXH_PREFETCH_DIST 384 -# endif -# endif /* __clang__ */ -#endif /* XXH_PREFETCH_DIST */ - -/* - * XXH3_accumulate() - * Loops over XXH3_accumulate_512(). - * Assumption: nbStripes will not overflow the secret size - */ -XXH_FORCE_INLINE void -XXH3_accumulate( xxh_u64* XXH_RESTRICT acc, - const xxh_u8* XXH_RESTRICT input, - const xxh_u8* XXH_RESTRICT secret, - size_t nbStripes, - XXH3_f_accumulate_512 f_acc512) -{ - size_t n; - for (n = 0; n < nbStripes; n++ ) { - const xxh_u8* const in = input + n*XXH_STRIPE_LEN; - XXH_PREFETCH(in + XXH_PREFETCH_DIST); - f_acc512(acc, - in, - secret + n*XXH_SECRET_CONSUME_RATE); - } -} - -XXH_FORCE_INLINE void -XXH3_hashLong_internal_loop(xxh_u64* XXH_RESTRICT acc, - const xxh_u8* XXH_RESTRICT input, size_t len, - const xxh_u8* XXH_RESTRICT secret, size_t secretSize, - XXH3_f_accumulate_512 f_acc512, - XXH3_f_scrambleAcc f_scramble) -{ - size_t const nbStripesPerBlock = (secretSize - XXH_STRIPE_LEN) / XXH_SECRET_CONSUME_RATE; - size_t const block_len = XXH_STRIPE_LEN * nbStripesPerBlock; - size_t const nb_blocks = (len - 1) / block_len; - - size_t n; - - XXH_ASSERT(secretSize >= XXH3_SECRET_SIZE_MIN); - - for (n = 0; n < nb_blocks; n++) { - XXH3_accumulate(acc, input + n*block_len, secret, nbStripesPerBlock, f_acc512); - f_scramble(acc, secret + secretSize - XXH_STRIPE_LEN); - } - - /* last partial block */ - XXH_ASSERT(len > XXH_STRIPE_LEN); - { size_t const nbStripes = ((len - 1) - (block_len * nb_blocks)) / XXH_STRIPE_LEN; - XXH_ASSERT(nbStripes <= (secretSize / XXH_SECRET_CONSUME_RATE)); - XXH3_accumulate(acc, input + nb_blocks*block_len, secret, nbStripes, f_acc512); - - /* last stripe */ - { const xxh_u8* const p = input + len - XXH_STRIPE_LEN; -#define XXH_SECRET_LASTACC_START 7 /* not aligned on 8, last secret is different from acc & scrambler */ - f_acc512(acc, p, secret + secretSize - XXH_STRIPE_LEN - XXH_SECRET_LASTACC_START); - } } -} - -XXH_FORCE_INLINE xxh_u64 -XXH3_mix2Accs(const xxh_u64* XXH_RESTRICT acc, const xxh_u8* XXH_RESTRICT secret) -{ - return XXH3_mul128_fold64( - acc[0] ^ XXH_readLE64(secret), - acc[1] ^ XXH_readLE64(secret+8) ); -} - -static XXH64_hash_t -XXH3_mergeAccs(const xxh_u64* XXH_RESTRICT acc, const xxh_u8* XXH_RESTRICT secret, xxh_u64 start) -{ - xxh_u64 result64 = start; - size_t i = 0; - - for (i = 0; i < 4; i++) { - result64 += XXH3_mix2Accs(acc+2*i, secret + 16*i); -#if defined(__clang__) /* Clang */ \ - && (defined(__arm__) || defined(__thumb__)) /* ARMv7 */ \ - && (defined(__ARM_NEON) || defined(__ARM_NEON__)) /* NEON */ \ - && !defined(XXH_ENABLE_AUTOVECTORIZE) /* Define to disable */ - /* - * UGLY HACK: - * Prevent autovectorization on Clang ARMv7-a. Exact same problem as - * the one in XXH3_len_129to240_64b. Speeds up shorter keys > 240b. - * XXH3_64bits, len == 256, Snapdragon 835: - * without hack: 2063.7 MB/s - * with hack: 2560.7 MB/s - */ - XXH_COMPILER_GUARD(result64); -#endif - } - - return XXH3_avalanche(result64); -} - -#define XXH3_INIT_ACC { XXH_PRIME32_3, XXH_PRIME64_1, XXH_PRIME64_2, XXH_PRIME64_3, \ - XXH_PRIME64_4, XXH_PRIME32_2, XXH_PRIME64_5, XXH_PRIME32_1 } - -XXH_FORCE_INLINE XXH64_hash_t -XXH3_hashLong_64b_internal(const void* XXH_RESTRICT input, size_t len, - const void* XXH_RESTRICT secret, size_t secretSize, - XXH3_f_accumulate_512 f_acc512, - XXH3_f_scrambleAcc f_scramble) -{ - XXH_ALIGN(XXH_ACC_ALIGN) xxh_u64 acc[XXH_ACC_NB] = XXH3_INIT_ACC; - - XXH3_hashLong_internal_loop(acc, (const xxh_u8*)input, len, (const xxh_u8*)secret, secretSize, f_acc512, f_scramble); - - /* converge into final hash */ - XXH_STATIC_ASSERT(sizeof(acc) == 64); - /* do not align on 8, so that the secret is different from the accumulator */ -#define XXH_SECRET_MERGEACCS_START 11 - XXH_ASSERT(secretSize >= sizeof(acc) + XXH_SECRET_MERGEACCS_START); - return XXH3_mergeAccs(acc, (const xxh_u8*)secret + XXH_SECRET_MERGEACCS_START, (xxh_u64)len * XXH_PRIME64_1); -} - -/* - * It's important for performance to transmit secret's size (when it's static) - * so that the compiler can properly optimize the vectorized loop. - * This makes a big performance difference for "medium" keys (<1 KB) when using AVX instruction set. - */ -XXH_FORCE_INLINE XXH64_hash_t -XXH3_hashLong_64b_withSecret(const void* XXH_RESTRICT input, size_t len, - XXH64_hash_t seed64, const xxh_u8* XXH_RESTRICT secret, size_t secretLen) -{ - (void)seed64; - return XXH3_hashLong_64b_internal(input, len, secret, secretLen, XXH3_accumulate_512, XXH3_scrambleAcc); -} - -/* - * It's preferable for performance that XXH3_hashLong is not inlined, - * as it results in a smaller function for small data, easier to the instruction cache. - * Note that inside this no_inline function, we do inline the internal loop, - * and provide a statically defined secret size to allow optimization of vector loop. - */ -XXH_NO_INLINE XXH64_hash_t -XXH3_hashLong_64b_default(const void* XXH_RESTRICT input, size_t len, - XXH64_hash_t seed64, const xxh_u8* XXH_RESTRICT secret, size_t secretLen) -{ - (void)seed64; (void)secret; (void)secretLen; - return XXH3_hashLong_64b_internal(input, len, XXH3_kSecret, sizeof(XXH3_kSecret), XXH3_accumulate_512, XXH3_scrambleAcc); -} - -/* - * XXH3_hashLong_64b_withSeed(): - * Generate a custom key based on alteration of default XXH3_kSecret with the seed, - * and then use this key for long mode hashing. - * - * This operation is decently fast but nonetheless costs a little bit of time. - * Try to avoid it whenever possible (typically when seed==0). - * - * It's important for performance that XXH3_hashLong is not inlined. Not sure - * why (uop cache maybe?), but the difference is large and easily measurable. - */ -XXH_FORCE_INLINE XXH64_hash_t -XXH3_hashLong_64b_withSeed_internal(const void* input, size_t len, - XXH64_hash_t seed, - XXH3_f_accumulate_512 f_acc512, - XXH3_f_scrambleAcc f_scramble, - XXH3_f_initCustomSecret f_initSec) -{ - if (seed == 0) - return XXH3_hashLong_64b_internal(input, len, - XXH3_kSecret, sizeof(XXH3_kSecret), - f_acc512, f_scramble); - { XXH_ALIGN(XXH_SEC_ALIGN) xxh_u8 secret[XXH_SECRET_DEFAULT_SIZE]; - f_initSec(secret, seed); - return XXH3_hashLong_64b_internal(input, len, secret, sizeof(secret), - f_acc512, f_scramble); - } -} - -/* - * It's important for performance that XXH3_hashLong is not inlined. - */ -XXH_NO_INLINE XXH64_hash_t -XXH3_hashLong_64b_withSeed(const void* input, size_t len, - XXH64_hash_t seed, const xxh_u8* secret, size_t secretLen) -{ - (void)secret; (void)secretLen; - return XXH3_hashLong_64b_withSeed_internal(input, len, seed, - XXH3_accumulate_512, XXH3_scrambleAcc, XXH3_initCustomSecret); -} - - -typedef XXH64_hash_t (*XXH3_hashLong64_f)(const void* XXH_RESTRICT, size_t, - XXH64_hash_t, const xxh_u8* XXH_RESTRICT, size_t); - -XXH_FORCE_INLINE XXH64_hash_t -XXH3_64bits_internal(const void* XXH_RESTRICT input, size_t len, - XXH64_hash_t seed64, const void* XXH_RESTRICT secret, size_t secretLen, - XXH3_hashLong64_f f_hashLong) -{ - XXH_ASSERT(secretLen >= XXH3_SECRET_SIZE_MIN); - /* - * If an action is to be taken if `secretLen` condition is not respected, - * it should be done here. - * For now, it's a contract pre-condition. - * Adding a check and a branch here would cost performance at every hash. - * Also, note that function signature doesn't offer room to return an error. - */ - if (len <= 16) - return XXH3_len_0to16_64b((const xxh_u8*)input, len, (const xxh_u8*)secret, seed64); - if (len <= 128) - return XXH3_len_17to128_64b((const xxh_u8*)input, len, (const xxh_u8*)secret, secretLen, seed64); - if (len <= XXH3_MIDSIZE_MAX) - return XXH3_len_129to240_64b((const xxh_u8*)input, len, (const xxh_u8*)secret, secretLen, seed64); - return f_hashLong(input, len, seed64, (const xxh_u8*)secret, secretLen); -} - - -/* === Public entry point === */ - -/*! @ingroup xxh3_family */ -XXH_PUBLIC_API XXH64_hash_t XXH3_64bits(const void* input, size_t len) -{ - return XXH3_64bits_internal(input, len, 0, XXH3_kSecret, sizeof(XXH3_kSecret), XXH3_hashLong_64b_default); -} - -/*! @ingroup xxh3_family */ -XXH_PUBLIC_API XXH64_hash_t -XXH3_64bits_withSecret(const void* input, size_t len, const void* secret, size_t secretSize) -{ - return XXH3_64bits_internal(input, len, 0, secret, secretSize, XXH3_hashLong_64b_withSecret); -} - -/*! @ingroup xxh3_family */ -XXH_PUBLIC_API XXH64_hash_t -XXH3_64bits_withSeed(const void* input, size_t len, XXH64_hash_t seed) -{ - return XXH3_64bits_internal(input, len, seed, XXH3_kSecret, sizeof(XXH3_kSecret), XXH3_hashLong_64b_withSeed); -} - -XXH_PUBLIC_API XXH64_hash_t -XXH3_64bits_withSecretandSeed(const void* input, size_t len, const void* secret, size_t secretSize, XXH64_hash_t seed) -{ - if (len <= XXH3_MIDSIZE_MAX) - return XXH3_64bits_internal(input, len, seed, XXH3_kSecret, sizeof(XXH3_kSecret), NULL); - return XXH3_hashLong_64b_withSecret(input, len, seed, (const xxh_u8*)secret, secretSize); -} - - -/* === XXH3 streaming === */ - -/* - * Malloc's a pointer that is always aligned to align. - * - * This must be freed with `XXH_alignedFree()`. - * - * malloc typically guarantees 16 byte alignment on 64-bit systems and 8 byte - * alignment on 32-bit. This isn't enough for the 32 byte aligned loads in AVX2 - * or on 32-bit, the 16 byte aligned loads in SSE2 and NEON. - * - * This underalignment previously caused a rather obvious crash which went - * completely unnoticed due to XXH3_createState() not actually being tested. - * Credit to RedSpah for noticing this bug. - * - * The alignment is done manually: Functions like posix_memalign or _mm_malloc - * are avoided: To maintain portability, we would have to write a fallback - * like this anyways, and besides, testing for the existence of library - * functions without relying on external build tools is impossible. - * - * The method is simple: Overallocate, manually align, and store the offset - * to the original behind the returned pointer. - * - * Align must be a power of 2 and 8 <= align <= 128. - */ -static void* XXH_alignedMalloc(size_t s, size_t align) -{ - XXH_ASSERT(align <= 128 && align >= 8); /* range check */ - XXH_ASSERT((align & (align-1)) == 0); /* power of 2 */ - XXH_ASSERT(s != 0 && s < (s + align)); /* empty/overflow */ - { /* Overallocate to make room for manual realignment and an offset byte */ - xxh_u8* base = (xxh_u8*)XXH_malloc(s + align); - if (base != NULL) { - /* - * Get the offset needed to align this pointer. - * - * Even if the returned pointer is aligned, there will always be - * at least one byte to store the offset to the original pointer. - */ - size_t offset = align - ((size_t)base & (align - 1)); /* base % align */ - /* Add the offset for the now-aligned pointer */ - xxh_u8* ptr = base + offset; - - XXH_ASSERT((size_t)ptr % align == 0); - - /* Store the offset immediately before the returned pointer. */ - ptr[-1] = (xxh_u8)offset; - return ptr; - } - return NULL; - } -} -/* - * Frees an aligned pointer allocated by XXH_alignedMalloc(). Don't pass - * normal malloc'd pointers, XXH_alignedMalloc has a specific data layout. - */ -static void XXH_alignedFree(void* p) -{ - if (p != NULL) { - xxh_u8* ptr = (xxh_u8*)p; - /* Get the offset byte we added in XXH_malloc. */ - xxh_u8 offset = ptr[-1]; - /* Free the original malloc'd pointer */ - xxh_u8* base = ptr - offset; - XXH_free(base); - } -} -/*! @ingroup xxh3_family */ -XXH_PUBLIC_API XXH3_state_t* XXH3_createState(void) -{ - XXH3_state_t* const state = (XXH3_state_t*)XXH_alignedMalloc(sizeof(XXH3_state_t), 64); - if (state==NULL) return NULL; - XXH3_INITSTATE(state); - return state; -} - -/*! @ingroup xxh3_family */ -XXH_PUBLIC_API XXH_errorcode XXH3_freeState(XXH3_state_t* statePtr) -{ - XXH_alignedFree(statePtr); - return XXH_OK; -} - -/*! @ingroup xxh3_family */ -XXH_PUBLIC_API void -XXH3_copyState(XXH3_state_t* dst_state, const XXH3_state_t* src_state) -{ - XXH_memcpy(dst_state, src_state, sizeof(*dst_state)); -} - -static void -XXH3_reset_internal(XXH3_state_t* statePtr, - XXH64_hash_t seed, - const void* secret, size_t secretSize) -{ - size_t const initStart = offsetof(XXH3_state_t, bufferedSize); - size_t const initLength = offsetof(XXH3_state_t, nbStripesPerBlock) - initStart; - XXH_ASSERT(offsetof(XXH3_state_t, nbStripesPerBlock) > initStart); - XXH_ASSERT(statePtr != NULL); - /* set members from bufferedSize to nbStripesPerBlock (excluded) to 0 */ - memset((char*)statePtr + initStart, 0, initLength); - statePtr->acc[0] = XXH_PRIME32_3; - statePtr->acc[1] = XXH_PRIME64_1; - statePtr->acc[2] = XXH_PRIME64_2; - statePtr->acc[3] = XXH_PRIME64_3; - statePtr->acc[4] = XXH_PRIME64_4; - statePtr->acc[5] = XXH_PRIME32_2; - statePtr->acc[6] = XXH_PRIME64_5; - statePtr->acc[7] = XXH_PRIME32_1; - statePtr->seed = seed; - statePtr->useSeed = (seed != 0); - statePtr->extSecret = (const unsigned char*)secret; - XXH_ASSERT(secretSize >= XXH3_SECRET_SIZE_MIN); - statePtr->secretLimit = secretSize - XXH_STRIPE_LEN; - statePtr->nbStripesPerBlock = statePtr->secretLimit / XXH_SECRET_CONSUME_RATE; -} - -/*! @ingroup xxh3_family */ -XXH_PUBLIC_API XXH_errorcode -XXH3_64bits_reset(XXH3_state_t* statePtr) -{ - if (statePtr == NULL) return XXH_ERROR; - XXH3_reset_internal(statePtr, 0, XXH3_kSecret, XXH_SECRET_DEFAULT_SIZE); - return XXH_OK; -} - -/*! @ingroup xxh3_family */ -XXH_PUBLIC_API XXH_errorcode -XXH3_64bits_reset_withSecret(XXH3_state_t* statePtr, const void* secret, size_t secretSize) -{ - if (statePtr == NULL) return XXH_ERROR; - XXH3_reset_internal(statePtr, 0, secret, secretSize); - if (secret == NULL) return XXH_ERROR; - if (secretSize < XXH3_SECRET_SIZE_MIN) return XXH_ERROR; - return XXH_OK; -} - -/*! @ingroup xxh3_family */ -XXH_PUBLIC_API XXH_errorcode -XXH3_64bits_reset_withSeed(XXH3_state_t* statePtr, XXH64_hash_t seed) -{ - if (statePtr == NULL) return XXH_ERROR; - if (seed==0) return XXH3_64bits_reset(statePtr); - if ((seed != statePtr->seed) || (statePtr->extSecret != NULL)) - XXH3_initCustomSecret(statePtr->customSecret, seed); - XXH3_reset_internal(statePtr, seed, NULL, XXH_SECRET_DEFAULT_SIZE); - return XXH_OK; -} - -/*! @ingroup xxh3_family */ -XXH_PUBLIC_API XXH_errorcode -XXH3_64bits_reset_withSecretandSeed(XXH3_state_t* statePtr, const void* secret, size_t secretSize, XXH64_hash_t seed64) -{ - if (statePtr == NULL) return XXH_ERROR; - if (secret == NULL) return XXH_ERROR; - if (secretSize < XXH3_SECRET_SIZE_MIN) return XXH_ERROR; - XXH3_reset_internal(statePtr, seed64, secret, secretSize); - statePtr->useSeed = 1; /* always, even if seed64==0 */ - return XXH_OK; -} - -/* Note : when XXH3_consumeStripes() is invoked, - * there must be a guarantee that at least one more byte must be consumed from input - * so that the function can blindly consume all stripes using the "normal" secret segment */ -XXH_FORCE_INLINE void -XXH3_consumeStripes(xxh_u64* XXH_RESTRICT acc, - size_t* XXH_RESTRICT nbStripesSoFarPtr, size_t nbStripesPerBlock, - const xxh_u8* XXH_RESTRICT input, size_t nbStripes, - const xxh_u8* XXH_RESTRICT secret, size_t secretLimit, - XXH3_f_accumulate_512 f_acc512, - XXH3_f_scrambleAcc f_scramble) -{ - XXH_ASSERT(nbStripes <= nbStripesPerBlock); /* can handle max 1 scramble per invocation */ - XXH_ASSERT(*nbStripesSoFarPtr < nbStripesPerBlock); - if (nbStripesPerBlock - *nbStripesSoFarPtr <= nbStripes) { - /* need a scrambling operation */ - size_t const nbStripesToEndofBlock = nbStripesPerBlock - *nbStripesSoFarPtr; - size_t const nbStripesAfterBlock = nbStripes - nbStripesToEndofBlock; - XXH3_accumulate(acc, input, secret + nbStripesSoFarPtr[0] * XXH_SECRET_CONSUME_RATE, nbStripesToEndofBlock, f_acc512); - f_scramble(acc, secret + secretLimit); - XXH3_accumulate(acc, input + nbStripesToEndofBlock * XXH_STRIPE_LEN, secret, nbStripesAfterBlock, f_acc512); - *nbStripesSoFarPtr = nbStripesAfterBlock; - } else { - XXH3_accumulate(acc, input, secret + nbStripesSoFarPtr[0] * XXH_SECRET_CONSUME_RATE, nbStripes, f_acc512); - *nbStripesSoFarPtr += nbStripes; - } -} - -#ifndef XXH3_STREAM_USE_STACK -# ifndef __clang__ /* clang doesn't need additional stack space */ -# define XXH3_STREAM_USE_STACK 1 -# endif -#endif -/* - * Both XXH3_64bits_update and XXH3_128bits_update use this routine. - */ -XXH_FORCE_INLINE XXH_errorcode -XXH3_update(XXH3_state_t* XXH_RESTRICT const state, - const xxh_u8* XXH_RESTRICT input, size_t len, - XXH3_f_accumulate_512 f_acc512, - XXH3_f_scrambleAcc f_scramble) -{ - if (input==NULL) { - XXH_ASSERT(len == 0); - return XXH_OK; - } - - XXH_ASSERT(state != NULL); - { const xxh_u8* const bEnd = input + len; - const unsigned char* const secret = (state->extSecret == NULL) ? state->customSecret : state->extSecret; -#if defined(XXH3_STREAM_USE_STACK) && XXH3_STREAM_USE_STACK >= 1 - /* For some reason, gcc and MSVC seem to suffer greatly - * when operating accumulators directly into state. - * Operating into stack space seems to enable proper optimization. - * clang, on the other hand, doesn't seem to need this trick */ - XXH_ALIGN(XXH_ACC_ALIGN) xxh_u64 acc[8]; memcpy(acc, state->acc, sizeof(acc)); -#else - xxh_u64* XXH_RESTRICT const acc = state->acc; -#endif - state->totalLen += len; - XXH_ASSERT(state->bufferedSize <= XXH3_INTERNALBUFFER_SIZE); - - /* small input : just fill in tmp buffer */ - if (state->bufferedSize + len <= XXH3_INTERNALBUFFER_SIZE) { - XXH_memcpy(state->buffer + state->bufferedSize, input, len); - state->bufferedSize += (XXH32_hash_t)len; - return XXH_OK; - } - - /* total input is now > XXH3_INTERNALBUFFER_SIZE */ - #define XXH3_INTERNALBUFFER_STRIPES (XXH3_INTERNALBUFFER_SIZE / XXH_STRIPE_LEN) - XXH_STATIC_ASSERT(XXH3_INTERNALBUFFER_SIZE % XXH_STRIPE_LEN == 0); /* clean multiple */ - - /* - * Internal buffer is partially filled (always, except at beginning) - * Complete it, then consume it. - */ - if (state->bufferedSize) { - size_t const loadSize = XXH3_INTERNALBUFFER_SIZE - state->bufferedSize; - XXH_memcpy(state->buffer + state->bufferedSize, input, loadSize); - input += loadSize; - XXH3_consumeStripes(acc, - &state->nbStripesSoFar, state->nbStripesPerBlock, - state->buffer, XXH3_INTERNALBUFFER_STRIPES, - secret, state->secretLimit, - f_acc512, f_scramble); - state->bufferedSize = 0; - } - XXH_ASSERT(input < bEnd); - - /* large input to consume : ingest per full block */ - if ((size_t)(bEnd - input) > state->nbStripesPerBlock * XXH_STRIPE_LEN) { - size_t nbStripes = (size_t)(bEnd - 1 - input) / XXH_STRIPE_LEN; - XXH_ASSERT(state->nbStripesPerBlock >= state->nbStripesSoFar); - /* join to current block's end */ - { size_t const nbStripesToEnd = state->nbStripesPerBlock - state->nbStripesSoFar; - XXH_ASSERT(nbStripes <= nbStripes); - XXH3_accumulate(acc, input, secret + state->nbStripesSoFar * XXH_SECRET_CONSUME_RATE, nbStripesToEnd, f_acc512); - f_scramble(acc, secret + state->secretLimit); - state->nbStripesSoFar = 0; - input += nbStripesToEnd * XXH_STRIPE_LEN; - nbStripes -= nbStripesToEnd; - } - /* consume per entire blocks */ - while(nbStripes >= state->nbStripesPerBlock) { - XXH3_accumulate(acc, input, secret, state->nbStripesPerBlock, f_acc512); - f_scramble(acc, secret + state->secretLimit); - input += state->nbStripesPerBlock * XXH_STRIPE_LEN; - nbStripes -= state->nbStripesPerBlock; - } - /* consume last partial block */ - XXH3_accumulate(acc, input, secret, nbStripes, f_acc512); - input += nbStripes * XXH_STRIPE_LEN; - XXH_ASSERT(input < bEnd); /* at least some bytes left */ - state->nbStripesSoFar = nbStripes; - /* buffer predecessor of last partial stripe */ - XXH_memcpy(state->buffer + sizeof(state->buffer) - XXH_STRIPE_LEN, input - XXH_STRIPE_LEN, XXH_STRIPE_LEN); - XXH_ASSERT(bEnd - input <= XXH_STRIPE_LEN); - } else { - /* content to consume <= block size */ - /* Consume input by a multiple of internal buffer size */ - if (bEnd - input > XXH3_INTERNALBUFFER_SIZE) { - const xxh_u8* const limit = bEnd - XXH3_INTERNALBUFFER_SIZE; - do { - XXH3_consumeStripes(acc, - &state->nbStripesSoFar, state->nbStripesPerBlock, - input, XXH3_INTERNALBUFFER_STRIPES, - secret, state->secretLimit, - f_acc512, f_scramble); - input += XXH3_INTERNALBUFFER_SIZE; - } while (inputbuffer + sizeof(state->buffer) - XXH_STRIPE_LEN, input - XXH_STRIPE_LEN, XXH_STRIPE_LEN); - } - } - - /* Some remaining input (always) : buffer it */ - XXH_ASSERT(input < bEnd); - XXH_ASSERT(bEnd - input <= XXH3_INTERNALBUFFER_SIZE); - XXH_ASSERT(state->bufferedSize == 0); - XXH_memcpy(state->buffer, input, (size_t)(bEnd-input)); - state->bufferedSize = (XXH32_hash_t)(bEnd-input); -#if defined(XXH3_STREAM_USE_STACK) && XXH3_STREAM_USE_STACK >= 1 - /* save stack accumulators into state */ - memcpy(state->acc, acc, sizeof(acc)); -#endif - } - - return XXH_OK; -} - -/*! @ingroup xxh3_family */ -XXH_PUBLIC_API XXH_errorcode -XXH3_64bits_update(XXH3_state_t* state, const void* input, size_t len) -{ - return XXH3_update(state, (const xxh_u8*)input, len, - XXH3_accumulate_512, XXH3_scrambleAcc); -} - - -XXH_FORCE_INLINE void -XXH3_digest_long (XXH64_hash_t* acc, - const XXH3_state_t* state, - const unsigned char* secret) -{ - /* - * Digest on a local copy. This way, the state remains unaltered, and it can - * continue ingesting more input afterwards. - */ - XXH_memcpy(acc, state->acc, sizeof(state->acc)); - if (state->bufferedSize >= XXH_STRIPE_LEN) { - size_t const nbStripes = (state->bufferedSize - 1) / XXH_STRIPE_LEN; - size_t nbStripesSoFar = state->nbStripesSoFar; - XXH3_consumeStripes(acc, - &nbStripesSoFar, state->nbStripesPerBlock, - state->buffer, nbStripes, - secret, state->secretLimit, - XXH3_accumulate_512, XXH3_scrambleAcc); - /* last stripe */ - XXH3_accumulate_512(acc, - state->buffer + state->bufferedSize - XXH_STRIPE_LEN, - secret + state->secretLimit - XXH_SECRET_LASTACC_START); - } else { /* bufferedSize < XXH_STRIPE_LEN */ - xxh_u8 lastStripe[XXH_STRIPE_LEN]; - size_t const catchupSize = XXH_STRIPE_LEN - state->bufferedSize; - XXH_ASSERT(state->bufferedSize > 0); /* there is always some input buffered */ - XXH_memcpy(lastStripe, state->buffer + sizeof(state->buffer) - catchupSize, catchupSize); - XXH_memcpy(lastStripe + catchupSize, state->buffer, state->bufferedSize); - XXH3_accumulate_512(acc, - lastStripe, - secret + state->secretLimit - XXH_SECRET_LASTACC_START); - } -} - -/*! @ingroup xxh3_family */ -XXH_PUBLIC_API XXH64_hash_t XXH3_64bits_digest (const XXH3_state_t* state) -{ - const unsigned char* const secret = (state->extSecret == NULL) ? state->customSecret : state->extSecret; - if (state->totalLen > XXH3_MIDSIZE_MAX) { - XXH_ALIGN(XXH_ACC_ALIGN) XXH64_hash_t acc[XXH_ACC_NB]; - XXH3_digest_long(acc, state, secret); - return XXH3_mergeAccs(acc, - secret + XXH_SECRET_MERGEACCS_START, - (xxh_u64)state->totalLen * XXH_PRIME64_1); - } - /* totalLen <= XXH3_MIDSIZE_MAX: digesting a short input */ - if (state->useSeed) - return XXH3_64bits_withSeed(state->buffer, (size_t)state->totalLen, state->seed); - return XXH3_64bits_withSecret(state->buffer, (size_t)(state->totalLen), - secret, state->secretLimit + XXH_STRIPE_LEN); -} - - - -/* ========================================== - * XXH3 128 bits (a.k.a XXH128) - * ========================================== - * XXH3's 128-bit variant has better mixing and strength than the 64-bit variant, - * even without counting the significantly larger output size. - * - * For example, extra steps are taken to avoid the seed-dependent collisions - * in 17-240 byte inputs (See XXH3_mix16B and XXH128_mix32B). - * - * This strength naturally comes at the cost of some speed, especially on short - * lengths. Note that longer hashes are about as fast as the 64-bit version - * due to it using only a slight modification of the 64-bit loop. - * - * XXH128 is also more oriented towards 64-bit machines. It is still extremely - * fast for a _128-bit_ hash on 32-bit (it usually clears XXH64). - */ - -XXH_FORCE_INLINE XXH128_hash_t -XXH3_len_1to3_128b(const xxh_u8* input, size_t len, const xxh_u8* secret, XXH64_hash_t seed) -{ - /* A doubled version of 1to3_64b with different constants. */ - XXH_ASSERT(input != NULL); - XXH_ASSERT(1 <= len && len <= 3); - XXH_ASSERT(secret != NULL); - /* - * len = 1: combinedl = { input[0], 0x01, input[0], input[0] } - * len = 2: combinedl = { input[1], 0x02, input[0], input[1] } - * len = 3: combinedl = { input[2], 0x03, input[0], input[1] } - */ - { xxh_u8 const c1 = input[0]; - xxh_u8 const c2 = input[len >> 1]; - xxh_u8 const c3 = input[len - 1]; - xxh_u32 const combinedl = ((xxh_u32)c1 <<16) | ((xxh_u32)c2 << 24) - | ((xxh_u32)c3 << 0) | ((xxh_u32)len << 8); - xxh_u32 const combinedh = XXH_rotl32(XXH_swap32(combinedl), 13); - xxh_u64 const bitflipl = (XXH_readLE32(secret) ^ XXH_readLE32(secret+4)) + seed; - xxh_u64 const bitfliph = (XXH_readLE32(secret+8) ^ XXH_readLE32(secret+12)) - seed; - xxh_u64 const keyed_lo = (xxh_u64)combinedl ^ bitflipl; - xxh_u64 const keyed_hi = (xxh_u64)combinedh ^ bitfliph; - XXH128_hash_t h128; - h128.low64 = XXH64_avalanche(keyed_lo); - h128.high64 = XXH64_avalanche(keyed_hi); - return h128; - } -} - -XXH_FORCE_INLINE XXH128_hash_t -XXH3_len_4to8_128b(const xxh_u8* input, size_t len, const xxh_u8* secret, XXH64_hash_t seed) -{ - XXH_ASSERT(input != NULL); - XXH_ASSERT(secret != NULL); - XXH_ASSERT(4 <= len && len <= 8); - seed ^= (xxh_u64)XXH_swap32((xxh_u32)seed) << 32; - { xxh_u32 const input_lo = XXH_readLE32(input); - xxh_u32 const input_hi = XXH_readLE32(input + len - 4); - xxh_u64 const input_64 = input_lo + ((xxh_u64)input_hi << 32); - xxh_u64 const bitflip = (XXH_readLE64(secret+16) ^ XXH_readLE64(secret+24)) + seed; - xxh_u64 const keyed = input_64 ^ bitflip; - - /* Shift len to the left to ensure it is even, this avoids even multiplies. */ - XXH128_hash_t m128 = XXH_mult64to128(keyed, XXH_PRIME64_1 + (len << 2)); - - m128.high64 += (m128.low64 << 1); - m128.low64 ^= (m128.high64 >> 3); - - m128.low64 = XXH_xorshift64(m128.low64, 35); - m128.low64 *= 0x9FB21C651E98DF25ULL; - m128.low64 = XXH_xorshift64(m128.low64, 28); - m128.high64 = XXH3_avalanche(m128.high64); - return m128; - } -} - -XXH_FORCE_INLINE XXH128_hash_t -XXH3_len_9to16_128b(const xxh_u8* input, size_t len, const xxh_u8* secret, XXH64_hash_t seed) -{ - XXH_ASSERT(input != NULL); - XXH_ASSERT(secret != NULL); - XXH_ASSERT(9 <= len && len <= 16); - { xxh_u64 const bitflipl = (XXH_readLE64(secret+32) ^ XXH_readLE64(secret+40)) - seed; - xxh_u64 const bitfliph = (XXH_readLE64(secret+48) ^ XXH_readLE64(secret+56)) + seed; - xxh_u64 const input_lo = XXH_readLE64(input); - xxh_u64 input_hi = XXH_readLE64(input + len - 8); - XXH128_hash_t m128 = XXH_mult64to128(input_lo ^ input_hi ^ bitflipl, XXH_PRIME64_1); - /* - * Put len in the middle of m128 to ensure that the length gets mixed to - * both the low and high bits in the 128x64 multiply below. - */ - m128.low64 += (xxh_u64)(len - 1) << 54; - input_hi ^= bitfliph; - /* - * Add the high 32 bits of input_hi to the high 32 bits of m128, then - * add the long product of the low 32 bits of input_hi and XXH_PRIME32_2 to - * the high 64 bits of m128. - * - * The best approach to this operation is different on 32-bit and 64-bit. - */ - if (sizeof(void *) < sizeof(xxh_u64)) { /* 32-bit */ - /* - * 32-bit optimized version, which is more readable. - * - * On 32-bit, it removes an ADC and delays a dependency between the two - * halves of m128.high64, but it generates an extra mask on 64-bit. - */ - m128.high64 += (input_hi & 0xFFFFFFFF00000000ULL) + XXH_mult32to64((xxh_u32)input_hi, XXH_PRIME32_2); - } else { - /* - * 64-bit optimized (albeit more confusing) version. - * - * Uses some properties of addition and multiplication to remove the mask: - * - * Let: - * a = input_hi.lo = (input_hi & 0x00000000FFFFFFFF) - * b = input_hi.hi = (input_hi & 0xFFFFFFFF00000000) - * c = XXH_PRIME32_2 - * - * a + (b * c) - * Inverse Property: x + y - x == y - * a + (b * (1 + c - 1)) - * Distributive Property: x * (y + z) == (x * y) + (x * z) - * a + (b * 1) + (b * (c - 1)) - * Identity Property: x * 1 == x - * a + b + (b * (c - 1)) - * - * Substitute a, b, and c: - * input_hi.hi + input_hi.lo + ((xxh_u64)input_hi.lo * (XXH_PRIME32_2 - 1)) - * - * Since input_hi.hi + input_hi.lo == input_hi, we get this: - * input_hi + ((xxh_u64)input_hi.lo * (XXH_PRIME32_2 - 1)) - */ - m128.high64 += input_hi + XXH_mult32to64((xxh_u32)input_hi, XXH_PRIME32_2 - 1); - } - /* m128 ^= XXH_swap64(m128 >> 64); */ - m128.low64 ^= XXH_swap64(m128.high64); - - { /* 128x64 multiply: h128 = m128 * XXH_PRIME64_2; */ - XXH128_hash_t h128 = XXH_mult64to128(m128.low64, XXH_PRIME64_2); - h128.high64 += m128.high64 * XXH_PRIME64_2; - - h128.low64 = XXH3_avalanche(h128.low64); - h128.high64 = XXH3_avalanche(h128.high64); - return h128; - } } -} - -/* - * Assumption: `secret` size is >= XXH3_SECRET_SIZE_MIN - */ -XXH_FORCE_INLINE XXH128_hash_t -XXH3_len_0to16_128b(const xxh_u8* input, size_t len, const xxh_u8* secret, XXH64_hash_t seed) -{ - XXH_ASSERT(len <= 16); - { if (len > 8) return XXH3_len_9to16_128b(input, len, secret, seed); - if (len >= 4) return XXH3_len_4to8_128b(input, len, secret, seed); - if (len) return XXH3_len_1to3_128b(input, len, secret, seed); - { XXH128_hash_t h128; - xxh_u64 const bitflipl = XXH_readLE64(secret+64) ^ XXH_readLE64(secret+72); - xxh_u64 const bitfliph = XXH_readLE64(secret+80) ^ XXH_readLE64(secret+88); - h128.low64 = XXH64_avalanche(seed ^ bitflipl); - h128.high64 = XXH64_avalanche( seed ^ bitfliph); - return h128; - } } -} - -/* - * A bit slower than XXH3_mix16B, but handles multiply by zero better. - */ -XXH_FORCE_INLINE XXH128_hash_t -XXH128_mix32B(XXH128_hash_t acc, const xxh_u8* input_1, const xxh_u8* input_2, - const xxh_u8* secret, XXH64_hash_t seed) -{ - acc.low64 += XXH3_mix16B (input_1, secret+0, seed); - acc.low64 ^= XXH_readLE64(input_2) + XXH_readLE64(input_2 + 8); - acc.high64 += XXH3_mix16B (input_2, secret+16, seed); - acc.high64 ^= XXH_readLE64(input_1) + XXH_readLE64(input_1 + 8); - return acc; -} - - -XXH_FORCE_INLINE XXH128_hash_t -XXH3_len_17to128_128b(const xxh_u8* XXH_RESTRICT input, size_t len, - const xxh_u8* XXH_RESTRICT secret, size_t secretSize, - XXH64_hash_t seed) -{ - XXH_ASSERT(secretSize >= XXH3_SECRET_SIZE_MIN); (void)secretSize; - XXH_ASSERT(16 < len && len <= 128); - - { XXH128_hash_t acc; - acc.low64 = len * XXH_PRIME64_1; - acc.high64 = 0; - if (len > 32) { - if (len > 64) { - if (len > 96) { - acc = XXH128_mix32B(acc, input+48, input+len-64, secret+96, seed); - } - acc = XXH128_mix32B(acc, input+32, input+len-48, secret+64, seed); - } - acc = XXH128_mix32B(acc, input+16, input+len-32, secret+32, seed); - } - acc = XXH128_mix32B(acc, input, input+len-16, secret, seed); - { XXH128_hash_t h128; - h128.low64 = acc.low64 + acc.high64; - h128.high64 = (acc.low64 * XXH_PRIME64_1) - + (acc.high64 * XXH_PRIME64_4) - + ((len - seed) * XXH_PRIME64_2); - h128.low64 = XXH3_avalanche(h128.low64); - h128.high64 = (XXH64_hash_t)0 - XXH3_avalanche(h128.high64); - return h128; - } - } -} - -XXH_NO_INLINE XXH128_hash_t -XXH3_len_129to240_128b(const xxh_u8* XXH_RESTRICT input, size_t len, - const xxh_u8* XXH_RESTRICT secret, size_t secretSize, - XXH64_hash_t seed) -{ - XXH_ASSERT(secretSize >= XXH3_SECRET_SIZE_MIN); (void)secretSize; - XXH_ASSERT(128 < len && len <= XXH3_MIDSIZE_MAX); - - { XXH128_hash_t acc; - int const nbRounds = (int)len / 32; - int i; - acc.low64 = len * XXH_PRIME64_1; - acc.high64 = 0; - for (i=0; i<4; i++) { - acc = XXH128_mix32B(acc, - input + (32 * i), - input + (32 * i) + 16, - secret + (32 * i), - seed); - } - acc.low64 = XXH3_avalanche(acc.low64); - acc.high64 = XXH3_avalanche(acc.high64); - XXH_ASSERT(nbRounds >= 4); - for (i=4 ; i < nbRounds; i++) { - acc = XXH128_mix32B(acc, - input + (32 * i), - input + (32 * i) + 16, - secret + XXH3_MIDSIZE_STARTOFFSET + (32 * (i - 4)), - seed); - } - /* last bytes */ - acc = XXH128_mix32B(acc, - input + len - 16, - input + len - 32, - secret + XXH3_SECRET_SIZE_MIN - XXH3_MIDSIZE_LASTOFFSET - 16, - 0ULL - seed); - - { XXH128_hash_t h128; - h128.low64 = acc.low64 + acc.high64; - h128.high64 = (acc.low64 * XXH_PRIME64_1) - + (acc.high64 * XXH_PRIME64_4) - + ((len - seed) * XXH_PRIME64_2); - h128.low64 = XXH3_avalanche(h128.low64); - h128.high64 = (XXH64_hash_t)0 - XXH3_avalanche(h128.high64); - return h128; - } - } -} - -XXH_FORCE_INLINE XXH128_hash_t -XXH3_hashLong_128b_internal(const void* XXH_RESTRICT input, size_t len, - const xxh_u8* XXH_RESTRICT secret, size_t secretSize, - XXH3_f_accumulate_512 f_acc512, - XXH3_f_scrambleAcc f_scramble) -{ - XXH_ALIGN(XXH_ACC_ALIGN) xxh_u64 acc[XXH_ACC_NB] = XXH3_INIT_ACC; - - XXH3_hashLong_internal_loop(acc, (const xxh_u8*)input, len, secret, secretSize, f_acc512, f_scramble); - - /* converge into final hash */ - XXH_STATIC_ASSERT(sizeof(acc) == 64); - XXH_ASSERT(secretSize >= sizeof(acc) + XXH_SECRET_MERGEACCS_START); - { XXH128_hash_t h128; - h128.low64 = XXH3_mergeAccs(acc, - secret + XXH_SECRET_MERGEACCS_START, - (xxh_u64)len * XXH_PRIME64_1); - h128.high64 = XXH3_mergeAccs(acc, - secret + secretSize - - sizeof(acc) - XXH_SECRET_MERGEACCS_START, - ~((xxh_u64)len * XXH_PRIME64_2)); - return h128; - } -} - -/* - * It's important for performance that XXH3_hashLong is not inlined. - */ -XXH_NO_INLINE XXH128_hash_t -XXH3_hashLong_128b_default(const void* XXH_RESTRICT input, size_t len, - XXH64_hash_t seed64, - const void* XXH_RESTRICT secret, size_t secretLen) -{ - (void)seed64; (void)secret; (void)secretLen; - return XXH3_hashLong_128b_internal(input, len, XXH3_kSecret, sizeof(XXH3_kSecret), - XXH3_accumulate_512, XXH3_scrambleAcc); -} - -/* - * It's important for performance to pass @secretLen (when it's static) - * to the compiler, so that it can properly optimize the vectorized loop. - */ -XXH_FORCE_INLINE XXH128_hash_t -XXH3_hashLong_128b_withSecret(const void* XXH_RESTRICT input, size_t len, - XXH64_hash_t seed64, - const void* XXH_RESTRICT secret, size_t secretLen) -{ - (void)seed64; - return XXH3_hashLong_128b_internal(input, len, (const xxh_u8*)secret, secretLen, - XXH3_accumulate_512, XXH3_scrambleAcc); -} - -XXH_FORCE_INLINE XXH128_hash_t -XXH3_hashLong_128b_withSeed_internal(const void* XXH_RESTRICT input, size_t len, - XXH64_hash_t seed64, - XXH3_f_accumulate_512 f_acc512, - XXH3_f_scrambleAcc f_scramble, - XXH3_f_initCustomSecret f_initSec) -{ - if (seed64 == 0) - return XXH3_hashLong_128b_internal(input, len, - XXH3_kSecret, sizeof(XXH3_kSecret), - f_acc512, f_scramble); - { XXH_ALIGN(XXH_SEC_ALIGN) xxh_u8 secret[XXH_SECRET_DEFAULT_SIZE]; - f_initSec(secret, seed64); - return XXH3_hashLong_128b_internal(input, len, (const xxh_u8*)secret, sizeof(secret), - f_acc512, f_scramble); - } -} - -/* - * It's important for performance that XXH3_hashLong is not inlined. - */ -XXH_NO_INLINE XXH128_hash_t -XXH3_hashLong_128b_withSeed(const void* input, size_t len, - XXH64_hash_t seed64, const void* XXH_RESTRICT secret, size_t secretLen) -{ - (void)secret; (void)secretLen; - return XXH3_hashLong_128b_withSeed_internal(input, len, seed64, - XXH3_accumulate_512, XXH3_scrambleAcc, XXH3_initCustomSecret); -} - -typedef XXH128_hash_t (*XXH3_hashLong128_f)(const void* XXH_RESTRICT, size_t, - XXH64_hash_t, const void* XXH_RESTRICT, size_t); - -XXH_FORCE_INLINE XXH128_hash_t -XXH3_128bits_internal(const void* input, size_t len, - XXH64_hash_t seed64, const void* XXH_RESTRICT secret, size_t secretLen, - XXH3_hashLong128_f f_hl128) -{ - XXH_ASSERT(secretLen >= XXH3_SECRET_SIZE_MIN); - /* - * If an action is to be taken if `secret` conditions are not respected, - * it should be done here. - * For now, it's a contract pre-condition. - * Adding a check and a branch here would cost performance at every hash. - */ - if (len <= 16) - return XXH3_len_0to16_128b((const xxh_u8*)input, len, (const xxh_u8*)secret, seed64); - if (len <= 128) - return XXH3_len_17to128_128b((const xxh_u8*)input, len, (const xxh_u8*)secret, secretLen, seed64); - if (len <= XXH3_MIDSIZE_MAX) - return XXH3_len_129to240_128b((const xxh_u8*)input, len, (const xxh_u8*)secret, secretLen, seed64); - return f_hl128(input, len, seed64, secret, secretLen); -} - - -/* === Public XXH128 API === */ - -/*! @ingroup xxh3_family */ -XXH_PUBLIC_API XXH128_hash_t XXH3_128bits(const void* input, size_t len) -{ - return XXH3_128bits_internal(input, len, 0, - XXH3_kSecret, sizeof(XXH3_kSecret), - XXH3_hashLong_128b_default); -} - -/*! @ingroup xxh3_family */ -XXH_PUBLIC_API XXH128_hash_t -XXH3_128bits_withSecret(const void* input, size_t len, const void* secret, size_t secretSize) -{ - return XXH3_128bits_internal(input, len, 0, - (const xxh_u8*)secret, secretSize, - XXH3_hashLong_128b_withSecret); -} - -/*! @ingroup xxh3_family */ -XXH_PUBLIC_API XXH128_hash_t -XXH3_128bits_withSeed(const void* input, size_t len, XXH64_hash_t seed) -{ - return XXH3_128bits_internal(input, len, seed, - XXH3_kSecret, sizeof(XXH3_kSecret), - XXH3_hashLong_128b_withSeed); -} - -/*! @ingroup xxh3_family */ -XXH_PUBLIC_API XXH128_hash_t -XXH3_128bits_withSecretandSeed(const void* input, size_t len, const void* secret, size_t secretSize, XXH64_hash_t seed) -{ - if (len <= XXH3_MIDSIZE_MAX) - return XXH3_128bits_internal(input, len, seed, XXH3_kSecret, sizeof(XXH3_kSecret), NULL); - return XXH3_hashLong_128b_withSecret(input, len, seed, secret, secretSize); -} - -/*! @ingroup xxh3_family */ -XXH_PUBLIC_API XXH128_hash_t -XXH128(const void* input, size_t len, XXH64_hash_t seed) -{ - return XXH3_128bits_withSeed(input, len, seed); -} - - -/* === XXH3 128-bit streaming === */ - -/* - * All initialization and update functions are identical to 64-bit streaming variant. - * The only difference is the finalization routine. - */ - -/*! @ingroup xxh3_family */ -XXH_PUBLIC_API XXH_errorcode -XXH3_128bits_reset(XXH3_state_t* statePtr) -{ - return XXH3_64bits_reset(statePtr); -} - -/*! @ingroup xxh3_family */ -XXH_PUBLIC_API XXH_errorcode -XXH3_128bits_reset_withSecret(XXH3_state_t* statePtr, const void* secret, size_t secretSize) -{ - return XXH3_64bits_reset_withSecret(statePtr, secret, secretSize); -} - -/*! @ingroup xxh3_family */ -XXH_PUBLIC_API XXH_errorcode -XXH3_128bits_reset_withSeed(XXH3_state_t* statePtr, XXH64_hash_t seed) -{ - return XXH3_64bits_reset_withSeed(statePtr, seed); -} - -/*! @ingroup xxh3_family */ -XXH_PUBLIC_API XXH_errorcode -XXH3_128bits_reset_withSecretandSeed(XXH3_state_t* statePtr, const void* secret, size_t secretSize, XXH64_hash_t seed) -{ - return XXH3_64bits_reset_withSecretandSeed(statePtr, secret, secretSize, seed); -} - -/*! @ingroup xxh3_family */ -XXH_PUBLIC_API XXH_errorcode -XXH3_128bits_update(XXH3_state_t* state, const void* input, size_t len) -{ - return XXH3_update(state, (const xxh_u8*)input, len, - XXH3_accumulate_512, XXH3_scrambleAcc); -} - -/*! @ingroup xxh3_family */ -XXH_PUBLIC_API XXH128_hash_t XXH3_128bits_digest (const XXH3_state_t* state) -{ - const unsigned char* const secret = (state->extSecret == NULL) ? state->customSecret : state->extSecret; - if (state->totalLen > XXH3_MIDSIZE_MAX) { - XXH_ALIGN(XXH_ACC_ALIGN) XXH64_hash_t acc[XXH_ACC_NB]; - XXH3_digest_long(acc, state, secret); - XXH_ASSERT(state->secretLimit + XXH_STRIPE_LEN >= sizeof(acc) + XXH_SECRET_MERGEACCS_START); - { XXH128_hash_t h128; - h128.low64 = XXH3_mergeAccs(acc, - secret + XXH_SECRET_MERGEACCS_START, - (xxh_u64)state->totalLen * XXH_PRIME64_1); - h128.high64 = XXH3_mergeAccs(acc, - secret + state->secretLimit + XXH_STRIPE_LEN - - sizeof(acc) - XXH_SECRET_MERGEACCS_START, - ~((xxh_u64)state->totalLen * XXH_PRIME64_2)); - return h128; - } - } - /* len <= XXH3_MIDSIZE_MAX : short code */ - if (state->seed) - return XXH3_128bits_withSeed(state->buffer, (size_t)state->totalLen, state->seed); - return XXH3_128bits_withSecret(state->buffer, (size_t)(state->totalLen), - secret, state->secretLimit + XXH_STRIPE_LEN); -} - -/* 128-bit utility functions */ - -#include /* memcmp, memcpy */ - -/* return : 1 is equal, 0 if different */ -/*! @ingroup xxh3_family */ -XXH_PUBLIC_API int XXH128_isEqual(XXH128_hash_t h1, XXH128_hash_t h2) -{ - /* note : XXH128_hash_t is compact, it has no padding byte */ - return !(memcmp(&h1, &h2, sizeof(h1))); -} - -/* This prototype is compatible with stdlib's qsort(). - * return : >0 if *h128_1 > *h128_2 - * <0 if *h128_1 < *h128_2 - * =0 if *h128_1 == *h128_2 */ -/*! @ingroup xxh3_family */ -XXH_PUBLIC_API int XXH128_cmp(const void* h128_1, const void* h128_2) -{ - XXH128_hash_t const h1 = *(const XXH128_hash_t*)h128_1; - XXH128_hash_t const h2 = *(const XXH128_hash_t*)h128_2; - int const hcmp = (h1.high64 > h2.high64) - (h2.high64 > h1.high64); - /* note : bets that, in most cases, hash values are different */ - if (hcmp) return hcmp; - return (h1.low64 > h2.low64) - (h2.low64 > h1.low64); -} - - -/*====== Canonical representation ======*/ -/*! @ingroup xxh3_family */ -XXH_PUBLIC_API void -XXH128_canonicalFromHash(XXH128_canonical_t* dst, XXH128_hash_t hash) -{ - XXH_STATIC_ASSERT(sizeof(XXH128_canonical_t) == sizeof(XXH128_hash_t)); - if (XXH_CPU_LITTLE_ENDIAN) { - hash.high64 = XXH_swap64(hash.high64); - hash.low64 = XXH_swap64(hash.low64); - } - XXH_memcpy(dst, &hash.high64, sizeof(hash.high64)); - XXH_memcpy((char*)dst + sizeof(hash.high64), &hash.low64, sizeof(hash.low64)); -} - -/*! @ingroup xxh3_family */ -XXH_PUBLIC_API XXH128_hash_t -XXH128_hashFromCanonical(const XXH128_canonical_t* src) -{ - XXH128_hash_t h; - h.high64 = XXH_readBE64(src); - h.low64 = XXH_readBE64(src->digest + 8); - return h; -} - - - -/* ========================================== - * Secret generators - * ========================================== - */ -#define XXH_MIN(x, y) (((x) > (y)) ? (y) : (x)) - -static void XXH3_combine16(void* dst, XXH128_hash_t h128) -{ - XXH_writeLE64( dst, XXH_readLE64(dst) ^ h128.low64 ); - XXH_writeLE64( (char*)dst+8, XXH_readLE64((char*)dst+8) ^ h128.high64 ); -} - -/*! @ingroup xxh3_family */ -XXH_PUBLIC_API XXH_errorcode -XXH3_generateSecret(void* secretBuffer, size_t secretSize, const void* customSeed, size_t customSeedSize) -{ - XXH_ASSERT(secretBuffer != NULL); - if (secretBuffer == NULL) return XXH_ERROR; - XXH_ASSERT(secretSize >= XXH3_SECRET_SIZE_MIN); - if (secretSize < XXH3_SECRET_SIZE_MIN) return XXH_ERROR; - if (customSeedSize == 0) { - customSeed = XXH3_kSecret; - customSeedSize = XXH_SECRET_DEFAULT_SIZE; - } - XXH_ASSERT(customSeed != NULL); - if (customSeed == NULL) return XXH_ERROR; - - /* Fill secretBuffer with a copy of customSeed - repeat as needed */ - { size_t pos = 0; - while (pos < secretSize) { - size_t const toCopy = XXH_MIN((secretSize - pos), customSeedSize); - memcpy((char*)secretBuffer + pos, customSeed, toCopy); - pos += toCopy; - } } - - { size_t const nbSeg16 = secretSize / 16; - size_t n; - XXH128_canonical_t scrambler; - XXH128_canonicalFromHash(&scrambler, XXH128(customSeed, customSeedSize, 0)); - for (n=0; n -#include -#include -#include - -using namespace core; - -namespace py = pybind11; -using namespace py::literals; - -static constexpr unsigned int RESERVED_INDICES[] = { 0 }; -static const py::bool_ py_true(true); - -template -inline const T& choice(const T* items, unsigned int length) { - return sample_uniform(items, length); -} - -inline unsigned int randrange(unsigned int start, unsigned int end) { - return start + sample_uniform(end - start); -} - -inline unsigned int randrange(unsigned int end) { - return sample_uniform(end); -} - -struct node { - unsigned int id; - core::array children; - core::array parents; - - node(unsigned int id) : id(id), children(8), parents(8) { } - - static inline void free(node& n) { - core::free(n.children); - core::free(n.parents); - } -}; - -inline bool init(node& n, unsigned int id) { - if (!array_init(n.children, 8)) { - return false; - } else if (!array_init(n.parents, 8)) { - core::free(n.children); - return false; - } - n.id = id; - return true; -} - -inline bool operator == (const node& first, const node& second) { - return first.id == second.id; -} - -bool generate_graph(array& vertices, node*& start, node*& end, unsigned int num_vertices, unsigned int max_num_parents, unsigned int max_vertex_id) -{ - if (!vertices.ensure_capacity(num_vertices)) - return false; - for (unsigned int i = 0; i < num_vertices; i++) - if (!init(vertices[i], i)) return false; - vertices.length = num_vertices; - - /* sample a random DAG */ - for (unsigned int i = 1; i < num_vertices; i++) { - /* sample the number of parent vertices */ - unsigned int num_parents; - if (randrange(2) == 0) - num_parents = 1; - else - num_parents = randrange(1, max_num_parents); - num_parents = min(num_parents, i); - - array available_parent_ids(i); - for (unsigned int j = 0; j < i; j++) - available_parent_ids.add(j); - for (unsigned int j = 0; j < num_parents; j++) { - unsigned int u = randrange(available_parent_ids.length); - unsigned int parent_id = available_parent_ids[u]; - vertices[parent_id].children.add(&vertices[i]); - vertices[i].parents.add(&vertices[parent_id]); - available_parent_ids.remove(u); - } - } - - /* remove any correlation between graph topology and vertex IDs by shuffling the vertices */ - unsigned int* new_indices = (unsigned int*) alloca(sizeof(unsigned int) * (max_vertex_id + 1)); - for (unsigned int i = 0; i < max_vertex_id + 1; i++) new_indices[i] = i; - shuffle(new_indices, max_vertex_id + 1); - unsigned int src_index = 0; - for (unsigned int i = 0; i < vertices.length; i++) { - bool is_reserved = false; - for (unsigned int j = 0; j < array_length(RESERVED_INDICES); j++) { - if (new_indices[src_index] == RESERVED_INDICES[j]) { - is_reserved = true; - break; - } - } - if (is_reserved) - src_index++; - vertices[i].id = new_indices[src_index]; - src_index++; - } - - /* randomly select two vertices */ - start = &vertices[randrange(vertices.length - 1)]; - do { - end = &vertices[randrange(vertices.length - 1)]; - } while (end == start); - return true; -} - -/* computes the number of lookahead steps to find the answer */ -unsigned int lookahead_depth(const node* vertex, const node* next_vertex, const node* goal) -{ - array> frontier(8); - for (const node* v : vertex->children) - frontier.add(make_pair(v, v)); - array visited(16); - visited.append(vertex->children.data, vertex->children.length); - unsigned int lookahead = 0; - while (frontier.length != 0) { - bool frontier_is_next_vertex = true; - for (pair& entry : frontier) { - if (entry.value != next_vertex) { - frontier_is_next_vertex = false; - break; - } - } - if (frontier_is_next_vertex) - return lookahead; - - lookahead++; - array> new_frontier(8); - for (const pair& entry : frontier) { - const node* v = entry.key; - const node* branch = entry.value; - - if (v == goal) - return lookahead; - for (const node* child : v->children) { - if (!visited.contains(child)) { - new_frontier.add(make_pair(child, branch)); - visited.add(child); - } else if (branch == next_vertex) { - for (unsigned int i = 0; i < new_frontier.length; i++) - if (new_frontier[i].key == child) - new_frontier[i].value = branch; - } - } - } - core::swap(frontier, new_frontier); - } - return lookahead; -} - -void get_descendants(const node& vertex, array& descendants) { - array queue(8); - array visited(16); - queue[0] = &vertex; - queue.length = 1; - while (queue.length != 0) { - const node* current = queue.pop(); - visited.add(current); - for (const node* child : current->children) { - if (!descendants.contains(child)) - descendants.add(child); - if (visited.contains(child)) - continue; - queue.add(child); - } - } -} - -bool has_cycles(array& vertices) { - for (const node& vertex : vertices) { - array descendants(8); - get_descendants(vertex, descendants); - if (descendants.contains(&vertex)) - return true; - } - return false; -} - -bool generate_graph_with_lookahead(array& vertices, node*& start, node*& end, unsigned int num_vertices, unsigned int max_num_parents, unsigned int max_vertex_id, unsigned int lookahead, unsigned int num_paths, unsigned int max_prefix_vertices) -{ - num_vertices = std::max(std::max(2u, num_vertices), 1 + num_paths * lookahead); - - if (!vertices.ensure_capacity(num_vertices)) - return false; - for (unsigned int i = 0; i < num_vertices; i++) - if (!init(vertices[i], i)) return false; - vertices.length = num_vertices; - - vertices[1].parents.add(&vertices[0]); - vertices[0].children.add(&vertices[1]); - for (unsigned int i = 1; i < lookahead; i++) { - vertices[1 + i].parents.add(&vertices[i]); - vertices[i].children.add(&vertices[1 + i]); - } - unsigned int index; - if (lookahead == 0) { - index = 2; - } else { - index = 1 + lookahead; - for (unsigned int j = 0; j < num_paths - 1; j++) { - vertices[index].parents.add(&vertices[0]); - vertices[0].children.add(&vertices[index]); - index++; - unsigned int other_branch_length = lookahead + randrange(std::min(2u, num_vertices - index - (num_paths - j - 1) * lookahead + 2)); - for (unsigned int i = 1; i < other_branch_length; i++) { - vertices[index].parents.add(&vertices[index - 1]); - vertices[index - 1].children.add(&vertices[index]); - index++; - } - } - } - - unsigned int num_prefix_vertices = randrange(min(max_prefix_vertices + 1, num_vertices - index + 1)); - node* prev_vertex = &vertices[0]; - for (unsigned int i = 0; i < num_prefix_vertices; i++) { - vertices[index].children.add(prev_vertex); - prev_vertex->parents.add(&vertices[index]); - prev_vertex = &vertices[index]; - index++; - } - - start = &vertices[0]; - end = &vertices[std::max(1u, lookahead)]; - - /* sample some parent/ancestor vertices */ - constexpr float ALPHA = 0.5f; - unsigned int* in_degrees = (unsigned int*) calloc(num_vertices, sizeof(unsigned int)); - unsigned int* out_degrees = (unsigned int*) calloc(num_vertices, sizeof(unsigned int)); - if (in_degrees == nullptr || out_degrees == nullptr) { - if (in_degrees != nullptr) free(in_degrees); - return false; - } - for (unsigned int i = 0; i < num_vertices; i++) { - in_degrees[i] = vertices[i].parents.length; - out_degrees[i] = vertices[i].children.length; - } - for (unsigned int i = index; i < num_vertices; i++) { - /* sample the number of child and parent vertices */ - unsigned int num_children = randrange(0, max_num_parents); - unsigned int num_parents = randrange(num_children == 0 ? 1 : 0, max_num_parents); - num_children = std::min(num_children, i); - num_parents = std::min(num_parents, i); - - /* sample the children of this new node */ - array probabilities(index); - float total_probability = 0.0f; - for (unsigned int j = 0; j < index; j++) { - probabilities[j] = ALPHA + in_degrees[j]; - total_probability += probabilities[j]; - } - probabilities.length = index; - - array sampled_children(std::max(1u, num_children)); - for (unsigned int j = 0; j < num_children; j++) { - unsigned int u = sample_categorical(probabilities.data, total_probability, probabilities.length); - sampled_children.add(u); - total_probability -= probabilities[u]; - probabilities[u] = 0.0f; - } - - for (unsigned int child_id : sampled_children) { - vertices[index].children.add(&vertices[child_id]); - vertices[child_id].parents.add(&vertices[index]); - in_degrees[child_id] += 1; - } - - /* sample the parents of this new node */ - total_probability = 0.0f; - for (unsigned int j = 0; j < index; j++) { - probabilities[j] = ALPHA + out_degrees[j]; - total_probability += probabilities[j]; - } - - /* to avoid creating a cycle, we have to remove any descendants from the possible parents */ - array descendants(8); - get_descendants(vertices[index], descendants); - for (const node* descendant : descendants) { - total_probability -= probabilities[descendant->id]; - probabilities[descendant->id] = 0.0f; - } - if (total_probability != 0.0f) { - num_parents = std::min(num_parents, index - (unsigned int) descendants.length); - - array sampled_parents(std::max(1u, num_parents)); - for (unsigned int j = 0; j < num_parents; j++) { - unsigned int u = sample_categorical(probabilities.data, total_probability, probabilities.length); - sampled_parents.add(u); - total_probability -= probabilities[u]; - probabilities[u] = 0.0f; - } - - for (unsigned int parent_id : sampled_parents) { - vertices[parent_id].children.add(&vertices[index]); - vertices[index].parents.add(&vertices[parent_id]); - out_degrees[parent_id] += 1; - } - } - index += 1; - } - free(in_degrees); - free(out_degrees); - - /* remove any correlation between graph topology and vertex IDs by shuffling the vertices */ - unsigned int* new_indices = (unsigned int*) alloca(sizeof(unsigned int) * (max_vertex_id + 1)); - for (unsigned int i = 0; i < max_vertex_id + 1; i++) new_indices[i] = i; - shuffle(new_indices, max_vertex_id + 1); - unsigned int src_index = 0; - for (unsigned int i = 0; i < vertices.length; i++) { - bool is_reserved = false; - for (unsigned int j = 0; j < array_length(RESERVED_INDICES); j++) { - if (new_indices[src_index] == RESERVED_INDICES[j]) { - is_reserved = true; - break; - } - } - if (is_reserved) - src_index++; - vertices[i].id = new_indices[src_index]; - src_index++; - } - return true; -} - -bool generate_example(array& vertices, node*& start, node*& end, array>& paths, unsigned int num_vertices, unsigned int max_num_parents, unsigned int max_vertex_id, bool get_shortest_paths, int lookahead, unsigned int num_paths, unsigned int max_prefix_vertices) -{ - if (lookahead == -1) { - if (!generate_graph(vertices, start, end, num_vertices, max_num_parents, max_vertex_id)) - return false; - } else { - if (!generate_graph_with_lookahead(vertices, start, end, num_vertices, max_num_parents, max_vertex_id, lookahead, num_paths, max_prefix_vertices)) - return false; - } - - /* find the shortest paths from `start` to `end` */ - array> queue(16); - queue[0].key = start; - queue[0].value = 0; - queue.length = 1; - array_map> reverse_pointers(16); - while (queue.length != 0) { - pair item = queue.pop(); - node* current = item.key; - unsigned int distance = item.value; - - for (node* child : current->children) { - if (!reverse_pointers.ensure_capacity(reverse_pointers.size + 1)) { - for (const auto& entry : reverse_pointers) core::free(entry.value); - return false; - } - bool contains; - array_map& value = reverse_pointers.get(child, contains); - if (!contains) { - if (!array_map_init(value, 4)) { - for (const auto& entry : reverse_pointers) core::free(entry.value); - return false; - } - reverse_pointers.keys[reverse_pointers.size++] = child; - } - - if (!value.ensure_capacity(value.size + 1)) { - for (const auto& entry : reverse_pointers) core::free(entry.value); - return false; - } - unsigned int& distance_value = value.get(current, contains); - if (!contains) { - distance_value = distance + 1; - value.keys[value.size++] = current; - } else if (distance_value > distance + 1) { - distance_value = distance + 1; - } else { - continue; - } - - bool found_child = false; - for (unsigned int j = 0; j < queue.length; j++) { - if (queue[j].key == child) { - queue[j].value = std::min(queue[j].value, distance + 1); - found_child = true; - break; - } - } - if (!found_child) - queue.add(make_pair(child, distance + 1)); - } - } - - if (!reverse_pointers.contains(end)) { - for (const auto& entry : reverse_pointers) core::free(entry.value); - return false; - } - - array_map> forward_pointers(16); - array dist_queue(16); - dist_queue[0] = end; - dist_queue.length = 1; - while (dist_queue.length != 0) { - node* current = dist_queue.pop(); - if (current == start) - continue; - - array prev_nodes(8); - const array_map& value = reverse_pointers.get(current); - if (get_shortest_paths) { - unsigned int min_distance = value.values[0]; - for (unsigned int i = 1; i < value.size; i++) - if (value.values[i] < min_distance) min_distance = value.values[i]; - - for (unsigned int i = 0; i < value.size; i++) - if (value.values[i] == min_distance) prev_nodes.add(value.keys[i]); - } else { - prev_nodes.append(value.keys, value.size); - } - for (const node* prev : prev_nodes) { - if (!forward_pointers.ensure_capacity(forward_pointers.size + 1)) { - for (const auto& entry : reverse_pointers) core::free(entry.value); - for (const auto& entry : forward_pointers) core::free(entry.value); - return false; - } - bool contains; - array& fptrs = forward_pointers.get(prev, contains); - if (!contains) { - if (!array_init(fptrs, 4)) { - for (const auto& entry : reverse_pointers) core::free(entry.value); - for (const auto& entry : forward_pointers) core::free(entry.value); - return false; - } - forward_pointers.keys[forward_pointers.size++] = prev; - } - fptrs.add(current); - } - dist_queue.append(prev_nodes.data, prev_nodes.length); - } - - /* free `reverse_pointers` */ - for (const auto& entry : reverse_pointers) core::free(entry.value); - - /* construct the shortest paths from the forward pointers */ - array> path_queue(8); - if (!array_init(path_queue[0], 1)) { - for (const auto& entry : forward_pointers) core::free(entry.value); - return false; - } - path_queue[0].add(start); - path_queue.length = 1; - while (path_queue.length != 0) { - array& partial_path = *((array*) alloca(sizeof(array))); - core::move(path_queue.last(), partial_path); - path_queue.length--; - - if (partial_path.last() == end) { - if (!paths.ensure_capacity(paths.length + 1)) { - for (const auto& entry : forward_pointers) core::free(entry.value); - for (auto& p : path_queue) core::free(p); - core::free(partial_path); - return false; - } - core::move(partial_path, paths[paths.length++]); - if (paths.length > 64) { - for (const auto& entry : forward_pointers) core::free(entry.value); - for (auto& p : path_queue) core::free(p); - return false; - } - continue; - } - for (node* next : forward_pointers.get(partial_path.last())) { - if (!path_queue.ensure_capacity(path_queue.length + 1) - || !array_init(path_queue[path_queue.length], partial_path.length + 1)) - { - for (const auto& entry : forward_pointers) core::free(entry.value); - for (auto& p : path_queue) core::free(p); - core::free(partial_path); - return false; - } - path_queue.length++; - path_queue.last().append(partial_path.data, partial_path.length); - path_queue.last().add(next); - } - core::free(partial_path); - } - - for (const auto& entry : forward_pointers) core::free(entry.value); - return true; -} - -bool has_path(const node* start, const node* end) -{ - array stack(8); - stack[0] = start; - stack.length = 1; - array visited(16); - while (stack.length != 0) { - const node* v = stack.pop(); - if (v == end) - return true; - for (node* child : v->children) { - if (!visited.contains(child)) { - visited.add(child); - stack.add(child); - } - } - } - return false; -} - -py::tuple generate_training_set(const unsigned int max_input_size, const uint64_t dataset_size, const int max_lookahead, const unsigned int max_edges, const py::object& reserved_inputs, const int distance_from_start, const int max_prefix_vertices, const bool quiet=false, const int num_paths=2) -{ - const unsigned int QUERY_PREFIX_TOKEN = (max_input_size-5) / 3 + 4; - const unsigned int PADDING_TOKEN = (max_input_size-5) / 3 + 3; - const unsigned int EDGE_PREFIX_TOKEN = (max_input_size-5) / 3 + 2; - const unsigned int PATH_PREFIX_TOKEN = (max_input_size-5) / 3 + 1; - - unsigned int num_generated = 0; - unsigned int num_collisions = 0; - unsigned int ntokens = (max_input_size - 5) / 3 + 5; - size_t input_shape[2]{dataset_size, max_input_size}; - size_t output_shape[2]{dataset_size, ntokens}; - size_t label_shape[1]{dataset_size}; - py::array_t inputs(input_shape); - py::array_t outputs(output_shape); - py::array_t labels(label_shape); - auto inputs_mem = inputs.mutable_unchecked<2>(); - auto outputs_mem = outputs.mutable_unchecked<2>(); - auto labels_mem = labels.mutable_unchecked<1>(); - unsigned int* lookahead_step_histogram = (unsigned int*) alloca(sizeof(unsigned int) * max_input_size); - unsigned int* path_length_histogram = (unsigned int*) alloca(sizeof(unsigned int) * max_input_size); - for (unsigned int i = 0; i < max_input_size; i++) { - lookahead_step_histogram[i] = 0; - path_length_histogram[i] = 0; - } - float* MAX_FREQS_PER_BUCKET = (float*) alloca(sizeof(float) * max_input_size); - if (max_lookahead == -1) { - for (unsigned int i = 0; i < max_input_size; i++) - MAX_FREQS_PER_BUCKET[i] = 1.0; - } else { - for (unsigned int i = 0; i < (unsigned) max_lookahead + 1; i++) - MAX_FREQS_PER_BUCKET[i] = 1.0 / (max_lookahead+1); - for (unsigned int i = max_lookahead + 1; i < max_input_size; i++) - MAX_FREQS_PER_BUCKET[i] = 0.0; - MAX_FREQS_PER_BUCKET[max_lookahead] += 0.05; - } - - unsigned int* potential_lookaheads = (unsigned int*) alloca(max((size_t) 1, sizeof(unsigned int) * (max_lookahead + 1))); - unsigned int potential_lookahead_count = 0; - while (num_generated < dataset_size) { - array g(32); - node* start; node* end; - array> paths(8); - while (true) { - if (max_lookahead == -1) { - unsigned int num_vertices = randrange(3, (max_input_size - 5) / 3); - if (!generate_example(g, start, end, paths, num_vertices, 4, (max_input_size - 5) / 3, true, -1, 0, max_prefix_vertices == -1 ? max_input_size : max_prefix_vertices)) { - for (node& n : g) core::free(n); - for (array& a : paths) core::free(a); - g.length = 0; paths.length = 0; - continue; - } - } else { - potential_lookahead_count = 0; - for (unsigned int i = 0; i < (unsigned) max_lookahead + 1; i++) - if (num_generated == 0 || lookahead_step_histogram[i] / num_generated < MAX_FREQS_PER_BUCKET[i]) - potential_lookaheads[potential_lookahead_count++] = i; -// unsigned int lookahead = choice(potential_lookaheads, potential_lookahead_count); - unsigned int lookahead = max_lookahead; - -// unsigned int num_paths; -// if (lookahead == 0) { -// num_paths = randrange(1, 3); -// } else { -// unsigned int max_num_paths = (max_edges - 1) / lookahead; -// num_paths = randrange(2, max_num_paths + 1); -// } - - unsigned int num_vertices = std::min(std::min(lookahead * num_paths + 1 + randrange(0, 6), (max_input_size-5) / 3), max_edges + 1); - if (!generate_example(g, start, end, paths, num_vertices, 4, (max_input_size - 5) / 3, true, lookahead, num_paths, max_prefix_vertices == -1 ? max_input_size : max_prefix_vertices)) { - for (node& n : g) core::free(n); - for (array& a : paths) core::free(a); - g.length = 0; paths.length = 0; - continue; - } - - unsigned int shortest_path_length = paths[0].length; - for (unsigned int i = 1; i < paths.length; i++) - if (paths[i].length < shortest_path_length) - shortest_path_length = paths[i].length; - - // Count how many paths have the shortest length. - unsigned int shortest_path_count = 0; - for (unsigned int i = 0; i < paths.length; i++) { - if (paths[i].length == shortest_path_length) - shortest_path_count++; - } - - // Only accept graphs with a unique (one) shortest path. - if (shortest_path_length <= 1 || shortest_path_count != 1) { - for (node& n : g) core::free(n); - for (array& a : paths) core::free(a); - g.length = 0; paths.length = 0; - continue; - } - } - unsigned int shortest_path_length = paths[0].length; - for (unsigned int i = 1; i < paths.length; i++) - if (paths[i].length < shortest_path_length) - shortest_path_length = paths[i].length; - if (shortest_path_length > 1) - break; - for (node& n : g) core::free(n); - for (array& a : paths) core::free(a); - g.length = 0; paths.length = 0; - } - - array> edges(8); - for (node& vertex : g) - for (node* child : vertex.children) - edges.add(make_pair(vertex.id, child->id)); - if (edges.length > max_edges || edges.length * 3 + 4 > max_input_size) { - for (node& n : g) core::free(n); - for (array& a : paths) core::free(a); - g.length = 0; paths.length = 0; - continue; - } - shuffle(edges); - - array prefix(max_input_size); - for (auto& entry : edges) { - prefix[prefix.length++] = EDGE_PREFIX_TOKEN; - prefix[prefix.length++] = entry.key; - prefix[prefix.length++] = entry.value; - } - prefix[prefix.length++] = QUERY_PREFIX_TOKEN; - prefix[prefix.length++] = start->id; - prefix[prefix.length++] = end->id; - prefix[prefix.length++] = PATH_PREFIX_TOKEN; - - for (const array& path : paths) { - if (path.length == 1) - continue; - for (unsigned int j = 1; j < path.length; j++) { - if (distance_from_start != -1 && j != (unsigned int) distance_from_start) - continue; - array example(prefix.length + j); - for (unsigned int i = 0; i < prefix.length; i++) - example[i] = prefix[i]; - for (unsigned int i = 0; i < j; i++) - example[prefix.length + i] = path[i]->id; - example.length = prefix.length + j; - if (example.length > max_input_size) - continue; - - unsigned int lookahead_steps = lookahead_depth(path[j-1], path[j], end); - array useful_steps(8); - for (node* v : path[j-1]->children) - if (has_path(v, end)) useful_steps.add(v); - - /* check if this input is reserved */ - py::object contains = reserved_inputs.attr("__contains__"); - py::tuple example_tuple(example.length); - for (unsigned int i = 0; i < example.length; i++) - example_tuple[i] = example[i]; - if (contains(example_tuple).is(py_true)) { - num_collisions += 1; - continue; - } - - if (num_generated != 0 && lookahead_step_histogram[lookahead_steps] / num_generated >= MAX_FREQS_PER_BUCKET[lookahead_steps]) - continue; - lookahead_step_histogram[lookahead_steps] += 1; - path_length_histogram[j] += 1; - - for (unsigned int i = 0; i < max_input_size - example.length; i++) - inputs_mem(num_generated, i) = PADDING_TOKEN; - for (unsigned int i = 0; i < example.length; i++) - inputs_mem(num_generated, max_input_size - example.length + i) = example[i]; - for (unsigned int i = 0; i < ntokens; i++) - outputs_mem(num_generated, i) = 0.0f; - for (unsigned int i = 0; i < useful_steps.length; i++) - outputs_mem(num_generated, useful_steps[i]->id) = 1.0f; - labels_mem(num_generated) = choice(useful_steps.data, useful_steps.length)->id; - num_generated++; - if (num_generated == dataset_size) - break; - } - if (num_generated == dataset_size) - break; - } - - if (!quiet && num_generated > 0 && (num_generated % 1000 == 0 || num_generated >= dataset_size)) { - printf("%d examples generated.\n", num_generated); - - printf("Lookahead steps histogram:\n"); - printf("["); - bool first = true; - for (unsigned int i = 0; i < max_input_size; i++) { - if (lookahead_step_histogram[i] == 0) - continue; - if (!first) printf(", "); - printf("%d:%.2f", i, (float) lookahead_step_histogram[i] / num_generated + 1e-9); - first = false; - } - printf("]\n"); - - printf("Path length histogram:\n"); - printf("["); - first = true; - for (unsigned int i = 0; i < max_input_size; i++) { - if (path_length_histogram[i] == 0) - continue; - if (!first) printf(", "); - printf("%d:%.2f", i, (float) path_length_histogram[i] / num_generated + 1e-9); - first = false; - } - printf("]\n"); - fflush(stdout); - } - - for (node& n : g) core::free(n); - for (array& a : paths) core::free(a); - g.length = 0; paths.length = 0; - continue; - } - - return py::make_tuple(inputs, outputs, labels, num_collisions); -} - -py::array_t lookahead_histogram(const unsigned int max_input_size, const uint64_t num_samples, const unsigned int max_edges, const int distance_from_start, const bool quiet=false) -{ - const unsigned int QUERY_PREFIX_TOKEN = (max_input_size-5) / 3 + 4; - const unsigned int EDGE_PREFIX_TOKEN = (max_input_size-5) / 3 + 2; - const unsigned int PATH_PREFIX_TOKEN = (max_input_size-5) / 3 + 1; - - unsigned int num_generated = 0; - unsigned int max_lookahead = ((max_input_size - 5) / 3 - 1) / 2; - size_t histogram_shape[1]{max_lookahead}; - py::array_t histogram(histogram_shape); - auto histogram_mem = histogram.mutable_unchecked<1>(); - for (unsigned int i = 0; i < max_lookahead; i++) - histogram_mem(i) = 0; - - while (num_generated < num_samples) { - array g(32); - node* start; node* end; - array> paths(8); - while (true) { - unsigned int num_vertices = randrange(3, (max_input_size - 5) / 3); - if (!generate_example(g, start, end, paths, num_vertices, 4, (max_input_size - 5) / 3, true, -1, 0, -1)) { - for (node& n : g) core::free(n); - for (array& a : paths) core::free(a); - g.length = 0; paths.length = 0; - continue; - } - unsigned int shortest_path_length = paths[0].length; - for (unsigned int i = 1; i < paths.length; i++) - if (paths[i].length < shortest_path_length) - shortest_path_length = paths[i].length; - if (shortest_path_length > 1) - break; - for (node& n : g) core::free(n); - for (array& a : paths) core::free(a); - g.length = 0; paths.length = 0; - } - - array> edges(8); - for (node& vertex : g) - for (node* child : vertex.children) - edges.add(make_pair(vertex.id, child->id)); - if (edges.length > max_edges || edges.length * 3 + 4 > max_input_size) { - for (node& n : g) core::free(n); - for (array& a : paths) core::free(a); - g.length = 0; paths.length = 0; - continue; - } - shuffle(edges); - - array prefix(max_input_size); - for (auto& entry : edges) { - prefix[prefix.length++] = EDGE_PREFIX_TOKEN; - prefix[prefix.length++] = entry.key; - prefix[prefix.length++] = entry.value; - } - prefix[prefix.length++] = QUERY_PREFIX_TOKEN; - prefix[prefix.length++] = start->id; - prefix[prefix.length++] = end->id; - prefix[prefix.length++] = PATH_PREFIX_TOKEN; - - for (const array& path : paths) { - if (path.length == 1) - continue; - for (unsigned int j = 1; j < path.length; j++) { - if (distance_from_start != -1 && j != (unsigned int) distance_from_start) - continue; - array example(prefix.length + j); - for (unsigned int i = 0; i < prefix.length; i++) - example[i] = prefix[i]; - for (unsigned int i = 0; i < j; i++) - example[prefix.length + i] = path[i]->id; - example.length = prefix.length + j; - if (example.length > max_input_size) - continue; - - unsigned int lookahead_steps = lookahead_depth(path[j-1], path[j], end); - histogram_mem(lookahead_steps) += 1; - num_generated++; - if (num_generated == num_samples) - break; - } - if (num_generated == num_samples) - break; - } - - if (!quiet && num_generated > 0 && (num_generated % 1000 == 0 || num_generated >= num_samples)) { - printf("%d examples generated.\n", num_generated); - - printf("Lookahead steps histogram:\n"); - printf("["); - bool first = true; - for (unsigned int i = 0; i < max_lookahead; i++) { - if (histogram_mem(i) == 0) - continue; - if (!first) printf(", "); - printf("%d:%.2f", i, (float) histogram_mem(i) / num_generated + 1e-9); - first = false; - } - printf("]\n"); - } - - for (node& n : g) core::free(n); - for (array& a : paths) core::free(a); - g.length = 0; paths.length = 0; - continue; - } - - return histogram; -} - -py::tuple generate_reachable_training_set(const unsigned int max_input_size, const uint64_t dataset_size, const unsigned int lookahead, const unsigned int max_edges, const py::object& reserved_inputs, const int distance_from_start, const int reachable_distance, const unsigned int start_vertex_index, const bool exclude_start_vertex) -{ - const unsigned int QUERY_PREFIX_TOKEN = (max_input_size-5) / 3 + 4; - const unsigned int PADDING_TOKEN = (max_input_size-5) / 3 + 3; - const unsigned int EDGE_PREFIX_TOKEN = (max_input_size-5) / 3 + 2; - const unsigned int PATH_PREFIX_TOKEN = (max_input_size-5) / 3 + 1; - - unsigned int num_generated = 0; - unsigned int num_collisions = 0; - size_t input_shape[2]{dataset_size, max_input_size}; - size_t output_shape[2]{dataset_size, max_input_size}; - py::array_t inputs(input_shape); - py::array_t outputs(output_shape); - auto inputs_mem = inputs.mutable_unchecked<2>(); - auto outputs_mem = outputs.mutable_unchecked<2>(); - py::list valid_outputs; - - unsigned int max_vertex_id = (max_input_size - 5) / 3; - while (num_generated < dataset_size) { - array g(32); - node* start; node* end; - array> paths(8); - while (true) { - unsigned int num_paths; - if (lookahead == 0) { - num_paths = randrange(1, 3); - } else { - unsigned int max_num_paths = (max_edges - 1) / lookahead; - num_paths = randrange(2, max_num_paths + 1); - } - - unsigned int num_vertices = std::min(std::min(lookahead * num_paths + 1 + randrange(0, 6), (max_input_size-5) / 3), max_edges + 1); - if (!generate_example(g, start, end, paths, num_vertices, 4, max_vertex_id, true, lookahead, num_paths, -1)) { - for (node& n : g) core::free(n); - for (array& a : paths) core::free(a); - g.length = 0; paths.length = 0; - continue; - } - unsigned int shortest_path_length = paths[0].length; - for (unsigned int i = 1; i < paths.length; i++) - if (paths[i].length < shortest_path_length) - shortest_path_length = paths[i].length; - if (shortest_path_length > 1) - break; - for (node& n : g) core::free(n); - for (array& a : paths) core::free(a); - g.length = 0; paths.length = 0; - } - - array> edges(8); - for (node& vertex : g) - for (node* child : vertex.children) - edges.add(make_pair(vertex.id, child->id)); - if (edges.length > max_edges || edges.length * 3 + 4 > max_input_size) { - for (node& n : g) core::free(n); - for (array& a : paths) core::free(a); - g.length = 0; paths.length = 0; - continue; - } - shuffle(edges); - - array prefix(max_input_size); - for (auto& entry : edges) { - prefix[prefix.length++] = EDGE_PREFIX_TOKEN; - prefix[prefix.length++] = entry.key; - prefix[prefix.length++] = entry.value; - } - prefix[prefix.length++] = QUERY_PREFIX_TOKEN; - prefix[prefix.length++] = start->id; - prefix[prefix.length++] = end->id; - prefix[prefix.length++] = PATH_PREFIX_TOKEN; - - for (const array& path : paths) { - if (path.length == 1) - continue; - for (unsigned int j = 1; j < path.length; j++) { - if (distance_from_start != -1 && j != (unsigned int) distance_from_start) - continue; - array example(prefix.length + j); - for (unsigned int i = 0; i < prefix.length; i++) - example[i] = prefix[i]; - for (unsigned int i = 0; i < j; i++) - example[prefix.length + i] = path[i]->id; - example.length = prefix.length + j; - if (example.length > max_input_size) - continue; - - /* compute the set of reachable vertices */ - node** vertex_id_map = (node**) calloc(max_vertex_id + 1, sizeof(node*)); - for (unsigned int i = 0; i < g.length; i++) - vertex_id_map[g[i].id] = &g[i]; - array reachable(16); - array> stack(16); - unsigned int start_vertex; - if (example.length < start_vertex_index) - start_vertex = start->id; - else start_vertex = example[example.length - start_vertex_index]; - stack.add(make_pair(start_vertex, 0u)); - while (stack.length != 0) { - pair entry = stack.pop(); - unsigned int current_vertex = entry.key; - unsigned int current_distance = entry.value; - if (!reachable.contains(current_vertex)) - reachable.add(current_vertex); - if (reachable_distance > 0 && current_distance + 1 <= (unsigned int) reachable_distance) { - for (node* child : vertex_id_map[current_vertex]->children) - stack.add(make_pair(child->id, current_distance + 1)); - } else if (reachable_distance < 0 && current_distance + 1 <= (unsigned int) -reachable_distance) { - for (node* parent : vertex_id_map[current_vertex]->parents) - stack.add(make_pair(parent->id, current_distance + 1)); - } - } - if (exclude_start_vertex) - reachable.remove(reachable.index_of(start_vertex)); - - array useful_steps(8); - for (node* v : path[j-1]->children) - if (has_path(v, end)) useful_steps.add(v); - - /* check if this input is reserved */ - py::object contains = reserved_inputs.attr("__contains__"); - py::tuple example_tuple(example.length); - for (unsigned int i = 0; i < example.length; i++) - example_tuple[i] = example[i]; - if (contains(example_tuple).is(py_true)) { - num_collisions += 1; - continue; - } - - for (unsigned int i = 0; i < max_input_size - example.length; i++) - inputs_mem(num_generated, i) = PADDING_TOKEN; - for (unsigned int i = 0; i < example.length; i++) - inputs_mem(num_generated, max_input_size - example.length + i) = example[i]; - for (unsigned int i = 0; i < max_input_size - example.length; i++) - outputs_mem(num_generated, i) = 0; - for (unsigned int i = 0; i < example.length; i++) - outputs_mem(num_generated, max_input_size - example.length + i) = reachable.contains(example[i]) ? 1 : 0; - num_generated++; - if (num_generated == dataset_size) - break; - } - if (num_generated == dataset_size) - break; - } - - for (node& n : g) core::free(n); - for (array& a : paths) core::free(a); - g.length = 0; paths.length = 0; - continue; - } - - return py::make_tuple(inputs, outputs, num_collisions); -} - -bool generate_dfs_example(array& vertices, const node*& start, const node*& end, unsigned int& current_node_index, array& path, unsigned int num_vertices, unsigned int max_num_parents, unsigned int max_vertex_id, const int requested_backtrack, unsigned int max_edges) -{ - if (!vertices.ensure_capacity(num_vertices)) - return false; - for (unsigned int i = 0; i < num_vertices; i++) - if (!init(vertices[i], i)) return false; - vertices.length = num_vertices; - - /* sample some parent/ancestor vertices */ - constexpr float ALPHA = 1.0f; - unsigned int* out_degrees = (unsigned int*) calloc(num_vertices, sizeof(unsigned int)); - if (out_degrees == nullptr) - return false; - for (unsigned int i = 1; i < num_vertices; i++) { - /* sample the number of child and parent vertices */ - unsigned int num_parents = randrange(1, max_num_parents); - num_parents = std::min(num_parents, i); - - /* sample the parents of this new node */ - float total_probability = 0.0f; - array probabilities(i); - for (unsigned int j = 0; j < i; j++) { - probabilities[j] = ALPHA + out_degrees[j]; - total_probability += probabilities[j]; - } - probabilities.length = i; - - array sampled_parents(std::max(1u, num_parents)); - for (unsigned int j = 0; j < num_parents; j++) { - unsigned int u = sample_categorical(probabilities.data, total_probability, probabilities.length); - sampled_parents.add(u); - total_probability -= probabilities[u]; - probabilities[u] = 0.0f; - } - - for (unsigned int parent_id : sampled_parents) { - vertices[parent_id].children.add(&vertices[i]); - vertices[i].parents.add(&vertices[parent_id]); - out_degrees[parent_id] += 1; - } - } - free(out_degrees); - - /* create a list of edges that we can remove */ - unsigned int new_edge_count = 0; - array> all_edges(16); - for (const node& vertex : vertices) - for (const node* child : vertex.children) - all_edges.add(make_pair(vertex.id, child->id)); - unsigned int total_edge_count = all_edges.length; - - array> removable_edges(16); - removable_edges.append(all_edges.data, all_edges.length); - while (true) { - /* select a start and goal vertex uniformly at random */ - unsigned int start_index; - unsigned int end_index; - if (requested_backtrack == 0 || requested_backtrack == -1) { - start_index = randrange(vertices.length - 1); - end_index = start_index + 1 + randrange(vertices.length - start_index - 1); - } else { - /* if `requested_backtrack` is specified, divide the vertices after the - start vertex into two subgraphs: the first subgraph will contain the goal, - and the second subgraph must have size at least `requested_backtrack` */ - const node* second_subgraph_root = &vertices[vertices.length - requested_backtrack]; - if (second_subgraph_root->parents.length == 1 && second_subgraph_root->parents[0]->id == second_subgraph_root->id - 1) { - start_index = randrange(vertices.length - requested_backtrack - 1); - } else { - do { - start_index = choice(second_subgraph_root->parents.data, second_subgraph_root->parents.length)->id; - } while (start_index == second_subgraph_root->id - 1); - } - end_index = start_index + 1 + randrange(vertices.length - start_index - requested_backtrack - 1); - } - start = &vertices[start_index]; - end = &vertices[end_index]; - - if (requested_backtrack != 0 && requested_backtrack != -1) { - /* if `requested_backtrack` is specified, make sure there is a path from the - start vertex to all vertices in the first and second subgraphs */ - for (unsigned int i = start_index + 1; i < vertices.length; i++) { - bool has_path_to_start = false; - for (const node* parent : vertices[i].parents) { - if (i >= vertices.length - requested_backtrack && (parent->id == start_index || parent->id >= vertices.length - requested_backtrack)) { - has_path_to_start = true; - break; - } else if (i < vertices.length - requested_backtrack && parent->id >= start_index) { - has_path_to_start = true; - break; - } - } - if (!has_path_to_start) { - vertices[start_index].children.add(&vertices[i]); - vertices[i].parents.add(&vertices[start_index]); - new_edge_count++; - total_edge_count++; - } - } - } - - /* perform DFS from the start vertex */ - array> queue(8); - if (requested_backtrack == 0 || requested_backtrack == -1) { - queue[0] = make_pair((const node*) nullptr, start); - queue.length = 1; - } else { - /* if `requested_backtrack` is specified, start the DFS in the second subgraph */ - path.add(start); - const node** first_subgraph_children = (const node**) malloc(sizeof(const node*) * start->children.length); - const node** second_subgraph_children = (const node**) malloc(sizeof(const node*) * start->children.length); - unsigned int num_first_subgraph_children = 0; - unsigned int num_second_subgraph_children = 0; - for (const node* child : start->children) { - if (child->id >= vertices.length - requested_backtrack) { - second_subgraph_children[num_second_subgraph_children++] = child; - } else { - first_subgraph_children[num_first_subgraph_children++] = child; - } - } - shuffle(first_subgraph_children, num_first_subgraph_children); - shuffle(second_subgraph_children, num_second_subgraph_children); - for (unsigned int i = 0; i < num_first_subgraph_children; i++) - queue.add(make_pair(start, first_subgraph_children[i])); - for (unsigned int i = 0; i < num_second_subgraph_children; i++) - queue.add(make_pair(start, second_subgraph_children[i])); - free(first_subgraph_children); - free(second_subgraph_children); - } - bool found_goal = false; - while (queue.length != 0) { - pair state = queue.pop(); - const node* parent = state.key; - const node* current = state.value; - if (path.contains(current)) - continue; - path.add(current); - if (parent != nullptr) { - unsigned int index = removable_edges.index_of(make_pair(parent->id, current->id)); - if (index < removable_edges.length) - removable_edges.remove(index); - } - - if (current == end) { - found_goal = true; - break; - } else if (current->children.contains(end)) { - found_goal = true; - path.add(end); - unsigned int index = removable_edges.index_of(make_pair(current->id, end->id)); - if (index < removable_edges.length) - removable_edges.remove(index); - break; - } - - if (current->children.length == 0) - continue; - const node** children = (const node**) malloc(sizeof(const node*) * current->children.length); - for (unsigned int i = 0; i < current->children.length; i++) - children[i] = current->children[i]; - shuffle(children, current->children.length); - for (unsigned int i = 0; i < current->children.length; i++) { - if (path.contains(children[i])) continue; - queue.add(make_pair(current, children[i])); - } - free(children); - } - - /* check if the goal vertex is reachable from the start vertex */ - if (found_goal) - break; - removable_edges.clear(); - removable_edges.append(all_edges.data, all_edges.length); - path.clear(); - } - - if (requested_backtrack == -1) - current_node_index = randrange(path.length - 1); - else - current_node_index = requested_backtrack; - - /* remove edges to avoid generating a graph with too many edges */ - unsigned int edges_to_remove; - if (total_edge_count > max_edges) - edges_to_remove = total_edge_count - max_edges; - else - edges_to_remove = new_edge_count; - edges_to_remove = min((unsigned int) removable_edges.length, edges_to_remove); - for (unsigned int i = 0; i < edges_to_remove; i++) { - unsigned int u = randrange(removable_edges.length); - unsigned int parent_id = removable_edges[u].key; - unsigned int child_id = removable_edges[u].value; - vertices[parent_id].children.remove(vertices[parent_id].children.index_of(&vertices[child_id])); - vertices[child_id].parents.remove(vertices[child_id].parents.index_of(&vertices[parent_id])); - removable_edges.remove(u); - } - - /* remove any correlation between graph topology and vertex IDs by shuffling the vertices */ - unsigned int* new_indices = (unsigned int*) alloca(sizeof(unsigned int) * (max_vertex_id + 1)); - for (unsigned int i = 0; i < max_vertex_id + 1; i++) new_indices[i] = i; - shuffle(new_indices, max_vertex_id + 1); - unsigned int src_index = 0; - for (unsigned int i = 0; i < vertices.length; i++) { - bool is_reserved = false; - for (unsigned int j = 0; j < array_length(RESERVED_INDICES); j++) { - if (new_indices[src_index] == RESERVED_INDICES[j]) { - is_reserved = true; - break; - } - } - if (is_reserved) - src_index++; - vertices[i].id = new_indices[src_index]; - src_index++; - } - - return true; -} - -py::tuple generate_dfs_training_set(const unsigned int max_input_size, const uint64_t dataset_size, const py::object& reserved_inputs, const int requested_backtrack, const bool random_padding, const bool uniform, const bool quiet=false) -{ - const unsigned int QUERY_PREFIX_TOKEN = (max_input_size-5) / 3 + 4; - const unsigned int PADDING_TOKEN = (max_input_size-5) / 3 + 3; - const unsigned int EDGE_PREFIX_TOKEN = (max_input_size-5) / 3 + 2; - const unsigned int PATH_PREFIX_TOKEN = (max_input_size-5) / 3 + 1; - - unsigned int longest_path_length = (max_input_size - 4) / 4; - - unsigned int num_generated = 0; - unsigned int num_collisions = 0; - unsigned int ntokens = (max_input_size - 5) / 3 + 5; - size_t input_shape[2]{dataset_size, max_input_size}; - size_t output_shape[2]{dataset_size, ntokens}; - size_t label_shape[1]{dataset_size}; - py::array_t inputs(input_shape); - py::array_t outputs(output_shape); - py::array_t labels(label_shape); - auto inputs_mem = inputs.mutable_unchecked<2>(); - auto outputs_mem = outputs.mutable_unchecked<2>(); - auto labels_mem = labels.mutable_unchecked<1>(); - - unsigned int* backtrack_distance_histogram = (unsigned int*) alloca(sizeof(unsigned int) * max_input_size); - for (unsigned int i = 0; i < max_input_size; i++) - backtrack_distance_histogram[i] = 0; - - float* MAX_FREQS_PER_BUCKET = (float*) alloca(sizeof(float) * max_input_size); - if (requested_backtrack == -1) { - for (unsigned int i = 0; i < max_input_size; i++) - MAX_FREQS_PER_BUCKET[i] = 1.0; - } else { - for (unsigned int i = 0; i < (unsigned) requested_backtrack + 1; i++) - MAX_FREQS_PER_BUCKET[i] = 1.0 / (requested_backtrack+1); - for (unsigned int i = requested_backtrack + 1; i < max_input_size; i++) - MAX_FREQS_PER_BUCKET[i] = 0.0; - MAX_FREQS_PER_BUCKET[requested_backtrack] += 0.05; - } - - array path(32); - unsigned int* potential_backtracks = (unsigned int*) alloca(max((size_t) 1, sizeof(unsigned int) * (requested_backtrack + 1))); - unsigned int potential_backtrack_count = 0; - unsigned int num_attempts = 0; - while (num_generated < dataset_size) { - if (num_attempts >= 10000000) - break; - num_attempts++; - array g(32); - const node* start; const node* end; - unsigned int current_node_index; - while (true) { - unsigned int num_vertices = std::max(2u, randrange(longest_path_length + 1)); - int backtrack = requested_backtrack; - if (requested_backtrack != -1) { - if (uniform) { - potential_backtrack_count = 0; - for (unsigned int i = 0; i < (unsigned) requested_backtrack + 1; i++) - if (num_generated == 0 || backtrack_distance_histogram[i] / num_generated < MAX_FREQS_PER_BUCKET[i]) - potential_backtracks[potential_backtrack_count++] = i; - backtrack = choice(potential_backtracks, potential_backtrack_count); - } - num_vertices = std::max((unsigned int) backtrack + 2, num_vertices); - } - if (!generate_dfs_example(g, start, end, current_node_index, path, num_vertices, max_input_size / 24 + 1, (max_input_size - 5) / 3, backtrack, longest_path_length)) { - for (node& n : g) core::free(n); - g.length = 0; path.length = 0; - continue; - } - break; - } - - array> edges(8); - for (node& vertex : g) - for (node* child : vertex.children) - edges.add(make_pair(vertex.id, child->id)); - if (edges.length > longest_path_length) { - for (node& n : g) core::free(n); - g.length = 0; path.length = 0; - continue; - } - shuffle(edges); - - array prefix(max_input_size); - for (auto& entry : edges) { - prefix[prefix.length++] = EDGE_PREFIX_TOKEN; - prefix[prefix.length++] = entry.key; - prefix[prefix.length++] = entry.value; - } - prefix[prefix.length++] = QUERY_PREFIX_TOKEN; - prefix[prefix.length++] = start->id; - prefix[prefix.length++] = end->id; - prefix[prefix.length++] = PATH_PREFIX_TOKEN; - - /* randomly select a vertex in the DFS trace */ - path.length = current_node_index + 1; - - array unvisited(4); - unsigned int backtrack_distance = max_input_size - 1; - for (unsigned int j = current_node_index + 1; j > 0; j--) { - for (const node* child : path[j-1]->children) { - if (!path.contains(child)) - unvisited.add(child); - } - if (unvisited.length != 0) { - backtrack_distance = current_node_index + 1 - j; - break; - } - } - - if (random_padding) { - /* sample padding randomly until the path is the correct length */ - unsigned int* padding_lengths = (unsigned int*) calloc(path.length + 1, sizeof(unsigned int)); - unsigned int available_padding = longest_path_length - path.length; - for (unsigned int j = path.length; j > 0; j--) { - padding_lengths[j] = randrange(available_padding + 1); - available_padding -= padding_lengths[j]; - } - padding_lengths[0] = available_padding; - - for (unsigned int j = 0; j < path.length; j++) { - for (unsigned int k = 0; k < padding_lengths[j]; k++) - prefix[prefix.length++] = PATH_PREFIX_TOKEN; - prefix[prefix.length++] = path[j]->id; - } - for (unsigned int k = 0; k < padding_lengths[path.length]; k++) - prefix[prefix.length++] = PATH_PREFIX_TOKEN; - unsigned int effective_backtrack = backtrack_distance; - for (unsigned int j = path.length; j > 0; j--) { - effective_backtrack += padding_lengths[j]; - if (path.length - j == backtrack_distance) - break; - } - backtrack_distance = effective_backtrack; - free(padding_lengths); - } else { - for (unsigned int j = path.length; j < longest_path_length; j++) - prefix[prefix.length++] = PATH_PREFIX_TOKEN; - for (unsigned int j = 0; j < path.length; j++) - prefix[prefix.length++] = path[j]->id; - } - - if ((requested_backtrack != -1 && !uniform && (unsigned int) requested_backtrack != backtrack_distance) - || (uniform && num_generated != 0 && backtrack_distance_histogram[backtrack_distance] / num_generated >= MAX_FREQS_PER_BUCKET[backtrack_distance])) - { - for (node& n : g) core::free(n); - g.length = 0; path.length = 0; - continue; - } - - /* check if this input is reserved */ - py::object contains = reserved_inputs.attr("__contains__"); - py::tuple example_tuple(prefix.length); - for (unsigned int i = 0; i < prefix.length; i++) - example_tuple[i] = prefix[i]; - if (contains(example_tuple).is(py_true)) { - for (node& n : g) core::free(n); - g.length = 0; path.length = 0; - num_collisions += 1; - continue; - } - - backtrack_distance_histogram[backtrack_distance]++; - - for (unsigned int i = 0; i < max_input_size - prefix.length; i++) - inputs_mem(num_generated, i) = PADDING_TOKEN; - for (unsigned int i = 0; i < prefix.length; i++) - inputs_mem(num_generated, max_input_size - prefix.length + i) = prefix[i]; - for (unsigned int i = 0; i < ntokens; i++) - outputs_mem(num_generated, i) = 0.0f; - for (unsigned int i = 0; i < unvisited.length; i++) - outputs_mem(num_generated, unvisited[i]->id) = 1.0f; - labels_mem(num_generated) = path[current_node_index+1]->id; - num_generated++; - - if (!quiet && num_generated > 0 && (num_generated % 1000 == 0 || num_generated >= dataset_size)) { - printf("%d examples generated.\n", num_generated); - fflush(stdout); - - printf("Backtrack distance histogram: (log frequencies)\n"); - printf("["); - bool first = true; - for (unsigned int i = 0; i < max_input_size; i++) { - if (backtrack_distance_histogram[i] == 0) - continue; - if (!first) printf(", "); - printf("%d:%.2f", i, log(backtrack_distance_histogram[i]) - log(num_generated)); - first = false; - } - printf("]\n"); - } - - for (node& n : g) core::free(n); - g.length = 0; path.length = 0; - continue; - } - - return py::make_tuple(inputs, outputs, labels, num_collisions); -} - -bool generate_si_example(array& vertices, const node*& start, const node*& end, unsigned int& current_node_index, array& path, unsigned int num_vertices, unsigned int max_num_parents, unsigned int max_vertex_id, unsigned int max_edges, unsigned int frontier_size, unsigned int branch_size) -{ - if (!vertices.ensure_capacity(num_vertices)) - return false; - for (unsigned int i = 0; i < num_vertices; i++) - if (!init(vertices[i], i)) return false; - vertices.length = num_vertices; - - /* the current vertex must have at least `branch_size` vertices following it */ - current_node_index = randrange(num_vertices - branch_size); - node* current_node = &vertices[current_node_index]; - - /* randomly sample the start and goal vertices */ - unsigned int start_index = randrange(min(current_node_index + 1, num_vertices - frontier_size)); - unsigned int end_index = randrange(max(current_node_index + 1, start_index + frontier_size), num_vertices); - start = &vertices[start_index]; end = &vertices[end_index]; - - /* sample some parent/ancestor vertices */ - constexpr float ALPHA = 1.0f; - unsigned int* out_degrees = (unsigned int*) calloc(num_vertices, sizeof(unsigned int)); - if (out_degrees == nullptr) - return false; - for (unsigned int i = 1; i < num_vertices; i++) { - /* sample the number of child and parent vertices */ - unsigned int num_parents = randrange(1, max_num_parents); - num_parents = std::min(num_parents, i); - - /* sample the parents of this new node */ - float total_probability = 0.0f; - array probabilities(i); - for (unsigned int j = 0; j < i; j++) { - probabilities[j] = ALPHA + out_degrees[j]; - total_probability += probabilities[j]; - } - probabilities.length = i; - if (current_node->children.length == branch_size) { - total_probability -= probabilities[current_node_index]; - probabilities[current_node_index] = 0.0f; - } - - array sampled_parents(std::max(1u, num_parents)); - for (unsigned int j = 0; j < num_parents; j++) { - unsigned int u = sample_categorical(probabilities.data, total_probability, probabilities.length); - sampled_parents.add(u); - total_probability -= probabilities[u]; - probabilities[u] = 0.0f; - if (total_probability < 1.0e-6f) - break; - } - - for (unsigned int parent_id : sampled_parents) { - vertices[parent_id].children.add(&vertices[i]); - vertices[i].parents.add(&vertices[parent_id]); - out_degrees[parent_id] += 1; - } - } - free(out_degrees); - - /* randomly sample the frontier (visited vertices with unvisited child nodes) */ - array frontier(8); - array available_nodes(8); - frontier.add(&vertices[start_index]); - if (start_index != current_node_index) - frontier.add(&vertices[current_node_index]); - for (unsigned int i = start_index + 1; i < vertices.length - 1; i++) - if (i != end_index && !frontier.contains(&vertices[i])) - available_nodes.add(&vertices[i]); - while (frontier.length < frontier_size) { - unsigned int r = randrange(available_nodes.length); - frontier.add(available_nodes[r]); - available_nodes.remove(r); - } - - unsigned int new_edge_count = 0; - array> removable_edges(16); - for (const node& vertex : vertices) - for (const node* child : vertex.children) - removable_edges.add(make_pair(vertex.id, child->id)); - unsigned int total_edge_count = removable_edges.length; - - /* construct a path such that all vertices in `frontier` are visited */ - array> queue(8); - for (node* child : start->children) - if (frontier.contains(child)) queue.add(make_pair(start, child)); - while (queue.length != 0) { - unsigned int r = randrange(queue.length); - const node* prev = queue[r].key; - const node* next = queue[r].value; - queue.remove(r); - if (path.contains(next)) - continue; - path.add(prev); - path.add(next); - removable_edges.remove(removable_edges.index_of(make_pair(prev->id, next->id))); - - for (node* child : next->children) - if (frontier.contains(child)) queue.add(make_pair(next, child)); - } - - /* get set of vertices reachable from `start` */ - array stack(8); - array reachable(8); - stack[stack.length++] = &vertices[start_index]; - while (stack.length != 0) { - node* current = stack.pop(); - if (reachable.contains(current)) - continue; - reachable.add(current); - - for (node* child : current->children) { - if (!reachable.contains(child) && !stack.contains(child)) - stack.add(child); - } - } - - /* for the unvisited vertices in `frontier`, move edges until they have an incoming edge from a visited vertex */ - for (node* frontier_vertex : frontier) { - if (frontier_vertex == start || path.index_of(frontier_vertex) % 2 == 1) - continue; - - /* get the most recent visited ancestors */ - array stack(8); - array visited_ancestors(8); - stack.append(frontier_vertex->parents.data, frontier_vertex->parents.length); - while (stack.length != 0) { - node* current = stack.pop(); - - bool has_visited_parent = false; - for (node* parent : current->parents) { - if (parent == start || path.contains(parent)) { - if (!visited_ancestors.contains(parent)) - visited_ancestors.add(parent); - has_visited_parent = true; - } - } - if (!has_visited_parent) - stack.append(current->parents.data, current->parents.length); - } - - if (current_node->children.length == branch_size && visited_ancestors.contains(current_node)) - visited_ancestors.remove(visited_ancestors.index_of(current_node)); - if (visited_ancestors.length != 0) { - node* new_parent = choice(visited_ancestors.data, visited_ancestors.length); - - /* replace an existing parent with `new_parent` */ - unsigned int r = randrange(frontier_vertex->parents.length); - node* old_parent = frontier_vertex->parents[r]; - frontier_vertex->parents[r] = new_parent; - old_parent->children.remove(old_parent->children.index_of(frontier_vertex)); - new_parent->children.add(frontier_vertex); - path.add(new_parent); - path.add(frontier_vertex); - removable_edges.remove(removable_edges.index_of(make_pair(old_parent->id, frontier_vertex->id))); - - /* move the new edge in `path` to a random position */ - unsigned int insert_path_index = path.index_of(new_parent); - if (path.length == 0) insert_path_index = 0; - else if (insert_path_index % 2 == 1) insert_path_index++; - unsigned int insert_index = randrange(insert_path_index, path.length); - insert_index -= insert_index % 2; - for (unsigned int j = path.length - 1; j >= insert_index + 2; j--) - path[j] = path[j - 2]; - path[insert_index] = new_parent; - path[insert_index + 1] = frontier_vertex; - } else { - /* there is no path from `start` to `frontier_vertex`, so add a new parent that is reachable from `start` */ - array candidate_parents(8); - for (node* n : reachable) { - if (n->id < frontier_vertex->id) - candidate_parents.add(n); - } - - if (current_node->children.length == branch_size && candidate_parents.contains(current_node) && candidate_parents.length != 1) - candidate_parents.remove(candidate_parents.index_of(current_node)); - - node* new_parent = choice(candidate_parents.data, candidate_parents.length); - frontier_vertex->parents.add(new_parent); - new_parent->children.add(frontier_vertex); - path.add(new_parent); - path.add(frontier_vertex); - - /* move the new edge in `path` to a random position */ - unsigned int insert_path_index = path.index_of(new_parent); - if (path.length == 0) insert_path_index = 0; - else if (insert_path_index % 2 == 1) insert_path_index++; - unsigned int insert_index = randrange(insert_path_index, path.length); - insert_index -= insert_index % 2; - for (unsigned int j = path.length - 1; j >= insert_index + 2; j--) - path[j] = path[j - 2]; - path[insert_index] = new_parent; - path[insert_index + 1] = frontier_vertex; - } - } - - /* next we have to make sure every vertex in `frontier` has an unvisited child */ - for (node* frontier_vertex : frontier) { - bool has_unvisited_child = false; - for (const node* child : frontier_vertex->children) { - if (!path.contains(child)) { - has_unvisited_child = true; - break; - } - } - if (has_unvisited_child) continue; - - /* randomly select an unvisited vertex after `frontier_vertex` */ - array unvisited(8); - for (unsigned int i = frontier_vertex->id + 1; i < vertices.length; i++) - if (!path.contains(&vertices[i])) unvisited.add(i); - node* new_child = &vertices[choice(unvisited.data, unvisited.length)]; - - /* see if we can replace an unvisited parent node of `new_child` */ - unvisited.clear(); - for (unsigned int i = 0; i < new_child->parents.length; i++) - if (new_child->parents[i] != start && !path.contains(new_child->parents[i])) - unvisited.add(i); - if (unvisited.length != 0) { - unsigned int r = choice(unvisited.data, unvisited.length); - node* old_parent = new_child->parents[r]; - old_parent->children.remove(old_parent->children.index_of(new_child)); - frontier_vertex->children.add(new_child); - new_child->parents[r] = frontier_vertex; - removable_edges.remove(removable_edges.index_of(make_pair(old_parent->id, new_child->id))); - } else { - /* we can't remove any parents from `new_child` so simply add an edge */ - frontier_vertex->children.add(new_child); - new_child->parents.add(frontier_vertex); - total_edge_count++; new_edge_count++; - } - } - - /* next, we want to make sure `vertices[current_node_index]` has an edge to `branch_size` vertices that follows it */ - unsigned int current_branch_size = 0; - available_nodes.clear(); - array available_frontier_branches(8); - for (unsigned int i = current_node_index + 1; i < vertices.length; i++) { - available_nodes.add(&vertices[i]); - - if (frontier.contains(&vertices[i])) - available_frontier_branches.add(&vertices[i]); - } - unsigned int num_frontier_branches = 0; - while (available_frontier_branches.length > 0 && num_frontier_branches < max(2*frontier_size + branch_size, max_edges) - max_edges) { - unsigned int r = randrange(available_frontier_branches.length); - node* new_child = available_frontier_branches[r]; - available_frontier_branches.remove(r); - available_nodes.remove(available_nodes.index_of(new_child)); - - /* find the edge in the path that leads to `new_child` */ - unsigned int path_index; - for (path_index = 0; path_index < path.length; path_index += 2) - if (path[path_index + 1] == new_child) break; - - if (path[path_index] == current_node) { - /* this vertex is already visited from current */ - unsigned int index = removable_edges.index_of(make_pair(current_node_index, new_child->id)); - if (index < removable_edges.length) - removable_edges.remove(index); - current_branch_size++; - num_frontier_branches++; - continue; - } - - /* connect this vertex directly to `current_node` */ - array visited(8); - node* old_parent = (node*) path[path_index]; - old_parent->children.remove(old_parent->children.index_of(new_child)); - if (current_node->children.contains(new_child)) { - new_child->parents.remove(new_child->parents.index_of(old_parent)); - } else { - current_node->children.add(new_child); - new_child->parents[new_child->parents.index_of(old_parent)] = current_node; - } - unsigned int index = removable_edges.index_of(make_pair(old_parent->id, new_child->id)); - if (index < removable_edges.length) - removable_edges.remove(index); - index = removable_edges.index_of(make_pair(current_node_index, new_child->id)); - if (index < removable_edges.length) - removable_edges.remove(index); - - /* replace the corresponding edge from `path` */ - path[path_index] = current_node; - - current_branch_size++; - num_frontier_branches++; - } - while (current_branch_size < branch_size) { - unsigned int r = randrange(available_nodes.length); - node* new_child = available_nodes[r]; - available_nodes.remove(r); - - if (current_node->children.contains(new_child)){ - unsigned int index = removable_edges.index_of(make_pair(current_node_index, new_child->id)); - if (index < removable_edges.length) - removable_edges.remove(index); - current_branch_size++; - continue; - } - - /* see if we can replace an unvisited parent node of `new_child` */ - array unvisited(8); - for (unsigned int i = 0; i < new_child->parents.length; i++) - if (new_child->parents[i] != start && !path.contains(new_child->parents[i])) - unvisited.add(i); - if (unvisited.length != 0) { - unsigned int r = choice(unvisited.data, unvisited.length); - node* old_parent = new_child->parents[r]; - old_parent->children.remove(old_parent->children.index_of(new_child)); - current_node->children.add(new_child); - new_child->parents[r] = current_node; - removable_edges.remove(removable_edges.index_of(make_pair(old_parent->id, new_child->id))); - } else { - /* we can't remove any parents from `new_child` so simply add an edge */ - current_node->children.add(new_child); - new_child->parents.add(current_node); - total_edge_count++; new_edge_count++; - } - current_branch_size++; - } - - /* visit some of the children of `current_node` */ - //unsigned int visited_branches = randrange(branch_size); - //array available_nodes(8); - //available_nodes.append(current_node->children.data, current_node->children.length); - //unsigned int index = available_nodes.index_of(end); - //if (index != available_nodes.length) - // available_nodes.remove(index); /* make sure we don't visit the goal vertex prematurely */ - //unsigned int insert_path_index = path.index_of(current_node) + 1; - //if (path.length == 0) insert_path_index = 0; - //else if (insert_path_index % 2 == 1) insert_path_index++; - //for (unsigned int i = 0; i < visited_branches && available_nodes.length != 0; i++) { - /* count the number of branches with unvisited children */ - // unsigned int r = randrange(available_nodes.length); - // node* child = available_nodes[r]; - // available_nodes.remove(r); - - // path.add(current_node); - // path.add(child); - // unsigned int insert_index = randrange(insert_path_index, path.length); - // insert_index -= insert_index % 2; - // for (unsigned int j = path.length - 1; j >= insert_index + 2; j--) - // path[j] = path[j - 2]; - // path[insert_index] = current_node; - // path[insert_index + 1] = child; - //} - - /* make sure `end` is reachable from `start` */ - array> reachability_stack(8); - for (node* child : start->children) - reachability_stack.add(make_pair(child, start)); - array_map reverse_ptrs(8); - while (reachability_stack.length != 0) { - pair entry = reachability_stack.pop(); - node* next = entry.key; - const node* parent = entry.value; - if (next->id > end_index || reverse_ptrs.contains(next)) - continue; - reverse_ptrs.put(next, parent); - for (node* child : next->children) - if (!reachability_stack.contains(make_pair(child, next))) - reachability_stack.add(make_pair(child, next)); - } - const node* current; - const node* parent; - if (reverse_ptrs.contains(&vertices[end_index])) { - current = &vertices[end_index]; - parent = reverse_ptrs.get(current); - } else { - reverse_ptrs.put(&vertices[start_index], nullptr); - node* new_parent = choice(reverse_ptrs.keys, reverse_ptrs.size); - new_parent->children.add(&vertices[end_index]); - vertices[end_index].parents.add(new_parent); - total_edge_count++; new_edge_count++; - - current = &vertices[end_index]; - parent = new_parent; - } - /* make sure none of the edges on the path from `start` to `end` are removable */ - while (true) { - pair entry = make_pair(parent->id, current->id); - unsigned int index = removable_edges.index_of(entry); - if (index != removable_edges.length) - removable_edges.remove(index); - if (parent == start) - break; - current = parent; - parent = reverse_ptrs.get(current); - } - - /* remove edges to avoid generating a graph with too many edges */ - unsigned int edges_to_remove; - if (total_edge_count > max_edges) - edges_to_remove = total_edge_count - max_edges; - else - edges_to_remove = new_edge_count; - for (unsigned int i = 0; i < edges_to_remove && removable_edges.length != 0; i++) { - unsigned int u = randrange(removable_edges.length); - unsigned int parent_id = removable_edges[u].key; - unsigned int child_id = removable_edges[u].value; - removable_edges.remove(u); - - /* make sure that removing this edge will not remove an unvisited child node from any frontier vertex */ - if (frontier.contains(&vertices[parent_id])) { - bool can_remove_edge = false; - for (const node* child : vertices[parent_id].children) { - if (child->id != child_id && !path.contains(child)) { - can_remove_edge = true; - break; - } - } - if (!can_remove_edge) { - i--; - continue; - } - } - - vertices[parent_id].children.remove(vertices[parent_id].children.index_of(&vertices[child_id])); - vertices[child_id].parents.remove(vertices[child_id].parents.index_of(&vertices[parent_id])); - } - - /* visit some additional vertices */ - array path_length_probabilities(8); - for (unsigned int i = path.length / 2; i < max_edges + 1; i++) - path_length_probabilities.add((float) i); - unsigned int target_path_length = path.length / 2 + sample_categorical(path_length_probabilities.data, path_length_probabilities.length); - available_nodes.clear(); - for (node* frontier_vertex : frontier) { - unsigned int unvisited_child_count = 0; - for (const node* child : frontier_vertex->children) - if (!path.contains(child)) unvisited_child_count++; - - if (unvisited_child_count > 1) - available_nodes.add(frontier_vertex); - } - while (available_nodes.length != 0 && path.length < target_path_length * 2) { - unsigned int u = randrange(available_nodes.length); - node* selected_vertex = available_nodes[u]; - available_nodes.remove(u); - - array unvisited_children(4); - for (node* child : selected_vertex->children) { - if (path.contains(child)) continue; - - /* make sure this is not the unvisited child of another frontier vertex */ - bool can_visit = true; - for (const node* parent : child->parents) { - if (!frontier.contains(parent)) continue; - unsigned int unvisited_child_count = 0; - for (const node* other_child : parent->children) - if (!path.contains(other_child)) unvisited_child_count++; - if (unvisited_child_count == 1) { - can_visit = false; - break; - } - } - if (can_visit) - unvisited_children.add(child); - } - if (unvisited_children.length == 0) - continue; - u = randrange(unvisited_children.length); - node* selected_child = unvisited_children[u]; - path.add(selected_vertex); - path.add(selected_child); - if (frontier.contains(selected_vertex) && unvisited_children.length - 1 > 1) - available_nodes.add(selected_vertex); - else if (!frontier.contains(selected_vertex) && unvisited_children.length - 1 > 0) - available_nodes.add(selected_vertex); - - /* check if `selected_child` has any unvisited child nodes */ - unsigned int unvisited_child_count = 0; - for (const node* child : selected_child->children) - if (!path.contains(child)) unvisited_child_count++; - if (unvisited_child_count > 0) - available_nodes.add(selected_child); - } - - /* remove any correlation between graph topology and vertex IDs by shuffling the vertices */ - unsigned int* new_indices = (unsigned int*) alloca(sizeof(unsigned int) * (max_vertex_id + 1)); - for (unsigned int i = 0; i < max_vertex_id + 1; i++) new_indices[i] = i; - shuffle(new_indices, max_vertex_id + 1); - unsigned int src_index = 0; - for (unsigned int i = 0; i < vertices.length; i++) { - bool is_reserved = false; - for (unsigned int j = 0; j < array_length(RESERVED_INDICES); j++) { - if (new_indices[src_index] == RESERVED_INDICES[j]) { - is_reserved = true; - break; - } - } - if (is_reserved) - src_index++; - vertices[i].id = new_indices[src_index]; - src_index++; - } - - return true; -} - -py::tuple generate_si_training_set(const unsigned int max_input_size, const uint64_t dataset_size, const py::object& reserved_inputs, const int requested_frontier_size, const int requested_branch_size, bool uniform, const bool quiet=false) -{ - const unsigned int QUERY_PREFIX_TOKEN = (max_input_size-5) / 3 + 4; - const unsigned int PADDING_TOKEN = (max_input_size-5) / 3 + 3; - const unsigned int EDGE_PREFIX_TOKEN = (max_input_size-5) / 3 + 2; - const unsigned int PATH_PREFIX_TOKEN = (max_input_size-5) / 3 + 1; - - unsigned int max_edges = (max_input_size - 2) / 6; - - unsigned int num_generated = 0; - unsigned int num_collisions = 0; - unsigned int ntokens = (max_input_size - 5) / 3 + 5; - size_t input_shape[2]{dataset_size, max_input_size}; - size_t output_shape[2]{dataset_size, ntokens}; - size_t label_shape[1]{dataset_size}; - py::array_t inputs(input_shape); - py::array_t outputs(output_shape); - py::array_t labels(label_shape); - auto inputs_mem = inputs.mutable_unchecked<2>(); - auto outputs_mem = outputs.mutable_unchecked<2>(); - auto labels_mem = labels.mutable_unchecked<1>(); - - unsigned int* frontier_branch_histogram = (unsigned int*) alloca(sizeof(unsigned int) * (max_edges + 1) * (max_edges + 1)); - for (unsigned int i = 0; i < (max_edges + 1) * (max_edges + 1); i++) - frontier_branch_histogram[i] = 0; - unsigned int* visited_edges_histogram = (unsigned int*) alloca(sizeof(unsigned int) * (max_edges + 1)); - for (unsigned int i = 0; i < max_edges + 1; i++) - visited_edges_histogram[i] = 0; - - float* MAX_FREQS_PER_BUCKET = (float*) alloca(sizeof(float) * (max_edges + 1) * (max_edges + 1)); - unsigned int nonzero_buckets = 0; - for (unsigned int i = 0; i < max_edges + 1; i++) { - /* frontier_size is i */ - for (unsigned int j = 0; j < max_edges + 1; j++) { - /* branch_size is j */ - if (i == 0 || j == 0 || i > (unsigned) requested_frontier_size || j > (unsigned) requested_branch_size || i + j > max_edges + 1) { - MAX_FREQS_PER_BUCKET[i*(max_edges+1) + j] = 0.0; - } else { - MAX_FREQS_PER_BUCKET[i*(max_edges+1) + j] = 1.0; - nonzero_buckets++; - } - } - } - for (unsigned int i = 0; i < (max_edges + 1) * (max_edges + 1); i++) - if (MAX_FREQS_PER_BUCKET[i] != 0.0) MAX_FREQS_PER_BUCKET[i] /= nonzero_buckets; - if ((unsigned) requested_branch_size < max_edges + 1 - requested_frontier_size) - MAX_FREQS_PER_BUCKET[requested_frontier_size*(max_edges+1) + requested_branch_size] += 0.05 / nonzero_buckets; - else - MAX_FREQS_PER_BUCKET[requested_frontier_size*(max_edges+1) + max_edges + 1 - requested_frontier_size] += 0.05 / nonzero_buckets; - - array path(32); - pair* potential_frontier_branches = (pair*) alloca(max((size_t) 1, sizeof(pair) * ((max_edges + 1) * (max_edges + 1)))); - unsigned int potential_frontier_branch_count = 0; - unsigned int num_attempts = 0; - while (num_generated < dataset_size) { - if (num_attempts >= 10000000) - break; - num_attempts++; - array g(32); - const node* start; const node* end; - unsigned int current_node_index; - while (true) { - unsigned int frontier_size; unsigned int branch_size; - if (uniform) { - potential_frontier_branch_count = 0; - for (unsigned int i = 1; i < max_edges + 1; i++) { - for (unsigned int j = 1; j < max_edges + 1; j++) { - if ((num_generated == 0 && MAX_FREQS_PER_BUCKET[i*(max_edges+1) + j] != 0.0) || (float) frontier_branch_histogram[i*(max_edges+1) + j] / num_generated < MAX_FREQS_PER_BUCKET[i*(max_edges+1) + j]) - potential_frontier_branches[potential_frontier_branch_count++] = make_pair(i, j); - } - } - pair frontier_branch = choice(potential_frontier_branches, potential_frontier_branch_count); - frontier_size = frontier_branch.key; - branch_size = frontier_branch.value; - } else { - frontier_size = requested_frontier_size; - branch_size = requested_branch_size; - } - unsigned int num_vertices = std::max(2u, randrange(max_edges - frontier_size + 1)); - num_vertices = std::max(frontier_size + branch_size + 1 - std::min(frontier_size, branch_size), num_vertices); - if (!generate_si_example(g, start, end, current_node_index, path, num_vertices, max_input_size / 24 + 1, (max_input_size - 2) / 6 + 2, max_edges, frontier_size, branch_size)) { - for (node& n : g) core::free(n); - g.length = 0; path.length = 0; - continue; - } - break; - } - - array> edges(8); - for (node& vertex : g) - for (node* child : vertex.children) - edges.add(make_pair(vertex.id, child->id)); - if (edges.length > max_edges) { - for (node& n : g) core::free(n); - g.length = 0; path.length = 0; - continue; - } - shuffle(edges); - - array prefix(max_input_size); - for (auto& entry : edges) { - prefix[prefix.length++] = EDGE_PREFIX_TOKEN; - prefix[prefix.length++] = entry.key; - prefix[prefix.length++] = entry.value; - } - prefix[prefix.length++] = QUERY_PREFIX_TOKEN; - prefix[prefix.length++] = start->id; - prefix[prefix.length++] = end->id; - prefix[prefix.length++] = PATH_PREFIX_TOKEN; - - /* compute the frontier size and branch size of this example */ - array frontier(8); - for (unsigned int i = 0; i < path.length; i++) { - bool has_unvisited_children = false; - for (const node* child : path[i]->children) { - if (!path.contains(child)) { - has_unvisited_children = true; - break; - } - } - if (has_unvisited_children && !frontier.contains(path[i])) - frontier.add(path[i]); - } - if (frontier.length == 0) - frontier.add(start); - const node* current_node = &g[current_node_index]; - unsigned int branch_size = current_node->children.length; - - bool is_selection_step; - if (path.length == 0) - is_selection_step = false; - else - is_selection_step = (randrange(2) == 0); - if (3*(path.length/2) + (is_selection_step ? 1 : 2) > 3*(max_edges - 1) + 1) { - /* we have just barely too many edges */ - for (node& n : g) core::free(n); - g.length = 0; path.length = 0; - continue; - } - for (unsigned int j = 3*(path.length/2) + (is_selection_step ? 1 : 2); j < 3*(max_edges - 1) + 1; j++) - prefix[prefix.length++] = PATH_PREFIX_TOKEN; - for (unsigned int j = 0; j < path.length; j += 2) { - prefix[prefix.length++] = PATH_PREFIX_TOKEN; - prefix[prefix.length++] = path[j]->id; - prefix[prefix.length++] = path[j+1]->id; - } - prefix[prefix.length++] = PATH_PREFIX_TOKEN; - if (!is_selection_step) - prefix[prefix.length++] = current_node->id; - - if ((requested_frontier_size != -1 && !uniform && (unsigned int) requested_frontier_size != (unsigned int) frontier.length) - || (requested_branch_size != -1 && !uniform && (unsigned int) requested_branch_size != branch_size) - || (uniform && num_generated != 0 && (float) frontier_branch_histogram[frontier.length*(max_edges+1) + branch_size] / num_generated >= MAX_FREQS_PER_BUCKET[frontier.length*(max_edges+1) + branch_size])) - { - for (node& n : g) core::free(n); - g.length = 0; path.length = 0; - continue; - } - - /* check if this input is reserved */ - py::object contains = reserved_inputs.attr("__contains__"); - py::tuple example_tuple(prefix.length); - for (unsigned int i = 0; i < prefix.length; i++) - example_tuple[i] = prefix[i]; - if (contains(example_tuple).is(py_true)) { - for (node& n : g) core::free(n); - g.length = 0; path.length = 0; - num_collisions += 1; - continue; - } - - frontier_branch_histogram[frontier.length*(max_edges+1) + branch_size]++; - visited_edges_histogram[path.length/2]++; - - for (unsigned int i = 0; i < max_input_size - prefix.length; i++) - inputs_mem(num_generated, i) = PADDING_TOKEN; - for (unsigned int i = 0; i < prefix.length; i++) - inputs_mem(num_generated, max_input_size - prefix.length + i) = prefix[i]; - for (unsigned int i = 0; i < ntokens; i++) - outputs_mem(num_generated, i) = 0.0f; - if (is_selection_step) { - for (const node* frontier_vertex : frontier) - outputs_mem(num_generated, frontier_vertex->id) = 1.0f; - labels_mem(num_generated) = choice(frontier.data, frontier.length)->id; - } else { - array correct_answers(8); - for (const node* child : current_node->children) { - if (!path.contains(child)) { - outputs_mem(num_generated, child->id) = 1.0f; - correct_answers.add(child->id); - } - } - labels_mem(num_generated) = choice(correct_answers.data, correct_answers.length); - } - num_generated++; - - if (!quiet && num_generated > 0 && (num_generated % 1000 == 0 || num_generated >= dataset_size)) { - printf("%d examples generated.\n", num_generated); - fflush(stdout); - - printf("Frontier-branch size histogram: (log frequencies)\n"); - printf("["); - bool first = true; - for (unsigned int i = 0; i < max_edges + 1; i++) { - for (unsigned int j = 0; j < max_edges + 1; j++) { - if (frontier_branch_histogram[i*(max_edges+1) + j] == 0) - continue; - if (!first) printf(", "); - printf("(%d,%d):%.2f", i, j, log(frontier_branch_histogram[i*(max_edges+1) + j]) - log(num_generated)); - first = false; - } - } - printf("]\n"); - - printf("Visited edge count histogram: (log frequencies)\n"); - printf("["); - first = true; - for (unsigned int i = 0; i < max_edges + 1; i++) { - if (visited_edges_histogram[i] == 0) - continue; - if (!first) printf(", "); - printf("%d:%.2f", i, log(visited_edges_histogram[i]) - log(num_generated)); - first = false; - } - printf("]\n"); - } - - for (node& n : g) core::free(n); - g.length = 0; path.length = 0; - continue; - } - - return py::make_tuple(inputs, outputs, labels, num_collisions); -} - -PYBIND11_MODULE(generator, m) { - m.def("generate_training_set", &generate_training_set); - m.def("generate_reachable_training_set", &generate_reachable_training_set); - m.def("generate_dfs_training_set", &generate_dfs_training_set); - m.def("generate_si_training_set", &generate_si_training_set); - m.def("lookahead_histogram", &lookahead_histogram); - m.def("set_seed", &core::set_seed); -} diff --git a/reasoners/data_preprocess/logic/graph_dataset_gen/jobfile.sub b/reasoners/data_preprocess/logic/graph_dataset_gen/jobfile.sub deleted file mode 100644 index 8dfaf35f8..000000000 --- a/reasoners/data_preprocess/logic/graph_dataset_gen/jobfile.sub +++ /dev/null @@ -1,11 +0,0 @@ -#!/bin/bash -# FILENAME: jobfile.sub -#SBATCH --nodes=1 -#SBATCH --time=4:00:00 -#SBATCH --job-name gen_graph_dataset -#SBATCH --gpus-per-node=1 -module load conda -source ~/.bashrc -conda activate llm_env -module list -python logic.py --num_samples 10000 --output_file graph_dataset.json \ No newline at end of file diff --git a/reasoners/data_preprocess/logic/graph_dataset_gen/logic.py b/reasoners/data_preprocess/logic/graph_dataset_gen/logic.py deleted file mode 100644 index 04aa92cf8..000000000 --- a/reasoners/data_preprocess/logic/graph_dataset_gen/logic.py +++ /dev/null @@ -1,359 +0,0 @@ -import os -import random -import sysconfig -from io import StringIO -from faker import Faker -from pybind11.__main__ import print_includes -import math -import argparse -import json - -def build_module(name): - import sys - old_stdout = sys.stdout - try: - sys.stdout = StringIO() - print_includes() - includes = sys.stdout.getvalue().strip() - sys.stdout.close() - sys.stdout = old_stdout - except Exception as e: - raise e - finally: - sys.stdout = old_stdout - - python_extension_suffix = sysconfig.get_config_var("EXT_SUFFIX") - - if sys.platform == "darwin": - # macOS command - command = ( - f"g++ -std=c++11 -Ofast -DNDEBUG -fno-stack-protector " - f"-Wall -Wpedantic -undefined dynamic_lookup -shared -fPIC " - f"{includes} -I. {name}.cpp -o {name}{python_extension_suffix}" - ) - else: - # Non-macOS command - command = ( - f"g++ -Ofast -std=c++11 -DNDEBUG -fno-stack-protector " - f"-Wall -Wpedantic -shared -fPIC " - f"{includes} -I. {name}.cpp -o {name}{python_extension_suffix}" - ) - - print(command) - if os.system(command) != 0: - print(f"ERROR: Unable to compile `{name}.cpp`.") - sys.exit(1) - -try: - from os.path import getmtime - from importlib.util import find_spec - - generator_spec = find_spec('generator') - if generator_spec == None: - raise ModuleNotFoundError - if getmtime(generator_spec.origin) < getmtime('generator.cpp'): - print("C++ module `generator` is out-of-date. Compiling from source...") - build_module("generator") - import generator -except ModuleNotFoundError: - print("C++ module `generator` not found. Compiling from source...") - build_module("generator") - import generator -except ImportError: - print("Error loading C++ module `generator`. Compiling from source...") - build_module("generator") - import generator -print("C++ module `generator` loaded.") - - -def generate_graph_text( - max_input_size, - num_vertices, - max_num_parents, - lookahead, - num_paths, -): - - # Get graph using the C++ generator - inputs, outputs, labels, num_collisions = generator.generate_training_set( - max_input_size, - 1, # batch size - lookahead, - max_num_parents * num_vertices, # max edges - set(), # reserved vertices - -1, # distance from start - 0, # max_prefix_vertices - True, # quiet - num_paths # number of paths - ) - - # If the random DAG fails to generate a valid example, return None - if inputs is None or outputs is None: - return None - - # Token IDs - PADDING_TOKEN = (max_input_size - 5) // 3 + 3 - EDGE_PREFIX_TOKEN = (max_input_size - 5) // 3 + 2 - PATH_PREFIX_TOKEN = (max_input_size - 5) // 3 + 1 - QUERY_PREFIX_TOKEN = (max_input_size - 5) // 3 + 4 - - # Function for mapping tokens to letters - def token_to_str(token): - if token == EDGE_PREFIX_TOKEN: - return "E" - elif token == QUERY_PREFIX_TOKEN: - return "Q" - elif token == PATH_PREFIX_TOKEN: - return "P" - # return "" - elif token == PADDING_TOKEN: - return "" - else: - return str(token) - - # Convert all tokens to letters - text_tokens = [token_to_str(t) for t in inputs[0] if token_to_str(t) != ""] - - final_str = " ".join(text_tokens) - return final_str, labels[0] - - -def generate_name(n): - """ - Generates n random names using the Faker library. - """ - fake = Faker() - # return [fake.unique.first_name() for i in range(n)] - return [fake.unique.name() for i in range(n)] - - -def random_syllable(): - """ - Generates a random syllable based on a randomly chosen pattern. - The pattern is built from 'C' (consonant) and 'V' (vowel). - """ - vowels = "aeiou" - consonants = "bcdfghjklmnpqrstvwxyz" - # Syllable patterns (e.g. "CV" means a consonant followed by a vowel) - patterns = ["CV", "VC", "CVC", "CVV", "CCV", "VCV", "VCC"] - pattern = random.choice(patterns) - - syllable = "" - for char in pattern: - if char == "C": - syllable += random.choice(consonants) - elif char == "V": - syllable += random.choice(vowels) - return syllable - - -def generate_fake_noun(): - """ - Generates a fake noun by concatenating two randomly generated syllables. - """ - return random_syllable() + random_syllable() - - -def generate_fake_nouns(n): - """ - Returns a list of n fake nouns. - """ - nouns = set() - while len(nouns) < n: - nouns.add(generate_fake_noun()) - return list(nouns) - - -# def generate_words(n): -# fake = Faker() -# return fake.words(nb=n) - - -def logic_paragraph_from_tokens(tokens_str: str, next_step: int, use_diff_names=True): - # Split string into tokens - tokens = tokens_str.strip().split() - - edges = [] - queries = [] - - # Keep track of nodes in a set - node_ids = set() - - i = 0 - while i < len(tokens): - t = tokens[i] - - if t == "E": - # Expect "E A B" - if i + 2 >= len(tokens): - break # malformed - A = int(tokens[i + 1]) - B = int(tokens[i + 2]) - edges.append((A, B)) - node_ids.update([A, B]) - i += 3 - - elif t == "Q": - # Expect "Q X Y" - if i + 2 >= len(tokens): - break - X = int(tokens[i + 1]) - Y = int(tokens[i + 2]) - queries.append((X, Y)) - node_ids.update([X, Y]) - i += 3 - - elif t == "P": - # "P" is just a path prefix token; ignore - i += 1 - - else: - # Possibly a stray token or something else (like a vertex ID alone). - i += 1 - - # Generate fake names and fake adjectives - num_nodes = len(node_ids) - all_names = generate_name(num_nodes) - all_adjs = generate_fake_nouns(num_nodes) - - sorted_nodes = sorted(node_ids) - id_to_pair = {} # node_id -> (name, adjective) - # Assign all nodes the same name if use_diff_names == False - for idx, node_id in enumerate(sorted_nodes): - id_to_pair[node_id] = (all_names[idx if use_diff_names else 0], all_adjs[idx]) - - # Lines of logic - lines = ["Given the following list of predicates:"] - - def get_logic_line(name_a: str, name_b: str, adj_a: str, adj_b: str) -> str: - choices = [ - f"If {name_a} is {adj_a}, then {name_b} is {adj_b}.", - f"{name_a} is {adj_a} implies {name_b} is {adj_b}.", - f"{name_b} is {adj_b} is true if {name_a} is {adj_a}.", - # f"{adj_b} is true if {adj_a} is true.", - # f"If {adj_a} then {adj_b} is true.", - # f"If {adj_a} is true then {adj_b}.", - f"Given {name_a} is {adj_a} then {name_b} is {adj_b}.", - ] - - sentence = random.choice(choices) - return sentence[0].upper() + sentence[1:] - - # For each edge: E A B => "If name(A) is adj(A), then name(A) is adj(B)." - for (A, B) in edges: - name_A, adj_A = id_to_pair[A] - name_B, adj_B = id_to_pair[B] - lines.append(get_logic_line(name_A, name_B, adj_A, adj_B)) - - # For each query: Q X Y => "If name(X) is adj(X), prove that name(X) is adj(Y)." - for (X, Y) in queries: - name_x, adj_x = id_to_pair[X] - name_y, adj_y = id_to_pair[Y] - lines.append(f"\n\nIf {name_x} is {adj_x}, what is the next step to prove that {name_y} is {adj_y}?") - - # Join all lines into one paragraph - paragraph = " ".join(lines) - - # Get the correct adjective for the next step - next_step_adj = id_to_pair[next_step] - next_step_adj = id_to_pair[next_step][1] - - return paragraph, next_step_adj - - - -def main(samples_per_test: int = 3, lookahead_range: list = range(1, 5), num_paths: int = 2, max_num_parents: int = 3, logic: bool = False, seed: int = None, verbose: bool = True, print_prompts: bool = False, output_file:str ="graph_search_data.json", model: str = "gpt-4o"): - if seed is not None: - random.seed(seed) - generator.set_seed(seed) - - graph_search_puzzles = [] - correct_responses = [] - look_ahead_values = [] - counter = -1 - for look_ahead in lookahead_range: - for _ in range(samples_per_test): - txt, next_step = generate_graph_text( - max_input_size=max_num_parents * look_ahead * num_paths * 4, - num_vertices=max_num_parents * look_ahead, - max_num_parents=max_num_parents, - lookahead=look_ahead, - num_paths=num_paths, - ) - counter += 1 - # Create the prompt for the graph search - prompt = (f"{txt}\nAbove is a representation of a directed graph search problem, " - f"where E A B represents an edge from A to B, and Q X Y represents starting from X and ending at Y, " - f"find the shortest path. The vertex after P indicates our current position. Respond with only the " - f"next vertex on the shortest path from X to Y and nothing else.") - - # Change the prompt to a logic puzzle if the logic option is enabled - if logic: - logic, next_step_adj = logic_paragraph_from_tokens(txt, next_step) - prompt = f"{logic} Respond with only the trait of the next step." - - - correct_response = next_step_adj if logic else next_step - puzzle = { - "id": counter, - "prompt": prompt, - "correct_response": correct_response, - "look_ahead": look_ahead - } - graph_search_puzzles.append(puzzle) - if print_prompts: - print(f"Prompt: {prompt}\n") - print(f"Correct: {next_step_adj if logic else next_step}\n") - - with open(output_file, "w") as json_file: - json.dump(graph_search_puzzles, json_file, indent=2) - - -def multiplicative_range(start, stop, step): - result = [] - current = start - while (step > 1 and current < stop) or (step < 1 and current > stop): - result.append(current) - current *= step - return result - -def exponential_range(base, start, end): - res = [round(base ** n) for n in range(start, end + 1)] - res = list(set(res)) - return sorted(res) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Generate graph search problems.") - parser.add_argument("--output_dir", type=str, default="../../../data/raw", help="output directory for generated problems") - parser.add_argument("--output_file", type=str, default="graph_search.json", help="output file for generated problems") - parser.add_argument("--num_samples", type=int, default=10000, help="number of samples to generate") - args = parser.parse_args() - # Set the output directory and file - # base_dir = os.path.dirname(os.path.abspath(__file__)) - output_dir = args.output_dir - # Create the output directory if it doesn't exist - if not os.path.exists(output_dir): - os.makedirs(output_dir) - output_file = os.path.join(output_dir, args.output_file) - # Set the random seed for reproducibility - seed = 42 - random.seed(seed) - look_ahead = exponential_range(1.3, 1, 24) - print(f"look_ahead range={look_ahead}") - samples_per_test = round(args.num_samples/len(look_ahead)) - print(f"samples_per_test={samples_per_test}") - main( - samples_per_test=samples_per_test, - lookahead_range=look_ahead, - num_paths=9, - logic=True, - verbose=False, - print_prompts=False, - output_file = output_file, - seed=seed, - ) - -# example usage: -# python graph_search.py --output_dir ../data/graph_dataset --output_file graph_search.json --num_samples 10000 - diff --git a/reasoners/data_preprocess/logic/process_graph_dataset.py b/reasoners/data_preprocess/logic/process_graph_dataset.py deleted file mode 100644 index 01279a61c..000000000 --- a/reasoners/data_preprocess/logic/process_graph_dataset.py +++ /dev/null @@ -1,159 +0,0 @@ -"""Downloads, processes, and saves graph logical reasoning dataset.""" - -import os -import datasets -import argparse -import re -from sklearn.model_selection import train_test_split -from transformers import AutoTokenizer - -from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset - -# Constants -INSTRUCTION_FOLLOW = "Please put your answer within and tags, for example fdebme ." - -def get_dataset(json_path): - """ - Load the graph dataset from a JSON file. - - Args: - json_path (str): Path to the JSON file - - Returns: - Dataset: Loaded dataset - """ - print(f"Loading the dataset from {json_path}...") - return datasets.load_dataset('json', data_files=json_path)['train'] - -def make_map_fn(split: str, data_source: str) -> callable: - """ - Create a mapping function for processing dataset examples. - - Args: - split (str): Data split ('train' or 'test') - data_source (str): Name of the data source - - Returns: - callable: Function to map over the dataset - """ - def process_fn(example, idx): - question = example['prompt'] - lookahead = example['look_ahead'] - answer = example['correct_response'] - - # Create user message with instructions - formatted_question = question + " " + INSTRUCTION_FOLLOW - - data = { - "data_source": data_source, - "prompt": [{ - "role": "user", - "content": formatted_question, - }], - "ability": "logical_reasoning", - "apply_chat_template": True, - "reward_model": { - "style": "rule", - "ground_truth": answer - }, - "extra_info": { - 'id': example['id'] if 'id' in example else str(idx), - 'lookahead': lookahead, - 'split': split - } - } - - if idx == 0: - print(f"data_source: {data_source}, split: {split}, idx: {idx}") - print("\n" + "=" * 100 + f"{data_source} {split} {idx}" + "=" * 10) - print(data) - - return data - - return process_fn - -if __name__ == '__main__': - """Main script execution: parse args, load, process, and save datasets.""" - parser = argparse.ArgumentParser() - parser.add_argument('--json_path', default='data/raw/graph_search.json', help='Path to json file') - parser.add_argument('--output_dir', default='data', help='Directory to save processed data') - parser.add_argument('--hdfs_dir', default=None, help='HDFS directory (optional)') - parser.add_argument('--train_size', type=float, default=0.8, help='Proportion of data for train set') - parser.add_argument('--test_size', type=float, default=0.2, help='Proportion of data for test set') - parser.add_argument('--data_source', default='graph_logical_dataset', help='Name of data source') - parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility when splitting data') - parser.add_argument('--train_sample_size', type=int, default=None, help='Number of samples to use from train. If None, use all.') - args = parser.parse_args() - - # Set random seed for reproducibility - set_seed(args.seed) - - # Validate train and test sizes - if args.train_size + args.test_size > 1.0: - raise ValueError(f"The sum of train_size ({args.train_size}) and test_size ({args.test_size}) cannot exceed 1.0") - - # Load the dataset - dataset = get_dataset(args.json_path) - - # Split dataset into train and test - train_indices, test_indices = train_test_split( - range(len(dataset)), - train_size=args.train_size, - test_size=args.test_size, - random_state=args.seed - ) - - print(f"Train set size: {len(train_indices)}, Test set size: {len(test_indices)}") - - # Process the datasets - process_train_fn = make_map_fn('train', args.data_source) - process_test_fn = make_map_fn('test', args.data_source) - - train_dataset = dataset.select(train_indices).map(function=process_train_fn, with_indices=True) - test_dataset = dataset.select(test_indices).map(function=process_test_fn, with_indices=True) - - # Store the original training dataset size - original_train_size = len(train_dataset) - - # Sample the training dataset if needed - train_dataset = sample_dataset(train_dataset, args.train_sample_size) - - # Create output directories - train_output_dir = os.path.join(args.output_dir, "train") - test_output_dir = os.path.join(args.output_dir, "test") - os.makedirs(train_output_dir, exist_ok=True) - os.makedirs(test_output_dir, exist_ok=True) - - # Save train dataset - train_output_path = save_dataset( - dataset=train_dataset, - output_dir=train_output_dir, - filename_prefix=f"logic__{args.data_source}", - sample_size=args.train_sample_size if args.train_sample_size else len(train_dataset) - ) - - # Save test dataset - test_output_path = save_dataset( - dataset=test_dataset, - output_dir=test_output_dir, - filename_prefix=f"logic__{args.data_source}", - sample_size=len(test_dataset) - ) - - # Copy to HDFS if specified - if args.hdfs_dir is not None: - try: - from verl.utils.hdfs_io import copy, makedirs - makedirs(args.hdfs_dir) - copy(src=args.output_dir, dst=args.hdfs_dir) - print(f"Data copied to HDFS: {args.hdfs_dir}") - except ImportError: - print("HDFS utilities not available. Install verl package for HDFS support.") - - print(f"Done! \n" - f"Train data saved to {train_output_path}\n" - f"Test data saved to {test_output_path}") - print(f"Original train set size: {original_train_size} examples") - print(f"Final train set size: {len(train_dataset)} examples") - print(f"Test set: {len(test_dataset)} examples") - diff --git a/reasoners/data_preprocess/logic/process_puzzles_dataset.py b/reasoners/data_preprocess/logic/process_puzzles_dataset.py deleted file mode 100644 index 21f8f8083..000000000 --- a/reasoners/data_preprocess/logic/process_puzzles_dataset.py +++ /dev/null @@ -1,198 +0,0 @@ -# Copyright 2024 -# -# 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. -""" -Preprocess a puzzle dataset to parquet format. -""" -import os -import datasets -import argparse -import re -from sklearn.model_selection import train_test_split - -from transformers import ( - LlamaForCausalLM, - LlamaTokenizer, - AutoTokenizer, -) - -from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset - -# Constants -INSTRUCTION_FOLLOW = "Please put your answer within and tags, for example ['pigeon', 'sparrow', 'quail'] ." - -def make_prefix(dp): - """ - Create a suitable prompt prefix. - - Args: - dp (dict): Data point containing input, instruction, etc. - - Returns: - str: Formatted prompt - """ - constraints = dp['input'] - result = dp['ground_truth'] - instruction = dp['instruction'] - prefix = f"{instruction}. The constraints are: {constraints}." - return prefix - - -def extract_from_ground_truth(text): - """ - Extract answer from the ground truth text. - - Args: - text: Ground truth answer, could be string or list - - Returns: - list: Extracted answer - """ - if isinstance(text, list): - return text - else: - return eval(text) - - -def make_map_fn(split: str, data_source: str) -> callable: - """ - Create a mapping function for processing dataset examples. - - Args: - split (str): Data split ('train' or 'test') - data_source (str): Name of the data source - - Returns: - callable: Function to map over the dataset - """ - def process_fn(example, idx): - # Generate the appropriate question format based on model type - question = make_prefix(example) - num_objects = example['num_objects'] - final_arrangement = extract_from_ground_truth(example['ground_truth']) - - data = { - "data_source": data_source, - "prompt": [{ - "role": "user", - "content": question + " " + INSTRUCTION_FOLLOW - }], - "ability": "logical_reasoning", - "reward_model": { - "style": "rule", - "ground_truth": final_arrangement - }, - "apply_chat_template": True, - "extra_info": { - 'id': example['id'] if 'id' in example else str(idx), - 'raw_instruction': example['instruction'], - 'raw_input': example['input'], - 'num_objects': num_objects, - 'split': split, - } - } - - if idx == 0: - print(f"data_source: {data_source}, split: {split}, idx: {idx}") - print("\n" + "=" * 100 + f"{data_source} {split} {idx}" + "=" * 10) - print(data) - return data - - return process_fn - - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument('--json_path', default='data/raw/puzzles_dataset.json', help='Path to json file') - parser.add_argument('--output_dir', default='data', help='Directory to save processed data') - parser.add_argument('--hdfs_dir', default=None, help='HDFS directory (optional)') - parser.add_argument('--train_size', type=float, default=0.8, help='Proportion of data for train set') - parser.add_argument('--test_size', type=float, default=0.2, help='Proportion of data for test set') - parser.add_argument('--data_source', default='ordering_puzzle_dataset', help='Name of data source') - parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility') - parser.add_argument('--train_sample_size', type=int, default=None, help='Number of samples to use from train. If None, use all.') - args = parser.parse_args() - - # Set seed for reproducibility - set_seed(args.seed) - - # Load dataset from JSON - dataset = datasets.load_dataset('json', data_files=args.json_path)['train'] - print(f"Loaded dataset with {len(dataset)} examples") - - # Validate train and test sizes - if args.train_size + args.test_size > 1.0: - raise ValueError(f"The sum of train_size ({args.train_size}) and test_size ({args.test_size}) cannot exceed 1.0") - - # Transform dataset - process_train_fn = make_map_fn('train', args.data_source) - processed_dataset = dataset.map(function=process_train_fn, with_indices=True) - - # Split dataset into train and test - train_indices, test_indices = train_test_split( - range(len(dataset)), - train_size=args.train_size, - test_size=args.test_size, - random_state=args.seed - ) - - print(f"Train set size: {len(train_indices)}, Test set size: {len(test_indices)}") - - # Create train and test datasets - train_dataset = processed_dataset.select(train_indices) - test_dataset = processed_dataset.select(test_indices) - - # Store the original training dataset size - original_train_size = len(train_dataset) - - # Sample the training dataset if needed - train_dataset = sample_dataset(train_dataset, args.train_sample_size) - - # Create output directories - train_output_dir = os.path.join(args.output_dir, "train") - test_output_dir = os.path.join(args.output_dir, "test") - os.makedirs(train_output_dir, exist_ok=True) - os.makedirs(test_output_dir, exist_ok=True) - - # Save train dataset - train_output_path = save_dataset( - dataset=train_dataset, - output_dir=train_output_dir, - filename_prefix=f"logic__{args.data_source}", - sample_size=args.train_sample_size if args.train_sample_size else len(train_dataset) - ) - - # Save test dataset - test_output_path = save_dataset( - dataset=test_dataset, - output_dir=test_output_dir, - filename_prefix=f"logic__{args.data_source}", - sample_size=len(test_dataset) - ) - - # Copy to HDFS if specified - if args.hdfs_dir is not None: - try: - from verl.utils.hdfs_io import copy, makedirs - makedirs(args.hdfs_dir) - copy(src=args.output_dir, dst=args.hdfs_dir) - print(f"Data copied to HDFS: {args.hdfs_dir}") - except ImportError: - print("HDFS utilities not available. Install verl package for HDFS support.") - - print(f"Done! \n" - f"Train data saved to {train_output_path}\n" - f"Test data saved to {test_output_path}") - print(f"Original train set size: {original_train_size} examples") - print(f"Final train set size: {len(train_dataset)} examples") - print(f"Test set: {len(test_dataset)} examples") \ No newline at end of file diff --git a/reasoners/data_preprocess/logic/process_zebrapuzzle_dataset.py b/reasoners/data_preprocess/logic/process_zebrapuzzle_dataset.py deleted file mode 100644 index 04c11d29d..000000000 --- a/reasoners/data_preprocess/logic/process_zebrapuzzle_dataset.py +++ /dev/null @@ -1,143 +0,0 @@ -import os -import datasets -import argparse -from sklearn.model_selection import train_test_split - -from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset - -# Instruction for zebra puzzle answers -InstructionFollow = "Output the grid in the form of a dictionary with keys as header containing a list of the attributes and rows denoting each row of the final grid. Please return the final answer in tags, for example {\"header\": [\"Position\", \"Nationality\", \"Job\"], \"rows\": [[\"1\", \"british\", \"plumber\"], [\"2\", \"polish\", \"carpenter\"]]} ." - -def make_prefix(dp): - clues = dp['clues'] - if isinstance(clues, list): - clues = "\n".join(clues) - result = dp['ground_truth'] - instruction = dp['instruction'] - prefix = f"{instruction} The clues are: {clues}." - return prefix - - -def extract_from_ground_truth(text): - if isinstance(text, dict): - return text - else: - return eval(text) - -def make_map_fn(split, data_source): - def process_fn(example, idx): - question = make_prefix(example) - grid_size = str(example['config']["cols"]) + "x" + str(example['config']["rows"]) - final_grid = extract_from_ground_truth(example['ground_truth']) - - data = { - "data_source": data_source, - "prompt": [{ - "role": "user", - "content": question + " " + InstructionFollow - }], - "ability": "logical_reasoning", - "reward_model": { - "style": "rule", - "ground_truth": final_grid, - }, - "apply_chat_template": True, - "extra_info": { - 'id': example['id'] if 'id' in example else str(idx), - "grid_size": grid_size, - 'raw_instruction': example['instruction'], - 'raw_input': example['clues'], - 'split': split, - } - } - - if idx == 0: - print(f"data_source: {data_source}, split: {split}, idx: {idx}") - print("\n" + "=" * 100 + f"{data_source} {split} {idx}" + "=" * 10) - print(data) - return data - - return process_fn - - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument('--json_path', default='data/raw/zebra_puzzles.json', help='Path to json file') - parser.add_argument('--output_dir', default='data', help='Directory to save processed data') - parser.add_argument('--hdfs_dir', default=None, help='HDFS directory (optional)') - parser.add_argument('--train_size', type=float, default=0.8, help='Proportion of data for train set') - parser.add_argument('--test_size', type=float, default=0.2, help='Proportion of data for test set') - parser.add_argument('--data_source', default='zebra_puzzle_dataset', help='Name of data source') - parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility') - parser.add_argument('--train_sample_size', type=int, default=None, help='Number of samples to use from train. If None, use all.') - args = parser.parse_args() - - # Set seed for reproducibility - set_seed(args.seed) - - # Load dataset from JSON - dataset = datasets.load_dataset('json', data_files=args.json_path)['train'] - - # Transform dataset - process_train_fn = make_map_fn('train', args.data_source) - processed_dataset = dataset.map(function=process_train_fn, with_indices=True) - - if args.train_size + args.test_size > 1.0: - raise ValueError(f"The sum of train_size ({args.train_size}) and test_size ({args.test_size}) cannot exceed 1.0") - - # Split dataset into train and test - train_indices, test_indices = train_test_split( - range(len(dataset)), - train_size=args.train_size, - test_size=args.test_size, - random_state=args.seed - ) - - # Create train and test datasets - train_dataset = processed_dataset.select(train_indices) - test_dataset = processed_dataset.select(test_indices) - - # Store the original training dataset size - original_train_size = len(train_dataset) - - # Sample the training dataset if needed - train_dataset = sample_dataset(train_dataset, args.train_sample_size) - - # Create output directories - train_output_dir = os.path.join(args.output_dir, "train") - test_output_dir = os.path.join(args.output_dir, "test") - os.makedirs(train_output_dir, exist_ok=True) - os.makedirs(test_output_dir, exist_ok=True) - - # Save train dataset - train_output_path = save_dataset( - dataset=train_dataset, - output_dir=train_output_dir, - filename_prefix=f"logic__{args.data_source}", - sample_size=args.train_sample_size if args.train_sample_size else len(train_dataset) - ) - - # Save test dataset - test_output_path = save_dataset( - dataset=test_dataset, - output_dir=test_output_dir, - filename_prefix=f"logic__{args.data_source}", - sample_size=len(test_dataset) - ) - - # Copy to HDFS if specified - if args.hdfs_dir is not None: - try: - from verl.utils.hdfs_io import copy, makedirs - makedirs(args.hdfs_dir) - copy(src=args.output_dir, dst=args.hdfs_dir) - print(f"Data copied to HDFS: {args.hdfs_dir}") - except ImportError: - print("HDFS utilities not available. Install verl package for HDFS support.") - - print(f"Done! \n" - f"Train data saved to {train_output_path}\n" - f"Test data saved to {test_output_path}") - print(f"Original train set size: {original_train_size} examples") - print(f"Final train set size: {len(train_dataset)} examples") - print(f"Test set: {len(test_dataset)} examples") \ No newline at end of file diff --git a/reasoners/data_preprocess/logic/puzzle_gen.py b/reasoners/data_preprocess/logic/puzzle_gen.py deleted file mode 100644 index 655ce7651..000000000 --- a/reasoners/data_preprocess/logic/puzzle_gen.py +++ /dev/null @@ -1,318 +0,0 @@ -import random -from itertools import permutations -from typing import List, Set, Dict, Tuple, Optional -from dataclasses import dataclass -import os -from datetime import datetime -import random -import gen_proof as gen_proof -import csv -import argparse -import json - -@dataclass -class Constraint: - text: str # Human-readable constraint - object1: Optional[int] # First object's encoded value (if applicable) - object2: Optional[int] # Second object's encoded value (if applicable) - constraint_type: str # Type of mathematical constraint - value: Optional[int] # Additional value needed for constraint (e.g., number of birds between) - - def evaluate(self, positions: List[int]) -> bool: - """Evaluate if the constraint is satisfied for given positions.""" - if self.constraint_type == "left_of": - return positions.index(self.object1) < positions.index(self.object2) - elif self.constraint_type == "right_of": - return positions.index(self.object1) > positions.index(self.object2) - elif self.constraint_type == "adjacent_left": - return positions.index(self.object1) + 1 == positions.index(self.object2) - elif self.constraint_type == "adjacent_right": - return positions.index(self.object1) - 1 == positions.index(self.object2) - elif self.constraint_type == "birds_between": - return abs(positions.index(self.object1) - positions.index(self.object2)) - 1 == self.value - elif self.constraint_type == "absolute_position_left": - return positions.index(self.object1) == self.value - 1 - elif self.constraint_type == "absolute_position_right": - return positions.index(self.object1) == len(positions) - self.value - return False - -class BirdPuzzleGenerator: - def __init__(self, objects: List[str]): - self.birds = objects - self.num_positions = len(objects) - # Encode birds as integers (0 to n-1) - self.bird_encoding = {bird: i for i, bird in enumerate(objects)} - self.bird_decoding = {i: bird for i, bird in enumerate(objects)} - - # Generate ground truth - self.ground_truth_birds = list(objects) - random.shuffle(self.ground_truth_birds) - # Store encoded ground truth positions - self.ground_truth = [self.bird_encoding[bird] for bird in self.ground_truth_birds] - - def generate_all_constraints(self) -> List[Constraint]: - """Generate all possible true constraints based on ground truth.""" - absolute_constraints = [] - relative_constraints = [] - # Generate all types of constraints - for i in range(self.num_positions): - bird1_code = self.ground_truth[i] - bird1_name = self.bird_decoding[bird1_code] - - # Absolute position constraints - pos_suffix = lambda n: 'th' if 11 <= n % 100 <= 13 else {1: 'st', 2: 'nd', 3: 'rd'}.get(n % 10, 'th') - left_pos = i + 1 - right_pos = self.num_positions - i - - absolute_constraints.append(Constraint( - text=f"The {bird1_name} is the {left_pos}{pos_suffix(left_pos)} from the left", - object1=bird1_code, - object2=None, - constraint_type="absolute_position_left", - value=left_pos - )) - - absolute_constraints.append(Constraint( - text=f"The {bird1_name} is the {right_pos}{pos_suffix(right_pos)} from the right", - object1=bird1_code, - object2=None, - constraint_type="absolute_position_right", - value=right_pos - )) - - # Relative position constraints - for j in range(i + 1, self.num_positions): - bird2_code = self.ground_truth[j] - bird2_name = self.bird_decoding[bird2_code] - - # Left/Right relationships - relative_constraints.append(Constraint( - text=f"The {bird1_name} is to the left of the {bird2_name}", - object1=bird1_code, - object2=bird2_code, - constraint_type="left_of", - value=None - )) - - relative_constraints.append(Constraint( - text=f"The {bird2_name} is to the right of the {bird1_name}", - object1=bird2_code, - object2=bird1_code, - constraint_type="right_of", - value=None - )) - - # Birds between - birds_between = j - i - 1 - if birds_between > 0: - plural = "s" if birds_between > 1 else "" - relative_constraints.append(Constraint( - text=f"There {['is', 'are'][birds_between > 1]} {birds_between} bird{plural} between the {bird1_name} and the {bird2_name}", - object1=bird1_code, - object2=bird2_code, - constraint_type="birds_between", - value=birds_between - )) - - # Adjacent relationships - if birds_between == 0: - relative_constraints.append(Constraint( - text=f"The {bird1_name} is immediately to the left of the {bird2_name}", - object1=bird1_code, - object2=bird2_code, - constraint_type="adjacent_left", - value=None - )) - - relative_constraints.append(Constraint( - text=f"The {bird2_name} is immediately to the right of the {bird1_name}", - object1=bird2_code, - object2=bird1_code, - constraint_type="adjacent_right", - value=None - )) - - return absolute_constraints, relative_constraints - - def verify_arrangement(self, constraints: List[Constraint], test_arrangement: List[int]) -> bool: - """Verify if a given arrangement satisfies all constraints using mathematical evaluation.""" - return all(constraint.evaluate(test_arrangement) for constraint in constraints) - - def find_minimal_constraints(self) -> List[Constraint]: - """Find a minimal set of constraints that uniquely determines the ground truth. - Uses incremental constraint addition until solution space reduces to one arrangement.""" - absolute_constraints, relative_constraints = self.generate_all_constraints() - # Shuffle constraints to randomize selection order - abs_constraints = list(absolute_constraints) - random.shuffle(abs_constraints) - rel_constraints = list(relative_constraints) - random.shuffle(rel_constraints) - - selected_constraints = [] - - # num_absolute = random.randint(0, int(self.num_positions*0.4)) - num_absolute = 0 - while len(selected_constraints) < num_absolute: - selected_constraints.append(abs_constraints.pop(0)) - - while rel_constraints and len(selected_constraints) < len(rel_constraints): - # Take the next constraint - current_constraint = rel_constraints.pop(0) - selected_constraints.append(current_constraint) - - # Find all valid arrangements given current constraints - valid_arrangements = [] - for perm in permutations(range(self.num_positions)): - if self.verify_arrangement(selected_constraints, list(perm)): - valid_arrangements.append(list(perm)) - # if len(valid_arrangements) > 1: - # break - - # If we've found a unique solution matching ground truth - if len(valid_arrangements) == 1: - if valid_arrangements[0] == self.ground_truth: - print("Found minimal constraint set!") - return selected_constraints - else: - # Our constraints led to a wrong unique solution - print("Warning: Constraints led to incorrect unique solution.") - # Remove the last constraint and try a different one - selected_constraints.pop() - continue - - # If we have no valid arrangements, the last constraint was incompatible - if len(valid_arrangements) == 0: - print("Warning: No valid arrangements with these constraints.") - # Remove the incompatible constraint and continue - selected_constraints.pop() - continue - - print("Warning: Could not find a minimal constraint set.") - return selected_constraints - - def generate_puzzle(self) -> Tuple[List[Constraint], List[List[int]]]: - """Generate a puzzle with constraints and wrong arrangements.""" - constraints = self.find_minimal_constraints() - return constraints - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Generate ordering puzzles for evaluating reasoning in LLMs") - parser.add_argument("--num_puzzles", type=int, default=10000, help="Number of puzzles to generate") - parser.add_argument("--output_dir", type=str, default="data/raw", help="Output directory for generated puzzles") - parser.add_argument("--output_file", type=str, default="puzzles_dataset.json", help="Output filename for generated puzzles") - parser.add_argument("--test", type=bool, default=False, help="Generate set without solution traces") - args = parser.parse_args() - - birds = ["hawk", "hummingbird", "quail", "owl", "crow", "robin", "cardinal", "sparrow", "bluejay", "pigeon", - "eagle", "falcon", "finch", "woodpecker", "duck", "goose", "swan", "seagull", "pelican", "flamingo", "parrot", - "macaw", "peacock", "penguin", "ostrich", "emu", "vulture", "raven", "magpie", "starling", "warbler", "chickadee", - "swallow", "martin", "oriole", "dove", "pheasant", "turkey", "toucan", "kingfisher", "heron", "egret", "stork", - "kiwi", "albatross", "condor"] - - mammals = ["lion", "tiger", "bear", "wolf", "fox", "deer", "elephant", "giraffe", "zebra", "rhino", - "hippo", "monkey", "gorilla", "chimpanzee", "kangaroo", "koala", "panda", "raccoon", "squirrel", - "rabbit", "mouse", "rat", "bat", "whale", "dolphin", "seal", "walrus", "otter", "beaver", - "hedgehog", "armadillo", "sloth", "opossum", "weasel", "badger", "skunk", "moose", "bison", - "camel", "llama", "alpaca", "horse", "donkey", "cow", "goat", "sheep", "pig"] - - fruits = ["apple", "orange", "banana", "grape", "strawberry", "blueberry", "raspberry", "blackberry", - "watermelon", "cantaloupe", "honeydew", "pineapple", "mango", "papaya", "kiwi", "pear", - "peach", "plum", "apricot", "cherry", "lemon", "lime", "grapefruit", "coconut", "fig", - "date", "guava", "passion fruit", "pomegranate", "dragonfruit", "lychee", "persimmon", - "avocado", "starfruit", "durian", "jackfruit", "nectarine", "tangerine", "clementine"] - - vehicles = ["car", "truck", "bus", "motorcycle", "scooter", "bicycle", "train", "tram", "subway", - "helicopter", "airplane", "jet", "rocket", "spaceship", "submarine", "ship", "yacht", - "sailboat", "speedboat", "canoe", "kayak", "van", "ambulance", "firetruck", "tractor", - "bulldozer", "crane", "forklift", "excavator", "tank", "snowmobile", "golf cart", - "limousine", "taxi", "trolley", "rickshaw", "hovercraft", "blimp", "hot air balloon"] - - instruments = ["guitar", "piano", "violin", "cello", "bass", "harp", "flute", "clarinet", "saxophone", - "trumpet", "trombone", "tuba", "oboe", "bassoon", "piccolo", "recorder", - "harmonica", "accordion", "banjo", "mandolin", "ukulele", "drums", "bongos", "congas", - "tambourine", "xylophone", "marimba", "vibraphone", "glockenspiel", "triangle", "cymbals", - "bagpipes", "sitar", "shamisen", "erhu", "theremin", "synthesizer", "organ", "harpsichord"] - - gemstones = ["diamond", "ruby", "emerald", "sapphire", "amethyst", "topaz", "opal", "pearl", "jade", - "garnet", "aquamarine", "turquoise", "moonstone", "amber", "onyx", "quartz", "citrine", - "peridot", "tanzanite", "agate", "alexandrite", "beryl", "malachite", "obsidian", "jasper", - "zircon", "spinel", "sunstone", "tourmaline", "lapis lazuli", "hematite", "pyrite", - "rhodochrosite", "fluorite", "calcite", "azurite", "chrysoberyl", "sodalite", "carnelian"] - - people = ["Emma", "Liam", "Wei", "Priya", "Mohammed", "Sofia", "Jamal", "Yuki", - "Isabella", "Raj", "Noah", "Ling", "Santiago", "Amara", "Daniel", - "Fatima", "Chen", "Miguel", "Aisha", "Sanjay", "Harper", "Takashi", - "Kwame", "Elena", "Omar", "Zoe", "Hiroshi", "Nia", "Gabriel", "Jin", - "Maya", "Mateo", "Layla", "Arjun", "Clara", "Ahmed", "Mei", "Luis", - "Imani", "Alejandro", "Divya", "Kenji", "Chioma", "Paulo", "Ananya", - "Hassan", "Valentina", "Kofi", "Sakura", "Amir"] - - random.seed(42) - num_puzzles = args.num_puzzles - base_dir = os.path.dirname(os.path.abspath(__file__)) - - # Create the "generated_QAs" folder at the same level as the Python file if it doesn't exist - # output_folder = os.path.join(base_dir, args.output_dir) - output_folder = args.output_dir - if not os.path.exists(output_folder): - os.makedirs(output_folder) - - output_file = os.path.join(output_folder, args.output_file) - - if args.test: - output_file_test = os.path.join(output_folder, args.output_file.replace(".json", "_test.json")) - else: - output_file_test = os.path.join(output_folder, args.output_file.replace(".json", "_unsolved.json")) - - puzzles_data = [] - puzzles_test_data = [] - object_types = { - 'birds': birds, - 'mammals': mammals, - 'fruits': fruits, - 'vehicles': vehicles, - 'instruments': instruments, - 'gemstones': gemstones, - 'people': people - } - for id in range(num_puzzles): - num_objects = random.randint(3, 7) - - type_name = random.choice(list(object_types.keys())) - type_objects = object_types[type_name] - curr_birds = random.sample(type_objects, num_objects) - birds_str = ", ".join(curr_birds) - instruction = f"Solve the following puzzle to determine the order of the {type_name} from left to right. The {type_name} are {birds_str}." - generator = BirdPuzzleGenerator(curr_birds) - - constraints = generator.generate_all_constraints() - constraints = generator.generate_puzzle() - input = "" - input_constraints = [] - for i, constraint in enumerate(constraints, 1): - input_constraints.append(constraint.text) - input += f"{i}. {constraint.text}\n" - if not args.test: - output = gen_proof.solve_puzzle(ground_truth = generator.ground_truth_birds, raw_constraints = input_constraints) - if output: - output = "\n".join(output) - puzzle_data = {'id': id, 'instruction': instruction, 'input': input, 'output': output, 'ground_truth': generator.ground_truth_birds, 'num_objects': num_objects} - puzzles_data.append(puzzle_data) - else: - puzzle_data = {'id': id, 'instruction':instruction, 'input': input, 'ground_truth': generator.ground_truth_birds, 'num_objects': num_objects} - puzzles_test_data.append(puzzle_data) - else: - puzzle_data = {'id': id, 'instruction':instruction, 'input': input, 'ground_truth': generator.ground_truth_birds, 'num_objects': num_objects} - puzzles_test_data.append(puzzle_data) - if args.test: - with open(output_file, 'w') as json_file: - json.dump(puzzles_test_data, json_file, indent=2) - else: - with open(output_file, 'w') as json_file: - json.dump(puzzles_data, json_file, indent=2) - with open(output_file_test, 'w') as json_file_test: - json.dump(puzzles_test_data, json_file_test, indent=2) - -# example usage -# python puzzle_gen.py --num_puzzles 10000 --output_dir data/puzzles_dataset --output_file puzzles_dataset.json --test True - diff --git a/reasoners/data_preprocess/logic/zebrapuzzle_gen/generator.py b/reasoners/data_preprocess/logic/zebrapuzzle_gen/generator.py deleted file mode 100644 index 6c7615701..000000000 --- a/reasoners/data_preprocess/logic/zebrapuzzle_gen/generator.py +++ /dev/null @@ -1,488 +0,0 @@ -import random -import collections -import time -from typing import Literal, List, Set, Tuple, Callable - - -def format_table(header: List[str], table: List[List[str]], - top_format='{:^{}}', left_format=' {:<{}}', cell_format='{:<{}}', - col_delim=' | ', row_delim='\n', prefix_format='|', postfix_format='|'): - table = [[''] + header + ['']] + [row for row in table] - table_format = [[prefix_format + left_format] + len(header) * [top_format]] \ - + len(table) * [[prefix_format + left_format] + len(header) * [cell_format]] - col_widths = [max(len(format.format(cell, 0)) - for format, cell in zip(col_format, col)) - for col_format, col in zip(zip(*table_format), zip(*table))] - return row_delim.join( - col_delim.join( - format.format(cell, width) - for format, cell, width in zip(row_format, row, col_widths)) + f" {postfix_format}" - for row_format, row in zip(table_format, table)) - - -def update_range(wns: List[str], rns: List[List[Set[str]]], cmp: Callable): - changed = False - for rn in rns: - classified_words = set() - for n_col, set_of_words in enumerate(rn): - if len(set_of_words) == 1: - classified_words.add(next(iter(set_of_words))) - word_to_cols = dict() - for n_col, set_of_words in enumerate(rn): - if len(set_of_words) != 1: - prev_length = len(set_of_words) - set_of_words.difference_update(classified_words) - changed |= prev_length != len(set_of_words) - for word in set_of_words: - word_to_cols.setdefault(word, set()).add(n_col) - for word, cols in word_to_cols.items(): - if len(cols) == 1: - x = rn[next(iter(cols))] - if len(x) != 1: - x.clear() - x.add(word) - changed = True - - new_rns = [[{x for x in xs if x != wn} for xs in rn] for wn, rn in zip(wns, rns)] - pairs = [] - for wn, rn in zip(wns, rns): - new_pairs = [] - break_condition = True - for cn, setn in enumerate(rn): - if wn in setn: - break_condition = False - if not pairs: - pairs = [[]] - for v in pairs: - new_pairs.append([*v, cn]) - pairs = new_pairs - if break_condition: - break - for pair in pairs: - if cmp(*pair): - for nrn, cn, wn in zip(new_rns, pair, wns): - nrn[cn].add(wn) - changed |= any(rn != new_rn for rn, new_rn in zip(rns, new_rns)) - if changed: - for rn, new_rn in zip(rns, new_rns): - for old, new in zip(rn, new_rn): - old.intersection_update(new) - return changed - - -def update_ranges(relations: List[Tuple[List[int], List[str], Callable, ...]], - ranges: List[List[Set[str]]]): - changed = False - for ins, wns, callable_object, *_ in relations: - changed |= update_range(wns, [ranges[i] for i in ins], callable_object) - return changed - - -def generate_puzzle(table: List[List[str]], *, - level: Literal[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20], - minimal_conditions: bool = False, max_seconds_for_minimizing: float = None, - tries: int = 1): - if level not in range(1, 20 + 1): - raise ValueError('level must be >= 1 and <= 20') - - table_wo_left = [row[1:] for row in table] - n_attributes = len(table_wo_left) - m_objects = len(table_wo_left[0]) - - if level >= 19 and m_objects == 2: - raise ValueError('too few objects for level >= 19') - elif m_objects <= 1: - raise ValueError('m_objects must be >= 2') - elif n_attributes <= 0: - raise ValueError('n_attributes must be >= 1') - - center = m_objects // 2 - except_flag = True - rules_for_relations = [ - (2, lambda j1, j2: j1 == j2, ['{0}:{1} == {2}:{3}', '{2}:{3} == {0}:{1}']), - (2, lambda j1, j2: j1 == j2 - 1, ['{0}:{1} is on the left of {2}:{3}']), - (2, lambda j1, j2: j1 == j2 + 1, ['{0}:{1} is on the right of {2}:{3}']), - (1, lambda j1: j1 == 0, ['{0}:{1} is on the far left']), - (1, lambda j1, last_index=m_objects - 1: j1 == last_index, ['{0}:{1} is on the far right']), - ] + (m_objects % 2 != 0) * [(1, lambda j1, mid=center: j1 == mid, ['{0}:{1} is in the middle'])] - if level >= 2: - rules_for_relations += [ - (3, lambda j1, j2, j3: j2 + 1 == j1 == j3 - 1 or j3 + 1 == j1 == j2 - 1, - ['{0}:{1} is between {2}:{3} and {4}:{5}', '{0}:{1} is between {4}:{5} and {2}:{3}']), - ] - if level >= 3: - rules_for_relations += [ - (2, lambda j1, j2: j1 == j2 - 1 or j1 == j2 + 1, - ['{0}:{1} is on the left or right of {2}:{3}']), - (1, lambda j1, last_index=m_objects - 1: j1 == 0 or j1 == last_index, - ['{0}:{1} is on the far left or far right']), - ] - if level >= 4: - rules_for_relations += [ - (1, lambda j1: (j1 + 1) % 2 != 0, ['{0}:{1} is in an odd position']), - (1, lambda j1: (j1 + 1) % 2 == 0, ['{0}:{1} is in an even position']), - ] - if level >= 5: - rules_for_relations += [ - (2, lambda j1, j2: j1 < j2, ['{0}:{1} is somewhere to the left of {2}:{3}']), - (2, lambda j1, j2: j1 > j2, ['{0}:{1} is somewhere to the right of {2}:{3}']), - ] - if level >= 6: - rules_for_relations += [ - (2, lambda j1, j2: j1 != j2, ['{0}:{1} != {2}:{3}', '{2}:{3} != {0}:{1}'], except_flag), - ] - if level >= 7: - rules_for_relations += [ - (3, lambda j1, j2, j3: j2 < j1 < j3 or j3 < j1 < j2, - ['{0}:{1} is somewhere between {2}:{3} and {4}:{5}', - '{0}:{1} is somewhere between {4}:{5} and {2}:{3}']), - ] - if level >= 8: - rules_for_relations += [ - (2, lambda j1, j2: j1 >= j2, ['{0}:{1} is not to the left of {2}:{3}']), - (2, lambda j1, j2: j1 <= j2, ['{0}:{1} is not to the right of {2}:{3}']), - ] - if level >= 9: - rules_for_relations += [ - (2, lambda j1, j2: j1 % 2 != j2 % 2, - ['{0}:{1} and {2}:{3} have different parity positions', - '{2}:{3} and {0}:{1} have different parity positions'], except_flag), - (2, lambda j1, j2: j1 % 2 == j2 % 2, - ['{0}:{1} and {2}:{3} have the same parity positions', - '{2}:{3} and {0}:{1} have the same parity positions'], except_flag), - ] - if level >= 10: - rules_for_relations += [ - (3, lambda j1, j2, j3: (j1 == j2 and j1 != j3) or (j1 != j2 and j1 == j3), - ['{0}:{1} == {2}:{3} or {0}:{1} == {4}:{5}, but not both', - '{0}:{1} == {4}:{5} or {0}:{1} == {2}:{3}, but not both'], except_flag), - (3, lambda j1, j2, j3: (j1 == j2 and j2 != j3) or (j1 != j2 and j2 == j3), - ['{0}:{1} == {2}:{3} or {2}:{3} == {4}:{5}, but not both', - '{2}:{3} == {4}:{5} or {0}:{1} == {2}:{3}, but not both'], except_flag), - ] - if level >= 11: - rules_for_relations += [ - (3, lambda j1, j2, j3: j1 == j2 or j1 == j3, - ['{0}:{1} == {2}:{3} or {0}:{1} == {4}:{5} or both', - '{0}:{1} == {4}:{5} or {0}:{1} == {2}:{3} or both'], except_flag), - (3, lambda j1, j2, j3: j1 == j2 or j2 == j3, - ['{0}:{1} == {2}:{3} or {2}:{3} == {4}:{5} or both', - '{2}:{3} == {4}:{5} or {0}:{1} == {2}:{3} or both'], except_flag), - ] - if level >= 12: - rules_for_relations += [ - (3, lambda j1, j2, j3: j1 != j2 or j1 != j3, - ['{0}:{1} != {2}:{3} or {0}:{1} != {4}:{5} or both', - '{0}:{1} != {4}:{5} or {0}:{1} != {2}:{3} or both'], except_flag), - (3, lambda j1, j2, j3: j1 != j2 or j2 != j3, - ['{0}:{1} != {2}:{3} or {2}:{3} != {4}:{5} or both', - '{2}:{3} != {4}:{5} or {0}:{1} != {2}:{3} or both'], except_flag), - ] - if level >= 13: - rules_for_relations.pop(0) # pop '==' - if level >= 14: - rules_for_relations.pop(0) # pop 'is on the left of' - rules_for_relations.pop(0) # pop 'is on the right of' - if level >= 15: - rules_for_relations.pop(0) # pop 'is on the far left' - rules_for_relations.pop(0) # pop 'is on the far right' - if m_objects % 2 != 0: - rules_for_relations.pop(0) # pop 'is in the middle' - if level >= 16: - rules_for_relations.pop(0) # pop 'is between' - if level >= 17: - rules_for_relations.pop(0) # pop 'is on the left or right of' - rules_for_relations.pop(0) # pop 'is on the far left or far right' - if level >= 18: - rules_for_relations.pop(0) # pop 'is in an odd position' - rules_for_relations.pop(0) # pop 'is in an even position' - if level >= 19: - rules_for_relations.pop(0) # pop 'is somewhere to the left of' - rules_for_relations.pop(0) # pop 'is somewhere to the right of' - if level >= 20: - rules_for_relations.pop(0) # pop '!=' - is_minimized = False - time_elapsed = False - min_relations = None - while True: - ranges = [[set(table_wo_left[i]) for _ in range(len(table_wo_left[i]))] for i in range(len(table_wo_left))] - relations = list() - fail = False - while not fail: - needs_clarification = list() - no_solutions = False - solved = True - for i, rng in enumerate(ranges): - for j, rs in enumerate(rng): - if len(rs) == 0: - no_solutions = True - solved = False - break - elif len(rs) > 1: - solved = False - needs_clarification.append((i, j)) - if no_solutions: - break - if solved or min_relations is not None and len(relations) >= len(min_relations): - tries -= 1 - if min_relations is None or len(relations) < len(min_relations): - min_relations = relations - if tries > 0: - fail = True - continue - if tries <= 0: - relations = min_relations - if not minimal_conditions: - break - number_of_relations_min = len(relations) - number_of_relations_before = len(relations) - start_time = time.monotonic() - main_q = collections.deque([relations]) - while main_q: - current_relations = main_q.popleft() - for k in range(len(current_relations)): - new_ranges = [[set(table_wo_left[i]) for _ in range(len(table_wo_left[i]))] - for i in range(len(table_wo_left))] - new_relations = current_relations.copy() - new_relations.pop(k) - changed = True - while changed: - changed = update_ranges(new_relations, new_ranges) - - q = collections.deque([new_ranges]) - possible_solutions = [] - while q: - current_ranges = q.popleft() - - no_solutions = False - solved = True - for rng in current_ranges: - for rs in rng: - if len(rs) == 0: - no_solutions = True - solved = False - break - elif len(rs) > 1: - solved = False - if no_solutions or not solved: - break - if no_solutions: - continue - if solved: - if current_ranges not in possible_solutions: - possible_solutions.append(current_ranges) - if len(possible_solutions) >= 2: - break - continue - - for n_group, rng in enumerate(current_ranges): - founded = False - for n_x, rs in enumerate(rng): - if len(rs) > 1: - founded = True - for r in rs: - new_ranges = [[x.copy() for x in row] for row in current_ranges] - new_ranges[n_group][n_x] = {r} - changed = True - while changed: - changed = update_ranges(new_relations, new_ranges) - q.appendleft(new_ranges) - break - if founded: - break - if len(possible_solutions) == 1: - number_of_relations_after = len(new_relations) - if number_of_relations_min > number_of_relations_after: - number_of_relations_min = number_of_relations_after - relations = new_relations - main_q.append(new_relations) - if max_seconds_for_minimizing is not None and \ - time.monotonic() >= start_time + max_seconds_for_minimizing: - time_elapsed = True - break - if time_elapsed: - break - is_minimized = number_of_relations_min < number_of_relations_before or not time_elapsed - break - if no_solutions or not needs_clarification: - fail = True - continue - - i, j = item = random.choice(needs_clarification) - next2_i, next2_j = None, None - if level >= 2 and len(needs_clarification) > 1: - needs_clarification.remove(item) - next2_i, next2_j = random.choice(needs_clarification) - - neighbours = [] - right_neighbours = [] - for dj in range(-1, 1 + 1): - if not (0 <= j + dj < m_objects): - continue - for new_i in range(0, n_attributes): - if new_i == i and dj == 0: - continue - new_item = (new_i, j + dj) - neighbours.append(new_item) - if level >= 2 and dj == 1: - right_neighbours.append(new_item) - if not neighbours: - continue - next_i, next_j = random.choice(neighbours) - if level >= 2 and next2_i is None and right_neighbours: - next2_i, next2_j = random.choice(right_neighbours) - - permutations3 = [ - ((i, j), (next_i, next_j), (next2_i, next2_j)), ((i, j), (next2_i, next2_j), (next_i, next_j)), - ((next_i, next_j), (i, j), (next2_i, next2_j)), ((next_i, next_j), (next2_i, next2_j), (i, j)), - ((next2_i, next2_j), (i, j), (next_i, next_j)), ((next2_i, next2_j), (next_i, next_j), (i, j)) - ] if next2_i is not None else [] - permutations2 = [ - ((i, j), (next_i, next_j)), ((next_i, next_j), (next2_i, next2_j)), ((i, j), (next2_i, next2_j)), - ((next_i, next_j), (i, j)), ((next2_i, next2_j), (next_i, next_j)), ((next2_i, next2_j), (i, j)), - ] if next2_i is not None else [ - ((i, j), (next_i, next_j)), ((next_i, next_j), (i, j)) - ] - possible_variants = [] - for (n_args, cmp_function, str_variants, *flags) in rules_for_relations: - if n_args == 3: - for items in permutations3: - (ti, tj), (t_next_i, t_next_j), (t_next2_i, t_next2_j) = items - if flags and flags[0] and (ti == t_next_i or ti == t_next2_i or t_next_i == t_next2_i): - continue - if cmp_function(tj, t_next_j, t_next2_j): - possible_variants.append((n_args, items, cmp_function, random.choice(str_variants))) - elif n_args == 2: - for items in permutations2: - (ti, tj), (t_next_i, t_next_j) = items - if flags and flags[0] and ti == t_next_i: - continue - if cmp_function(tj, t_next_j): - possible_variants.append((n_args, items, cmp_function, random.choice(str_variants))) - elif n_args == 1 and cmp_function(j): - possible_variants.append((n_args, [(i, j)], cmp_function, random.choice(str_variants))) - if not possible_variants: - continue - - n_args, list_of_ij, cmp_function, string_format = random.choice(possible_variants) - list_for_format = [] - ins, wns = [], [] - for i, j in list_of_ij: - list_for_format.extend([table[i][0], table_wo_left[i][j]]) - ins.append(i) - wns.append(table_wo_left[i][j]) - relations.append((ins, wns, cmp_function, string_format.format(*list_for_format))) - - changed = True - while changed: - changed = update_ranges(relations, ranges) - - if not fail: - if minimal_conditions and not is_minimized and not time_elapsed: - continue - break - - premises = [t[-1] for t in relations] - random.shuffle(premises) - return premises - -def table_to_dict(table: List[List[str]]) -> dict: - res = {} - number_of_rows = len(table[0])-1 - res["header"] = ["Position"] - res["rows"] = [[str(i+1)] for i in range(number_of_rows)] - for i, row in enumerate(table): - res['header'].append(row[0]) - for j in range(1, len(row)): - res['rows'][j-1].append(row[j]) - return res - - - -def generate_table(n_attributes: int, m_objects: int) -> List[List[str]]: - kinds_dict = { - "Nationality": { - "american", "argentine", "australian", "brazilian", "british", - "canadian", "chinese", "colombian", "dutch", "egyptian", - "french", "german", "indian", "indonesian", "italian", - "japanese", "malaysian", "mexican", "nigerian", "pakistani", - "polish", "russian", "spanish", "thai", "turkish", - }, - "Food": { - "apple", "apricot", "artichoke", "asparagus", "avocado", - "banana", "blueberry", "broccoli", "cabbage", "carrot", - "cauliflower", "cherry", "corn", "cranberry", "cucumber", - "eggplant", "garlic", "grapefruit", "grapes", "kale", - "kiwi", "lemon", "lettuce", "lime", "mango", - "nectarine", "onion", "orange", "papaya", "peach", - "pear", "peas", "pepper", "pineapple", "plum", - "pomegranate", "potato", "pumpkin", "radish", "raspberry", - "spinach", "strawberry", "tomato", "watermelon", "zucchini", - }, - "Pet": { - "bird", "cat", "chinchilla", "dog", "ferret", - "fish", "frog", "goat", "goldfish", "guinea-pig", - "hamster", "hedgehog", "horse", "lizard", "mouse", - "pony", "rabbit", "rat", "snake", "turtle", - }, - "Job": { - "accountant", "analyst", "architect", "bartender", "chef", - "coach", "dancer", "designer", "doctor", "dressmaker", - "electrician", "engineer", "entrepreneur", "firefighter", "fisherman", - "freelancer", "journalist", "lawyer", "librarian", "manager", - "mechanic", "musician", "nurse", "paramedic", "photographer", - "pilot", "police-officer", "project-manager", "scientist", "security-guard", - "social-worker", "software-developer", "teacher", "videographer", "writer", - }, - "Beverage": { - "7up", "almond-milk", "coffee", "cola", "fanta", - "hot-chocolate", "iced-tea", "juice", "lemonade", "milk", - "mirinda", "soy-milk", "sprite", "tea", "water", - }, - "Transport": { - "airplane", "bike", "boat", "bus", "car", - "helicopter", "jet-ski", "motorbike", "quad-bike", "roller", - "scooter", "ship", "skateboard", "snowmobile", - "subway", "taxi", "train", "tram", "trike", "van", - }, - "Music-Genre": { - "ambient", "blues", "classical", "country", "d&b", - "disco", "dubstep", "electronic", "folk", "funk", - "gospel", "hip-hop", "house", "indie", "jazz", - "metal", "pop", "punk", "r&b", "reggae", - "rock", "salsa", "soul", "techno", "trance", - }, - "Movie-Genre": { - "action", "adventure", "animation", "comedy", "crime", - "disaster", "documentary", "drama", "epic", "family", - "fantasy", "horror", "martial-arts", "musical", "mystery", - "romance", "satire", "scientific", "sports", "spy", - "superhero", "thriller", "time-travel", "western", "zombie", - }, - "Sport": { - "badminton", "baseball", "basketball", "biathlon", "climbing", - "cricket", "cycling", "golf", "handball", "ice-hockey", - "lacrosse", "parkour", "rowing", "rugby", "sailing", - "skateboarding", "skiing", "snowboarding", "soccer", "surfing", - "swimming", "tennis", "volleyball", "water-polo", "weightlifting", - }, - "Hobby": { - "baking", "board-games", "camping", "card-games", "chess", - "collecting", "cooking", "dancing", "drawing", "filmmaking", - "fishing", "gardening", "hiking", "magic-tricks", "photography", - "puzzles", "reading", "rock-climbing", "singing", "skydiving", - "sudoku", "traveling", "video-games", "woodworking", "writing", - } - } - kinds = sorted(kinds_dict) - - # Check - assert n_attributes <= len(kinds_dict),\ - f'Not enough attributes: actual {len(kinds_dict)}, expected {n_attributes}' - assert all(m_objects <= len(v) for k, v in kinds_dict.items()), 'Not enough objects: ' +\ - f'actual {next(f"{k}={len(v)}" for k, v in kinds_dict.items() if m_objects > len(v))}, expected {m_objects}' - - chosen_kinds = sorted(random.sample(kinds, k=n_attributes)) - table = [[kind] + random.sample(sorted(kinds_dict[kind]), k=m_objects) for kind in chosen_kinds] - return table - diff --git a/reasoners/data_preprocess/logic/zebrapuzzle_gen/puzzle_generator.py b/reasoners/data_preprocess/logic/zebrapuzzle_gen/puzzle_generator.py deleted file mode 100644 index 1c47964d3..000000000 --- a/reasoners/data_preprocess/logic/zebrapuzzle_gen/puzzle_generator.py +++ /dev/null @@ -1,202 +0,0 @@ -import threading -import queue -import time -import random -from typing import List, Literal -import concurrent.futures -from generator import generate_puzzle, generate_table, table_to_dict -import multiprocessing -import argparse - - -# diverse set of configs, rows = num_attributes, cols = num_objects -# The level is a number from 1 to 20, where 1 is the easiest and 20 is the hardest -# minimal_conditions is a boolean indicating whether to use minimal conditions -PUZZLE_CONFIGS = [ - {"rows": 2, "cols": 2, "level": 5, "minimal_conditions": True}, - {"rows": 2, "cols": 4, "level": 12, "minimal_conditions": True}, - {"rows": 2, "cols": 5, "level": 3, "minimal_conditions": False}, - {"rows": 2, "cols": 6, "level": 15, "minimal_conditions": True}, - {"rows": 2, "cols": 8, "level": 7, "minimal_conditions": False}, - {"rows": 2, "cols": 10, "level": 18, "minimal_conditions": True}, - {"rows": 2, "cols": 12, "level": 10, "minimal_conditions": False}, - {"rows": 2, "cols": 15, "level": 20, "minimal_conditions": True}, - - {"rows": 3, "cols": 2, "level": 14, "minimal_conditions": False}, - {"rows": 3, "cols": 3, "level": 6, "minimal_conditions": True}, - {"rows": 3, "cols": 5, "level": 19, "minimal_conditions": False}, - {"rows": 3, "cols": 7, "level": 11, "minimal_conditions": True}, - {"rows": 3, "cols": 9, "level": 4, "minimal_conditions": False}, - {"rows": 3, "cols": 12, "level": 16, "minimal_conditions": True}, - {"rows": 3, "cols": 15, "level": 9, "minimal_conditions": False}, - - {"rows": 4, "cols": 2, "level": 17, "minimal_conditions": True}, - {"rows": 4, "cols": 4, "level": 1, "minimal_conditions": False}, - {"rows": 4, "cols": 6, "level": 13, "minimal_conditions": True}, - {"rows": 4, "cols": 8, "level": 20, "minimal_conditions": False}, - {"rows": 4, "cols": 10, "level": 5, "minimal_conditions": True}, - {"rows": 4, "cols": 13, "level": 15, "minimal_conditions": False}, - - {"rows": 5, "cols": 3, "level": 8, "minimal_conditions": True}, - {"rows": 5, "cols": 5, "level": 18, "minimal_conditions": False}, - {"rows": 5, "cols": 7, "level": 3, "minimal_conditions": True}, - {"rows": 5, "cols": 9, "level": 12, "minimal_conditions": False}, - {"rows": 5, "cols": 11, "level": 7, "minimal_conditions": True}, - {"rows": 5, "cols": 14, "level": 19, "minimal_conditions": False}, - - {"rows": 6, "cols": 2, "level": 10, "minimal_conditions": False}, - {"rows": 6, "cols": 4, "level": 6, "minimal_conditions": True}, - {"rows": 6, "cols": 6, "level": 16, "minimal_conditions": False}, - {"rows": 6, "cols": 8, "level": 2, "minimal_conditions": True}, - {"rows": 6, "cols": 10, "level": 14, "minimal_conditions": False}, - {"rows": 6, "cols": 12, "level": 9, "minimal_conditions": True}, - {"rows": 6, "cols": 15, "level": 4, "minimal_conditions": False}, - - {"rows": 7, "cols": 3, "level": 11, "minimal_conditions": True}, - {"rows": 7, "cols": 5, "level": 5, "minimal_conditions": False}, - {"rows": 7, "cols": 7, "level": 17, "minimal_conditions": True}, - {"rows": 7, "cols": 9, "level": 8, "minimal_conditions": False}, - {"rows": 7, "cols": 11, "level": 20, "minimal_conditions": True}, - {"rows": 7, "cols": 13, "level": 3, "minimal_conditions": False}, - - {"rows": 8, "cols": 2, "level": 15, "minimal_conditions": True}, - {"rows": 8, "cols": 4, "level": 7, "minimal_conditions": False}, - {"rows": 8, "cols": 6, "level": 19, "minimal_conditions": True}, - {"rows": 8, "cols": 8, "level": 10, "minimal_conditions": False}, - {"rows": 8, "cols": 10, "level": 1, "minimal_conditions": True}, - {"rows": 8, "cols": 12, "level": 13, "minimal_conditions": False}, - {"rows": 8, "cols": 14, "level": 6, "minimal_conditions": True}, - - {"rows": 9, "cols": 3, "level": 18, "minimal_conditions": False}, - {"rows": 9, "cols": 5, "level": 4, "minimal_conditions": True}, - {"rows": 9, "cols": 7, "level": 16, "minimal_conditions": False}, - {"rows": 9, "cols": 9, "level": 9, "minimal_conditions": True}, - {"rows": 9, "cols": 11, "level": 2, "minimal_conditions": False}, - {"rows": 9, "cols": 13, "level": 14, "minimal_conditions": True}, - - {"rows": 10, "cols": 2, "level": 8, "minimal_conditions": False}, - {"rows": 10, "cols": 4, "level": 20, "minimal_conditions": True}, - {"rows": 10, "cols": 6, "level": 5, "minimal_conditions": False}, - {"rows": 10, "cols": 8, "level": 17, "minimal_conditions": True}, - {"rows": 10, "cols": 10, "level": 10, "minimal_conditions": False}, - {"rows": 10, "cols": 12, "level": 3, "minimal_conditions": True}, - {"rows": 10, "cols": 14, "level": 18, "minimal_conditions": False}, - {"rows": 10, "cols": 15, "level": 20, "minimal_conditions": True}, -] - -INSTRUCTION = """Solve the following puzzle where you are given {{n_objects}} objects and their attributes like job, nationality etc. -The list of attributes and their unique values are: {{table}}. The goal is to determine the correct attribute for each object based on the provided clues.""" - - -def worker_function(config_and_id): - config, worker_id = config_and_id - rows, cols = config["rows"], config["cols"] - level = config["level"] - minimal_conditions = config["minimal_conditions"] - instance_id = config.get("instance_id", 0) - - print(f"Worker {worker_id}: Starting puzzle generation with {rows}x{cols} table, level {level}, instance {instance_id}") - - # Create a table for the puzzle - table = generate_table(rows, cols) - - try: - # Call the generate_puzzle function with our table - premises = generate_puzzle( - table, - level=level, - minimal_conditions=minimal_conditions, - max_seconds_for_minimizing=30.0, - tries=3 - ) - - ground_truth = table_to_dict(table) - # zebralogic convention and ground truth in format of n_objects x n_attributes - puzzle_key = f"puzzle_{config['cols']}x{config['rows']}_level{config['level']}_instance{instance_id}" - list_attributes = "" - for row in table: - list_attributes += f"\n{row[0]}: " + ', '.join(sorted(row[1:])) - instruction = INSTRUCTION.replace("{{table}}", list_attributes) - instruction = instruction.replace("{{n_objects}}", str(cols)) - result = { - "worker_id": worker_id, - "puzzle_id": puzzle_key, - "config": config, - "instruction": instruction, - "clues": premises, - "ground_truth": ground_truth, - } - print(f"Worker {worker_id}: Successfully generated puzzle") - return result - - except Exception as e: - print(f"Worker {worker_id}: Error generating puzzle: {e}") - return None - -def run_parallel_generation(total_puzzles = None, num_processes=None): - """Run puzzle generation in parallel using multiple threads""" - num_configs = len(PUZZLE_CONFIGS) - - if total_puzzles is None: - total_puzzles = num_configs - - puzzles_per_config = total_puzzles // num_configs - - extra_puzzles = total_puzzles % num_configs - expanded_configs = [] - for i, config in enumerate(PUZZLE_CONFIGS): - repeats = puzzles_per_config + (1 if i < extra_puzzles else 0) - for j in range(repeats): - # Create a copy with a unique instance ID - config_copy = config.copy() - config_copy['instance_id'] = j - expanded_configs.append(config_copy) - - if num_processes is None: - num_processes = min(len(expanded_configs), - concurrent.futures.cpu_count()) - - print(f"Generating {total_puzzles} puzzles across {num_configs} configurations using {num_processes} processes") - - worker_args = [(config, i) for i, config in enumerate(expanded_configs)] - - with concurrent.futures.ProcessPoolExecutor(max_workers=num_processes) as executor: - results = list(executor.map(worker_function, worker_args)) - - # Collect all results from the queue - results = [r for r in results if r is not None] - - return results - -def save_puzzles(results, output_dir="./data/raw", output_file="zebra_puzzles.json"): - import os - import json - - # Create output directory if it doesn't exist - os.makedirs(output_dir, exist_ok=True) - - output_path = os.path.join(output_dir, output_file) - with open(output_path, "w") as f: - json.dump(results, f, indent=4) - -if __name__ == "__main__": - # Parse command-line arguments - parser = argparse.ArgumentParser(description="Generate puzzles in parallel.") - parser.add_argument("--num_puzzles", type=int, default=100, help="Total number of puzzles to generate") - parser.add_argument("--num_processes", type=int, default=16, help="Number of processes to use for generation") - parser.add_argument("--output_dir", type=str, default="data/raw", help="Directory to save generated puzzles") - parser.add_argument("--output_file", type=str, default="zebra_puzzles.json", help="Filename for the generated puzzles") - args = parser.parse_args() - multiprocessing.freeze_support() - start_time = time.time() - print("Starting parallel puzzle generation...") - - # Run parallel generation with 4 threads (or adjust as needed) - results = run_parallel_generation(total_puzzles = args.num_puzzles, num_processes=args.num_processes) - - # Save the generated puzzles - save_puzzles(results, output_dir=args.output_dir, output_file=args.output_file) - - end_time = time.time() - print(f"Completed in {end_time - start_time:.2f} seconds") - print(f"Generated {len(results)} puzzles") \ No newline at end of file diff --git a/reasoners/data_preprocess/math/bigmath_preview_filtered_mar21.py b/reasoners/data_preprocess/math/bigmath_preview_filtered_mar21.py deleted file mode 100644 index a5a2a6956..000000000 --- a/reasoners/data_preprocess/math/bigmath_preview_filtered_mar21.py +++ /dev/null @@ -1,159 +0,0 @@ -"""Downloads, processes, and saves math datasets.""" - -import os -import datasets -from typing import Dict, List, Optional, Any, Union -import enum -import argparse -import pandas as pd -import json -import random -import transformers - -from verl.utils.reward_score.math import remove_boxed, last_boxed_only_string -from verl.utils.data_process.prompt import build_zero_style_prompt -from verl.utils.data_process.utils import add_suffix, set_seed, sample_dataset, save_dataset - -def get_datasets(cache_dir: str): - """ - Load the math datasets from Hugging Face Hub. - - Args: - train_data_source (str): Source for training data - test_data_sources (list): List of sources for test data - - Returns: - tuple: (train_dataset, test_datasets) as Dataset objects - """ - train_data_source = "SDSB/big_math_partial_mar21_filtered_basic" - test_data_sources = [ # hard-coded for now as only bigmath handles all the test datasets - "nanoverl/minerva", - "SDSB/aime_repeated_8x", - "SDSB/amc_repeated_4x", - "nanoverl/olympiad_bench", - "nanoverl/math", - ] - - print(f"Loading the {train_data_source} dataset...") - train_dataset = datasets.load_dataset( - train_data_source, trust_remote_code=True, split="train", cache_dir=cache_dir - ) - - print(f"Loading the test datasets...") - test_datasets = [ - datasets.load_dataset(test_data_source, trust_remote_code=True, split="test", cache_dir=cache_dir) - for test_data_source in test_data_sources - ] - - return train_dataset, test_datasets - - -def make_map_fn(split: str, data_source: str, prompt_style: str="zero_style", reward_metric: str=None) -> callable: - def process_fn(example: Dict[str, Any], idx: int) -> Dict[str, Any]: - question = example.pop("problem") - answer = example.pop("answer") - data = { - "data_source": data_source, - "prompt": [ - {"role": "user", "content": question + " Please output the final answer within \\boxed{}."}, - ], - "ability": "math", - "apply_chat_template": True, - "reward_model": {"style": "rule", "ground_truth": answer}, - "extra_info": { - "split": split, - "index": idx, - "reward_metric": reward_metric, - "original_question": question, - }, - } - - if idx == 0 or idx == 1: - print(f"data_source: {data_source}, split: {split}, idx: {idx}") - print("\n" + "=" * 100 + f"{data_source} {split} {idx}" + "=" * 10) - print(data) - return data - - return process_fn - - -if __name__ == "__main__": - """Main script execution: parse args, load, process, and save datasets.""" - parser = argparse.ArgumentParser() - parser.add_argument( - "--data-dir", - default="data", - help="Directory to save the processed data files. Will be modified based on other parameters.", - ) - parser.add_argument( - "--domain", - default="math", - type=str, - help="Data domain", - ) - parser.add_argument( - "--name", - default="bigmath_preview_filtered_mar21", - type=str, - help="Dataset name", - ) - parser.add_argument( - "--train-sample-size", - type=int, - default=None, - help="Number of samples to use from the training dataset. If None, use all samples.", - ) - parser.add_argument( - "--train-reward-metric", - type=str, - default=None, - help="Reward metric to use for training. If None, use the naive_dapo.compute_score.", - ) - parser.add_argument( - "--test-reward-metric", - type=str, - default=None, - help="Reward metric to use for testing. If None, use the naive_dapo.compute_score.", - ) - parser.add_argument( - "--suffix", - type=str, - default=None, - help="Suffix to add to the dataset name.", - ) - parser.add_argument( - "--seed", - type=int, - default=42, - help="Random seed for reproducibility when sampling data.", - ) - args = parser.parse_args() - set_seed(args.seed) - - # Download the datasets from Hugging Face Hub - cache_dir = datasets.config.HF_DATASETS_CACHE - train_dataset, _ = get_datasets(cache_dir) - data_source = f"{args.domain}__{args.name}" - - # Process train dataset - process_train_fn = make_map_fn("train", data_source, args.train_reward_metric) - train_data = train_dataset.map(function=process_train_fn, with_indices=True) - # Sample - train_data = sample_dataset(train_data, args.train_sample_size) - # Save - train_output_dir = os.path.join(args.data_dir, "train") - args.train_sample_size = len(train_dataset) if args.train_sample_size is None else args.train_sample_size - train_output_path = save_dataset( - dataset=train_data, - output_dir=train_output_dir, - filename_prefix=data_source + args.suffix if args.suffix else data_source, - sample_size=args.train_sample_size - ) - - print(f"Done! \n" - f"Train data saved to {train_output_path}") - - # python bigmath_preview_filtered_mar21.py - - # with llm judge - # python bigmath_preview_filtered_mar21.py --train-reward-metric math_llm_judge --suffix _llm_judge \ No newline at end of file diff --git a/reasoners/data_preprocess/math/dapo_or1_merge_dedup_apr30.py b/reasoners/data_preprocess/math/dapo_or1_merge_dedup_apr30.py deleted file mode 100644 index 4adc216d7..000000000 --- a/reasoners/data_preprocess/math/dapo_or1_merge_dedup_apr30.py +++ /dev/null @@ -1,183 +0,0 @@ -"""Downloads, processes, and saves math datasets.""" - -import os -import datasets -from typing import Dict, List, Optional, Any, Union -import enum -import argparse -import pandas as pd -import json -import random -import transformers - -from verl.utils.reward_score.math import remove_boxed, last_boxed_only_string -from verl.utils.data_process.prompt import build_zero_style_prompt -from verl.utils.data_process.utils import add_suffix, set_seed, sample_dataset, save_dataset - -def get_datasets(cache_dir: str): - """ - Load the math datasets from Hugging Face Hub. - - Args: - train_data_source (str): Source for training data - test_data_sources (list): List of sources for test data - - Returns: - tuple: (train_dataset, test_datasets) as Dataset objects - """ - train_data_source = "SDSB/merged_deduped_dapo_or1_dataset" - test_data_sources = [ - "nanoverl/minerva", - "SDSB/aime_repeated_8x", - "SDSB/amc_repeated_4x", - "nanoverl/olympiad_bench", - "nanoverl/math", - "nanoverl/aime2025_repeated_8x", - ] - - print(f"Loading the {train_data_source} dataset...") - train_dataset = datasets.load_dataset( - train_data_source, trust_remote_code=True, split="train", cache_dir=cache_dir - ) - - print(f"Loading the test datasets...") - test_datasets = { - os.path.basename(test_data_source.lower()): - datasets.load_dataset(test_data_source, trust_remote_code=True, split="test", cache_dir=cache_dir) - for test_data_source in test_data_sources - } - - return train_dataset, test_datasets - - -def make_map_fn(split: str, data_source: str, reward_metric: str='default') -> callable: - def process_fn(example: Dict[str, Any], idx: int) -> Dict[str, Any]: - # if original_question in example, use it as the question - question = example.pop("problem") - # question = example.pop("problem") - answer = example.pop("answer") - if isinstance(answer, list): - answer = answer[0] - data = { - "data_source": data_source, - "prompt": [ - {"role": "user", "content": question + " Please output the final answer within \\boxed{}."}, - ], - "ability": "math", - "apply_chat_template": True, - "reward_model": {"style": "rule", "ground_truth": answer}, - "extra_info": { - "split": split, - "index": idx, - "reward_metric": reward_metric, - "original_question": question, - }, - } - - if idx == 0 or idx == 1: - print(f"data_source: {data_source}, split: {split}, idx: {idx}") - print("\n" + "=" * 100 + f"{data_source} {split} {idx}" + "=" * 10) - print(data) - return data - - return process_fn - - -if __name__ == "__main__": - """Main script execution: parse args, load, process, and save datasets.""" - parser = argparse.ArgumentParser() - parser.add_argument( - "--data-dir", - default="data", - help="Directory to save the processed data files. Will be modified based on other parameters.", - ) - parser.add_argument( - "--domain", - default="math", - type=str, - help="Data domain", - ) - parser.add_argument( - "--name", - default="merged_deduped_dapo_or1_dataset", - type=str, - help="Dataset name", - ) - parser.add_argument( - "--train-sample-size", - type=int, - default=None, - help="Number of samples to use from the training dataset. If None, use all samples.", - ) - parser.add_argument( - "--train-reward-metric", - type=str, - default="default", - help="Reward metric to use for training. If None, use the naive_dapo.compute_score.", - ) - parser.add_argument( - "--test-reward-metric", - type=str, - default="default", - help="Reward metric to use for testing. If None, use the naive_dapo.compute_score.", - ) - parser.add_argument( - "--suffix", - type=str, - default=None, - help="Suffix to add to the dataset name.", - ) - parser.add_argument( - "--seed", - type=int, - default=42, - help="Random seed for reproducibility when sampling data.", - ) - args = parser.parse_args() - set_seed(args.seed) - - # Download the datasets from Hugging Face Hub - cache_dir = datasets.config.HF_DATASETS_CACHE - train_dataset, test_datasets = get_datasets(cache_dir) - data_source = f"{args.domain}__{args.name}" - - # Process train dataset - process_train_fn = make_map_fn("train", data_source, args.train_reward_metric) - train_data = train_dataset.map(function=process_train_fn, with_indices=True) - # Sample - train_data = sample_dataset(train_data, args.train_sample_size) - # Save - train_output_dir = os.path.join(args.data_dir, "train") - args.train_sample_size = len(train_dataset) if args.train_sample_size is None else args.train_sample_size - train_output_path = save_dataset( - dataset=train_data, - output_dir=train_output_dir, - filename_prefix=data_source + args.suffix if args.suffix else data_source, - sample_size=args.train_sample_size - ) - - # Process test datasets - test_output_dir = os.path.join(args.data_dir, "test") - test_output_paths = [] - for test_data_source, test_data in test_datasets.items(): - test_data_source = f"{args.domain}__{test_data_source}" - process_fn = make_map_fn("test", test_data_source, args.test_reward_metric) - test_data = test_data.map(process_fn, with_indices=True) - test_output_path = save_dataset( - dataset=test_data, - output_dir=test_output_dir, - filename_prefix=test_data_source, - sample_size=None - ) - test_output_paths.append(test_output_path) - print(f"test_data_source: {test_data_source}") - print(f"Test data saved to {test_output_path}") - - print(f"Done! \n" - f"Train data saved to {train_output_path}\n" - f"Test data saved to {test_output_paths}") - - # python dapo_or1_merge_dedup_apr30.py - - # with llm judge - # python dapo_or1_merge_dedup_apr30.py --train-reward-metric math_llm_judge --suffix _llm_judge \ No newline at end of file diff --git a/reasoners/data_preprocess/math/deepscaler_preview.py b/reasoners/data_preprocess/math/deepscaler_preview.py deleted file mode 100644 index ef37ed0bd..000000000 --- a/reasoners/data_preprocess/math/deepscaler_preview.py +++ /dev/null @@ -1,143 +0,0 @@ -# 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. -""" -Preprocess the https://huggingface.co/datasets/agentica-org/DeepScaleR-Preview-Dataset to parquet format -Code referring to https://github.com/agentica-project/deepscaler/blob/main/scripts/data/deepscaler_dataset.py -""" - -import os -import datasets -from typing import Dict, Any -import argparse - -from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset -from verl.utils.data_process.prompt import build_zero_style_prompt -from verl.utils.reward_score.math import remove_boxed, last_boxed_only_string - -def get_datasets(cache_dir: str): - """ - Load the math datasets from Hugging Face Hub. - - Args: - train_data_source (str): Source for training data - test_data_sources (list): List of sources for test data - - Returns: - tuple: (train_dataset, test_datasets) as Dataset objects - """ - train_data_source = "agentica-org/DeepScaleR-Preview-Dataset" - print(f"Loading the {train_data_source} dataset...") - train_dataset = datasets.load_dataset( - train_data_source, trust_remote_code=True, split="train", cache_dir=cache_dir - ) - - return train_dataset, None - - -def make_map_fn(split: str, data_source: str) -> callable: - """ - Create a mapping function for processing dataset examples. - - Args: - split (str): Dataset split ('train' or 'test') - data_source (str): Source of the data - - Returns: - callable: Function to process individual examples - """ - def process_fn(example: Dict[str, Any], idx: int) -> Dict[str, Any]: - question = example.pop("problem") - answer = example.pop("answer") - - if isinstance(answer, list): - answer = answer[0] # minerva, olympiad_bench - - data = { - "data_source": data_source, - "prompt": [ - {"role": "user", "content": question + " Please output the final answer within \\boxed{}."}, - ], - "ability": "math", - "apply_chat_template": True, - "reward_model": {"style": "rule", "ground_truth": answer}, - "extra_info": {"split": split, "index": idx, "original_question": question}, - } - if idx == 0 or idx == 1: - print(f"data_source: {data_source}, split: {split}, idx: {idx}") - print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) - print(data) - return data - - return process_fn - - -if __name__ == "__main__": - """Main script execution: parse args, load, process, and save datasets.""" - parser = argparse.ArgumentParser() - parser.add_argument( - "--data-dir", - default="data", - help="Directory to save the processed data files. Will be modified based on other parameters.", - ) - parser.add_argument( - "--domain", - default="math", - type=str, - help="Data domain", - ) - parser.add_argument( - "--name", - default="deepscaler_preview", - type=str, - help="Dataset name", - ) - parser.add_argument( - "--train-sample-size", - type=int, - default=None, - help="Number of samples to use from the training dataset. If None, use all samples.", - ) - parser.add_argument( - "--seed", - type=int, - default=42, - help="Random seed for reproducibility when sampling data.", - ) - args = parser.parse_args() - set_seed(args.seed) - data_source = f"{args.domain}__{args.name}" - - # Configure data sources - cache_dir = datasets.config.HF_DATASETS_CACHE - # Download the datasets from Hugging Face Hub - train_dataset, test_datasets = get_datasets(cache_dir) - - # Process train dataset - process_train_fn = make_map_fn("train", data_source) - train_data = train_dataset.map(function=process_train_fn, with_indices=True) - - # Sample - train_data = sample_dataset(train_data, args.train_sample_size) - - # Save - train_output_dir = os.path.join(args.data_dir, "train") - args.train_sample_size = len(train_dataset) if args.train_sample_size is None else args.train_sample_size - train_output_path = save_dataset( - dataset=train_data, - output_dir=train_output_dir, - filename_prefix=data_source, - sample_size=args.train_sample_size - ) - print(f"Done! \n" - f"Train data saved to {train_output_path}") \ No newline at end of file diff --git a/reasoners/data_preprocess/norm_json.py b/reasoners/data_preprocess/norm_json.py deleted file mode 100644 index fd43c2902..000000000 --- a/reasoners/data_preprocess/norm_json.py +++ /dev/null @@ -1,203 +0,0 @@ -#!/usr/bin/env python3 -""" -Normalise every .json file in a directory so that each record contains -only the allowed keys; anything else is moved under `extra_info` as a -compressed Base64‐encoded JSON string split into 10 parts. - -If a record already has `extra_info` or `reward_model` from a previous run, -we first de-serialise, merge with any new stray keys, then re-serialise, -compress, and Base64-encode. - -Canonical keys: - ability - apply_chat_template - data_source - extra_info (dict → JSON → gzip → Base64 split into parts) - prompt - qwen2.5_7b_pass_rate - qwen3_30b_pass_rate - reward_model (dict → JSON → gzip → Base64) -""" - -from pathlib import Path -from typing import Any, Dict, List, Union -import json -import gzip -import base64 -import sys -import argparse -# --------------------------------------------------------------------------- -# Configuration -# --------------------------------------------------------------------------- - -KEEP_KEYS = { - "ability", - "apply_chat_template", - "data_source", - "extra_info", # handled specially - "prompt", - "qwen2.5_7b_pass_rate", - "qwen3_30b_pass_rate", - "reward_model", -} - -Json = Union[Dict[str, Any], List[Any]] - -# --------------------------------------------------------------------------- -# Serialisation & Compression helpers -# --------------------------------------------------------------------------- - -def _serialise_extra(extra: Dict[str, Any]) -> str: - """Return a compact JSON string for `extra`.""" - return json.dumps(extra, ensure_ascii=False, separators=(',', ':')) - -def _deserialise_extra(extra_str: str) -> Dict[str, Any]: - """Load JSON string if valid, else return empty dict.""" - try: - val = json.loads(extra_str) - return val if isinstance(val, dict) else {} - except Exception: - return {} - -def _compress_str(s: str) -> str: - """Gzip-compress then Base64-encode a JSON string.""" - gz = gzip.compress(s.encode('utf-8')) - return base64.b64encode(gz).decode('ascii') - -def _decompress_str(b64: str) -> str: - """Base64-decode then gzip-decompress to original JSON string.""" - gz = base64.b64decode(b64.encode('ascii')) - return gzip.decompress(gz).decode('utf-8') - -# --------------------------------------------------------------------------- -# JSON/NDJSON loader -# --------------------------------------------------------------------------- - -def load_json_or_ndjson(path: Path) -> Json: - text = path.read_text(encoding='utf-8').strip() - try: - return json.loads(text) - except json.JSONDecodeError: - pass - records: List[Dict[str, Any]] = [] - for line in text.splitlines(): - ln = line.strip() - if ln: - records.append(json.loads(ln)) - return records - -# --------------------------------------------------------------------------- -# Record normalisation (with 10-way split for extra_info) -# --------------------------------------------------------------------------- - -def normalise_record(record: Dict[str, Any]) -> Dict[str, Any]: - # 1) extract existing extra_info dict (if any) - extra_raw = record.get('extra_info', {}) - if isinstance(extra_raw, str): - nested = json.loads(extra_raw).get('compressed', '') - extra = _deserialise_extra(_decompress_str(nested)) - elif isinstance(extra_raw, dict) and 'compressed' in extra_raw: - extra = _deserialise_extra(_decompress_str(extra_raw['compressed'])) - elif isinstance(extra_raw, dict): - extra = extra_raw.copy() - else: - extra = {} - - # 2) extract existing reward_model dict (if any) - rm_raw = record.get('reward_model', {}) - if isinstance(rm_raw, str): - reward_model = _deserialise_extra(rm_raw) - elif isinstance(rm_raw, dict) and 'compressed' in rm_raw: - reward_model = _deserialise_extra(_decompress_str(rm_raw['compressed'])) - elif isinstance(rm_raw, dict): - reward_model = rm_raw.copy() - else: - reward_model = {} - - # 3) copy canonical keys, gather others into extra - out: Dict[str, Any] = {} - for k, v in record.items(): - if k in KEEP_KEYS and k not in ('extra_info', 'reward_model'): - out[k] = v - elif k not in KEEP_KEYS: - extra[k] = v - - # 4) serialise + compress + split extra_info into 10 parts - if extra: - ser = _serialise_extra(extra) - comp = _compress_str(ser) - total_len = len(comp) - chunk_size = total_len // 10 - parts: Dict[str, str] = {} - for i in range(10): - start = i * chunk_size - end = start + chunk_size if i < 9 else total_len - parts[f'part_{i+1}'] = comp[start:end] - out['extra_info'] = parts - - # 5) serialise + compress reward_model if present - if reward_model: - ser_rm = _serialise_extra(reward_model) - comp_rm = _compress_str(ser_rm) - out['reward_model'] = {"ground_truth": {"compressed": comp_rm}} - - return out - -# --------------------------------------------------------------------------- -# Directory processing (parallel) -# --------------------------------------------------------------------------- - -from multiprocessing import Pool, cpu_count - -def _process_file(args): - """ - Worker to normalise one JSON file and write out result. - args: tuple(src_path_str, dst_dir_str) - """ - path_str, dst_dir_str = args - src_path = Path(path_str) - dst_dir = Path(dst_dir_str) - - try: - data = load_json_or_ndjson(src_path) - if isinstance(data, list): - out_data = [normalise_record(obj) for obj in data] - else: - out_data = normalise_record(data) - - out_path = dst_dir / src_path.name - with out_path.open('w', encoding='utf-8') as f: - json.dump(out_data, f, ensure_ascii=False, indent=2) - return f"✔ {src_path.name} → {out_path.name}" - except Exception as e: - return f"✖ {src_path.name} failed: {e}" - -def process_dir_parallel(src_dir: Path, dst_dir: Path, workers: int = None) -> None: - """ - Process all *.json files in src_dir in parallel, writing to dst_dir. - workers: number of parallel worker processes. Defaults to number of CPU cores minus one. - """ - dst_dir.mkdir(parents=True, exist_ok=True) - if workers is None: - workers = max(1, cpu_count() - 1) - - args_list = [(str(path), str(dst_dir)) for path in src_dir.glob('*.json')] - with Pool(processes=workers) as pool: - for result in pool.imap_unordered(_process_file, args_list): - print(result) - -if __name__ == '__main__': - parser = argparse.ArgumentParser(description="Normalize JSON files by moving non-canonical keys to extra_info") - parser.add_argument("src", type=str, help="Source directory containing JSON files") - parser.add_argument("dst", type=str, help="Destination directory for normalized JSON files") - parser.add_argument("-w", "--workers", type=int, help="Number of worker processes (default: CPU count - 1)", default=None) - args = parser.parse_args() - - src = Path(args.src).expanduser().resolve() - dst = Path(args.dst).expanduser().resolve() - - if not src.is_dir(): - print(f"Error: {src} is not a directory") - sys.exit(1) - - process_dir_parallel(src, dst, workers=args.workers) diff --git a/reasoners/data_preprocess/ood/ifeval.py b/reasoners/data_preprocess/ood/ifeval.py deleted file mode 100644 index b02436ccd..000000000 --- a/reasoners/data_preprocess/ood/ifeval.py +++ /dev/null @@ -1,139 +0,0 @@ -import argparse -import json -import os - -import datasets -import transformers -from datasets import load_dataset -from tqdm import tqdm - -from verl.utils.data_process.filter import LengthFilter -from verl.utils.data_process.utils import (sample_dataset, save_dataset, - set_seed) - -""" -python data_preprocess/ood/ifeval.py -""" - - -def get_datasets(cache_dir: str): - """ - Loads the IFEval dataset. - """ - try: - dataset = load_dataset("google/IFEval", cache_dir=cache_dir)["train"] # they updated the dataset under train, but its still called test - print(f"IFEval dataset: {len(dataset)} examples") - return None, dataset - except Exception as e: - print(f"Error loading dataset: {e}") - return None, None - - -non_numeric_answer = 0 -non_arithmetic_question = 0 - -PromptTemplate = """{{context}}""" - - -def make_map_fn(split: str, data_source: str) -> callable: - - def process_fn(example, idx): - prompt = example["prompt"] - instruction_id_list = example["instruction_id_list"] - kwargs = example["kwargs"] - - data = { - "data_source": data_source, - "prompt": [ - { - "role": "user", - "content": PromptTemplate.replace("{{context}}", prompt) - } - ], - "ability": "ood", - "apply_chat_template": True, - "reward_model": { - "style": "rule", - "ground_truth": kwargs, - }, - "extra_info": { - "instruction_id_list": instruction_id_list, - "prompt": prompt, - } - } - - if idx == 0 or idx == 1: - print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) - print(data) - - return data - - return process_fn - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="Download, process, and save OOD datasets." - ) - parser.add_argument( - "--data-dir", - default="data", - help="Base directory to save the processed data files.", - ) - parser.add_argument("--domain", default="ood", help="Domain of the dataset.") - parser.add_argument("--name", default="ifeval", help="Name of the dataset.") - parser.add_argument( - "--sample-size", - type=int, - default=None, - help="Number of samples to use from dataset. If None, use all samples.", - ) - parser.add_argument( - "--seed", type=int, default=42, help="Random seed for reproducibility" - ) - - args = parser.parse_args() - - # Config - set_seed(args.seed) - data_source = f"{args.domain}__{args.name}" - test_output_dir = os.path.join(args.data_dir, "test") - - # Download dataset - cache_dir = datasets.config.HF_DATASETS_CACHE - _, dataset = get_datasets(cache_dir) - - # Process datasets - process_fn = make_map_fn("test", data_source) - - dataset = dataset.map(function=process_fn, with_indices=True) - - # Filter dataset - try: - # length filter - tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B") - length_filter = LengthFilter(tokenizer=tokenizer, max_length=4096) - dataset = dataset.filter(lambda x: length_filter.check(x)) - - # non-numeric answer filter - dataset = dataset.filter(lambda x: x["reward_model"]["ground_truth"] is not None) - except Exception as e: - print(f"Warning: Could not perform length filtering. Error: {e}") - print("Proceeding without length filtering.") - - # Sample the dataset - dataset = sample_dataset(dataset, args.sample_size) - - # Save the dataset to test directory - test_output_path = save_dataset( - dataset=dataset, - output_dir=test_output_dir, - filename_prefix=data_source, - sample_size=len(dataset), - ) - - print( - f"\nDone! \n" - f"Data source: {data_source}\n" - f"Test data saved to {test_output_path} ({len(dataset)} samples)" - ) diff --git a/reasoners/data_preprocess/ood/livebench.py b/reasoners/data_preprocess/ood/livebench.py deleted file mode 100644 index e972c7673..000000000 --- a/reasoners/data_preprocess/ood/livebench.py +++ /dev/null @@ -1,160 +0,0 @@ -import argparse -import json -import os -import datetime - -import datasets -import transformers -from datasets import load_dataset -from tqdm import tqdm - -from verl.utils.data_process.filter import LengthFilter -from verl.utils.data_process.utils import (sample_dataset, save_dataset, - set_seed) - -""" -python data_preprocess/ood/livebench.py -""" - - -def get_datasets(cache_dir: str): - """ - Loads the LiveBench dataset. - """ - - test_data_sources = [ - "livebench/reasoning", - "livebench/data_analysis", - # "livebench/coding", - # "livebench/instruction_following", - # "livebench/math", - "livebench/language", - ] - - print(f"Loading the test datasets...") - test_datasets = { - os.path.basename(test_data_source.lower()): - datasets.load_dataset(test_data_source, trust_remote_code=True, split="test", cache_dir=cache_dir) - for test_data_source in test_data_sources - } - - return test_datasets - - - -non_numeric_answer = 0 -non_arithmetic_question = 0 - -PromptTemplate = """{{context}}""" - - -def make_map_fn(split: str, data_source: str) -> callable: - - def process_fn(example, idx): - turns = example["turns"] - if len(turns) != 1: - print(f"⚠️ Warning: {len(turns)} turns found in {example['task']}") - ground_truth = example["ground_truth"] - date = example["livebench_release_date"] - - data = { - "data_source": data_source, - "prompt": [ - { - "role": "user", - "content": PromptTemplate.replace("{{context}}", turns[0]) - } - ], - "ability": "ood", - "apply_chat_template": True, - "reward_model": { - "style": "rule", - "ground_truth": ground_truth, - }, - "extra_info": { - "category": example["category"], - "task": example["task"], - "prompt": turns, - "date": date, - } - } - - if idx == 0 or idx == 1: - print(data["prompt"][0]["content"]) - print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) - print(data) - - return data - - return process_fn - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="Download, process, and save OOD datasets." - ) - parser.add_argument( - "--data-dir", - default="data", - help="Base directory to save the processed data files.", - ) - parser.add_argument("--domain", default="ood", help="Domain of the dataset.") - parser.add_argument("--name", default="livebench", help="Name of the dataset.") - parser.add_argument( - "--sample-size", - type=int, - default=None, - help="Number of samples to use from dataset. If None, use all samples.", - ) - parser.add_argument( - "--seed", type=int, default=42, help="Random seed for reproducibility" - ) - - args = parser.parse_args() - - # Config - set_seed(args.seed) - data_source = f"{args.domain}__{args.name}" - test_output_dir = os.path.join(args.data_dir, "test") - - # Download dataset - cache_dir = datasets.config.HF_DATASETS_CACHE - test_datasets = get_datasets(cache_dir) - - # Process datasets - process_fn = make_map_fn("test", data_source) - for test_data_source, test_data in test_datasets.items(): - dataset = test_data.map(function=process_fn, with_indices=True) - # filter date before 2024-6 and after 2024-07 - dataset = dataset.filter(lambda x: x["extra_info"]["date"] >= datetime.datetime(2024, 6, 1) and x["extra_info"]["date"] <= datetime.datetime(2024, 7, 31)) - # Filter dataset - try: - # length filter - tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B") - length_filter = LengthFilter(tokenizer=tokenizer, max_length=4096) - dataset = dataset.filter(lambda x: length_filter.check(x)) - - # null answer filter - dataset = dataset.filter(lambda x: x["reward_model"]["ground_truth"] is not None) - except Exception as e: - print(f"Warning: Could not perform length filtering. Error: {e}") - print("Proceeding without length filtering.") - - # Sample the dataset - dataset = sample_dataset(dataset, args.sample_size) - - # Save the dataset to test directory - print(test_data_source, data_source) - file_prefix = test_data_source.replace("/", "_") - test_output_path = save_dataset( - dataset=dataset, - output_dir=test_output_dir, - filename_prefix=f"{data_source}_{file_prefix}", - sample_size=len(dataset), - ) - - print( - f"\nDone! \n" - f"Data source: {data_source}\n" - f"Test data saved to {test_output_path} ({len(dataset)} samples)" - ) diff --git a/reasoners/data_preprocess/parquet2json.py b/reasoners/data_preprocess/parquet2json.py deleted file mode 100644 index 89f969cc7..000000000 --- a/reasoners/data_preprocess/parquet2json.py +++ /dev/null @@ -1,53 +0,0 @@ -#!/usr/bin/env python3 -""" -Convert every .parquet in a folder to a single JSON array per file, -using Polars + the standard json module. - -Usage: - python parquet2json_array.py -""" - -from pathlib import Path -import sys -import polars as pl # pip install -U polars -import json - -def parquet_to_json_array(src_dir: Path, dst_dir: Path) -> None: - dst_dir.mkdir(parents=True, exist_ok=True) - - for pq_path in src_dir.glob("*.parquet"): - # 1) Read the Parquet file into a Polars DataFrame - df = pl.read_parquet(pq_path) - - # 2) Convert to a list of dicts (one dict per row) - records = df.to_dicts() - - # 3) Write out as one big JSON array, letting json.dumps call str() on datetimes - out_path = dst_dir / f"{pq_path.stem}.json" - with out_path.open("w", encoding="utf-8") as f: - json.dump( - records, - f, - ensure_ascii=False, - indent=2, - default=str # ← this makes datetime → ISO string - ) - - print(f"✔ {pq_path.name} → {out_path.name}") - -if __name__ == "__main__": - import argparse - - parser = argparse.ArgumentParser(description="Convert Parquet files to JSON arrays") - parser.add_argument("src", type=str, help="Source directory containing .parquet files") - parser.add_argument("dst", type=str, help="Destination directory for output .json files") - args = parser.parse_args() - - src = Path(args.src).expanduser().resolve() - dst = Path(args.dst).expanduser().resolve() - - if not src.is_dir(): - print(f"Error: {src} is not a directory") - sys.exit(1) - - parquet_to_json_array(src, dst) diff --git a/reasoners/data_preprocess/sft_rlhf/full_hh_rlhf.py b/reasoners/data_preprocess/sft_rlhf/full_hh_rlhf.py deleted file mode 100644 index 07e0884cb..000000000 --- a/reasoners/data_preprocess/sft_rlhf/full_hh_rlhf.py +++ /dev/null @@ -1,146 +0,0 @@ -# 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. -""" -- Preprocess data and split the training set into 75% for training RM and 25% for validting RM. -- All the training data is used to train SFT and RL. -- Both chosen and rejected is used to train SFT -""" -import argparse -import os - -import pandas as pd -from datasets import load_dataset - -from tqdm.auto import tqdm - -from verl.utils.fs import copy, makedirs - - -def generate_sft_dataset(target_hdfs_path_dir, local_dir='~/data/full_hh_rlh/sft'): - dataset = load_dataset('Dahoas/full-hh-rlhf') - output = {'prompt': [], 'response': []} - for data in tqdm(dataset['train']): - # add chosen - output['prompt'].append(data['prompt']) - output['response'].append(data['chosen']) - - # add rejection - output['prompt'].append(data['prompt']) - output['response'].append(data['rejected']) - - df = pd.DataFrame(output) - - local_dir = os.path.expanduser(local_dir) - os.makedirs(local_dir, exist_ok=True) - - local_path = os.path.join(local_dir, 'train.parquet') - - df.to_parquet(path=local_path) - - if target_hdfs_path_dir is not None: - hdfs_dir = target_hdfs_path_dir + '/' + 'train.parquet' - makedirs(hdfs_dir) - - copy(local_path, hdfs_dir) - - -def generate_rm_dataset(target_hdfs_path_dir, local_dir='~/data/full_hh_rlh/rm'): - train_dataset = load_dataset('Dahoas/full-hh-rlhf', split='train[:75%]') - test_dataset = load_dataset('Dahoas/full-hh-rlhf', split='train[-25%:]') - - local_dir = os.path.expanduser(local_dir) - os.makedirs(local_dir, exist_ok=True) - - for dataset, name in zip([train_dataset, test_dataset], ['train', 'test']): - output = {'prompt': [], 'chosen': [], 'rejected': []} - for data in tqdm(dataset): - # add chosen - output['prompt'].append(data['prompt']) - output['chosen'].append(data['chosen']) - output['rejected'].append(data['rejected']) - - df = pd.DataFrame(output) - - local_path = os.path.join(local_dir, name + '.parquet') - - df.to_parquet(path=local_path) - - if target_hdfs_path_dir is not None: - hdfs_dir = target_hdfs_path_dir + '/' + name + '.parquet' - makedirs(hdfs_dir) - - copy(local_path, hdfs_dir) - - -def generate_rl_dataset(target_hdfs_path_dir, local_dir='~/data/full_hh_rlhf/rl'): - dataset = load_dataset('Dahoas/full-hh-rlhf') - train_dataset = dataset['train'] - - data_source = 'Dahoas/full-hh-rlhf' - - # add a row to each data item that represents a unique id - def make_map_fn(split): - - def process_fn(example, idx): - prompt = example.pop('prompt') - response = example.pop('response') - - data = { - "data_source": data_source, - "prompt": [{ - "role": "user", - "content": prompt - }], - "ability": "alignment", - "reward_model": { - "style": "model", - "ground_truth": response # should not be used - }, - "extra_info": { - 'split': split, - 'index': idx - } - } - return data - - return process_fn - - train_dataset = train_dataset.map(function=make_map_fn('train'), with_indices=True) - local_dir = os.path.expanduser(local_dir) - local_path = os.path.join(local_dir, 'train.parquet') - train_dataset.to_parquet(local_path) - - if target_hdfs_path_dir is not None: - hdfs_dir = target_hdfs_path_dir + '/' + 'train.parquet' - makedirs(hdfs_dir) - - copy(local_path, hdfs_dir) - - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument('--split', type=str, choices=['sft', 'rm', 'rl'], required=True) - parser.add_argument('--local_dir', type=str, default='~/data/full_hh_rlhf') - parser.add_argument('--hdfs_dir', type=str, required=False, default=None) - - args = parser.parse_args() - - if args.split == 'sft': - generate_sft_dataset(args.hdfs_dir, os.path.join(args.local_dir, args.split)) - elif args.split == 'rm': - generate_rm_dataset(args.hdfs_dir, os.path.join(args.local_dir, args.split)) - elif args.split == 'rl': - generate_rl_dataset(args.hdfs_dir, os.path.join(args.local_dir, args.split)) - else: - raise NotImplementedError diff --git a/reasoners/data_preprocess/sft_rlhf/limo.py b/reasoners/data_preprocess/sft_rlhf/limo.py deleted file mode 100644 index b9ee0dd3f..000000000 --- a/reasoners/data_preprocess/sft_rlhf/limo.py +++ /dev/null @@ -1,93 +0,0 @@ -# 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. -""" -Preprocess the GSM8k dataset to parquet format -""" - -import re -import os -import datasets - -from verl.utils.hdfs_io import copy, makedirs -import argparse - - -def extract_solution(solution_str): - solution = re.search(r"\\boxed\{(.*?)\}", solution_str) - assert solution is not None - final_solution = solution.group(1) - return final_solution - - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument('--local_dir', default='~/data/limo') - parser.add_argument('--hdfs_dir', default=None) - - args = parser.parse_args() - - data_source = 'GAIR/LIMO' - - dataset = datasets.load_dataset(data_source) - - train_dataset = dataset['train'] - test_dataset = dataset['train'] - - instruction_following = "Let's think step by step and output the final answer in \\boxed\{\}." - - # add a row to each data item that represents a unique id - def make_map_fn(split): - - def process_fn(example, idx): - question_raw = example.pop('question') - - question = question_raw + ' ' + instruction_following - - solution = example.pop('solution') - exact_answer = example.pop('answer') - data = { - "data_source": data_source, - "prompt": [{ - "role": "user", - "content": question, - }], - "ability": "math", - "reward_model": { - "style": "rule", - "ground_truth": exact_answer - }, - "extra_info": { - 'split': split, - 'index': idx, - 'answer': solution, - "question": question, # TODO: use prompt \\boxed later - } - } - return data - - return process_fn - - train_dataset = train_dataset.map(function=make_map_fn('train'), with_indices=True) - test_dataset = test_dataset.map(function=make_map_fn('test'), with_indices=True) - - local_dir = args.local_dir - hdfs_dir = args.hdfs_dir - - train_dataset.to_parquet(os.path.join(local_dir, 'train_w_prompt.parquet')) - test_dataset.to_parquet(os.path.join(local_dir, 'test_w_prompt.parquet')) - - if hdfs_dir is not None: - makedirs(hdfs_dir) - - copy(src=local_dir, dst=hdfs_dir) diff --git a/reasoners/data_preprocess/simulation/arcagi.py b/reasoners/data_preprocess/simulation/arcagi.py deleted file mode 100644 index a862d8289..000000000 --- a/reasoners/data_preprocess/simulation/arcagi.py +++ /dev/null @@ -1,183 +0,0 @@ -"""Downloads, processes, and saves Abstraction and Reasoning Corpus for Artificial General Intelligence v2 (ARC-AGI-2) datasets.""" - -import os -import argparse -import json -import datasets -from datasets import Dataset -import subprocess - -from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset -from verl.utils.data_process.filter import LengthFilter - -import transformers - - - -RawARCPrompt = """You are a world-class puzzle solver with exceptional pattern recognition skills. Your task is to analyze puzzles, spot patterns, and provide direct solutions. -Given input-output grid pairs as reference examples, carefully observe the patterns to predict the output grid for new test input. Each pair follows the same transformation rule. Grids are 2D arrays. Here are the input and output grids for the reference examples: ----------------------------------------- -{{training_data}} ----------------------------------------- -Now, solve the following puzzle based on its input grid by applying the rules you have learned from the training data. ----------------------------------------- -{{input_test_data}} ----------------------------------------- -What is the output grid? Please put your answer within and tags, your final answer should be only the output grid (2d array). -""" - - -def get_datasets(cache_dir, name, download=False): - """ - Downloads (if specified) and loads the ARC-AGI-1 and ARC-AGI-2 datasets. - """ - # Download the dataset - if name == "arcagi1": - dir_path = os.path.join(cache_dir, "ARC-AGI-1/") - url = "https://github.com/fchollet/ARC-AGI.git" - elif name == "arcagi2": - dir_path = os.path.join(cache_dir, "ARC-AGI-2/") - url = "https://github.com/arcprize/ARC-AGI-2.git" - if download: - if os.path.isdir(dir_path): - pass - else: - try: - subprocess.run(["git", "clone", url, str(dir_path)], check=True) - except subprocess.CalledProcessError as e: - print("Error occurred while cloning the repo:", e) - - # Build the dataset. - train_dataset, test_dataset = [], [] - data_dir = dir_path - train_dir = os.path.join(data_dir, "data/training/") - for data_path in os.listdir(train_dir): - with open(os.path.join(train_dir, data_path), 'r') as f: - data = json.load(f) - train_dataset.append(data) - test_dir = os.path.join(data_dir, "data/evaluation/") - for data_path in os.listdir(test_dir): - with open(os.path.join(test_dir, data_path), 'r') as f: - data = json.load(f) - test_dataset.append(data) - - print(f"Training {len(train_dataset)} samples.") - print(f"Evaluation {len(test_dataset)} samples.") - return Dataset.from_list(train_dataset), Dataset.from_list(test_dataset) - - - -def make_map_fn(split: str, data_source: str, prompt_style: str="zero_style") -> callable: - """ - Creates a mapping function to process individual examples of the dataset. - - Args: - split: The dataset split ('train' or 'test'). - data_source: Identifier for the data source. - prompt_style: The style of prompt to use (e.g., 'zero_style'). - - Returns: - A callable function that takes an example and index, and returns the processed data. - """ - def process_fn(example, idx): - max_examples = 6 # max number of examples to show in the prompt - train_data = example.pop("train") - N_pairs = len(train_data) - training_data_prompt = "" - for i in range(min(N_pairs, max_examples)): - pair = train_data[i] - training_data_prompt += f"Example {i+1}\n" - training_data_prompt += f"Input: {pair['input']}\nOutput: {pair['output']}\n\n" - raw_prompt = RawARCPrompt.replace("{{training_data}}", training_data_prompt) - test_data = example.pop("test") - raw_prompt = raw_prompt.replace("{{input_test_data}}", str(test_data[0]["input"])) - answer = test_data[0]["output"] - data = { - "data_source": data_source, - "prompt": [{ - "role": "user", - "content": raw_prompt, - }], - "ability": "reasoning", - "apply_chat_template": True, - "reward_model": {"style": "rule", "ground_truth": answer}, - "extra_info": {"split": split, - "index": idx, - }, - } - if idx == 0 or idx == 1: - print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) - print(data) - return data - - return process_fn - - - -if __name__ == "__main__": - """Main script execution: parse args, load, process, and save datasets.""" - parser = argparse.ArgumentParser(description="Download, process, and save ARC-AGI-2 datasets.") - parser.add_argument('--data-dir', default='data', - help='Base directory to save the processed data files.') - parser.add_argument('--domain', default="simulation", help='Domain of the dataset.') - parser.add_argument('--name', default="arcagi1", choices=['arcagi1', 'arcagi2'], help='Name of the dataset.') - parser.add_argument('--train-sample-size', type=int, default=None, - help='Number of samples to use from training dataset. If None, use all samples.') - parser.add_argument('--test-sample-size', type=int, default=None, - help='Number of samples to use from test dataset. If None, use all samples.') - parser.add_argument('--prompt-style', type=str, choices=['zero_style'], default='zero_style', - help='Prompt style to use (currently only zero_style supported).') - parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility') - - args = parser.parse_args() - - set_seed(args.seed) - - data_source = f"{args.domain}__{args.name}" - train_output_dir = os.path.join(args.data_dir, 'train') - test_output_dir = os.path.join(args.data_dir, 'test') - os.makedirs(train_output_dir, exist_ok=True) - os.makedirs(test_output_dir, exist_ok=True) - - # Download the dataset from Github, but still saving to HF_datasets_cache - cache_dir = datasets.config.HF_DATASETS_CACHE - train_dataset, test_dataset = get_datasets(cache_dir, args.name, download=True) - - # Process the dataset - process_train_fn = make_map_fn('train', data_source, args.prompt_style) - process_test_fn = make_map_fn('test', data_source, args.prompt_style) - train_dataset = train_dataset.map(function=process_train_fn, with_indices=True) - test_dataset = test_dataset.map(function=process_test_fn, with_indices=True) - - # Length filter - try: - tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B") - length_filter = LengthFilter(tokenizer=tokenizer, max_length=4096) - train_dataset = train_dataset.filter(lambda x: length_filter.check(x), num_proc=64) - test_dataset = test_dataset.filter(lambda x: length_filter.check(x), num_proc=64) - except Exception as e: - print(f"Warning: Could not perform length filtering. Error: {e}") - print("Proceeding without length filtering.") - - # Sample the datasets using utility function - train_dataset = sample_dataset(train_dataset, args.train_sample_size) - test_dataset = sample_dataset(test_dataset, args.test_sample_size) - - # Save the datasets using utility function - train_output_path = save_dataset( - dataset=train_dataset, - output_dir=train_output_dir, - filename_prefix=data_source, - sample_size=len(train_dataset) - ) - test_output_path = save_dataset( - dataset=test_dataset, - output_dir=test_output_dir, - filename_prefix=data_source, - sample_size=len(test_dataset) - ) - - print(f"\nDone! \n" - f"Data source: {data_source}\n" - f"Train data saved to {train_output_path} ({len(train_dataset)} samples)\n" - f"Test data saved to {test_output_path} ({len(test_dataset)} samples)") diff --git a/reasoners/data_preprocess/simulation/barc.py b/reasoners/data_preprocess/simulation/barc.py deleted file mode 100644 index 902645730..000000000 --- a/reasoners/data_preprocess/simulation/barc.py +++ /dev/null @@ -1,166 +0,0 @@ -"""Downloads, processes, and saves BARC (Bootstrapping ARC: Synthetic Problem Generation for ARC Visual Reasoning Tasks) datasets.""" - -import os -import argparse -import datasets -import numpy as np -from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset -from verl.utils.data_process.filter import LengthFilter - -import transformers - - -RawARCPrompt = """You are a world-class puzzle solver with exceptional pattern recognition skills. Your task is to analyze puzzles, spot patterns, and provide direct solutions. -Given input-output grid pairs as reference examples, carefully observe the patterns to predict the output grid for new test input. Each pair follows the same transformation rule. Grids are 2D arrays. Here are the input and output grids for the reference examples: ----------------------------------------- -{{training_data}} ----------------------------------------- -Now, solve the following puzzle based on its input grid by applying the rules you have learned from the training data. ----------------------------------------- -{{input_test_data}} ----------------------------------------- -What is the output grid? Please put your answer within and tags, your final answer should be only the output grid (2d array). -""" - - -def get_datasets(cache_dir, N_samples=10000): - """ - Downloads (if specified) and loads the ARC-AGI-1 and ARC-AGI-2 datasets. - """ - # Load the dataset - ds = datasets.load_dataset("barc0/200k_HEAVY_gpt4o-description-gpt4omini-code_generated_problems", - split="train", cache_dir=cache_dir) # only have train split, no test split, max 200k samples - - ds = ds.select(range(min(N_samples, len(ds)))) - split_dataset = ds.train_test_split(test_size=0.2) - train_dataset = split_dataset["train"] - test_dataset = split_dataset["test"] - - return train_dataset, test_dataset - - - - -def make_map_fn(split: str, data_source: str, prompt_style: str="zero_style") -> callable: - """ - Creates a mapping function to process individual examples of the dataset. - - Args: - split: The dataset split ('train' or 'test'). - data_source: Identifier for the data source. - prompt_style: The style of prompt to use (e.g., 'zero_style'). - - Returns: - A callable function that takes an example and index, and returns the processed data. - """ - def process_fn(example, idx): - """ - BARC dataset contains multiple examples, we only use the maximum 6 examples as the training examples. - """ - max_examples = 6 # max number of examples to show in the prompt - train_data = example.pop("examples") - N_pairs = min(len(train_data)-1, max_examples) - training_data_prompt = "" - for i in range(N_pairs): - pair = train_data[i] - training_data_prompt += f"Example {i+1}\n" - input_data = [[int(float(x)) for x in row] for row in pair[0]] - output_data = [[int(float(x)) for x in row] for row in pair[1]] - training_data_prompt += f"Input: {input_data}\nOutput: {output_data}\n\n" - raw_prompt = RawARCPrompt.replace("{{training_data}}", training_data_prompt) - test_pair = train_data[N_pairs] - input_test_data = [[int(float(x)) for x in row] for row in test_pair[0]] - output_test_data = [[int(float(x)) for x in row] for row in test_pair[1]] - raw_prompt = raw_prompt.replace("{{input_test_data}}", str(input_test_data)) - answer = np.array(output_test_data) - data = { - "data_source": data_source, - "prompt": [{ - "role": "user", - "content": raw_prompt, - }], - "ability": "reasoning", - "apply_chat_template": True, - "reward_model": {"style": "rule", "ground_truth": answer}, - "extra_info": {"split": split, - "index": idx, - }, - } - if idx == 0 or idx == 1: - print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) - print(data) - return data - - return process_fn - - - - - -if __name__ == "__main__": - """Main script execution: parse args, load, process, and save datasets.""" - parser = argparse.ArgumentParser(description="Download, process, and save ARC-AGI-2 datasets.") - parser.add_argument('--data-dir', default='data', - help='Base directory to save the processed data files.') - parser.add_argument('--domain', default="simulation", help='Domain of the dataset.') - parser.add_argument('--name', default="barc", help='Name of the dataset.') - parser.add_argument('--train-sample-size', type=int, default=None, - help='Number of samples to use from training dataset. If None, use all samples.') - parser.add_argument('--test-sample-size', type=int, default=None, - help='Number of samples to use from test dataset. If None, use all samples.') - parser.add_argument('--prompt-style', type=str, choices=['zero_style'], default='zero_style', - help='Prompt style to use (currently only zero_style supported).') - parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility') - - args = parser.parse_args() - - set_seed(args.seed) - - data_source = f"{args.domain}__{args.name}" - train_output_dir = os.path.join(args.data_dir, 'train') - test_output_dir = os.path.join(args.data_dir, 'test') - os.makedirs(train_output_dir, exist_ok=True) - os.makedirs(test_output_dir, exist_ok=True) - - # Download the dataset from Github, but still saving to HF_datasets_cache - cache_dir = datasets.config.HF_DATASETS_CACHE - train_dataset, test_dataset = get_datasets(cache_dir) - - # Process the dataset - process_train_fn = make_map_fn('train', data_source, args.prompt_style) - process_test_fn = make_map_fn('test', data_source, args.prompt_style) - train_dataset = train_dataset.map(function=process_train_fn, with_indices=True) - test_dataset = test_dataset.map(function=process_test_fn, with_indices=True) - - # Length filter - try: - tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B") - length_filter = LengthFilter(tokenizer=tokenizer, max_length=4096) - train_dataset = train_dataset.filter(lambda x: length_filter.check(x), num_proc=64) - test_dataset = test_dataset.filter(lambda x: length_filter.check(x), num_proc=64) - except Exception as e: - print(f"Warning: Could not perform length filtering. Error: {e}") - print("Proceeding without length filtering.") - - # Sample the datasets using utility function - train_dataset = sample_dataset(train_dataset, args.train_sample_size) - test_dataset = sample_dataset(test_dataset, args.test_sample_size) - - # Save the datasets using utility function - train_output_path = save_dataset( - dataset=train_dataset, - output_dir=train_output_dir, - filename_prefix=data_source, - sample_size=len(train_dataset) - ) - test_output_path = save_dataset( - dataset=test_dataset, - output_dir=test_output_dir, - filename_prefix=data_source, - sample_size=len(test_dataset) - ) - - print(f"\nDone! \n" - f"Data source: {data_source}\n" - f"Train data saved to {train_output_path} ({len(train_dataset)} samples)\n" - f"Test data saved to {test_output_path} ({len(test_dataset)} samples)") diff --git a/reasoners/data_preprocess/simulation/codeio.py b/reasoners/data_preprocess/simulation/codeio.py deleted file mode 100644 index 97f5d097b..000000000 --- a/reasoners/data_preprocess/simulation/codeio.py +++ /dev/null @@ -1,193 +0,0 @@ -"""Downloads, processes, and saves CodeIO PyEdu Reasoning datasets.""" - -import os -import argparse -import json -import random -import time -import transformers -import datasets -from datasets import Dataset - -from verl.utils.data_process.prompt import build_zero_style_prompt -from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset -from verl.utils.data_process.filter import LengthFilter - -InstructionFollow = "Please output the final answer within ```json```" -RawInputPredictionPrompt = """You are given a question that requires some input and output variables as follows: -{{problem_description}} -The input and output requirements are as follows: -{{io_requirements}} -Given the following {{given_type}}: -{{given}} -Can you predict a feasible input without writing any code? Please reason and put your final answer in the following json format: "input": , where should be a dictionary, even if the there is only one input variable, with keys strictly match the input variables' names as specified. Please put your answer in \\boxed{} tags. -Tip: Here is a reference code snippet for this question. You can refer to this code tov guide your reasoning but not copy spans of code directly. -{{refcode}} - -""" -RawOutputPredictionPrompt = """You are given a question that requires some input and output variables as follows: -{{problem_description}} -The input and output requirements are as follows: -{{io_requirements}} -Given the following {{given_type}}: -{{given}} -Can you predict the output without writing any code? Please reason and put your final answer in the following json format: "output": , where should strictly match the the output requirement as specified. Please put your answer in \\boxed{} tags. -Tip: Here is a reference code snippet for this question. You can refer to this code tov guide your reasoning but not copy spans of code directly. -{{refcode}} - -""" -AnswerTemplate = """"{{predict_type}}": {{sol}}""" - - -def get_datasets(cache_dir, download=False): - """ - Downloads (if specified) and loads the CodeIO PyEdu Reasoning dataset. - """ - # Download the dataset - data_path = os.path.join(cache_dir, "PythonEdu-Reasoning.jsonl") - if download: - if os.path.exists(data_path): - pass - else: - url = "https://huggingface.co/datasets/hkust-nlp/CodeIO-PyEdu-Reasoning-Raw/resolve/main/0_368500_filtered_v2_ds25.sced.jsonl" - os.system(f'wget -O {data_path} {url}') - dataset = [] - N_code = 0 - if not os.path.exists(data_path): - time.sleep(5) - - # Build the dataset. Only keep the first "input/output" pair for each code sample for diversity. - max_entry_to_load = 200000 # hard-code to save processing time - with open(data_path, "r") as f: - for line in f: - if N_code >= max_entry_to_load: - break - N_code += 1 - data = json.loads(line) - common_fields = {k: v for k, v in data.items() if k != "ios"} - if data["ios"]: - io = data["ios"][0] - if random.random() < 0.5: - dataset.append({**common_fields, "input": json.dumps(io["input"]), "output": json.dumps(io["output"]), "given_type": "input", "predict_type": "output"}) - else: - dataset.append({**common_fields, "input": json.dumps(io["input"]), "output": json.dumps(io["output"]), "given_type": "output", "predict_type": "input"}) - - N = len(dataset) - N_train = int(N * 0.1) - train_dataset = dataset[:N_train] - test_dataset = dataset[N_train:] - print(f"Total {N_code} code samples, {N} I/O samples.") - return Dataset.from_list(train_dataset), Dataset.from_list(test_dataset) - - -def make_map_fn(split: str, data_source: str) -> callable: - """ - Creates a mapping function to process individual examples of the dataset. - - Args: - split: The dataset split ('train' or 'test'). - data_source: Identifier for the data source. - - Returns: - A callable function that takes an example and index, and returns the processed data. - """ - def process_fn(example, idx): - given_type = example.pop("given_type") - predict_type = example.pop("predict_type") - if predict_type == "input": - PromptTemplate = RawInputPredictionPrompt - else: - PromptTemplate = RawOutputPredictionPrompt - prompt = PromptTemplate.replace("{{given_type}}", given_type) + InstructionFollow - for key in ["problem_description", "io_requirements", given_type, "refcode"]: - feature = example.pop(key) - if key in ["input", "output"]: - prompt = prompt.replace("{{given}}", str(feature)) - else: - prompt = prompt.replace(f"{{{{{key}}}}}", str(feature)) - - sol = example.pop(predict_type) - answer = AnswerTemplate.replace("{{predict_type}}", predict_type) - answer = answer.replace("{{sol}}", sol) - data = { - "data_source": data_source, - "prompt": [{"role": "user", "content": prompt}], - "ability": "coding-inference", - "apply_chat_template": True, - "reward_model": {"style": "rule", "ground_truth": answer}, - "extra_info": {"split": split, - "index": idx, - }, - } - if idx == 0 or idx == 1: - print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) - print(data) - return data - - return process_fn - -if __name__ == '__main__': - """Main script execution: parse args, load, process, and save datasets.""" - parser = argparse.ArgumentParser(description="Download, process, and save CodeIO PyEdu Reasoning datasets.") - parser.add_argument('--data-dir', default='data', - help='Base directory to save the processed data files.') - parser.add_argument('--domain', default="simulation", help='Domain of the dataset.') - parser.add_argument('--name', default="codeio", help='Name of the dataset.') - parser.add_argument('--train-sample-size', type=int, default=None, - help='Number of samples to use from training dataset. If None, use all samples.') - parser.add_argument('--test-sample-size', type=int, default=None, - help='Number of samples to use from test dataset. If None, use all samples.') - parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility') - - args = parser.parse_args() - - set_seed(args.seed) - - data_source = f"{args.domain}__{args.name}" - train_output_dir = os.path.join(args.data_dir, 'train') - test_output_dir = os.path.join(args.data_dir, 'test') - os.makedirs(train_output_dir, exist_ok=True) - os.makedirs(test_output_dir, exist_ok=True) - - # Download the dataset - cache_dir = datasets.config.HF_DATASETS_CACHE - train_dataset, test_dataset = get_datasets(cache_dir, download=True) - - # Process the dataset - process_train_fn = make_map_fn('train', data_source) - process_test_fn = make_map_fn('test', data_source) - train_dataset = train_dataset.map(function=process_train_fn, with_indices=True) - test_dataset = test_dataset.map(function=process_test_fn, with_indices=True) - - # Length filter - try: - tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B") - length_filter = LengthFilter(tokenizer=tokenizer, max_length=2048) - train_dataset = train_dataset.filter(lambda x: length_filter.check(x), num_proc=64) - test_dataset = test_dataset.filter(lambda x: length_filter.check(x), num_proc=64) - except Exception as e: - print(f"Warning: Could not perform length filtering. Error: {e}") - print("Proceeding without length filtering.") - - # Sample the datasets using utility function - train_dataset = sample_dataset(train_dataset, args.train_sample_size) - test_dataset = sample_dataset(test_dataset, args.test_sample_size) - - # Save the datasets using utility function - train_output_path = save_dataset( - dataset=train_dataset, - output_dir=train_output_dir, - filename_prefix=data_source, - sample_size=len(train_dataset) - ) - test_output_path = save_dataset( - dataset=test_dataset, - output_dir=test_output_dir, - filename_prefix=data_source, - sample_size=len(test_dataset) - ) - - print(f"\nDone! \n" - f"Data source: {data_source}\n" - f"Train data saved to {train_output_path} ({len(train_dataset)} samples)\n" - f"Test data saved to {test_output_path} ({len(test_dataset)} samples)") diff --git a/reasoners/data_preprocess/simulation/cruxeval-i.py b/reasoners/data_preprocess/simulation/cruxeval-i.py deleted file mode 100644 index 44d54095a..000000000 --- a/reasoners/data_preprocess/simulation/cruxeval-i.py +++ /dev/null @@ -1,136 +0,0 @@ -"""Downloads, processes, and saves HumanEval dataset.""" - -import os -import argparse -import json -import transformers -import datasets -from datasets import load_dataset, Dataset - -from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset -from verl.utils.data_process.filter import LengthFilter -from verl.utils.reward_score.coder1 import code_exec - - -def get_datasets(cache_dir: str): - """ - Loads the CruxEval dataset. - """ - try: - dataset = load_dataset("cruxeval-org/cruxeval", cache_dir=cache_dir)["test"] - print(f"CruxEval dataset: {len(dataset)} examples") - return None, dataset - except Exception as e: - print(f"Error loading dataset: {e}") - return None, None - - -def make_map_fn(split: str, data_source: str) -> callable: - def process_fn(example, idx): - - # Extract task ID and prompt - task_id = example["id"] - - prompt = ( - "You will be given a function f and an output in the form f(??) == output. Your task is to find any input such that executing f on the input leads to the given output. There may be multiple answers, but only output one. First, think step by step. You MUST surround the answer with ```python and ``` tags. Express your answer as a passing assertion containing the input and the given output.\n\n" - "```python\n" - f"{example['code'].strip()}\n" - f"assert f(??) == {example['output']}\n" - "```\n" - ) - - data = { - "data_source": data_source, - "prompt": [{"role": "user", "content": prompt}], - "ability": "simulation", - "apply_chat_template": True, - "reward_model": { - "style": "rule", - "ground_truth": json.dumps({"functional": example["code"].strip()}), - }, - "extra_info": { - "split": split, - "index": idx, - "reference": example[ - "output" - ], # Include the canonical solution as reference - "original_prompt": prompt, - "dataset": "cruxeval-i", - "task_id": task_id, - }, - } - - if idx == 0 or idx == 1: - print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) - print(data) - - return data - - return process_fn - - -if __name__ == "__main__": - """Main script execution: parse args, load, process, and save datasets.""" - parser = argparse.ArgumentParser(description="Process and save HumanEval dataset.") - parser.add_argument( - "--data-dir", - default="data", - help="Base directory to save the processed data files.", - ) - parser.add_argument("--domain", default="simulation", help="Domain of the dataset.") - parser.add_argument("--name", default="cruxeval-i", help="Name of the dataset.") - parser.add_argument( - "--sample-size", - type=int, - default=None, - help="Number of samples to use from dataset. If None, use all samples.", - ) - parser.add_argument( - "--seed", type=int, default=42, help="Random seed for reproducibility" - ) - - args = parser.parse_args() - - set_seed(args.seed) - - data_source = f"{args.domain}__{args.name}" - test_output_dir = os.path.join(args.data_dir, "test") - os.makedirs(test_output_dir, exist_ok=True) - - # Load the dataset - cache_dir = datasets.config.HF_DATASETS_CACHE - _, dataset = get_datasets(cache_dir=cache_dir) - - # Process the dataset - process_fn = make_map_fn("test", data_source) - - dataset = dataset.map(function=process_fn, with_indices=True) - - # Filter out examples where processing failed - dataset = dataset.filter(lambda x: x["data_source"] == data_source) - - # Length filter - try: - tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B") - length_filter = LengthFilter(tokenizer=tokenizer, max_length=4096) - dataset = dataset.filter(lambda x: length_filter.check(x)) - except Exception as e: - print(f"Warning: Could not perform length filtering. Error: {e}") - print("Proceeding without length filtering.") - - # Sample the dataset - dataset = sample_dataset(dataset, args.sample_size) - - # Save the dataset to test directory - test_output_path = save_dataset( - dataset=dataset, - output_dir=test_output_dir, - filename_prefix=data_source, - sample_size=len(dataset), - ) - - print( - f"\nDone! \n" - f"Data source: {data_source}\n" - f"Test data saved to {test_output_path} ({len(dataset)} samples)" - ) diff --git a/reasoners/data_preprocess/simulation/cruxeval-o.py b/reasoners/data_preprocess/simulation/cruxeval-o.py deleted file mode 100644 index d34618145..000000000 --- a/reasoners/data_preprocess/simulation/cruxeval-o.py +++ /dev/null @@ -1,136 +0,0 @@ -"""Downloads, processes, and saves HumanEval dataset.""" - -import os -import argparse -import json -import transformers -import datasets -from datasets import load_dataset, Dataset - -from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset -from verl.utils.data_process.filter import LengthFilter -from verl.utils.reward_score.coder1 import code_exec - - -def get_datasets(cache_dir: str): - """ - Loads the CruxEval dataset. - """ - try: - dataset = load_dataset("cruxeval-org/cruxeval", cache_dir=cache_dir)["test"] - print(f"CruxEval dataset: {len(dataset)} examples") - return None, dataset - except Exception as e: - print(f"Error loading dataset: {e}") - return None, None - - -def make_map_fn(split: str, data_source: str) -> callable: - def process_fn(example, idx): - - # Extract task ID and prompt - task_id = example["id"] - - prompt = ( - "You will be given a Python function and an assertion containing an input to the function. Complete the assertion with a literal (no unsimplified expressions, no function calls) containing the output when executing the provided code on the given input, even if the function is incorrect or incomplete. Do NOT output any extra information. Execute the program step by step before arriving at an answer, and provide the full assertion with the correct output in ```python and ``` tags, following the examples.\n\n" - "```python\n" - f"{example['code'].strip()}\n" - f"assert f({example['input']}) == ??\n" - "```\n" - ) - - data = { - "data_source": data_source, - "prompt": [{"role": "user", "content": prompt}], - "ability": "simulation", - "apply_chat_template": True, - "reward_model": { - "style": "rule", - "ground_truth": json.dumps({"functional": example["code"].strip()}), - }, - "extra_info": { - "split": split, - "index": idx, - "reference": example[ - "output" - ], # Include the canonical solution as reference - "original_prompt": prompt, - "dataset": "cruxeval-o", - "task_id": task_id, - }, - } - - if idx == 0 or idx == 1: - print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) - print(data) - - return data - - return process_fn - - -if __name__ == "__main__": - """Main script execution: parse args, load, process, and save datasets.""" - parser = argparse.ArgumentParser(description="Process and save HumanEval dataset.") - parser.add_argument( - "--data-dir", - default="data", - help="Base directory to save the processed data files.", - ) - parser.add_argument("--domain", default="simulation", help="Domain of the dataset.") - parser.add_argument("--name", default="cruxeval-o", help="Name of the dataset.") - parser.add_argument( - "--sample-size", - type=int, - default=None, - help="Number of samples to use from dataset. If None, use all samples.", - ) - parser.add_argument( - "--seed", type=int, default=42, help="Random seed for reproducibility" - ) - - args = parser.parse_args() - - set_seed(args.seed) - - data_source = f"{args.domain}__{args.name}" - test_output_dir = os.path.join(args.data_dir, "test") - os.makedirs(test_output_dir, exist_ok=True) - - # Load the dataset - cache_dir = datasets.config.HF_DATASETS_CACHE - _, dataset = get_datasets(cache_dir=cache_dir) - - # Process the dataset - process_fn = make_map_fn("test", data_source) - - dataset = dataset.map(function=process_fn, with_indices=True) - - # Filter out examples where processing failed - dataset = dataset.filter(lambda x: x["data_source"] == data_source) - - # Length filter - try: - tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B") - length_filter = LengthFilter(tokenizer=tokenizer, max_length=4096) - dataset = dataset.filter(lambda x: length_filter.check(x)) - except Exception as e: - print(f"Warning: Could not perform length filtering. Error: {e}") - print("Proceeding without length filtering.") - - # Sample the dataset - dataset = sample_dataset(dataset, args.sample_size) - - # Save the dataset to test directory - test_output_path = save_dataset( - dataset=dataset, - output_dir=test_output_dir, - filename_prefix=data_source, - sample_size=len(dataset), - ) - - print( - f"\nDone! \n" - f"Data source: {data_source}\n" - f"Test data saved to {test_output_path} ({len(dataset)} samples)" - ) diff --git a/reasoners/data_preprocess/stem/bbeh.py b/reasoners/data_preprocess/stem/bbeh.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/reasoners/data_preprocess/stem/gpqa.py b/reasoners/data_preprocess/stem/gpqa.py deleted file mode 100644 index c06607cc1..000000000 --- a/reasoners/data_preprocess/stem/gpqa.py +++ /dev/null @@ -1,137 +0,0 @@ -""" -Preprocess the GPQA dataset to parquet format -""" - -import os -import re -import argparse -import random -from datasets import load_dataset, Dataset - -from verl.utils.hdfs_io import copy, makedirs -from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset - - -def get_datasets(): - """ - Loads the GPQA dataset. - """ - try: - dataset = load_dataset("Idavidrein/gpqa", "gpqa_main")["train"] - print(f"GPQA dataset: {len(dataset)} examples") - return None, dataset - except Exception as e: - print(f"Error loading dataset: {e}") - return None, None - - -# adopted from https://github.com/EleutherAI/lm-evaluation-harness/blob/main/lm_eval/tasks/gpqa/zeroshot/utils.py -def preprocess(text): - if text is None: - return " " - text = text.strip() - text = text.replace(" [title]", ". ") - text = re.sub("\\[.*?\\]", "", text) - text = text.replace(" ", " ") - return text - - -def make_map_fn(split: str, data_source: str) -> callable: - def process_fn(example, idx): - # Create a default "skip" response with all required fields - question = example["Question"].strip() - correct = preprocess(example["Correct Answer"]) - incorrect1 = preprocess(example["Incorrect Answer 1"]) - incorrect2 = preprocess(example["Incorrect Answer 2"]) - incorrect3 = preprocess(example["Incorrect Answer 3"]) - - all_choices = [incorrect1, incorrect2, incorrect3, correct] - random.shuffle(all_choices) - - correct_index = all_choices.index(correct) - correct_letter = chr(65 + correct_index) - - formatted_choices = "" - for i, choice in enumerate(all_choices): - letter = chr(65 + i) - formatted_choices += f"{letter}) {choice}\n" - - # deepseek uses OpenAI's simple-eval for GPQA-Diamond, so we adopt prompts from here: https://github.com/openai/simple-evals/blob/main/gpqa_eval.py - prompt = ( - f"Answer the following multiple choice question. The last line of your response should be of the following format: 'Answer: $LETTER' (without quotes) where LETTER is one of ABCD. Think step by step before answering." - f"\n{question}\n" - f"{formatted_choices}" - ) - - - data = { - "data_source": data_source, - "prompt": [ - {"role": "user", "content": prompt} - ], - "ability": "stem", - "apply_chat_template": True, - "reward_model": { - "style": "rule", - "ground_truth": correct_letter, - }, - "extra_info": { - "split": split, - "index": idx, - "original_prompt": prompt, - "dataset": "Idavidrein/gpqa", - }, - } - - if idx == 0 or idx == 1: - print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) - print(data) - print(f'\none prompt example is \n{prompt}') - - return data - - return process_fn - -if __name__ == '__main__': - """Main script execution: parse args, load, process, and save datasets.""" - parser = argparse.ArgumentParser(description="Process and save GPQA dataset.") - parser.add_argument('--data-dir', default='data', - help='Base directory to save the processed data files.') - parser.add_argument('--domain', default="stem", - help='Domain of the dataset.') - parser.add_argument('--name', default="gpqa", - help='Name of the dataset.') - parser.add_argument('--sample-size', type=int, default=None, - help='Number of samples to use from dataset. If None, use all samples.') - parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility') - - args = parser.parse_args() - - set_seed(args.seed) - - data_source = f"{args.domain}__{args.name}" - test_output_dir = os.path.join(args.data_dir, 'test') - os.makedirs(test_output_dir, exist_ok=True) - - # Load the dataset - _, dataset = get_datasets() - - # Process the dataset - process_fn = make_map_fn('test', data_source) - - dataset = dataset.map(function=process_fn, with_indices=True) - - # Sample the dataset - dataset = sample_dataset(dataset, args.sample_size) - - # Save the dataset to test directory - test_output_path = save_dataset( - dataset=dataset, - output_dir=test_output_dir, - filename_prefix=data_source, - sample_size=len(dataset) - ) - - print(f"\nDone! \n" - f"Data source: {data_source}\n" - f"Test data saved to {test_output_path} ({len(dataset)} samples)") \ No newline at end of file diff --git a/reasoners/data_preprocess/stem/gpqa_diamond.py b/reasoners/data_preprocess/stem/gpqa_diamond.py deleted file mode 100644 index c1dbad484..000000000 --- a/reasoners/data_preprocess/stem/gpqa_diamond.py +++ /dev/null @@ -1,142 +0,0 @@ -""" -Preprocess the GPQA - Diamond dataset to parquet format -""" - -import os -import re -import argparse -import random -from datasets import load_dataset, Dataset - -from verl.utils.hdfs_io import copy, makedirs -from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset - - -def get_datasets(): - """ - Loads the GPQA - Diamond dataset. - """ - try: - dataset = load_dataset("Idavidrein/gpqa", "gpqa_diamond")["train"] - print(f"GPQA - Diamond dataset: {len(dataset)} examples") - return None, dataset - except Exception as e: - print(f"Error loading dataset: {e}") - return None, None - - -# adopted from https://github.com/EleutherAI/lm-evaluation-harness/blob/main/lm_eval/tasks/gpqa/zeroshot/utils.py -def preprocess(text): - if text is None: - return " " - text = text.strip() - text = text.replace(" [title]", ". ") - text = re.sub("\\[.*?\\]", "", text) - text = text.replace(" ", " ") - return text - - -def make_map_fn(split: str, data_source: str) -> callable: - def process_fn(example, idx): - # Create a default "skip" response with all required fields - question = example["Question"].strip() - correct = preprocess(example["Correct Answer"]) - incorrect1 = preprocess(example["Incorrect Answer 1"]) - incorrect2 = preprocess(example["Incorrect Answer 2"]) - incorrect3 = preprocess(example["Incorrect Answer 3"]) - - all_choices = [incorrect1, incorrect2, incorrect3, correct] - random.shuffle(all_choices) - - correct_index = all_choices.index(correct) - correct_letter = chr(65 + correct_index) - - formatted_choices = "" - for i, choice in enumerate(all_choices): - letter = chr(65 + i) - formatted_choices += f"{letter}) {choice}\n" - - # # deepseek uses OpenAI's simple-eval for GPQA-Diamond, so we adopt prompts from here: https://github.com/openai/simple-evals/blob/main/gpqa_eval.py - # prompt = ( - # f"Answer the following multiple choice question. The last line of your response should be of the following format: 'Answer: $LETTER' (without quotes) where LETTER is one of ABCD. Think step by step before answering." - # f"\n{question}\n" - # f"{formatted_choices}" - # ) - prompt = ( - f"{question}\n" - f"{formatted_choices}" - "Please reason step by step, and put your final answer option within \\boxed{}. Only put the letter in the box, e.g. \\boxed{A}. There is only one correct answer." - ) - - - data = { - "data_source": data_source, - "prompt": [ - {"role": "user", "content": prompt} - ], - "ability": "stem", - "apply_chat_template": True, - "reward_model": { - "style": "rule", - "ground_truth": correct_letter, - }, - "extra_info": { - "split": split, - "index": idx, - "original_prompt": prompt, - "dataset": "Idavidrein/gpqa", - }, - } - - if idx == 0 or idx == 1: - print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) - print(data["prompt"][0]["content"]) - # print(f'\none prompt example is \n{prompt}') - - return data - - return process_fn - -if __name__ == '__main__': - """Main script execution: parse args, load, process, and save datasets.""" - parser = argparse.ArgumentParser(description="Process and save GPQA - Diamond dataset.") - parser.add_argument('--data-dir', default='data', - help='Base directory to save the processed data files.') - parser.add_argument('--domain', default="stem", - help='Domain of the dataset.') - parser.add_argument('--name', default="gpqa_diamond_no_box", - help='Name of the dataset.') - parser.add_argument('--sample-size', type=int, default=None, - help='Number of samples to use from dataset. If None, use all samples.') - parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility') - - args = parser.parse_args() - - set_seed(args.seed) - - data_source = f"{args.domain}__{args.name}" - test_output_dir = os.path.join(args.data_dir, 'test') - os.makedirs(test_output_dir, exist_ok=True) - - # Load the dataset - _, dataset = get_datasets() - - # Process the dataset - process_fn = make_map_fn('test', data_source) - - dataset = dataset.map(function=process_fn, with_indices=True) - - # Sample the dataset - dataset = sample_dataset(dataset, args.sample_size) - - # Save the dataset to test directory - test_output_path = save_dataset( - dataset=dataset, - output_dir=test_output_dir, - filename_prefix=data_source, - sample_size=len(dataset) - ) - - print(f"\nDone! \n" - f"Data source: {data_source}\n" - f"Test data saved to {test_output_path} ({len(dataset)} samples)") \ No newline at end of file diff --git a/reasoners/data_preprocess/stem/supergpqa.py b/reasoners/data_preprocess/stem/supergpqa.py deleted file mode 100644 index e954ef43e..000000000 --- a/reasoners/data_preprocess/stem/supergpqa.py +++ /dev/null @@ -1,118 +0,0 @@ -""" -Preprocess the GPQA dataset to parquet format -""" - -import os -import argparse -import random -from datasets import load_dataset, Dataset - -from verl.utils.hdfs_io import copy, makedirs -from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset - - -def get_datasets(): - """ - Loads the SuperGPQA dataset. - """ - try: - dataset = load_dataset("m-a-p/SuperGPQA")["train"] - print(f"SuperGPQA dataset: {len(dataset)} examples") - return None, dataset - except Exception as e: - print(f"Error loading dataset: {e}") - return None, None - - - -def make_map_fn(split: str, data_source: str) -> callable: - def process_fn(example, idx): - def form_options(options: list): - option_str = 'Options are:\n' - opts = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J'] - for opt, o in zip(options, opts): - option_str += f'({o}): {opt}\n' - return option_str - - question = example["question"].strip() - options = form_options(example["options"]) - query = question + '\n' + options + '\n' - - # prompt format is adopted from "General Reasoner" in https://github.com/TIGER-AI-Lab/General-Reasoner/blob/main/evaluation/eval_supergpqa.py - prompt = ( - f"{query}" - "Please reason step by step, and put your final answer option within \\boxed{}. Only put the letter in the box, e.g. \\boxed{A}. There is only one correct answer." - ) - - - data = { - "data_source": data_source, - "prompt": [ - {"role": "user", "content": prompt} - ], - "ability": "stem", - "apply_chat_template": True, - "reward_model": { - "style": "rule", - "ground_truth": example["answer_letter"], - }, - "extra_info": { - "split": split, - "index": idx, - "original_prompt": prompt, - "dataset": "m-a-p/SuperGPQA", - }, - } - - if idx == 0 or idx == 1: - print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) - print(data) - print(f'\none prompt example is \n{prompt}') - - return data - - return process_fn - -if __name__ == '__main__': - """Main script execution: parse args, load, process, and save datasets.""" - parser = argparse.ArgumentParser(description="Process and save GPQA dataset.") - parser.add_argument('--data-dir', default='data', - help='Base directory to save the processed data files.') - parser.add_argument('--domain', default="stem", - help='Domain of the dataset.') - parser.add_argument('--name', default="supergpqa", - help='Name of the dataset.') - parser.add_argument('--sample-size', type=int, default=1000, - help='Number of samples to use from dataset. If None, use all samples.') - parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility') - - args = parser.parse_args() - - set_seed(args.seed) - - data_source = f"{args.domain}__{args.name}" - test_output_dir = os.path.join(args.data_dir, 'test') - os.makedirs(test_output_dir, exist_ok=True) - - # Load the dataset - _, dataset = get_datasets() - - # Process the dataset - process_fn = make_map_fn('test', data_source) - - dataset = dataset.map(function=process_fn, with_indices=True) - - # Sample the dataset - dataset = sample_dataset(dataset, args.sample_size) - - # Save the dataset to test directory - test_output_path = save_dataset( - dataset=dataset, - output_dir=test_output_dir, - filename_prefix=data_source, - sample_size=len(dataset) - ) - - print(f"\nDone! \n" - f"Data source: {data_source}\n" - f"Test data saved to {test_output_path} ({len(dataset)} samples)") \ No newline at end of file diff --git a/reasoners/data_preprocess/stem_web/README.md b/reasoners/data_preprocess/stem_web/README.md deleted file mode 100644 index 5f2a389a4..000000000 --- a/reasoners/data_preprocess/stem_web/README.md +++ /dev/null @@ -1,20 +0,0 @@ -# STEM Web Dataset Preprocessing - -This directory contains scripts for preprocessing the STEM web dataset and running the verifier judge. - -## Running the Verifier Judge - -To run the verifier judge service: - -1. Submit the job using sbatch: - ```bash - sbatch run_verifier_judge.sh - ``` - -2. The script will: - - Detect the node's IP address automatically - - Launch a vLLM server running the TIGER-Lab/general-verifier model - -3. The service will be available at `http://:8000` - -4. User need to set a environment variable using following format: *export STEM_LLM_JUDGE_URL="http://{node_ip}:8000"* \ No newline at end of file diff --git a/reasoners/data_preprocess/stem_web/json2_parquet.py b/reasoners/data_preprocess/stem_web/json2_parquet.py deleted file mode 100644 index 0a5528122..000000000 --- a/reasoners/data_preprocess/stem_web/json2_parquet.py +++ /dev/null @@ -1,109 +0,0 @@ -#!/usr/bin/env python3 -# make_webinstruct_le30_parquet.py -""" -Convert samples_le30.json into a chat-style template and save as Parquet. -""" - -import argparse, json, os -from typing import Dict, Any, List -import datasets -from transformers import AutoTokenizer -import pprint -import random -from verl.utils.data_process.utils import set_seed, save_dataset - -# ------------------ prompt template definition ------------------ # -SYSTEM_MSG = ( - "You are a knowledgeable assistant. " - "Answer the following questions and think step by step. Please output the final answer within \\boxed{}. " -) - -def make_prompt(question: str) -> List[Dict[str, str]]: - """ - Chat-style message list: - - system: fixed instruction - - user : original question - """ - return [ - {"role": "user", "content": SYSTEM_MSG + question.strip()} - ] - -# --------------------------- main ------------------------------- # -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("--in", dest="in_file", - default="samples_le30.json", - help="Input JSON file created earlier.") - parser.add_argument("--out", dest="out_file", - default="webinstruct_le30.parquet", - help="Output Parquet filename.") - parser.add_argument("--tokenizer", - default="Qwen/Qwen3-8B", - help="Name or path of tokenizer used for stats.") - args = parser.parse_args() - - # 1. load json - with open(args.in_file, encoding="utf-8") as f: - raw: List[Dict[str, Any]] = json.load(f) - - # 2. wrap into new template - processed = [] - for ex in raw: - question = ex["question"] - answer = ex["answer"] - processed.append({ - "data_source": "WebInstruct-le30", - "prompt": make_prompt(question), - "raw_prompt": question, - "ability": "QA", - "apply_chat_template": True, - "response": answer, - "reward_model": {"ground_truth": answer}, - "extra_info": { - "category": ex.get("category"), - "difficulty": ex.get("difficulty"), - "answer_type": ex.get("answer_type"), - "token_length": ex["token_length"] - }, - }) - # ---------- NEW: print a few samples ----------------- - print("\n*** Example prompts (3 random rows) ***") - for sample in random.sample(processed, k=min(3, len(processed))): - pprint.pprint( - { - "prompt": sample["prompt"], - "response": sample["response"], - "extra_info": sample["extra_info"] - }, - compact=True, - width=120 - ) - print("-" * 80) - # ------------------------------------------------------ - - # 3. build Dataset and save parquet - ds = datasets.Dataset.from_list(processed) - save_path = save_dataset(dataset=ds, - output_dir="./", - filename_prefix="stem", - sample_size=len(ds)) - - print(f"Saved {len(ds)} rows to {save_path}") - - - # 4. (optional) token-length stats over *prompt* messages - tok = AutoTokenizer.from_pretrained(args.tokenizer, trust_remote_code=True) - lengths = [ - len(tok.encode(msg["content"], add_special_tokens=False)) - for row in ds - for msg in row["prompt"] - if msg["role"] == "user" - ] - import numpy as np - lengths = np.asarray(lengths) - print(f"User-prompt token stats — min:{lengths.min()}, " - f"max:{lengths.max()}, mean:{lengths.mean():.1f}, " - f"std:{lengths.std():.1f}") - -if __name__ == "__main__": - main() diff --git a/reasoners/data_preprocess/stem_web/simple_filter.py b/reasoners/data_preprocess/stem_web/simple_filter.py deleted file mode 100644 index 1c2f50075..000000000 --- a/reasoners/data_preprocess/stem_web/simple_filter.py +++ /dev/null @@ -1,35 +0,0 @@ -import os -import json -from datasets import load_dataset -from transformers import AutoTokenizer - -dataset = load_dataset("TIGER-Lab/WebInstruct-verified", split="train") -def keep_example(example): - return ( - example.get("category") != "Mathematics" and - example.get("difficulty") in ["University", "PhD"] and - example.get("answer_type") not in ["Boolean", "Multiple Choice"] - ) -filtered = dataset.filter(keep_example) - -tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-8B", trust_remote_code=True) -answers = filtered["answer"] -encodings = tokenizer(answers, padding=False, truncation=False) -token_lengths = [len(ids) for ids in encodings["input_ids"]] - -selected_indices = [i for i, length in enumerate(token_lengths) if length <= 30] -print(f"Number of examples with answer token length ≤ 30: {len(selected_indices)}") - -samples = [] -for idx in selected_indices: - example = filtered[idx] - ex_dict = dict(example) - ex_dict["token_length"] = token_lengths[idx] - samples.append(ex_dict) - -os.makedirs("output", exist_ok=True) -output_path = os.path.join("output", "samples_le30.json") -with open(output_path, "w", encoding="utf-8") as f: - json.dump(samples, f, ensure_ascii=False, indent=2) - -print(f"Saved {len(samples)} examples to {output_path}") \ No newline at end of file diff --git a/reasoners/data_preprocess/stem_web/stem.py b/reasoners/data_preprocess/stem_web/stem.py deleted file mode 100644 index da91e179e..000000000 --- a/reasoners/data_preprocess/stem_web/stem.py +++ /dev/null @@ -1,133 +0,0 @@ -#!/usr/bin/env python3 -import os -import json -import argparse -import random -import pprint -import re - -from datasets import load_dataset, Dataset -from transformers import AutoTokenizer -import numpy as np - -# ------------------ prompt template ------------------ # -SYSTEM_MSG = ( - "You are a knowledgeable assistant. " - "Answer the following questions and think step by step. Please output the final answer within \\boxed{}. " -) - -def make_prompt(question: str): - return [ - {"role": "user", "content": SYSTEM_MSG + question.strip()} - ] - -# ------------------ filter predicates ------------------ # -def keep_example(example, max_answer_len): - if len(example.get("answer", "")) == 1 and example.get("answer", "").isalpha(): - return False - return ( - example.get("category") != "Mathematics" - and example.get("difficulty") in ["University", "PhD"] - and example.get("answer_type") not in ["Boolean", "Multiple Choice"] - and len(example.get("answer", "").split()) > 0 - ) - -def has_too_many_decimal_places(s: str, max_decimals: int = 3) -> bool: - pattern = re.compile(r"\d+\.(\d{" + str(max_decimals+1) + r",})") - return bool(pattern.search(s)) - -# --------------------------- main ------------------------------- # -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("--json_out", default="stem_web.json") - parser.add_argument("--parquet_out", default="stem_web.parquet") - parser.add_argument("--dataset", default="TIGER-Lab/WebInstruct-verified") - parser.add_argument("--split", default="train") - parser.add_argument("--tokenizer", default="Qwen/Qwen3-8B") - parser.add_argument("--max_answer_len", type=int, default=30) - args = parser.parse_args() - - # 1) load and master filter - ds = load_dataset(args.dataset, split=args.split) - - def master_filter(ex): - ans = ex.get("answer", "") - if not keep_example(ex, args.max_answer_len): - return False - if has_too_many_decimal_places(ans, max_decimals=3): - return False - return True - - ds = ds.filter(master_filter) - print(f"Filtered down to {len(ds)} examples after applying all criteria") - - # 2) tokenize answers and select by token length - tok = AutoTokenizer.from_pretrained(args.tokenizer, trust_remote_code=True) - enc = tok(ds["answer"], padding=False, truncation=False) - lengths = [len(ids) for ids in enc["input_ids"]] - keep_idxs = [i for i, L in enumerate(lengths) if L <= args.max_answer_len] - print(f"{len(keep_idxs)} examples with answer ≤ {args.max_answer_len} tokens") - - # 3) build JSON samples - samples_q2idx = {} - samples = [] - for i in keep_idxs: - ex = dict(ds[i]) - ex["token_length"] = lengths[i] - prompt_str = make_prompt(ex["question"])[0]["content"] - samples_q2idx[prompt_str] = ex - samples.append(ex) - - os.makedirs(os.path.dirname(args.json_out) or ".", exist_ok=True) - with open(args.json_out, "w", encoding="utf-8") as f: - json.dump(samples_q2idx, f, ensure_ascii=False, indent=2) - print(f"Saved {len(samples)} samples to {args.json_out}") - - # 4) wrap into chat‐template records - processed = [] - for ex in samples: - processed.append({ - "data_source": "stem_web", - "prompt": make_prompt(ex["question"]), - "raw_prompt": ex["question"], - "ability": "QA", - "apply_chat_template": True, - "response": ex["answer"], - "reward_model": {"ground_truth": ex["answer"]}, - "extra_info": { - "idx": ex.get("id"), - "question": ex["question"], - "category": ex.get("category"), - "difficulty": ex.get("difficulty"), - "answer_type": ex.get("answer_type"), - "token_length": ex["token_length"], - }, - }) - - # 5) show a few examples - print("\n*** Example prompts (3 random rows) ***") - for row in random.sample(processed, k=min(3, len(processed))): - pprint.pprint({ - "prompt": row["prompt"], - "response": row["response"], - "extra_info": row["extra_info"], - }, compact=True, width=120) - print("-" * 80) - - # 6) save to Parquet - hf_ds = Dataset.from_list(processed) - hf_ds.to_parquet(args.parquet_out) - print(f"Saved {len(hf_ds)} rows to {args.parquet_out}") - - # 7) (optional) prompt‐length stats - user_lens = [] - for row in processed: - for msg in row["prompt"]: - if msg["role"] == "user": - user_lens.append(len(tok.encode(msg["content"], add_special_tokens=False))) - arr = np.array(user_lens) - print(f"User‐prompt token stats — " - f"min:{arr.min()}, max:{arr.max()}, mean:{arr.mean():.1f}, std:{arr.std():.1f}") - -if __name__ == "__main__": - main() \ No newline at end of file diff --git a/reasoners/data_preprocess/table/finqa.py b/reasoners/data_preprocess/table/finqa.py deleted file mode 100644 index 16d307a1a..000000000 --- a/reasoners/data_preprocess/table/finqa.py +++ /dev/null @@ -1,170 +0,0 @@ -import argparse -import json -import os - -import datasets -import transformers -from datasets import load_dataset -from tqdm import tqdm - -from verl.utils.data_process.filter import LengthFilter -from verl.utils.data_process.utils import (sample_dataset, save_dataset, - set_seed) - -""" -python data_preprocess/table/finqa.py -""" - - -def get_datasets(cache_dir: str): - """ - Loads the FinQA dataset. - """ - try: - dataset = load_dataset("nanoverl/finqa", cache_dir=cache_dir)["test"] - print(f"FinQA dataset: {len(dataset)} examples") - return None, dataset - except Exception as e: - print(f"Error loading dataset: {e}") - return None, None - - -non_numeric_answer = 0 -non_arithmetic_question = 0 - -InstructionFollow = "Please output the final numeric answer within \\boxed{}." -PromptTemplate = """You are given some text and tables. Use the information in the given context to answer the following question. -{{context}} -The question is: -{{question}} -""" - - -def make_map_fn(split: str, data_source: str) -> callable: - - def process_fn(example, idx): - pre_text = example["pre_text"] - post_text = example["post_text"] - table = example["table"] - question = example["question"] - answer = example["answer"] - - # custom table - tablerows = [] - for idx, row in enumerate(table): - if idx == 0: - tablerows.append("| " + " | ".join(row) + " |") - tablerows.append("| " + " | ".join(["---"] * len(row)) + " |") - else: - tablerows.append("| " + " | ".join(row) + " |") - - table_str = "\n".join(tablerows) - context_str = ( - "\n".join(pre_text) + "\n" + table_str + "\n" + "\n".join(post_text) - ) - - # filter answer which is not numeric - nanswer = ( - answer.replace(",", "") - .replace("%", " / 100") - .replace("$", "") - .replace(":", "/") - ) - - try: - nanswer = float(eval(nanswer)) - # print(nanswer) - except: - nanswer = None - - data = { - "data_source": data_source, - "prompt": [ - { - "role": "user", - "content": PromptTemplate.replace("{{context}}", context_str).replace("{{question}}", question) + InstructionFollow - } - ], - "ability": "table", - "apply_chat_template": True, - "reward_model": { - "style": "rule", - "ground_truth": nanswer, - }, - } - - if idx == 0 or idx == 1: - print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) - print(data) - - return data - - return process_fn - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="Download, process, and save MultiHierTT table reasoning datasets." - ) - parser.add_argument( - "--data-dir", - default="data", - help="Base directory to save the processed data files.", - ) - parser.add_argument("--domain", default="table", help="Domain of the dataset.") - parser.add_argument("--name", default="finqa", help="Name of the dataset.") - parser.add_argument( - "--sample-size", - type=int, - default=None, - help="Number of samples to use from dataset. If None, use all samples.", - ) - parser.add_argument( - "--seed", type=int, default=42, help="Random seed for reproducibility" - ) - - args = parser.parse_args() - - # Config - set_seed(args.seed) - data_source = f"{args.domain}__{args.name}" - test_output_dir = os.path.join(args.data_dir, "test") - - # Download dataset - cache_dir = datasets.config.HF_DATASETS_CACHE - _, dataset = get_datasets(cache_dir) - - # Process datasets - process_fn = make_map_fn("test", data_source) - - dataset = dataset.map(function=process_fn, with_indices=True) - - # Filter dataset - try: - # length filter - tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B") - length_filter = LengthFilter(tokenizer=tokenizer, max_length=4096) - dataset = dataset.filter(lambda x: length_filter.check(x)) - - # non-numeric answer filter - dataset = dataset.filter(lambda x: x["reward_model"]["ground_truth"] is not None) - except Exception as e: - print(f"Warning: Could not perform length filtering. Error: {e}") - print("Proceeding without length filtering.") - - # Sample the dataset - dataset = sample_dataset(dataset, args.sample_size) - - # Save the dataset to test directory - test_output_path = save_dataset( - dataset=dataset, - output_dir=test_output_dir, - filename_prefix=data_source, - sample_size=len(dataset), - ) - - print( - f"\nDone! \n" - f"Data source: {data_source}\n" - f"Test data saved to {test_output_path} ({len(dataset)} samples)" - ) diff --git a/reasoners/data_preprocess/table/hitab.py b/reasoners/data_preprocess/table/hitab.py deleted file mode 100644 index 367db29ed..000000000 --- a/reasoners/data_preprocess/table/hitab.py +++ /dev/null @@ -1,263 +0,0 @@ -"""Downloads, processes, and saves HiTab table reasoning datasets.""" - -import os -import json -import argparse -import random -import transformers -from datasets import Dataset -import datasets -import requests -import zipfile - -from verl.utils.data_process.filter import LengthFilter -from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset -from verl.utils.data_process.prompt import build_zero_style_prompt - - -InstructionFollow = "Please output the final answer within \\boxed{}. If there are multiple answers, please output them separated by |." -PromptTemplate = """You are given one or more tables. Use the information in the table to answer the following question. -{{tables}} -The question is: -{{question}} -""" - - -class HiTabFilter(LengthFilter): - """Filter for HiTab dataset to keep only arithmetic questions without text evidence.""" - - def __init__(self, tokenizer, max_length=2048): - super().__init__(tokenizer=tokenizer, max_length=max_length) - - def check(self, example): - # Only keep questions requiring operations on the table - aggregation = example["aggregation"] - # TODO: Now comment to include non-operation questions to enlarge the dataset - # if 'none' in aggregation: - # return False - - # Ensure the prompt length is within the specified range - length_check = super().check(example) - if not length_check: - return False - - return True - - -def get_datasets(cache_dir, download=False): - """ - Load the HiTab dataset from cache or download it if necessary. - - Args: - cache_dir (str): Directory to cache the dataset - download (bool): Whether to download the dataset if not found - - Returns: - tuple: (train_dataset, test_dataset) as Dataset objects - """ - data_path = os.path.join(cache_dir, "hitab/") - if download: - if os.path.exists(data_path): - pass - else: - # Base URL for raw content - def download_file(url, output_path): - """Download a file from GitHub raw content URL.""" - response = requests.get(url) - if response.status_code == 200: - os.makedirs(os.path.dirname(output_path), exist_ok=True) - with open(output_path, 'wb') as f: - f.write(response.content) - print(f"Successfully downloaded: {output_path}") - else: - print(f"Failed to download: {url}") - base_url = "https://raw.githubusercontent.com/microsoft/HiTab/main/data" - - # Files to download - files = [ - "tables.zip", - "train_samples.jsonl", - "test_samples.jsonl" - ] - - # Download each file - for file in files: - url = f"{base_url}/{file}" - output_path = os.path.join(data_path, file) - download_file(url, output_path) - # Unzip the tables.zip file - zip_path = os.path.join(data_path, "tables.zip") - with zipfile.ZipFile(zip_path, 'r') as zip_ref: - zip_ref.extractall(data_path) - - train_path = os.path.join(data_path, "train_samples.jsonl") - test_path = os.path.join(data_path, "test_samples.jsonl") - - def _preprocess_hitab_entry(entry): - new_entry = {} - new_entry['id'] = entry['id'] - new_entry['question'] = entry['question'] - new_entry['answer'] = '|'.join([str(answer) for answer in entry['answer']]) - new_entry['aggregation'] = entry['aggregation'] - table_id = entry['table_id'] - table_path = os.path.join(data_path, "tables/raw", f"{table_id}.json") - with open(table_path, "r") as f: - table = json.load(f) - new_entry['table'] = table["texts"] - return new_entry - - with open(train_path, "r") as f: - train_dataset = [json.loads(line) for line in f] - train_dataset = [_preprocess_hitab_entry(entry) for entry in train_dataset] - for idx, entry in enumerate(train_dataset): - print(entry) - if idx > 5: - break - - with open(test_path, "r") as f: - test_dataset = [json.loads(line) for line in f] - test_dataset = [_preprocess_hitab_entry(entry) for entry in test_dataset] - - print("Total train samples:", len(train_dataset)) - print("Total test samples:", len(test_dataset)) - - return Dataset.from_list(train_dataset), Dataset.from_list(test_dataset) - - -def table_matrix_to_markdown(table_matrix: list[list[str]]) -> str: - """ - Convert table matrix to markdown format. - - Args: - table_matrix (list[list[str]]): Table matrix - - Returns: - str: Markdown representation of the table - """ - if not table_matrix or not table_matrix[0]: - return "" - - rows = [] - - # Header row - header = "| " + " | ".join(str(cell) for cell in table_matrix[0]) + " |" - rows.append(header) - - # Separator row - separator = "|" + "|".join("-" * 3 for _ in table_matrix[0]) + "|" - rows.append(separator) - - # Data rows - for row in table_matrix[1:]: - formatted_row = "| " + " | ".join(str(cell) for cell in row) + " |" - rows.append(formatted_row) - - return "\n".join(rows) - - -def make_map_fn(split: str, data_source: str): - def process_fn(example, idx): - try: - table = example.pop("table") - table_string = table_matrix_to_markdown(table) - question = example["question"] - answer = example["answer"] - - prompt = PromptTemplate.replace("{{tables}}", table_string).replace("{{question}}", question) + InstructionFollow - - except Exception as e: - print(e) - print(table) - exit() - - data = { - "data_source": data_source, - "prompt": [ - { - "role": "user", - "content": prompt - } - ], - "ability": "table", - "apply_chat_template": True, - "reward_model": {"style": "rule", "ground_truth": answer}, - "extra_info": { - "split": split, - "index": idx, - }, - } - if idx == 0 or idx == 1: - print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) - print(data) - return data - - return process_fn - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Download, process, and save Hitab table reasoning datasets.") - parser.add_argument('--data-dir', default='data', - help='Base directory to save the processed data files.') - parser.add_argument('--domain', default="table", help='Domain of the dataset.') - parser.add_argument('--name', default="hitab", help='Name of the dataset.') - parser.add_argument('--train-sample-size', type=int, default=200, - help='Number of samples to use from training dataset. If None, use all samples.') - parser.add_argument('--test-sample-size', type=int, default=None, - help='Number of samples to use from test dataset. If None, use all samples.') - parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility') - - args = parser.parse_args() - - # Config - set_seed(args.seed) - data_source = f"{args.domain}__{args.name}" - train_output_dir = os.path.join(args.data_dir, 'train') - test_output_dir = os.path.join(args.data_dir, 'test') - - # Download dataset - cache_dir = datasets.config.HF_DATASETS_CACHE - train_dataset, test_dataset = get_datasets(cache_dir, download=True) - - # Process datasets - process_train_fn = make_map_fn('train', data_source) - process_test_fn = make_map_fn('test', data_source) - train_dataset = train_dataset.map(function=process_train_fn, with_indices=True) - test_dataset = test_dataset.map(function=process_test_fn, with_indices=True) - - print("Total train samples:", len(train_dataset)) - print("Total test samples:", len(test_dataset)) - for idx, example in enumerate(train_dataset): - print(example) - if idx > 5: - break - - # Filter dataset - tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B") - table_filter = HiTabFilter(tokenizer=tokenizer, max_length=4096) - train_dataset = train_dataset.filter(lambda x: table_filter.check(x), num_proc=64) - test_dataset = test_dataset.filter(lambda x: table_filter.check(x), num_proc=64) - args.train_sample_size = len(train_dataset) if args.train_sample_size is None else args.train_sample_size - args.test_sample_size = len(test_dataset) if args.test_sample_size is None else args.test_sample_size - - # Sample datasets if specified - train_dataset = sample_dataset(train_dataset, args.train_sample_size) - test_dataset = sample_dataset(test_dataset, args.test_sample_size) - - # Save datasets - train_output_path = save_dataset( - dataset=train_dataset, - output_dir=train_output_dir, - filename_prefix=data_source, - sample_size=args.train_sample_size - ) - test_output_path = save_dataset( - dataset=test_dataset, - output_dir=test_output_dir, - filename_prefix=data_source, - sample_size=args.test_sample_size - ) - - print(f"\nDone! \n" - f"Data source: {data_source}\n" - f"Train data saved to {train_output_path} ({len(train_dataset)} samples)\n" - f"Test data saved to {test_output_path} ({len(test_dataset)} samples)") diff --git a/reasoners/data_preprocess/table/multihier.py b/reasoners/data_preprocess/table/multihier.py deleted file mode 100644 index 26a15d32c..000000000 --- a/reasoners/data_preprocess/table/multihier.py +++ /dev/null @@ -1,274 +0,0 @@ -"""Downloads, processes, and saves MultiHierTT table reasoning datasets.""" - -import os -import json -import argparse -import random -import transformers -from datasets import Dataset -import datasets - -from verl.utils.data_process.filter import LengthFilter -from verl.utils.data_process.utils import set_seed, sample_dataset, save_dataset -from verl.utils.data_process.prompt import build_zero_style_prompt - - -InstructionFollow = "Please output the final answer within \\boxed{}." -PromptTemplate = """You are given one or more tables. Use the information in the tables to answer the following question. -{{tables}} -The question is: -{{question}} -""" - - -class MultiHierTTFilter(LengthFilter): - """Filter for MultiHierTT dataset to keep only arithmetic questions without text evidence.""" - - def __init__(self, tokenizer, max_length=2048): - super().__init__(tokenizer=tokenizer, max_length=max_length) - - def check(self, example): - # Only keep arithmetic questions - # TODO: Now comment to include non-operation questions to enlarge the dataset - # question_type = example["qa"]["question_type"] - # if question_type != "arithmetic": - # return False - # Filter out questions that need text evidence - text_evidence = example["qa"]["text_evidence"] - if text_evidence != []: - return False - # Ensure the prompt length is within the specified range - length_check = super().check(example) - if not length_check: - return False - - return True - - -def get_datasets(cache_dir, download=False): - """ - Load the MultiHierTT dataset from cache or download it if necessary. - - Args: - cache_dir (str): Directory to cache the dataset - download (bool): Whether to download the dataset if not found - - Returns: - tuple: (train_dataset, test_dataset) as Dataset objects - """ - data_path = os.path.join(cache_dir, "multihier/") - if download: - if os.path.exists(data_path): - pass - else: - import gdown - - # Replace with your folder ID or full URL - url = "https://drive.google.com/drive/folders/1ituEWZ5F7G9T9AZ0kzZZLrHNhRigHCZJ" - - # Download the folder to the cache directory - gdown.download_folder(url, output=data_path, quiet=False, use_cookies=False) - - train_path = os.path.join(data_path, "MultiHiertt_dataset/train.json") - test_path = os.path.join(data_path, "MultiHiertt_dataset/dev.json") - - def _align_format(dataset): - for entry in dataset: - entry["qa"]["answer"] = str(entry["qa"]["answer"]) - return dataset - - with open(train_path, "r") as f: - train_dataset = json.load(f) - train_dataset = _align_format(train_dataset) - with open(test_path, "r") as f: - test_dataset = json.load(f) - test_dataset = _align_format(test_dataset) - print("Total train samples:", len(train_dataset)) - print("Total test samples:", len(test_dataset)) - - return Dataset.from_list(train_dataset), Dataset.from_list(test_dataset) - - -def html_table_to_markdown(table): - """ - Convert HTML table to markdown format. - Handles hierarchical tables with colspan attributes. - - Args: - table (str): HTML table string - - Returns: - str: Markdown representation of the table - """ - import re - from bs4 import BeautifulSoup - - # Parse HTML - soup = BeautifulSoup(table, 'html.parser') - - # Get all rows - rows = soup.find_all('tr') - if not rows: - return "" - - # Process each row to determine column structure - table_data = [] - max_cols = 0 - - for row in rows: - row_data = [] - col_idx = 0 - - for cell in row.find_all(['td', 'th']): - # Get cell content - content = cell.get_text().strip() - - # Get colspan (default to 1) - colspan = int(cell.get('colspan', 1)) - - # Add cell with colspan info - row_data.append({ - 'content': content, - 'colspan': colspan, - 'col_idx': col_idx - }) - - col_idx += colspan - - max_cols = max(max_cols, col_idx) - table_data.append(row_data) - - # Create markdown table - markdown_rows = [] - - # Process each row - for i, row in enumerate(table_data): - md_row = [''] * max_cols - - # Fill in cells - for cell in row: - content = cell['content'] - col_idx = cell['col_idx'] - colspan = cell.get('colspan', 1) - - # For cells with colspan > 1, center the content - if colspan > 1: - md_row[col_idx] = content - # Fill in empty cells for the span - for j in range(1, colspan): - if col_idx + j < max_cols: - md_row[col_idx + j] = '' - else: - md_row[col_idx] = content - - # Join cells with pipe separator - markdown_rows.append('| ' + ' | '.join(md_row) + ' |') - - # Add header separator after first row - if i == 0: - separator = '| ' + ' | '.join(['---'] * max_cols) + ' |' - markdown_rows.append(separator) - - return '\n'.join(markdown_rows) - - -def make_map_fn(split: str, data_source: str): - def process_fn(example, idx): - try: - tables = example.pop("tables") - paragraphs = example.pop("paragraphs") - question = example["qa"]["question"] - answer = example["qa"]["answer"] - table_string = "\n".join([html_table_to_markdown(table) for table in tables]) - - prompt = PromptTemplate.replace("{{tables}}", table_string).replace("{{question}}", question) + InstructionFollow - - except Exception as e: - print(e) - print(tables) - exit() - - data = { - "data_source": data_source, - "prompt": [ - { - "role": "user", - "content": prompt - } - ], - "ability": "table", - "apply_chat_template": True, - "reward_model": {"style": "rule", "ground_truth": answer}, - "extra_info": { - "split": split, - "index": idx, - }, - } - if idx == 0 or idx == 1: - print("\n" + "=" * 10 + f"{data_source} {split} {idx}" + "=" * 10) - print(data) - return data - - return process_fn - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Download, process, and save MultiHierTT table reasoning datasets.") - parser.add_argument('--data-dir', default='data', - help='Base directory to save the processed data files.') - parser.add_argument('--domain', default="table", help='Domain of the dataset.') - parser.add_argument('--name', default="multihier", help='Name of the dataset.') - parser.add_argument('--train-sample-size', type=int, default=None, - help='Number of samples to use from training dataset. If None, use all samples.') - parser.add_argument('--test-sample-size', type=int, default=500, - help='Number of samples to use from test dataset. If None, use all samples.') - parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility') - - args = parser.parse_args() - - # Config - set_seed(args.seed) - data_source = f"{args.domain}__{args.name}" - train_output_dir = os.path.join(args.data_dir, 'train') - test_output_dir = os.path.join(args.data_dir, 'test') - - # Download dataset - cache_dir = datasets.config.HF_DATASETS_CACHE - train_dataset, test_dataset = get_datasets(cache_dir, download=True) - - # Process datasets - process_train_fn = make_map_fn('train', data_source) - process_test_fn = make_map_fn('test', data_source) - train_dataset = train_dataset.map(function=process_train_fn, with_indices=True) - test_dataset = test_dataset.map(function=process_test_fn, with_indices=True) - - # Filter dataset - tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B") - table_filter = MultiHierTTFilter(tokenizer=tokenizer, max_length=4096) - train_dataset = train_dataset.filter(lambda x: table_filter.check(x), num_proc=64) - test_dataset = test_dataset.filter(lambda x: table_filter.check(x), num_proc=64) - args.train_sample_size = len(train_dataset) if args.train_sample_size is None else args.train_sample_size - args.test_sample_size = len(test_dataset) if args.test_sample_size is None else args.test_sample_size - - # Sample datasets if specified - train_dataset = sample_dataset(train_dataset, args.train_sample_size) - test_dataset = sample_dataset(test_dataset, args.test_sample_size) - - # Save datasets - train_output_path = save_dataset( - dataset=train_dataset, - output_dir=train_output_dir, - filename_prefix=data_source, - sample_size=args.train_sample_size - ) - test_output_path = save_dataset( - dataset=test_dataset, - output_dir=test_output_dir, - filename_prefix=data_source, - sample_size=args.test_sample_size - ) - - print(f"\nDone! \n" - f"Data source: {data_source}\n" - f"Train data saved to {train_output_path} ({len(train_dataset)} samples)\n" - f"Test data saved to {test_output_path} ({len(test_dataset)} samples)") From 9a8d8b8e4a1fcf2718ac9ebbd5e360dfd4ecd216 Mon Sep 17 00:00:00 2001 From: AndreasXie Date: Wed, 4 Jun 2025 05:03:55 +0000 Subject: [PATCH 4/5] add examples --- examples/Guru/model_filtering.sh | 11 +++ examples/Guru/reward_fn_import.sh | 114 ++++++++++++++++++++++++++++++ 2 files changed, 125 insertions(+) create mode 100644 examples/Guru/model_filtering.sh create mode 100644 examples/Guru/reward_fn_import.sh diff --git a/examples/Guru/model_filtering.sh b/examples/Guru/model_filtering.sh new file mode 100644 index 000000000..005728353 --- /dev/null +++ b/examples/Guru/model_filtering.sh @@ -0,0 +1,11 @@ +#! /bin/bash +TOTAL_SIZE=15000 + +# math +MATH_SIZE=2500 +python reasoners/tools/model_filtering/run_sample_and_postprocess.py \ + --input_data_dir data/train_filtered \ + --input_data_names math__merged \ + --output_data_dir data/train_guru15k \ + --target_sample_size $MATH_SIZE \ + --domain math > math_sample.log diff --git a/examples/Guru/reward_fn_import.sh b/examples/Guru/reward_fn_import.sh new file mode 100644 index 000000000..90f95240a --- /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 From 3269cc3747b02cab0279efce06e07cec68e68dd8 Mon Sep 17 00:00:00 2001 From: AndreasXie Date: Thu, 5 Jun 2025 21:28:54 +0000 Subject: [PATCH 5/5] seperate dirs and update readme --- examples/Guru/README.md | 170 ++ examples/Guru/model_filtering.sh | 11 - examples/model_filtering/README.md | 433 +++++ .../model_filtering/run_add_pr.py | 0 .../model_filtering/run_filter.py | 0 .../model_filtering/run_inference.py | 4 +- .../model_filtering/run_reward.py | 4 +- .../run_sample_and_postprocess.py | 9 +- reasoners/reward/README.md | 1 - reasoners/reward/__init__.py | 1 - reasoners/reward/coder1/__init__.py | 14 +- reasoners/reward/cruxeval/__init__.py | 2 +- reasoners/reward/cruxeval/cruxeval.py | 0 reasoners/reward/geo3k.py | 31 - reasoners/reward/gsm8k.py | 63 - reasoners/reward/ifeval/__init__.py | 47 - reasoners/reward/ifeval/instructions.py | 1611 ---------------- .../reward/ifeval/instructions_registry.py | 168 -- reasoners/reward/ifeval/instructions_util.py | 1701 ----------------- reasoners/reward/livebench/__init__.py | 68 - .../livebench/data_analysis/cta/utils.py | 32 - .../data_analysis/tablejoin/utils.py | 82 - .../data_analysis/tablereformat/utils.py | 240 --- .../reasoning/house_traversal/utils.py | 41 - .../livebench/reasoning/spatial/utils.py | 61 - .../reasoning/web_of_lies_v2/utils.py | 79 - .../reasoning/web_of_lies_v3/utils.py | 87 - .../livebench/reasoning/zebra_puzzle/utils.py | 128 -- reasoners/reward/livebench/util.py | 46 - .../livebench/writing/connections/utils.py | 108 -- .../writing/plot_unscrambling/utils.py | 66 - .../reward/livebench/writing/typos/utils.py | 59 - reasoners/reward/math_batch.py | 25 - reasoners/reward/math_dapo.py | 294 --- reasoners/reward/math_verify.py | 33 - reasoners/reward/orz/__init__.py | 3 - reasoners/reward/orz/math_utils.py | 448 ----- reasoners/reward/orz/math_utils_sync.py | 442 ----- reasoners/reward/prime_code/__init__.py | 73 - reasoners/reward/prime_code/testing_util.py | 721 ------- reasoners/reward/prime_code/utils.py | 59 - reasoners/reward/tablereason.py | 4 +- reasoners/tools/model_filtering/README.md | 194 -- reasoners/tools/model_filtering/filter.py | 40 + .../tools/model_filtering/test_parquet.py | 113 -- reasoners/tools/model_filtering/utils.py | 55 + 46 files changed, 716 insertions(+), 7155 deletions(-) create mode 100644 examples/Guru/README.md delete mode 100644 examples/Guru/model_filtering.sh create mode 100644 examples/model_filtering/README.md rename {reasoners/tools => examples}/model_filtering/run_add_pr.py (100%) rename {reasoners/tools => examples}/model_filtering/run_filter.py (100%) rename {reasoners/tools => examples}/model_filtering/run_inference.py (98%) rename {reasoners/tools => examples}/model_filtering/run_reward.py (98%) rename {reasoners/tools => examples}/model_filtering/run_sample_and_postprocess.py (95%) delete mode 100644 reasoners/reward/README.md delete mode 100644 reasoners/reward/cruxeval/cruxeval.py delete mode 100644 reasoners/reward/geo3k.py delete mode 100644 reasoners/reward/gsm8k.py delete mode 100644 reasoners/reward/ifeval/__init__.py delete mode 100644 reasoners/reward/ifeval/instructions.py delete mode 100644 reasoners/reward/ifeval/instructions_registry.py delete mode 100644 reasoners/reward/ifeval/instructions_util.py delete mode 100644 reasoners/reward/livebench/__init__.py delete mode 100644 reasoners/reward/livebench/data_analysis/cta/utils.py delete mode 100644 reasoners/reward/livebench/data_analysis/tablejoin/utils.py delete mode 100644 reasoners/reward/livebench/data_analysis/tablereformat/utils.py delete mode 100644 reasoners/reward/livebench/reasoning/house_traversal/utils.py delete mode 100644 reasoners/reward/livebench/reasoning/spatial/utils.py delete mode 100644 reasoners/reward/livebench/reasoning/web_of_lies_v2/utils.py delete mode 100644 reasoners/reward/livebench/reasoning/web_of_lies_v3/utils.py delete mode 100644 reasoners/reward/livebench/reasoning/zebra_puzzle/utils.py delete mode 100644 reasoners/reward/livebench/util.py delete mode 100644 reasoners/reward/livebench/writing/connections/utils.py delete mode 100644 reasoners/reward/livebench/writing/plot_unscrambling/utils.py delete mode 100644 reasoners/reward/livebench/writing/typos/utils.py delete mode 100644 reasoners/reward/math_batch.py delete mode 100644 reasoners/reward/math_dapo.py delete mode 100644 reasoners/reward/math_verify.py delete mode 100644 reasoners/reward/orz/__init__.py delete mode 100644 reasoners/reward/orz/math_utils.py delete mode 100644 reasoners/reward/orz/math_utils_sync.py delete mode 100644 reasoners/reward/prime_code/__init__.py delete mode 100644 reasoners/reward/prime_code/testing_util.py delete mode 100644 reasoners/reward/prime_code/utils.py delete mode 100644 reasoners/tools/model_filtering/README.md create mode 100644 reasoners/tools/model_filtering/filter.py delete mode 100644 reasoners/tools/model_filtering/test_parquet.py diff --git a/examples/Guru/README.md b/examples/Guru/README.md new file mode 100644 index 000000000..ebf9445c4 --- /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/model_filtering.sh b/examples/Guru/model_filtering.sh deleted file mode 100644 index 005728353..000000000 --- a/examples/Guru/model_filtering.sh +++ /dev/null @@ -1,11 +0,0 @@ -#! /bin/bash -TOTAL_SIZE=15000 - -# math -MATH_SIZE=2500 -python reasoners/tools/model_filtering/run_sample_and_postprocess.py \ - --input_data_dir data/train_filtered \ - --input_data_names math__merged \ - --output_data_dir data/train_guru15k \ - --target_sample_size $MATH_SIZE \ - --domain math > math_sample.log diff --git a/examples/model_filtering/README.md b/examples/model_filtering/README.md new file mode 100644 index 000000000..bfe82c87a --- /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/reasoners/tools/model_filtering/run_add_pr.py b/examples/model_filtering/run_add_pr.py similarity index 100% rename from reasoners/tools/model_filtering/run_add_pr.py rename to examples/model_filtering/run_add_pr.py diff --git a/reasoners/tools/model_filtering/run_filter.py b/examples/model_filtering/run_filter.py similarity index 100% rename from reasoners/tools/model_filtering/run_filter.py rename to examples/model_filtering/run_filter.py diff --git a/reasoners/tools/model_filtering/run_inference.py b/examples/model_filtering/run_inference.py similarity index 98% rename from reasoners/tools/model_filtering/run_inference.py rename to examples/model_filtering/run_inference.py index b2dbe78d3..04a8b0279 100644 --- a/reasoners/tools/model_filtering/run_inference.py +++ b/examples/model_filtering/run_inference.py @@ -11,8 +11,8 @@ from datetime import datetime from rich.panel import Panel -from model_filtering.utils import console -from model_filtering.pipeline import DifficultyFilterPipeline +from reasoners.tools.model_filtering.utils import console +from reasoners.tools.model_filtering.pipeline import DifficultyFilterPipeline from datasets import Dataset # --------------------------------------------------------------------------- # diff --git a/reasoners/tools/model_filtering/run_reward.py b/examples/model_filtering/run_reward.py similarity index 98% rename from reasoners/tools/model_filtering/run_reward.py rename to examples/model_filtering/run_reward.py index 48298d441..7710186a1 100644 --- a/reasoners/tools/model_filtering/run_reward.py +++ b/examples/model_filtering/run_reward.py @@ -19,8 +19,8 @@ from tqdm import tqdm # Local imports -from verl.utils.reward_score import _default_compute_score -from model_filtering.utils import console, json_default +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 # diff --git a/reasoners/tools/model_filtering/run_sample_and_postprocess.py b/examples/model_filtering/run_sample_and_postprocess.py similarity index 95% rename from reasoners/tools/model_filtering/run_sample_and_postprocess.py rename to examples/model_filtering/run_sample_and_postprocess.py index 3240c6fed..8c834959e 100644 --- a/reasoners/tools/model_filtering/run_sample_and_postprocess.py +++ b/examples/model_filtering/run_sample_and_postprocess.py @@ -4,9 +4,8 @@ import json import transformers import numpy as np -from verl.utils.data_process.utils import save_dataset -from verl.utils.data_process.filter import LengthFilter - +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) @@ -123,9 +122,9 @@ def sample_dataset(dataset, target_sample_size) -> datasets.Dataset: if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument("--input_data_dir", type=str, default="/mnt/weka/home/zhuojun.cheng/leo/Reasoning360/data/train_filtered") + 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, default="/mnt/weka/home/zhuojun.cheng/leo/Reasoning360/data/train_guru15k") + 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) diff --git a/reasoners/reward/README.md b/reasoners/reward/README.md deleted file mode 100644 index 6422b56e0..000000000 --- a/reasoners/reward/README.md +++ /dev/null @@ -1 +0,0 @@ -to be done \ No newline at end of file diff --git a/reasoners/reward/__init__.py b/reasoners/reward/__init__.py index 11ed8844a..fa356829c 100644 --- a/reasoners/reward/__init__.py +++ b/reasoners/reward/__init__.py @@ -11,7 +11,6 @@ # 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. -# from . import gsm8k, math, prime_math, prime_code import json from typing import Dict, Any from reward_score.utils import _deserialise_extra, _decompress_str diff --git a/reasoners/reward/coder1/__init__.py b/reasoners/reward/coder1/__init__.py index a1818ac93..347963ca7 100644 --- a/reasoners/reward/coder1/__init__.py +++ b/reasoners/reward/coder1/__init__.py @@ -8,7 +8,7 @@ import numpy as np -from reward_score.coder1.utils import _ERROR_MSG_PREFIX +from reasoners.reward.coder1.utils import _ERROR_MSG_PREFIX _MAX_CHAR_DISPLAY = 2048 @@ -16,27 +16,27 @@ CODER1_EXEC = os.environ.get("CODER1_EXEC", "unsafe_local") if CODER1_EXEC == "docker": - from .docker_exec import code_exec_docker + from reasoners.reward.coder1.docker_exec import code_exec_docker code_exec = code_exec_docker elif CODER1_EXEC == "firejail": - from .firejail_exec import code_exec_firejail + from reasoners.reward.coder1.firejail_exec import code_exec_firejail code_exec = code_exec_firejail elif CODER1_EXEC == "ces": - from .ces_exec import remote_code_exec_ces + from reasoners.reward.coder1.ces_exec import remote_code_exec_ces code_exec = remote_code_exec_ces elif CODER1_EXEC == "kira": - from .kira_exec import remote_code_exec_kira + from reasoners.reward.coder1.kira_exec import remote_code_exec_kira code_exec = remote_code_exec_kira elif CODER1_EXEC == "bwrap": - from .bwrap_exec import code_exec_bwrap + from reasoners.reward.coder1.bwrap_exec import code_exec_bwrap code_exec = code_exec_bwrap elif CODER1_EXEC == "unsafe_local": - from .unsafe_local_exec import code_exec_local + from reasoners.reward.coder1.unsafe_local_exec import code_exec_local code_exec = code_exec_local else: diff --git a/reasoners/reward/cruxeval/__init__.py b/reasoners/reward/cruxeval/__init__.py index dc8e95bbd..9455fc00f 100644 --- a/reasoners/reward/cruxeval/__init__.py +++ b/reasoners/reward/cruxeval/__init__.py @@ -1,5 +1,5 @@ import re -from verl.utils.reward_score.cruxeval.utils import check_correctness +from reasoners.reward.cruxeval.utils import check_correctness def compute_score(data_source: str, diff --git a/reasoners/reward/cruxeval/cruxeval.py b/reasoners/reward/cruxeval/cruxeval.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/reasoners/reward/geo3k.py b/reasoners/reward/geo3k.py deleted file mode 100644 index 92c455469..000000000 --- a/reasoners/reward/geo3k.py +++ /dev/null @@ -1,31 +0,0 @@ -# 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 -from mathruler.grader import extract_boxed_content, grade_answer - - -def format_reward(predict_str: str) -> float: - pattern = re.compile(r'.*.*\\boxed\{.*\}.*', re.DOTALL) - match_result = re.fullmatch(pattern, predict_str) - return 1.0 if match_result else 0.0 - - -def acc_reward(predict_str: str, ground_truth: str) -> float: - answer = extract_boxed_content(predict_str) - return 1.0 if grade_answer(answer, ground_truth) else 0.0 - - -def compute_score(predict_str: str, ground_truth: str) -> float: - return 0.9 * acc_reward(predict_str, ground_truth) + 0.1 * format_reward(predict_str) diff --git a/reasoners/reward/gsm8k.py b/reasoners/reward/gsm8k.py deleted file mode 100644 index 709103764..000000000 --- a/reasoners/reward/gsm8k.py +++ /dev/null @@ -1,63 +0,0 @@ -# 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'): - assert method in ['strict', 'flexible'] - - if method == 'strict': - # this also tests the formatting of the model - solution = re.search("#### (\\-?[0-9\\.\\,]+)", solution_str) - if solution is None: - final_answer = None - else: - final_answer = solution.group(0) - final_answer = final_answer.split('#### ')[1].replace(',', '').replace('$', '') - elif method == 'flexible': - answer = re.findall("(\\-?[0-9\\.\\,]+)", solution_str) - final_answer = None - if len(answer) == 0: - # no reward is there is no answer - pass - else: - invalid_str = ['', '.'] - # find the last number that is not '.' - 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.): - """The scoring function for GSM8k. - - 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 - """ - 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/ifeval/__init__.py b/reasoners/reward/ifeval/__init__.py deleted file mode 100644 index 3f6c312fb..000000000 --- a/reasoners/reward/ifeval/__init__.py +++ /dev/null @@ -1,47 +0,0 @@ -from verl.utils.reward_score.ifeval import instructions_registry -import numpy as np -def compute_score(solution_str, ground_truth, extra_info): - """The scoring function for IFEval. - - 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 "" in solution_str: - answer = solution_str.split("")[1] - else: - answer = solution_str - is_following_list = [] - for index, instruction_id in enumerate(extra_info["instruction_id_list"]): - instruction_cls = instructions_registry.INSTRUCTION_DICT[instruction_id] - instruction = instruction_cls(instruction_id) - - # Remove None values from kwargs to avoid unexpected keyword argument errors in build_description method. - # print(ground_truth) - # print(extra_info["instruction_id_list"]) - # print(len(extra_info["instruction_id_list"])) - # if v is , turn it into a list - # if v is float, turn it into an int - kwargs = { - k: int(v) if isinstance(v, float) else v.tolist() if isinstance(v, np.ndarray) else v for k, v in ground_truth[index].items() if v is not None - } - - instruction.build_description(**kwargs) - args = instruction.get_instruction_args() - if args and "prompt" in args: - instruction.build_description(prompt=extra_info["prompt"]) - - if answer.strip() and instruction.check_following(answer): - is_following_list.append(True) - else: - is_following_list.append(False) - - return { - "score": all(is_following_list), - "acc": all(is_following_list), - } diff --git a/reasoners/reward/ifeval/instructions.py b/reasoners/reward/ifeval/instructions.py deleted file mode 100644 index 995536ec1..000000000 --- a/reasoners/reward/ifeval/instructions.py +++ /dev/null @@ -1,1611 +0,0 @@ -# Copyright 2023 The Google Research Authors. -# -# 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. - -"""Library of instructions.""" - -import collections -import json -import logging -import random -import re -import string -from typing import Dict, Optional, Sequence, Union - -import langdetect - -from verl.utils.reward_score.ifeval import instructions_util - - -logger = logging.getLogger(__name__) - -_InstructionArgsDtype = Optional[Dict[str, Union[int, str, Sequence[str]]]] - -_LANGUAGES = instructions_util.LANGUAGE_CODES - -# The relational operation for comparison. -_COMPARISON_RELATION = ("less than", "at least") - -# The maximum number of sentences. -_MAX_NUM_SENTENCES = 20 - -# The number of placeholders. -_NUM_PLACEHOLDERS = 4 - -# The number of bullet lists. -_NUM_BULLETS = 5 - -# The options of constrained response. -_CONSTRAINED_RESPONSE_OPTIONS = ( - "My answer is yes.", - "My answer is no.", - "My answer is maybe.", -) - -# The options of starter keywords. -_STARTER_OPTIONS = ( - "I would say", - "My answer is", - "I believe", - "In my opinion", - "I think", - "I reckon", - "I feel", - "From my perspective", - "As I see it", - "According to me", - "As far as I'm concerned", - "To my understanding", - "In my view", - "My take on it is", - "As per my perception", -) - -# The options of ending keywords. -# TODO(jeffreyzhou) add more ending options -_ENDING_OPTIONS = ("Any other questions?", "Is there anything else I can help with?") - -# The number of highlighted sections. -_NUM_HIGHLIGHTED_SECTIONS = 4 - -# The section splitter. -_SECTION_SPLITER = ("Section", "SECTION") - -# The number of sections. -_NUM_SECTIONS = 5 - -# The number of paragraphs. -_NUM_PARAGRAPHS = 5 - -# The postscript marker. -_POSTSCRIPT_MARKER = ("P.S.", "P.P.S") - -# The number of keywords. -_NUM_KEYWORDS = 2 - -# The occurrences of a single keyword. -_KEYWORD_FREQUENCY = 3 - -# The occurrences of a single letter. -_LETTER_FREQUENCY = 10 - -# The occurrences of words with all capital letters. -_ALL_CAPITAL_WORD_FREQUENCY = 20 - -# The number of words in the response. -_NUM_WORDS_LOWER_LIMIT = 100 -_NUM_WORDS_UPPER_LIMIT = 500 - - -class Instruction: - """An instruction template.""" - - def __init__(self, instruction_id): - self.id = instruction_id - - def build_description(self, **kwargs): - raise NotImplementedError("`build_description` not implemented.") - - def get_instruction_args(self): - raise NotImplementedError("`get_instruction_args` not implemented.") - - def get_instruction_args_keys(self): - raise NotImplementedError("`get_instruction_args_keys` not implemented.") - - def check_following(self, value): - raise NotImplementedError("`check_following` not implemented.") - - -class ResponseLanguageChecker(Instruction): - """Check the language of the entire response.""" - - def build_description(self, *, language=None): - """Build the instruction description. - - Args: - language: A string representing the expected language of the response. The - language has to comply to the 97 types defined in - `langid.py` (https://pypi.org/project/langid/1.1.5/), which follows - ISO 639-1 codes (https://en.wikipedia.org/wiki/List_of_ISO_639-1_codes); - for example, `en` for English, `zh` for Chinese, `fr` for French. - - Returns: - A string representing the instruction description. - """ - self._language = language - if self._language is None: - self._language = random.choice(list(_LANGUAGES.keys())) - # TODO(tianjianlu): opens the description generation to more choices. - self._description_pattern = ( - "Your ENTIRE response should be in {language} language, no other " - + "language is allowed." - ) - return self._description_pattern.format(language=_LANGUAGES[self._language]) - - def get_instruction_args(self): - """Returns the keyword args of `build_description`.""" - return {"language": self._language} - - def get_instruction_args_keys(self): - """Returns the args keys of `build_description`.""" - return ["language"] - - def check_following(self, value): - """Check if the language of the entire response follows the instruction. - - Args: - value: A string representing the response. - - Returns: - True if the language of `value` follows instruction; otherwise False. - """ - assert isinstance(value, str) - - try: - return langdetect.detect(value) == self._language - except langdetect.LangDetectException as e: - # Count as instruction is followed. - logging.error( - "Unable to detect language for text %s due to %s", value, e - ) # refex: disable=pytotw.037 - return True - - -class NumberOfSentences(Instruction): - """Check the number of sentences.""" - - def build_description(self, *, num_sentences=None, relation=None): - """Build the instruction description. - - Args: - num_sentences: An integer specifying the number of sentences as a - threshold. - relation: A string in (`less than`, `at least`), defining the relational - operator for comparison. - Two relational comparisons are supported for now: - if 'less than', the actual number of sentences < the threshold; - if 'at least', the actual number of sentences >= the threshold. - - Returns: - A string representing the instruction description. - """ - # The number of sentences as a threshold for comparison. - self._num_sentences_threshold = num_sentences - if self._num_sentences_threshold is None or self._num_sentences_threshold < 0: - self._num_sentences_threshold = random.randint(1, _MAX_NUM_SENTENCES) - - if relation is None: - self._comparison_relation = random.choice(_COMPARISON_RELATION) - elif relation not in _COMPARISON_RELATION: - raise ValueError( - "The supported relation for comparison must be in " - f"{_COMPARISON_RELATION}, but {relation} is given." - ) - else: - self._comparison_relation = relation - - self._description_pattern = ( - "Your response should contain {relation} {num_sentences} sentences." - ) - return self._description_pattern.format( - relation=self._comparison_relation, - num_sentences=self._num_sentences_threshold, - ) - - def get_instruction_args(self): - """Returns the keyword args of `build_description`.""" - return { - "num_sentences": self._num_sentences_threshold, - "relation": self._comparison_relation, - } - - def get_instruction_args_keys(self): - """Returns the args keys of `build_description`.""" - return ["num_sentences", "relation"] - - def check_following(self, value): - """Check if the number of sentences follows the instruction. - - Args: - value: A string representing the response. - - Returns: - True if the response follows the instruction. - - Raise: - ValueError if the string in `instruction_args` is not in - [`less_than`, `at_least`]. - """ - num_sentences = instructions_util.count_sentences(value) - if self._comparison_relation == _COMPARISON_RELATION[0]: - return num_sentences < self._num_sentences_threshold - elif self._comparison_relation == _COMPARISON_RELATION[1]: - return num_sentences >= self._num_sentences_threshold - - -class PlaceholderChecker(Instruction): - """Check the placeholders in template writing.""" - - def build_description(self, *, num_placeholders=None): - """Build the instruction description. - - Args: - num_placeholders: An integer denoting the minimum number of - placeholders required in the response. - - Returns: - A string representing the instruction description. - """ - self._num_placeholders = num_placeholders - if self._num_placeholders is None or self._num_placeholders < 0: - self._num_placeholders = random.randint(1, _NUM_PLACEHOLDERS) - self._description_pattern = ( - "The response must contain at least {num_placeholders} placeholders " - + "represented by square brackets, such as [address]." - ) - return self._description_pattern.format(num_placeholders=self._num_placeholders) - - def get_instruction_args(self): - """Returns the keyword args of `build_description`.""" - return {"num_placeholders": self._num_placeholders} - - def get_instruction_args_keys(self): - """Returns the args keys of `build_description`.""" - return ["num_placeholders"] - - def check_following(self, value): - """Check if the number of placeholders follows the instruction. - - Args: - value: A string representing the response. - - Returns: - True if the actual number of placeholders in the response is greater than - or equal to `num_placeholders`; otherwise, False. - """ - placeholders = re.findall(r"\[.*?\]", value) - num_placeholders = len(placeholders) - return num_placeholders >= self._num_placeholders - - -class BulletListChecker(Instruction): - """Checks the bullet list in the prompt.""" - - def build_description(self, *, num_bullets=None): - """Build the instruction description. - - Args: - num_bullets: An integer specifying the exact number of bullet lists - that is required to appear in the response. - - Returns: - A string representing the instruction description. - """ - self._num_bullets = num_bullets - if self._num_bullets is None or self._num_bullets < 0: - self._num_bullets = random.randint(1, _NUM_BULLETS) - self._description_pattern = ( - "Your answer must contain exactly {num_bullets} bullet points. " - + "Use the markdown bullet points such as:\n" - + "* This is point 1. \n" - + "* This is point 2" - ) - return self._description_pattern.format(num_bullets=self._num_bullets) - - def get_instruction_args(self): - """Returns the keyword args of `build_description`.""" - return {"num_bullets": self._num_bullets} - - def get_instruction_args_keys(self): - """Returns the args keys of `build_description`.""" - return ["num_bullets"] - - def check_following(self, value): - r"""Check if the number of bullet lists meets the requirement. - - Args: - value: A string representing the response. The response is expected to - contain some bullet lists that start with `\*`. - - Returns: - True if the actual number of bullet lists in the response meets the - requirement. - """ - bullet_lists = re.findall(r"^\s*\*[^\*].*$", value, flags=re.MULTILINE) - bullet_lists_2 = re.findall(r"^\s*-.*$", value, flags=re.MULTILINE) - num_bullet_lists = len(bullet_lists) + len(bullet_lists_2) - return num_bullet_lists == self._num_bullets - - -class ConstrainedResponseChecker(Instruction): - """Checks the constrained response.""" - - def build_description(self): - """Build the instruction description.""" - # A sequence of string(s) representing the options of the expected response. - self._constrained_responses = _CONSTRAINED_RESPONSE_OPTIONS - self._description_pattern = ( - "Answer with one of the following options: {response_options}" - ) - return self._description_pattern.format( - response_options=self._constrained_responses - ) - - def get_instruction_args(self): - """Returns the keyword args of `build_description`.""" - return None - - def get_instruction_args_keys(self): - """Returns the args keys of `build_description`.""" - return [] - - def check_following(self, value): - """Checks if the response matches the constrained options. - - Args: - value: A string representing the response. - - Returns: - True if the actual response contains one of the options in the constrained - responses; otherwise False. - """ - value = value.strip() - for constrained_response in self._constrained_responses: - if constrained_response in value: - return True - return False - - -class ConstrainedStartChecker(Instruction): - """Checks the response start.""" - - def build_description(self, *, starter=None): - """Build the instruction description. - - Args: - starter: A string representing the keyword that the response should start - with. - - Returns: - A string representing the instruction description. - """ - self._starter = starter.strip() if isinstance(starter, str) else starter - if self._starter is None: - self._starter = random.choice(_STARTER_OPTIONS) - self._description_pattern = ( - "During the conversation, when it is your turn, " - + "please always start with {starter}" - ) - return self._description_pattern.format(starter=self._starter) - - def get_instruction_args(self): - """Returns the keyword args of `build_description`.""" - return {"starter": self._starter} - - def get_instruction_args_keys(self): - """Returns the args keys of `build_description`.""" - return ["starter"] - - def check_following(self, value): - """Checks if the response starts with the constrained keyword or phrase. - - Args: - value: A string representing the response. - - Returns: - True if the response starts with the given phrase or keyword that is - contained in `instruction_args`; otherwise, False. - """ - response_pattern = r"^\s*" + self._starter + r".*$" - response_with_constrained_start = re.search( - response_pattern, value, flags=re.MULTILINE - ) - return True if response_with_constrained_start else False - - -class HighlightSectionChecker(Instruction): - """Checks the highlighted section.""" - - def build_description(self, *, num_highlights=None): - """Build the instruction description. - - Args: - num_highlights: An integer specifying the minimum number of highlighted - sections. - - Returns: - A string representing the instruction description. - """ - self._num_highlights = num_highlights - if self._num_highlights is None or self._num_highlights < 0: - self._num_highlights = random.randint(1, _NUM_HIGHLIGHTED_SECTIONS) - - self._description_pattern = ( - "Highlight at least {num_highlights} sections in your answer with " - + "markdown, i.e. *highlighted section*." - ) - - return self._description_pattern.format(num_highlights=self._num_highlights) - - def get_instruction_args(self): - """Returns the keyword args of `build_description`.""" - return {"num_highlights": self._num_highlights} - - def get_instruction_args_keys(self): - """Returns the args keys of `build_description`.""" - return ["num_highlights"] - - def check_following(self, value): - """Checks if the number of highlighted sections meets the requirement. - - Args: - value: a string representing the response. The response is expected to - contain highlighted sections in the format of *highlighted*. - - Returns: - True if the actual number of highlighted sections in the format of - *highlighted sections* meets the minimum requirement; otherwise False. - """ - num_highlights = 0 - highlights = re.findall(r"\*[^\n\*]*\*", value) - double_highlights = re.findall(r"\*\*[^\n\*]*\*\*", value) - for highlight in highlights: - if highlight.strip("*").strip(): - num_highlights += 1 - for highlight in double_highlights: - if highlight.removeprefix("**").removesuffix("**").strip(): - num_highlights += 1 - - return num_highlights >= self._num_highlights - - -class SectionChecker(Instruction): - """Checks the sections.""" - - def build_description(self, *, section_spliter=None, num_sections=None): - """Build the instruction description. - - Args: - section_spliter: A string represents the section spliter keyword that - marks a new section, i.e., `Section` or `SECTION`. - num_sections: An integer specifying the number of sections. - - Returns: - A string representing the instruction description. - """ - self._section_spliter = ( - section_spliter.strip() - if isinstance(section_spliter, str) - else section_spliter - ) - if self._section_spliter is None: - self._section_spliter = random.choice(_SECTION_SPLITER) - - self._num_sections = num_sections - if self._num_sections is None or self._num_sections < 0: - self._num_sections = random.randint(1, _NUM_SECTIONS) - - self._description_pattern = ( - "Your response must have {num_sections} sections. Mark the beginning " - + "of each section with {section_spliter} X, such as:\n" - + "{section_spliter} 1\n" - + "[content of section 1]\n" - + "{section_spliter} 2\n" - + "[content of section 2]" - ) - - return self._description_pattern.format( - num_sections=self._num_sections, section_spliter=self._section_spliter - ) - - def get_instruction_args(self): - """Returns the keyword args of `build_description`.""" - return { - "section_spliter": self._section_spliter, - "num_sections": self._num_sections, - } - - def get_instruction_args_keys(self): - """Returns the args keys of `build_description`.""" - return ["section_spliter", "num_sections"] - - def check_following(self, value): - """Checks the response contains multiple sections. - - Args: - value: A string representing the response. The response is expected - to contain multiple sections (number of sections is greater than 1). - A new section starts with `Section 1`, where the number denotes the - section index. - - Returns: - True if the number of sections in the response is greater than or equal to - the minimum number of sections; otherwise, False. - """ - section_splitter_patten = r"\s?" + self._section_spliter + r"\s?\d+\s?" - sections = re.split(section_splitter_patten, value) - num_sections = len(sections) - 1 - return num_sections >= self._num_sections - - -class ParagraphChecker(Instruction): - """Checks the paragraphs.""" - - def build_description(self, *, num_paragraphs=None): - """Build the instruction description. - - Args: - num_paragraphs: An integer specifying the number of paragraphs. - - Returns: - A string representing the instruction description. - """ - self._num_paragraphs = num_paragraphs - if self._num_paragraphs is None or self._num_paragraphs < 0: - self._num_paragraphs = random.randint(1, _NUM_PARAGRAPHS) - - self._description_pattern = ( - "There should be {num_paragraphs} paragraphs. " - + "Paragraphs are separated with the markdown divider: ***" - ) - - return self._description_pattern.format(num_paragraphs=self._num_paragraphs) - - def get_instruction_args(self): - """Returns the keyword args of `build_description`.""" - return {"num_paragraphs": self._num_paragraphs} - - def get_instruction_args_keys(self): - """Returns the args keys of `build_description`.""" - return ["num_paragraphs"] - - def check_following(self, value): - """Checks the response contains required number of paragraphs. - - Args: - value: A string representing the response. The response may contain - paragraphs that are separated by the markdown divider: `***`. - - Returns: - True if the actual number of paragraphs is the same as required; - otherwise, False. - """ - paragraphs = re.split(r"\s?\*\*\*\s?", value) - num_paragraphs = len(paragraphs) - - for index, paragraph in enumerate(paragraphs): - if not paragraph.strip(): - if index == 0 or index == len(paragraphs) - 1: - num_paragraphs -= 1 - else: - return False - - return num_paragraphs == self._num_paragraphs - - -class PostscriptChecker(Instruction): - """Checks the postscript.""" - - def build_description(self, *, postscript_marker=None): - """Build the instruction description. - - Args: - postscript_marker: A string containing the keyword that marks the start - of the postscript section. - - Returns: - A string representing the instruction description. - """ - self._postscript_marker = ( - postscript_marker.strip() - if isinstance(postscript_marker, str) - else postscript_marker - ) - if self._postscript_marker is None: - self._postscript_marker = random.choice(_POSTSCRIPT_MARKER) - - self._description_pattern = ( - "At the end of your response, please explicitly add a postscript " - + "starting with {postscript}" - ) - - return self._description_pattern.format(postscript=self._postscript_marker) - - def get_instruction_args(self): - """Returns the keyword args of `build_description`.""" - return {"postscript_marker": self._postscript_marker} - - def get_instruction_args_keys(self): - """Returns the args keys of `build_description`.""" - return ["postscript_marker"] - - def check_following(self, value): - """Checks if the response follows the postscript format. - - Args: - value: a string representing the response. The response is expected to - contain a postscript section. - - Returns: - True if the response contains a postscript section starting with - the keyword containing in the `instruction_args`; otherwise False. - """ - value = value.lower() - if self._postscript_marker == "P.P.S": - postscript_pattern = r"\s*p\.\s?p\.\s?s.*$" - elif self._postscript_marker == "P.S.": - postscript_pattern = r"\s*p\.\s?s\..*$" - else: - postscript_pattern = r"\s*" + self._postscript_marker.lower() + r".*$" - postscript = re.findall(postscript_pattern, value, flags=re.MULTILINE) - return True if postscript else False - - -class RephraseChecker(Instruction): - """Checks the rephrase.""" - - def build_description(self, *, original_message): - """Build the instruction description. - - Args: - original_message: A string representing the original message. The - rephrased response should only change its words/sentences in between - its two asterisks, for example, *change me*. Both original and rephrased - messages should contain the changes in the form of *change me*. - - Returns: - A string representing the instruction description. - """ - if not self.is_change(original_message): - raise ValueError( - f"Message {original_message} does not contain changes " - "in the form of *change me*." - ) - - self._reference_without_change = original_message - self._description = ( - "Rephrasing: Your rephrased response should only" - + "change the words/sentences in between two asterisks" - + "such as *change me*." - ) - return self._description - - def get_instruction_args(self): - """Returns the keyword args of `build_description`.""" - return {"original_message": self._reference_without_change} - - def get_instruction_args_keys(self): - """Returns the args keys of `build_description`.""" - return ["original_message"] - - def check_following(self, value): - r"""Checks if the rephrasing follows the instruction. - - Args: - value: A string representing the response, which is expected to rephras - the string of `instruction_args`. - - Returns: - True if `value` and `instruction_args` only differ by the words/sentences - in between two asterisks such as *change me*; otherwise, False. - """ - - if not self.is_change(value): - raise ValueError( - f"value {value} does not contain changes in the form of *change me*." - ) - - response_without_changes = self.strip_changes(value) - reference_without_changes = self.strip_changes(self._reference_without_change) - - return response_without_changes == reference_without_changes - - def is_change(self, response): - """Check if there is change in the response in the form of *change me*.""" - return re.search(r"\*.*\*", response) - - def strip_changes(self, response): - """Strips off the changes.""" - return re.sub(r"\*.*\*", "", response) - - -class KeywordChecker(Instruction): - """Check the exisitence of certain keywords.""" - - def build_description(self, *, keywords=None): - """Build the instruction description. - - Args: - keywords: A sequence of strings representing the keywords that are - expected in the response. - - Returns: - A string representing the instruction description. - """ - if not keywords: - self._keywords = instructions_util.generate_keywords( - num_keywords=_NUM_KEYWORDS - ) - else: - self._keywords = keywords - self._keywords = sorted(self._keywords) - - self._description_pattern = "Include keywords {keywords} in the response." - - return self._description_pattern.format(keywords=self._keywords) - - def get_instruction_args(self): - """Returns the keyword args of `build_description`.""" - return {"keywords": self._keywords} - - def get_instruction_args_keys(self): - """Returns the args keys of `build_description`.""" - return ["keywords"] - - def check_following(self, value): - """Check if the response contain the expected keywords.""" - for keyword in self._keywords: - if not re.search(keyword, value, flags=re.IGNORECASE): - return False - return True - - -class KeywordFrequencyChecker(Instruction): - """Check the keyword frequency.""" - - def build_description(self, *, keyword=None, frequency=None, relation=None): - """Build the instruction description. - - Args: - keyword: A string representing a keyword that is expected in the response. - frequency: An integer specifying the number of times `keyword` is expected - to appear in the response. - relation: A string in (`less than`, `at least`), defining the relational - operator for comparison. - Two relational comparisons are supported for now: - if 'less than', the actual number of occurrences < frequency; - if 'at least', the actual number of occurrences >= frequency. - - Returns: - A string representing the instruction description. - """ - if not keyword: - self._keyword = instructions_util.generate_keywords(num_keywords=1)[0] - else: - self._keyword = keyword.strip() - - self._frequency = frequency - if self._frequency is None or self._frequency < 0: - self._frequency = random.randint(1, _KEYWORD_FREQUENCY) - - if relation is None: - self._comparison_relation = random.choice(_COMPARISON_RELATION) - elif relation not in _COMPARISON_RELATION: - raise ValueError( - "The supported relation for comparison must be in " - f"{_COMPARISON_RELATION}, but {relation} is given." - ) - else: - self._comparison_relation = relation - - self._description_pattern = ( - "In your response, the word {keyword} should appear {relation} " - + "{frequency} times." - ) - - return self._description_pattern.format( - keyword=self._keyword, - relation=self._comparison_relation, - frequency=self._frequency, - ) - - def get_instruction_args(self): - """Returns the keyword args of `build_description`.""" - return { - "keyword": self._keyword, - "frequency": self._frequency, - "relation": self._comparison_relation, - } - - def get_instruction_args_keys(self): - """Returns the args keys of `build_description`.""" - return ["keyword", "frequency", "relation"] - - def check_following(self, value): - """Checks if the response contain the keyword with required frequency.""" - actual_occurrences = len(re.findall(self._keyword, value, flags=re.IGNORECASE)) - - if self._comparison_relation == _COMPARISON_RELATION[0]: - return actual_occurrences < self._frequency - elif self._comparison_relation == _COMPARISON_RELATION[1]: - return actual_occurrences >= self._frequency - - -class NumberOfWords(Instruction): - """Checks the number of words.""" - - def build_description(self, *, num_words=None, relation=None): - """Build the instruction description. - - Args: - num_words: An integer specifying the number of words contained in the - response. - relation: A string in (`less than`, `at least`), defining the relational - operator for comparison. - Two relational comparisons are supported for now: - if 'less than', the actual number of words < num_words; - if 'at least', the actual number of words >= num_words. - - Returns: - A string representing the instruction description. - """ - - self._num_words = num_words - if self._num_words is None or self._num_words < 0: - self._num_words = random.randint( - _NUM_WORDS_LOWER_LIMIT, _NUM_WORDS_UPPER_LIMIT - ) - - if relation is None: - self._comparison_relation = random.choice(_COMPARISON_RELATION) - elif relation not in _COMPARISON_RELATION: - raise ValueError( - "The supported relation for comparison must be in " - f"{_COMPARISON_RELATION}, but {relation} is given." - ) - else: - self._comparison_relation = relation - - self._description_pattern = "Answer with {relation} {num_words} words." - - return self._description_pattern.format( - relation=self._comparison_relation, num_words=self._num_words - ) - - def get_instruction_args(self): - """Returns the keyword args of `build_description`.""" - return {"num_words": self._num_words, "relation": self._comparison_relation} - - def get_instruction_args_keys(self): - """Returns the args keys of `build_description`.""" - return ["num_words", "relation"] - - def check_following(self, value): - """Checks if the response contains the expected number of words.""" - num_words = instructions_util.count_words(value) - - if self._comparison_relation == _COMPARISON_RELATION[0]: - return num_words < self._num_words - elif self._comparison_relation == _COMPARISON_RELATION[1]: - return num_words >= self._num_words - - -class JsonFormat(Instruction): - """Check the Json format.""" - - def build_description(self): - self._description_pattern = ( - "Entire output should be wrapped in JSON format. You can use markdown" - " ticks such as ```." - ) - return self._description_pattern - - def get_instruction_args(self): - """Returns the keyword args of `build_description`.""" - return None - - def get_instruction_args_keys(self): - """Returns the args keys of `build_description`.""" - return [] - - def check_following(self, value): - value = ( - value.strip() - .removeprefix("```json") - .removeprefix("```Json") - .removeprefix("```JSON") - .removeprefix("```") - .removesuffix("```") - .strip() - ) - try: - json.loads(value) - except ValueError: - return False - return True - - -class ParagraphFirstWordCheck(Instruction): - """Check the paragraph and the first word of the nth paragraph.""" - - def build_description( - self, num_paragraphs=None, nth_paragraph=None, first_word=None - ): - r"""Build the instruction description. - - Args: - num_paragraphs: An integer indicating the number of paragraphs expected - in the response. A paragraph is a subset of the string that is - expected to be separated by '\n\n'. - nth_paragraph: An integer indicating the paragraph number that we look at. - Note that n starts from 1. - first_word: A string that represent the first word of the bth paragraph. - - Returns: - A string representing the instruction description. - """ - self._num_paragraphs = num_paragraphs - if self._num_paragraphs is None or self._num_paragraphs < 0: - self._num_paragraphs = random.randint(1, _NUM_PARAGRAPHS) - - self._nth_paragraph = nth_paragraph - if ( - self._nth_paragraph is None - or self._nth_paragraph <= 0 - or self._nth_paragraph > self._num_paragraphs - ): - self._nth_paragraph = random.randint(1, self._num_paragraphs + 1) - - self._first_word = first_word - if self._first_word is None: - self._first_word = instructions_util.generate_keywords(num_keywords=1)[0] - self._first_word = self._first_word.lower() - - self._description_pattern = ( - "There should be {num_paragraphs} paragraphs. " - + "Paragraphs and only paragraphs are separated with each other by two " - + "new lines as if it was '\\n\\n' in python. " - + "Paragraph {nth_paragraph} must start with word {first_word}." - ) - - return self._description_pattern.format( - num_paragraphs=self._num_paragraphs, - nth_paragraph=self._nth_paragraph, - first_word=self._first_word, - ) - - def get_instruction_args(self): - """Returns the keyword args of `build_description`.""" - return { - "num_paragraphs": self._num_paragraphs, - "nth_paragraph": self._nth_paragraph, - "first_word": self._first_word, - } - - def get_instruction_args_keys(self): - """Returns the args keys of `build_description`.""" - return ["num_paragraphs", "nth_paragraph", "first_word"] - - def check_following(self, value): - """Checks for required number of paragraphs and correct first word. - - Args: - value: a string representing the response. The response may contain - paragraphs that are separated by two new lines and the first word of - the nth paragraph will have to match a specified word. - - Returns: - True if the number of paragraphs is the same as required and the first - word of the specified paragraph is the same as required. Otherwise, false. - """ - - paragraphs = re.split(r"\n\n", value) - num_paragraphs = len(paragraphs) - - for paragraph in paragraphs: - if not paragraph.strip(): - num_paragraphs -= 1 - - # check that index doesn't go out of bounds - if self._nth_paragraph <= num_paragraphs: - paragraph = paragraphs[self._nth_paragraph - 1].strip() - if not paragraph: - return False - else: - return False - - first_word = "" - punctuation = {".", ",", "?", "!", "'", '"'} - - # get first word and remove punctuation - word = paragraph.split()[0].strip() - # TODO(jeffrey): make more complex? - word = word.lstrip("'") - word = word.lstrip('"') - - for letter in word: - if letter in punctuation: - break - first_word += letter.lower() - - return num_paragraphs == self._num_paragraphs and first_word == self._first_word - - -# TODO(jeffrey) add relation - at least/at most? -class KeySentenceChecker(Instruction): - """Check the existence of certain key sentences.""" - - def build_description(self, key_sentences=None, num_sentences=None): - """Build the instruction description. - - Args: - key_sentences: A sequences of strings representing the key sentences that - are expected in the response. - num_sentences: The number of key sentences that are expected to be seen in - the response. - - Returns: - A string representing the instruction description. - """ - - if not key_sentences: - # TODO(jeffrey) make a generate sentences function? wonderwords package - self._key_sentences = set(["For now, this is fine."]) - else: - self._key_sentences = key_sentences - - if not num_sentences: - self._num_sentences = random.randint(1, len(self._key_sentences)) - else: - self._num_sentences = num_sentences - - self._description_pattern = ( - "Include {num_sentences} of the following sentences {key_sentences}" - ) - - return self._description_pattern.format( - num_sentences=self._num_sentences, key_sentences=self._key_sentences - ) - - def get_instruction_args(self): - """Returns the keyword args of `build_description`.""" - return { - "num_sentences": self._num_sentences, - "key_sentences": list(self._key_sentences), - } - - def get_instruction_args_keys(self): - """Returns the args keys of `build_description`.""" - return ["num_sentences", "key_sentences"] - - def check_following(self, value): - """Checks if the response contains the expected key sentences.""" - count = 0 - sentences = instructions_util.split_into_sentences(value) - for sentence in self._key_sentences: - if sentence in sentences: - count += 1 - - return count == self._num_sentences - - -class ForbiddenWords(Instruction): - """Checks that specified words are not used in response.""" - - def build_description(self, forbidden_words=None): - """Build the instruction description. - - Args: - forbidden_words: A sequences of strings representing words that are not - allowed in the response. - - Returns: - A string representing the instruction description. - """ - - if not forbidden_words: - self._forbidden_words = instructions_util.generate_keywords( - num_keywords=_NUM_KEYWORDS - ) - else: - self._forbidden_words = list(set(forbidden_words)) - self._forbidden_words = sorted(self._forbidden_words) - self._description_pattern = ( - "Do not include keywords {forbidden_words} in the response." - ) - - return self._description_pattern.format(forbidden_words=self._forbidden_words) - - def get_instruction_args(self): - """Returns the keyword args of `build_description`.""" - return {"forbidden_words": self._forbidden_words} - - def get_instruction_args_keys(self): - """Returns the args keys of `build_description`.""" - return ["forbidden_words"] - - def check_following(self, value): - """Check if the response does not contain the expected keywords.""" - for word in self._forbidden_words: - if re.search(r"\b" + word + r"\b", value, flags=re.IGNORECASE): - return False - return True - - -class RephraseParagraph(Instruction): - """Checks that the paragraph is rephrased.""" - - def build_description(self, *, original_paragraph, low, high): - """Builds the instruction description. - - Args: - original_paragraph: A string presenting the original paragraph. The - rephrases response should have betweeb low-high words in common. - low: An integer presenting the lower bound of similar words. - high: An integer representing the upper bound of similar words. - - Returns: - A string representing the instruction description. - """ - # TODO(jeffrey) make more encompassing - self._original_paragraph = original_paragraph - self._low = low - self._high = high - - self._description = ( - "Rephrase the following paragraph: " - + "{original_paragraph}\nYour response should have " - + "between {low} and {high} of the same words. " - + "Words are the same if and only if all of the " - + "letters, ignoring cases, are the same. For " - + "example, 'run' is the same as 'Run' but different " - + "to 'ran'." - ) - - return self._description.format( - original_paragraph=original_paragraph, low=self._low, high=self._high - ) - - def get_instruction_args(self): - """Returns the keyword args of `build_description`.""" - return { - "original_paragraph": self._original_paragraph, - "low": self._low, - "high": self._high, - } - - def get_instruction_args_keys(self): - """Returns the args keys of `build_description`.""" - return ["original_paragraph", "low", "high"] - - def check_following(self, value): - val_words = re.findall(r"\w+", value.lower()) - original_words = re.findall(r"\w+", self._original_paragraph.lower()) - similar_words = 0 - - dict_val = collections.Counter(val_words) - dict_original = collections.Counter(original_words) - - for word in dict_original: - similar_words += min(dict_original[word], dict_val[word]) - - return similar_words >= self._low and similar_words <= self._high - - -class TwoResponsesChecker(Instruction): - """Check that two responses were given.""" - - def build_description(self): - """Build the instruction description.""" - self._description_pattern = ( - "Give two different responses. Responses and only responses should" - " be separated by 6 asterisk symbols: ******." - ) - return self._description_pattern - - def get_instruction_args(self): - """Returns the keyword args of `build_description`.""" - return None - - def get_instruction_args_keys(self): - """Returns the args keys of `build_description`.""" - return [] - - def check_following(self, value): - """Checks if the response has two different answers. - - Args: - value: A string representing the response. - - Returns: - True if two responses are detected and false otherwise. - """ - valid_responses = list() - responses = value.split("******") - for index, response in enumerate(responses): - if not response.strip(): - if index != 0 and index != len(responses) - 1: - return False - else: - valid_responses.append(response) - return ( - len(valid_responses) == 2 - and valid_responses[0].strip() != valid_responses[1].strip() - ) - - -class RepeatPromptThenAnswer(Instruction): - """Checks that Prompt is first repeated then answered.""" - - def build_description(self, *, prompt_to_repeat=None): - """Build the instruction description. - - Args: - prompt_to_repeat: The prompt that is meant to be repeated. - - Returns: - A string representing the instruction description. - """ - if not prompt_to_repeat: - raise ValueError("prompt_to_repeat must be set.") - else: - self._prompt_to_repeat = prompt_to_repeat - self._description_pattern = ( - "First repeat the request word for word without change," - " then give your answer (1. do not say any words or characters" - " before repeating the request; 2. the request you need to repeat" - " does not include this sentence)" - ) - return self._description_pattern - - def get_instruction_args(self): - return {"prompt_to_repeat": self._prompt_to_repeat} - - def get_instruction_args_keys(self): - """Returns the args keys of `build_description`.""" - return ["prompt_to_repeat"] - - def check_following(self, value): - if value.strip().lower().startswith(self._prompt_to_repeat.strip().lower()): - return True - return False - - -class EndChecker(Instruction): - """Checks that the prompt ends with a given phrase.""" - - def build_description(self, *, end_phrase=None): - """Build the instruction description. - - Args: - end_phrase: A string representing the phrase the response should end with. - - Returns: - A string representing the instruction description. - """ - self._end_phrase = ( - end_phrase.strip() if isinstance(end_phrase, str) else end_phrase - ) - if self._end_phrase is None: - self._end_phrase = random.choice(_ENDING_OPTIONS) - self._description_pattern = ( - "Finish your response with this exact phrase {ender}. " - "No other words should follow this phrase." - ) - return self._description_pattern.format(ender=self._end_phrase) - - def get_instruction_args(self): - return {"end_phrase": self._end_phrase} - - def get_instruction_args_keys(self): - """Returns the args keys of `build_description`.""" - return ["end_phrase"] - - def check_following(self, value): - """Checks if the response ends with the expected phrase.""" - value = value.strip().strip('"').lower() - self._end_phrase = self._end_phrase.strip().lower() - return value.endswith(self._end_phrase) - - -class TitleChecker(Instruction): - """Checks the response for a title.""" - - def build_description(self): - """Build the instruction description.""" - self._description_pattern = ( - "Your answer must contain a title, wrapped in double angular brackets," - " such as <>." - ) - return self._description_pattern - - def get_instruction_args(self): - return None - - def get_instruction_args_keys(self): - """Returns the args keys of `build_description`.""" - return [] - - def check_following(self, value): - """Checks if the response contains a title.""" - pattern = r"<<[^\n]+>>" - re_pattern = re.compile(pattern) - titles = re.findall(re_pattern, value) - - for title in titles: - if title.lstrip("<").rstrip(">").strip(): - return True - return False - - -class LetterFrequencyChecker(Instruction): - """Checks letter frequency.""" - - def build_description(self, *, letter=None, let_frequency=None, let_relation=None): - """Build the instruction description. - - Args: - letter: A string representing a letter that is expected in the response. - let_frequency: An integer specifying the number of times `keyword` is - expected to appear in the response. - let_relation: A string in (`less than`, `at least`), defining the - relational operator for comparison. Two relational comparisons are - supported for now; if 'less than', the actual number of - occurrences < frequency; if 'at least', the actual number of - occurrences >= frequency. - - Returns: - A string representing the instruction description. - """ - if ( - not letter - or len(letter) > 1 - or ord(letter.lower()) < 97 - or ord(letter.lower()) > 122 - ): - self._letter = random.choice(list(string.ascii_letters)) - else: - self._letter = letter.strip() - self._letter = self._letter.lower() - - self._frequency = let_frequency - if self._frequency is None or self._frequency < 0: - self._frequency = random.randint(1, _LETTER_FREQUENCY) - - if let_relation is None: - self._comparison_relation = random.choice(_COMPARISON_RELATION) - elif let_relation not in _COMPARISON_RELATION: - raise ValueError( - "The supported relation for comparison must be in " - f"{_COMPARISON_RELATION}, but {let_relation} is given." - ) - else: - self._comparison_relation = let_relation - - self._description_pattern = ( - "In your response, the letter {letter} should appear {let_relation}" - " {let_frequency} times." - ) - - return self._description_pattern.format( - letter=self._letter, - let_frequency=self._frequency, - let_relation=self._comparison_relation, - ) - - def get_instruction_args(self): - """Returns the keyword args of build description.""" - return { - "letter": self._letter, - "let_frequency": self._frequency, - "let_relation": self._comparison_relation, - } - - def get_instruction_args_keys(self): - """Returns the args keys of `build_description`.""" - return ["letter", "let_frequency", "let_relation"] - - def check_following(self, value): - """Checks that the response contains the letter at the right frequency.""" - value = value.lower() - letters = collections.Counter(value) - - if self._comparison_relation == _COMPARISON_RELATION[0]: - return letters[self._letter] < self._frequency - else: - return letters[self._letter] >= self._frequency - - -class CapitalLettersEnglishChecker(Instruction): - """Checks that the response is in english and is in all capital letters.""" - - def build_description(self): - """Build the instruction description.""" - self._description_pattern = ( - "Your entire response should be in English, and in all capital letters." - ) - return self._description_pattern - - def get_instruction_args(self): - return None - - def get_instruction_args_keys(self): - """Returns the args keys of `build_description`.""" - return [] - - def check_following(self, value): - """Checks that the response is in English and in all capital letters.""" - assert isinstance(value, str) - - try: - return value.isupper() and langdetect.detect(value) == "en" - except langdetect.LangDetectException as e: - # Count as instruction is followed. - logging.error( - "Unable to detect language for text %s due to %s", value, e - ) # refex: disable=pytotw.037 - return True - - -class LowercaseLettersEnglishChecker(Instruction): - """Checks that the response is in english and is in all lowercase letters.""" - - def build_description(self): - """Build the instruction description.""" - self._description_pattern = ( - "Your entire response should be in English, and in all lowercase" - " letters. No capital letters are allowed." - ) - return self._description_pattern - - def get_instruction_args(self): - return None - - def get_instruction_args_keys(self): - """Returns the args keys of `build_description`.""" - return [] - - def check_following(self, value): - """Checks that the response is in English and in all lowercase letters.""" - assert isinstance(value, str) - - try: - return value.islower() and langdetect.detect(value) == "en" - except langdetect.LangDetectException as e: - # Count as instruction is followed. - logging.error( - "Unable to detect language for text %s due to %s", value, e - ) # refex: disable=pytotw.037 - return True - - -class CommaChecker(Instruction): - """Checks the response for no commas.""" - - def build_description(self): - """Build the instruction description.""" - self._description_pattern = ( - "In your entire response, refrain from the use of any commas." - ) - return self._description_pattern - - def get_instruction_args(self): - return None - - def get_instruction_args_keys(self): - """Returns the args keys of `build_description`.""" - return [] - - def check_following(self, value): - """Checks that the response does not contain commas.""" - return not re.search(r"\,", value) - - -class CapitalWordFrequencyChecker(Instruction): - """Checks frequency of words with all capital letters.""" - - def build_description( - self, - capital_frequency=None, - capital_relation=None, - ): - """Build the instruction description. - - Args: - capital_frequency: An integer that represents the number of words that - should be in all capital letters. - capital_relation: A string that is 'at least' or 'at most' that refers to - the frequency. - - Returns: - A string representing the instruction description. - """ - self._frequency = capital_frequency - if self._frequency is None: - self._frequency = random.randint(1, _ALL_CAPITAL_WORD_FREQUENCY) - - self._comparison_relation = capital_relation - if capital_relation is None: - self._comparison_relation = random.choice(_COMPARISON_RELATION) - elif capital_relation not in _COMPARISON_RELATION: - raise ValueError( - "The supported relation for comparison must be in " - f"{_COMPARISON_RELATION}, but {capital_relation} is given." - ) - - self._description_pattern = ( - "In your response, words with all capital letters should appear" - " {relation} {frequency} times." - ) - - return self._description_pattern.format( - frequency=self._frequency, relation=self._comparison_relation - ) - - def get_instruction_args(self): - """Returns the keyword args of build description.""" - return { - "capital_frequency": self._frequency, - "capital_relation": self._comparison_relation, - } - - def get_instruction_args_keys(self): - """Returns the args keys of `build_description`.""" - return ["capital_frequency", "capital_relation"] - - def check_following(self, value): - """Checks the frequency of words with all capital letters.""" - # Hyphenated words will count as one word - words = instructions_util.nltk.word_tokenize(value) - capital_words = [word for word in words if word.isupper()] - - capital_words = len(capital_words) - - if self._comparison_relation == _COMPARISON_RELATION[0]: - return capital_words < self._frequency - else: - return capital_words >= self._frequency - - -class QuotationChecker(Instruction): - """Checks response is wrapped with double quotation marks.""" - - def build_description(self): - """Build the instruction description.""" - self._description_pattern = ( - "Wrap your entire response with double quotation marks." - ) - return self._description_pattern - - def get_instruction_args(self): - """Returns the keyword args of build description.""" - return None - - def get_instruction_args_keys(self): - """Returns the args keys of `build_description`.""" - return [] - - def check_following(self, value): - """Checks if the response is wrapped with double quotation marks.""" - value = value.strip() - return len(value) > 1 and value[0] == '"' and value[-1] == '"' \ No newline at end of file diff --git a/reasoners/reward/ifeval/instructions_registry.py b/reasoners/reward/ifeval/instructions_registry.py deleted file mode 100644 index cde7fb810..000000000 --- a/reasoners/reward/ifeval/instructions_registry.py +++ /dev/null @@ -1,168 +0,0 @@ -# Copyright 2023 The Google Research Authors. -# -# 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. - -"""Registry of all instructions.""" - -from verl.utils.reward_score.ifeval import instructions - - -_KEYWORD = "keywords:" - -_LANGUAGE = "language:" - -_LENGTH = "length_constraints:" - -_CONTENT = "detectable_content:" - -_FORMAT = "detectable_format:" - -_MULTITURN = "multi-turn:" - -_COMBINATION = "combination:" - -_STARTEND = "startend:" - -_CHANGE_CASES = "change_case:" - -_PUNCTUATION = "punctuation:" - -INSTRUCTION_DICT = { - _KEYWORD + "existence": instructions.KeywordChecker, - _KEYWORD + "frequency": instructions.KeywordFrequencyChecker, - # TODO(jeffreyzhou): make a proper set of sentences to choose from - # _KEYWORD + "key_sentences": instructions.KeySentenceChecker, - _KEYWORD + "forbidden_words": instructions.ForbiddenWords, - _KEYWORD + "letter_frequency": instructions.LetterFrequencyChecker, - _LANGUAGE + "response_language": instructions.ResponseLanguageChecker, - _LENGTH + "number_sentences": instructions.NumberOfSentences, - _LENGTH + "number_paragraphs": instructions.ParagraphChecker, - _LENGTH + "number_words": instructions.NumberOfWords, - _LENGTH + "nth_paragraph_first_word": instructions.ParagraphFirstWordCheck, - _CONTENT + "number_placeholders": instructions.PlaceholderChecker, - _CONTENT + "postscript": instructions.PostscriptChecker, - _FORMAT + "number_bullet_lists": instructions.BulletListChecker, - # TODO(jeffreyzhou): Pre-create paragraph or use prompt to replace - # _CONTENT + "rephrase_paragraph": instructions.RephraseParagraph, - _FORMAT + "constrained_response": instructions.ConstrainedResponseChecker, - _FORMAT + "number_highlighted_sections": (instructions.HighlightSectionChecker), - _FORMAT + "multiple_sections": instructions.SectionChecker, - # TODO(tianjianlu): Re-enable rephrasing with preprocessing the message. - # _FORMAT + "rephrase": instructions.RephraseChecker, - _FORMAT + "json_format": instructions.JsonFormat, - _FORMAT + "title": instructions.TitleChecker, - # TODO(tianjianlu): Re-enable with specific prompts. - # _MULTITURN + "constrained_start": instructions.ConstrainedStartChecker, - _COMBINATION + "two_responses": instructions.TwoResponsesChecker, - _COMBINATION + "repeat_prompt": instructions.RepeatPromptThenAnswer, - _STARTEND + "end_checker": instructions.EndChecker, - _CHANGE_CASES + "capital_word_frequency": instructions.CapitalWordFrequencyChecker, - _CHANGE_CASES + "english_capital": instructions.CapitalLettersEnglishChecker, - _CHANGE_CASES + "english_lowercase": instructions.LowercaseLettersEnglishChecker, - _PUNCTUATION + "no_comma": instructions.CommaChecker, - _STARTEND + "quotation": instructions.QuotationChecker, -} - -INSTRUCTION_CONFLICTS = { - _KEYWORD + "existence": {_KEYWORD + "existence"}, - _KEYWORD + "frequency": {_KEYWORD + "frequency"}, - # TODO(jeffreyzhou): make a proper set of sentences to choose from - # _KEYWORD + "key_sentences": instructions.KeySentenceChecker, - _KEYWORD + "forbidden_words": {_KEYWORD + "forbidden_words"}, - _KEYWORD + "letter_frequency": {_KEYWORD + "letter_frequency"}, - _LANGUAGE + "response_language": { - _LANGUAGE + "response_language", - _FORMAT + "multiple_sections", - _KEYWORD + "existence", - _KEYWORD + "frequency", - _KEYWORD + "forbidden_words", - _STARTEND + "end_checker", - _CHANGE_CASES + "english_capital", - _CHANGE_CASES + "english_lowercase", - }, - _LENGTH + "number_sentences": {_LENGTH + "number_sentences"}, - _LENGTH + "number_paragraphs": { - _LENGTH + "number_paragraphs", - _LENGTH + "nth_paragraph_first_word", - _LENGTH + "number_sentences", - _LENGTH + "nth_paragraph_first_word", - }, - _LENGTH + "number_words": {_LENGTH + "number_words"}, - _LENGTH + "nth_paragraph_first_word": { - _LENGTH + "nth_paragraph_first_word", - _LENGTH + "number_paragraphs", - }, - _CONTENT + "number_placeholders": {_CONTENT + "number_placeholders"}, - _CONTENT + "postscript": {_CONTENT + "postscript"}, - _FORMAT + "number_bullet_lists": {_FORMAT + "number_bullet_lists"}, - # TODO(jeffreyzhou): Pre-create paragraph or use prompt to replace - # _CONTENT + "rephrase_paragraph": instructions.RephraseParagraph, - _FORMAT + "constrained_response": set(INSTRUCTION_DICT.keys()), - _FORMAT + "number_highlighted_sections": {_FORMAT + "number_highlighted_sections"}, - _FORMAT + "multiple_sections": { - _FORMAT + "multiple_sections", - _LANGUAGE + "response_language", - _FORMAT + "number_highlighted_sections", - }, - # TODO(tianjianlu): Re-enable rephrasing with preprocessing the message. - # _FORMAT + "rephrase": instructions.RephraseChecker, - _FORMAT + "json_format": set(INSTRUCTION_DICT.keys()).difference( - {_KEYWORD + "forbidden_words", _KEYWORD + "existence"} - ), - _FORMAT + "title": {_FORMAT + "title"}, - # TODO(tianjianlu): Re-enable with specific prompts. - # _MULTITURN + "constrained_start": instructions.ConstrainedStartChecker, - _COMBINATION + "two_responses": set(INSTRUCTION_DICT.keys()).difference( - { - _KEYWORD + "forbidden_words", - _KEYWORD + "existence", - _LANGUAGE + "response_language", - _FORMAT + "title", - _PUNCTUATION + "no_comma", - } - ), - _COMBINATION + "repeat_prompt": set(INSTRUCTION_DICT.keys()).difference( - {_KEYWORD + "existence", _FORMAT + "title", _PUNCTUATION + "no_comma"} - ), - _STARTEND + "end_checker": {_STARTEND + "end_checker"}, - _CHANGE_CASES + "capital_word_frequency": { - _CHANGE_CASES + "capital_word_frequency", - _CHANGE_CASES + "english_lowercase", - _CHANGE_CASES + "english_capital", - }, - _CHANGE_CASES + "english_capital": {_CHANGE_CASES + "english_capital"}, - _CHANGE_CASES + "english_lowercase": { - _CHANGE_CASES + "english_lowercase", - _CHANGE_CASES + "english_capital", - }, - _PUNCTUATION + "no_comma": {_PUNCTUATION + "no_comma"}, - _STARTEND + "quotation": {_STARTEND + "quotation", _FORMAT + "title"}, -} - - -def conflict_make(conflicts): - """Makes sure if A conflicts with B, B will conflict with A. - - Args: - conflicts: Dictionary of potential conflicts where key is instruction id - and value is set of instruction ids that it conflicts with. - - Returns: - Revised version of the dictionary. All instructions conflict with - themselves. If A conflicts with B, B will conflict with A. - """ - for key in conflicts: - for k in conflicts[key]: - conflicts[k].add(key) - conflicts[key].add(key) - return conflicts \ No newline at end of file diff --git a/reasoners/reward/ifeval/instructions_util.py b/reasoners/reward/ifeval/instructions_util.py deleted file mode 100644 index 33e0a0a00..000000000 --- a/reasoners/reward/ifeval/instructions_util.py +++ /dev/null @@ -1,1701 +0,0 @@ -# Copyright 2023 The Google Research Authors. -# -# 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. - -"""Utility library of instructions.""" - -import functools -import os -import random -import re -from importlib.metadata import version - -import immutabledict -import nltk -from packaging.version import parse as parse_version - - -# Downloading 'punkt' with nltk<3.9 has a remote code vuln. -# see https://github.com/EleutherAI/lm-evaluation-harness/issues/2210 -# and https://github.com/nltk/nltk/issues/3266 -# for more information. -NLTK_MIN_VERSION = "3.9.1" -RANK = os.environ.get("LOCAL_RANK", "0") - - -def download_nltk_resources(): - """Download 'punkt' if not already installed""" - assert (nltk_version := parse_version(version("nltk"))) >= parse_version( - NLTK_MIN_VERSION - ), ( - f"`nltk` version {nltk_version} is not >= {NLTK_MIN_VERSION}. Please update `nltk` before proceeding--older versions are vulnerable to a remote code execution vulnerability." - ) - - try: - nltk.data.find("tokenizers/punkt_tab") - except LookupError: - if RANK == "0": - nltk.download("punkt_tab") - print("Downloaded punkt_tab on rank 0") - - -download_nltk_resources() - -WORD_LIST = [ - "western", - "sentence", - "signal", - "dump", - "spot", - "opposite", - "bottom", - "potato", - "administration", - "working", - "welcome", - "morning", - "good", - "agency", - "primary", - "wish", - "responsibility", - "press", - "problem", - "president", - "steal", - "brush", - "read", - "type", - "beat", - "trainer", - "growth", - "lock", - "bone", - "case", - "equal", - "comfortable", - "region", - "replacement", - "performance", - "mate", - "walk", - "medicine", - "film", - "thing", - "rock", - "tap", - "total", - "competition", - "ease", - "south", - "establishment", - "gather", - "parking", - "world", - "plenty", - "breath", - "claim", - "alcohol", - "trade", - "dear", - "highlight", - "street", - "matter", - "decision", - "mess", - "agreement", - "studio", - "coach", - "assist", - "brain", - "wing", - "style", - "private", - "top", - "brown", - "leg", - "buy", - "procedure", - "method", - "speed", - "high", - "company", - "valuable", - "pie", - "analyst", - "session", - "pattern", - "district", - "pleasure", - "dinner", - "swimming", - "joke", - "order", - "plate", - "department", - "motor", - "cell", - "spend", - "cabinet", - "difference", - "power", - "examination", - "engine", - "horse", - "dimension", - "pay", - "toe", - "curve", - "literature", - "bother", - "fire", - "possibility", - "debate", - "activity", - "passage", - "hello", - "cycle", - "background", - "quiet", - "author", - "effect", - "actor", - "page", - "bicycle", - "error", - "throat", - "attack", - "character", - "phone", - "tea", - "increase", - "outcome", - "file", - "specific", - "inspector", - "internal", - "potential", - "staff", - "building", - "employer", - "shoe", - "hand", - "direction", - "garden", - "purchase", - "interview", - "study", - "recognition", - "member", - "spiritual", - "oven", - "sandwich", - "weird", - "passenger", - "particular", - "response", - "reaction", - "size", - "variation", - "a", - "cancel", - "candy", - "exit", - "guest", - "condition", - "fly", - "price", - "weakness", - "convert", - "hotel", - "great", - "mouth", - "mind", - "song", - "sugar", - "suspect", - "telephone", - "ear", - "roof", - "paint", - "refrigerator", - "organization", - "jury", - "reward", - "engineering", - "day", - "possession", - "crew", - "bar", - "road", - "description", - "celebration", - "score", - "mark", - "letter", - "shower", - "suggestion", - "sir", - "luck", - "national", - "progress", - "hall", - "stroke", - "theory", - "offer", - "story", - "tax", - "definition", - "history", - "ride", - "medium", - "opening", - "glass", - "elevator", - "stomach", - "question", - "ability", - "leading", - "village", - "computer", - "city", - "grand", - "confidence", - "candle", - "priest", - "recommendation", - "point", - "necessary", - "body", - "desk", - "secret", - "horror", - "noise", - "culture", - "warning", - "water", - "round", - "diet", - "flower", - "bus", - "tough", - "permission", - "week", - "prompt", - "connection", - "abuse", - "height", - "save", - "corner", - "border", - "stress", - "drive", - "stop", - "rip", - "meal", - "listen", - "confusion", - "girlfriend", - "living", - "relation", - "significance", - "plan", - "creative", - "atmosphere", - "blame", - "invite", - "housing", - "paper", - "drink", - "roll", - "silver", - "drunk", - "age", - "damage", - "smoke", - "environment", - "pack", - "savings", - "influence", - "tourist", - "rain", - "post", - "sign", - "grandmother", - "run", - "profit", - "push", - "clerk", - "final", - "wine", - "swim", - "pause", - "stuff", - "singer", - "funeral", - "average", - "source", - "scene", - "tradition", - "personal", - "snow", - "nobody", - "distance", - "sort", - "sensitive", - "animal", - "major", - "negotiation", - "click", - "mood", - "period", - "arrival", - "expression", - "holiday", - "repeat", - "dust", - "closet", - "gold", - "bad", - "sail", - "combination", - "clothes", - "emphasis", - "duty", - "black", - "step", - "school", - "jump", - "document", - "professional", - "lip", - "chemical", - "front", - "wake", - "while", - "inside", - "watch", - "row", - "subject", - "penalty", - "balance", - "possible", - "adult", - "aside", - "sample", - "appeal", - "wedding", - "depth", - "king", - "award", - "wife", - "blow", - "site", - "camp", - "music", - "safe", - "gift", - "fault", - "guess", - "act", - "shame", - "drama", - "capital", - "exam", - "stupid", - "record", - "sound", - "swing", - "novel", - "minimum", - "ratio", - "machine", - "shape", - "lead", - "operation", - "salary", - "cloud", - "affair", - "hit", - "chapter", - "stage", - "quantity", - "access", - "army", - "chain", - "traffic", - "kick", - "analysis", - "airport", - "time", - "vacation", - "philosophy", - "ball", - "chest", - "thanks", - "place", - "mountain", - "advertising", - "red", - "past", - "rent", - "return", - "tour", - "house", - "construction", - "net", - "native", - "war", - "figure", - "fee", - "spray", - "user", - "dirt", - "shot", - "task", - "stick", - "friend", - "software", - "promotion", - "interaction", - "surround", - "block", - "purpose", - "practice", - "conflict", - "routine", - "requirement", - "bonus", - "hole", - "state", - "junior", - "sweet", - "catch", - "tear", - "fold", - "wall", - "editor", - "life", - "position", - "pound", - "respect", - "bathroom", - "coat", - "script", - "job", - "teach", - "birth", - "view", - "resolve", - "theme", - "employee", - "doubt", - "market", - "education", - "serve", - "recover", - "tone", - "harm", - "miss", - "union", - "understanding", - "cow", - "river", - "association", - "concept", - "training", - "recipe", - "relationship", - "reserve", - "depression", - "proof", - "hair", - "revenue", - "independent", - "lift", - "assignment", - "temporary", - "amount", - "loss", - "edge", - "track", - "check", - "rope", - "estimate", - "pollution", - "stable", - "message", - "delivery", - "perspective", - "mirror", - "assistant", - "representative", - "witness", - "nature", - "judge", - "fruit", - "tip", - "devil", - "town", - "emergency", - "upper", - "drop", - "stay", - "human", - "neck", - "speaker", - "network", - "sing", - "resist", - "league", - "trip", - "signature", - "lawyer", - "importance", - "gas", - "choice", - "engineer", - "success", - "part", - "external", - "worker", - "simple", - "quarter", - "student", - "heart", - "pass", - "spite", - "shift", - "rough", - "lady", - "grass", - "community", - "garage", - "youth", - "standard", - "skirt", - "promise", - "blind", - "television", - "disease", - "commission", - "positive", - "energy", - "calm", - "presence", - "tune", - "basis", - "preference", - "head", - "common", - "cut", - "somewhere", - "presentation", - "current", - "thought", - "revolution", - "effort", - "master", - "implement", - "republic", - "floor", - "principle", - "stranger", - "shoulder", - "grade", - "button", - "tennis", - "police", - "collection", - "account", - "register", - "glove", - "divide", - "professor", - "chair", - "priority", - "combine", - "peace", - "extension", - "maybe", - "evening", - "frame", - "sister", - "wave", - "code", - "application", - "mouse", - "match", - "counter", - "bottle", - "half", - "cheek", - "resolution", - "back", - "knowledge", - "make", - "discussion", - "screw", - "length", - "accident", - "battle", - "dress", - "knee", - "log", - "package", - "it", - "turn", - "hearing", - "newspaper", - "layer", - "wealth", - "profile", - "imagination", - "answer", - "weekend", - "teacher", - "appearance", - "meet", - "bike", - "rise", - "belt", - "crash", - "bowl", - "equivalent", - "support", - "image", - "poem", - "risk", - "excitement", - "remote", - "secretary", - "public", - "produce", - "plane", - "display", - "money", - "sand", - "situation", - "punch", - "customer", - "title", - "shake", - "mortgage", - "option", - "number", - "pop", - "window", - "extent", - "nothing", - "experience", - "opinion", - "departure", - "dance", - "indication", - "boy", - "material", - "band", - "leader", - "sun", - "beautiful", - "muscle", - "farmer", - "variety", - "fat", - "handle", - "director", - "opportunity", - "calendar", - "outside", - "pace", - "bath", - "fish", - "consequence", - "put", - "owner", - "go", - "doctor", - "information", - "share", - "hurt", - "protection", - "career", - "finance", - "force", - "golf", - "garbage", - "aspect", - "kid", - "food", - "boot", - "milk", - "respond", - "objective", - "reality", - "raw", - "ring", - "mall", - "one", - "impact", - "area", - "news", - "international", - "series", - "impress", - "mother", - "shelter", - "strike", - "loan", - "month", - "seat", - "anything", - "entertainment", - "familiar", - "clue", - "year", - "glad", - "supermarket", - "natural", - "god", - "cost", - "conversation", - "tie", - "ruin", - "comfort", - "earth", - "storm", - "percentage", - "assistance", - "budget", - "strength", - "beginning", - "sleep", - "other", - "young", - "unit", - "fill", - "store", - "desire", - "hide", - "value", - "cup", - "maintenance", - "nurse", - "function", - "tower", - "role", - "class", - "camera", - "database", - "panic", - "nation", - "basket", - "ice", - "art", - "spirit", - "chart", - "exchange", - "feedback", - "statement", - "reputation", - "search", - "hunt", - "exercise", - "nasty", - "notice", - "male", - "yard", - "annual", - "collar", - "date", - "platform", - "plant", - "fortune", - "passion", - "friendship", - "spread", - "cancer", - "ticket", - "attitude", - "island", - "active", - "object", - "service", - "buyer", - "bite", - "card", - "face", - "steak", - "proposal", - "patient", - "heat", - "rule", - "resident", - "broad", - "politics", - "west", - "knife", - "expert", - "girl", - "design", - "salt", - "baseball", - "grab", - "inspection", - "cousin", - "couple", - "magazine", - "cook", - "dependent", - "security", - "chicken", - "version", - "currency", - "ladder", - "scheme", - "kitchen", - "employment", - "local", - "attention", - "manager", - "fact", - "cover", - "sad", - "guard", - "relative", - "county", - "rate", - "lunch", - "program", - "initiative", - "gear", - "bridge", - "breast", - "talk", - "dish", - "guarantee", - "beer", - "vehicle", - "reception", - "woman", - "substance", - "copy", - "lecture", - "advantage", - "park", - "cold", - "death", - "mix", - "hold", - "scale", - "tomorrow", - "blood", - "request", - "green", - "cookie", - "church", - "strip", - "forever", - "beyond", - "debt", - "tackle", - "wash", - "following", - "feel", - "maximum", - "sector", - "sea", - "property", - "economics", - "menu", - "bench", - "try", - "language", - "start", - "call", - "solid", - "address", - "income", - "foot", - "senior", - "honey", - "few", - "mixture", - "cash", - "grocery", - "link", - "map", - "form", - "factor", - "pot", - "model", - "writer", - "farm", - "winter", - "skill", - "anywhere", - "birthday", - "policy", - "release", - "husband", - "lab", - "hurry", - "mail", - "equipment", - "sink", - "pair", - "driver", - "consideration", - "leather", - "skin", - "blue", - "boat", - "sale", - "brick", - "two", - "feed", - "square", - "dot", - "rush", - "dream", - "location", - "afternoon", - "manufacturer", - "control", - "occasion", - "trouble", - "introduction", - "advice", - "bet", - "eat", - "kill", - "category", - "manner", - "office", - "estate", - "pride", - "awareness", - "slip", - "crack", - "client", - "nail", - "shoot", - "membership", - "soft", - "anybody", - "web", - "official", - "individual", - "pizza", - "interest", - "bag", - "spell", - "profession", - "queen", - "deal", - "resource", - "ship", - "guy", - "chocolate", - "joint", - "formal", - "upstairs", - "car", - "resort", - "abroad", - "dealer", - "associate", - "finger", - "surgery", - "comment", - "team", - "detail", - "crazy", - "path", - "tale", - "initial", - "arm", - "radio", - "demand", - "single", - "draw", - "yellow", - "contest", - "piece", - "quote", - "pull", - "commercial", - "shirt", - "contribution", - "cream", - "channel", - "suit", - "discipline", - "instruction", - "concert", - "speech", - "low", - "effective", - "hang", - "scratch", - "industry", - "breakfast", - "lay", - "join", - "metal", - "bedroom", - "minute", - "product", - "rest", - "temperature", - "many", - "give", - "argument", - "print", - "purple", - "laugh", - "health", - "credit", - "investment", - "sell", - "setting", - "lesson", - "egg", - "middle", - "marriage", - "level", - "evidence", - "phrase", - "love", - "self", - "benefit", - "guidance", - "affect", - "you", - "dad", - "anxiety", - "special", - "boyfriend", - "test", - "blank", - "payment", - "soup", - "obligation", - "reply", - "smile", - "deep", - "complaint", - "addition", - "review", - "box", - "towel", - "minor", - "fun", - "soil", - "issue", - "cigarette", - "internet", - "gain", - "tell", - "entry", - "spare", - "incident", - "family", - "refuse", - "branch", - "can", - "pen", - "grandfather", - "constant", - "tank", - "uncle", - "climate", - "ground", - "volume", - "communication", - "kind", - "poet", - "child", - "screen", - "mine", - "quit", - "gene", - "lack", - "charity", - "memory", - "tooth", - "fear", - "mention", - "marketing", - "reveal", - "reason", - "court", - "season", - "freedom", - "land", - "sport", - "audience", - "classroom", - "law", - "hook", - "win", - "carry", - "eye", - "smell", - "distribution", - "research", - "country", - "dare", - "hope", - "whereas", - "stretch", - "library", - "if", - "delay", - "college", - "plastic", - "book", - "present", - "use", - "worry", - "champion", - "goal", - "economy", - "march", - "election", - "reflection", - "midnight", - "slide", - "inflation", - "action", - "challenge", - "guitar", - "coast", - "apple", - "campaign", - "field", - "jacket", - "sense", - "way", - "visual", - "remove", - "weather", - "trash", - "cable", - "regret", - "buddy", - "beach", - "historian", - "courage", - "sympathy", - "truck", - "tension", - "permit", - "nose", - "bed", - "son", - "person", - "base", - "meat", - "usual", - "air", - "meeting", - "worth", - "game", - "independence", - "physical", - "brief", - "play", - "raise", - "board", - "she", - "key", - "writing", - "pick", - "command", - "party", - "yesterday", - "spring", - "candidate", - "physics", - "university", - "concern", - "development", - "change", - "string", - "target", - "instance", - "room", - "bitter", - "bird", - "football", - "normal", - "split", - "impression", - "wood", - "long", - "meaning", - "stock", - "cap", - "leadership", - "media", - "ambition", - "fishing", - "essay", - "salad", - "repair", - "today", - "designer", - "night", - "bank", - "drawing", - "inevitable", - "phase", - "vast", - "chip", - "anger", - "switch", - "cry", - "twist", - "personality", - "attempt", - "storage", - "being", - "preparation", - "bat", - "selection", - "white", - "technology", - "contract", - "side", - "section", - "station", - "till", - "structure", - "tongue", - "taste", - "truth", - "difficulty", - "group", - "limit", - "main", - "move", - "feeling", - "light", - "example", - "mission", - "might", - "wait", - "wheel", - "shop", - "host", - "classic", - "alternative", - "cause", - "agent", - "consist", - "table", - "airline", - "text", - "pool", - "craft", - "range", - "fuel", - "tool", - "partner", - "load", - "entrance", - "deposit", - "hate", - "article", - "video", - "summer", - "feature", - "extreme", - "mobile", - "hospital", - "flight", - "fall", - "pension", - "piano", - "fail", - "result", - "rub", - "gap", - "system", - "report", - "suck", - "ordinary", - "wind", - "nerve", - "ask", - "shine", - "note", - "line", - "mom", - "perception", - "brother", - "reference", - "bend", - "charge", - "treat", - "trick", - "term", - "homework", - "bake", - "bid", - "status", - "project", - "strategy", - "orange", - "let", - "enthusiasm", - "parent", - "concentrate", - "device", - "travel", - "poetry", - "business", - "society", - "kiss", - "end", - "vegetable", - "employ", - "schedule", - "hour", - "brave", - "focus", - "process", - "movie", - "illegal", - "general", - "coffee", - "ad", - "highway", - "chemistry", - "psychology", - "hire", - "bell", - "conference", - "relief", - "show", - "neat", - "funny", - "weight", - "quality", - "club", - "daughter", - "zone", - "touch", - "tonight", - "shock", - "burn", - "excuse", - "name", - "survey", - "landscape", - "advance", - "satisfaction", - "bread", - "disaster", - "item", - "hat", - "prior", - "shopping", - "visit", - "east", - "photo", - "home", - "idea", - "father", - "comparison", - "cat", - "pipe", - "winner", - "count", - "lake", - "fight", - "prize", - "foundation", - "dog", - "keep", - "ideal", - "fan", - "struggle", - "peak", - "safety", - "solution", - "hell", - "conclusion", - "population", - "strain", - "alarm", - "measurement", - "second", - "train", - "race", - "due", - "insurance", - "boss", - "tree", - "monitor", - "sick", - "course", - "drag", - "appointment", - "slice", - "still", - "care", - "patience", - "rich", - "escape", - "emotion", - "royal", - "female", - "childhood", - "government", - "picture", - "will", - "sock", - "big", - "gate", - "oil", - "cross", - "pin", - "improvement", - "championship", - "silly", - "help", - "sky", - "pitch", - "man", - "diamond", - "most", - "transition", - "work", - "science", - "committee", - "moment", - "fix", - "teaching", - "dig", - "specialist", - "complex", - "guide", - "people", - "dead", - "voice", - "original", - "break", - "topic", - "data", - "degree", - "reading", - "recording", - "bunch", - "reach", - "judgment", - "lie", - "regular", - "set", - "painting", - "mode", - "list", - "player", - "bear", - "north", - "wonder", - "carpet", - "heavy", - "officer", - "negative", - "clock", - "unique", - "baby", - "pain", - "assumption", - "disk", - "iron", - "bill", - "drawer", - "look", - "double", - "mistake", - "finish", - "future", - "brilliant", - "contact", - "math", - "rice", - "leave", - "restaurant", - "discount", - "sex", - "virus", - "bit", - "trust", - "event", - "wear", - "juice", - "failure", - "bug", - "context", - "mud", - "whole", - "wrap", - "intention", - "draft", - "pressure", - "cake", - "dark", - "explanation", - "space", - "angle", - "word", - "efficiency", - "management", - "habit", - "star", - "chance", - "finding", - "transportation", - "stand", - "criticism", - "flow", - "door", - "injury", - "insect", - "surprise", - "apartment", -] # pylint: disable=line-too-long - -# ISO 639-1 codes to language names. -LANGUAGE_CODES = immutabledict.immutabledict( - { - "en": "English", - "es": "Spanish", - "pt": "Portuguese", - "ar": "Arabic", - "hi": "Hindi", - "fr": "French", - "ru": "Russian", - "de": "German", - "ja": "Japanese", - "it": "Italian", - "bn": "Bengali", - "uk": "Ukrainian", - "th": "Thai", - "ur": "Urdu", - "ta": "Tamil", - "te": "Telugu", - "bg": "Bulgarian", - "ko": "Korean", - "pl": "Polish", - "he": "Hebrew", - "fa": "Persian", - "vi": "Vietnamese", - "ne": "Nepali", - "sw": "Swahili", - "kn": "Kannada", - "mr": "Marathi", - "gu": "Gujarati", - "pa": "Punjabi", - "ml": "Malayalam", - "fi": "Finnish", - } -) - -_ALPHABETS = "([A-Za-z])" -_PREFIXES = "(Mr|St|Mrs|Ms|Dr)[.]" -_SUFFIXES = "(Inc|Ltd|Jr|Sr|Co)" -_STARTERS = r"(Mr|Mrs|Ms|Dr|Prof|Capt|Cpt|Lt|He\s|She\s|It\s|They\s|Their\s|Our\s|We\s|But\s|However\s|That\s|This\s|Wherever)" -_ACRONYMS = "([A-Z][.][A-Z][.](?:[A-Z][.])?)" -_WEBSITES = "[.](com|net|org|io|gov|edu|me)" -_DIGITS = "([0-9])" -_MULTIPLE_DOTS = r"\.{2,}" - - -def split_into_sentences(text): - """Split the text into sentences. - - Args: - text: A string that consists of more than or equal to one sentences. - - Returns: - A list of strings where each string is a sentence. - """ - text = " " + text + " " - text = text.replace("\n", " ") - text = re.sub(_PREFIXES, "\\1", text) - text = re.sub(_WEBSITES, "\\1", text) - text = re.sub(_DIGITS + "[.]" + _DIGITS, "\\1\\2", text) - text = re.sub( - _MULTIPLE_DOTS, - lambda match: "" * len(match.group(0)) + "", - text, - ) - if "Ph.D" in text: - text = text.replace("Ph.D.", "PhD") - text = re.sub(r"\s" + _ALPHABETS + "[.] ", " \\1 ", text) - text = re.sub(_ACRONYMS + " " + _STARTERS, "\\1 \\2", text) - text = re.sub( - _ALPHABETS + "[.]" + _ALPHABETS + "[.]" + _ALPHABETS + "[.]", - "\\1\\2\\3", - text, - ) - text = re.sub(_ALPHABETS + "[.]" + _ALPHABETS + "[.]", "\\1\\2", text) - text = re.sub(" " + _SUFFIXES + "[.] " + _STARTERS, " \\1 \\2", text) - text = re.sub(" " + _SUFFIXES + "[.]", " \\1", text) - text = re.sub(" " + _ALPHABETS + "[.]", " \\1", text) - if "”" in text: - text = text.replace(".”", "”.") - if '"' in text: - text = text.replace('."', '".') - if "!" in text: - text = text.replace('!"', '"!') - if "?" in text: - text = text.replace('?"', '"?') - text = text.replace(".", ".") - text = text.replace("?", "?") - text = text.replace("!", "!") - text = text.replace("", ".") - sentences = text.split("") - sentences = [s.strip() for s in sentences] - if sentences and not sentences[-1]: - sentences = sentences[:-1] - return sentences - - -def count_words(text): - """Counts the number of words.""" - tokenizer = nltk.tokenize.RegexpTokenizer(r"\w+") - tokens = tokenizer.tokenize(text) - num_words = len(tokens) - return num_words - - -@functools.lru_cache(maxsize=None) -def _get_sentence_tokenizer(): - return nltk.data.load("nltk:tokenizers/punkt/english.pickle") - - -def count_sentences(text): - """Count the number of sentences.""" - tokenizer = _get_sentence_tokenizer() - tokenized_sentences = tokenizer.tokenize(text) - return len(tokenized_sentences) - - -def generate_keywords(num_keywords): - """Randomly generates a few keywords.""" - return random.sample(WORD_LIST, k=num_keywords) diff --git a/reasoners/reward/livebench/__init__.py b/reasoners/reward/livebench/__init__.py deleted file mode 100644 index 908100cf0..000000000 --- a/reasoners/reward/livebench/__init__.py +++ /dev/null @@ -1,68 +0,0 @@ -import re -import numpy as np -import ast -from ast import literal_eval -import json - -# data analysis split -from verl.utils.reward_score.livebench.data_analysis.tablereformat.utils import table_process_results -from verl.utils.reward_score.livebench.data_analysis.cta.utils import cta_process_results -from verl.utils.reward_score.livebench.data_analysis.tablejoin.utils import joinmap_process_results - -# reasoning split -from verl.utils.reward_score.livebench.reasoning.web_of_lies_v2.utils import web_of_lies_process_results -from verl.utils.reward_score.livebench.reasoning.web_of_lies_v3.utils import web_of_lies_v3_process_results -from verl.utils.reward_score.livebench.reasoning.house_traversal.utils import house_traversal_process_results -from verl.utils.reward_score.livebench.reasoning.zebra_puzzle.utils import zebra_puzzle_process_results_old # get_zebra_puzzle_evaluator -from verl.utils.reward_score.livebench.reasoning.spatial.utils import spatial_process_results - -# language split -from verl.utils.reward_score.livebench.writing.plot_unscrambling.utils import plot_unscrambling_process_results -from verl.utils.reward_score.livebench.writing.typos.utils import typos_process_results -from verl.utils.reward_score.livebench.writing.connections.utils import connections_process_results_old # get_connections_puzzle_evaluator - -# excluded splits -# from livebench.process_results.math.math_competitions.utils import mathcontest_process_results,aime_process_results -# from livebench.process_results.math.olympiad.utils import proof_rearrangement_process_results -# from livebench.process_results.math.AMPS_Hard.utils import amps_hard_process_results -# from livebench.process_results.coding.utils import LCB_generation_process_results, code_generation_process_results -# from livebench.process_results.instruction_following.utils import instruction_following_process_results - - - -def compute_score(solution_str, ground_truth, extra_info): - """ - Compute the score for the solution string. - """ - score = 0 - - # data analysis split - if extra_info["task"] == "cta": - score = cta_process_results(ground_truth, solution_str) - elif extra_info["task"] == "tablejoin": - score = joinmap_process_results(ground_truth, solution_str) - elif extra_info["task"] == "tableformat": - score = table_process_results(ground_truth, solution_str) - # reasoning split - elif extra_info["task"] == "web_of_lies_v2": - score = web_of_lies_process_results(ground_truth, solution_str) - elif extra_info["task"] == "web_of_lies_v3": - score = web_of_lies_v3_process_results(ground_truth, solution_str) - elif extra_info["task"] == "house_traversal": - score = house_traversal_process_results(ground_truth, solution_str) - elif extra_info["task"] == "zebra_puzzle": - score = zebra_puzzle_process_results_old(ground_truth, solution_str) - elif extra_info["task"] == "spatial": - score = spatial_process_results(ground_truth, solution_str) - # language split - elif extra_info["task"] == "plot_unscrambling": - score = plot_unscrambling_process_results(ground_truth, solution_str) - elif extra_info["task"] == "typos": - score = typos_process_results(ground_truth, solution_str) - elif extra_info["task"] == "connections": - score = connections_process_results_old(ground_truth, solution_str) - - return { - "score": score, - "acc": score, - } diff --git a/reasoners/reward/livebench/data_analysis/cta/utils.py b/reasoners/reward/livebench/data_analysis/cta/utils.py deleted file mode 100644 index c96f20bd8..000000000 --- a/reasoners/reward/livebench/data_analysis/cta/utils.py +++ /dev/null @@ -1,32 +0,0 @@ -import re -from verl.utils.reward_score.livebench.util import last_boxed_only_string, remove_boxed - - -def clean_text(text): - text = text.lower().strip() - text = re.sub(r'[^\w]', '', text) - return text - - -def cta_process_results(ground_truth: str, llm_answer: str, debug=False) -> int: - - parsed_answer = llm_answer - - if '\\boxed{' in parsed_answer: - parsed_answer = last_boxed_only_string(parsed_answer) - parsed_answer = remove_boxed(parsed_answer) - parsed_answer = parsed_answer.replace('\\text{', '').replace('}', '').replace('\\', '') - - if clean_text(ground_truth) == clean_text(parsed_answer): - return 1 - elif clean_text(ground_truth) == clean_text(parsed_answer)[-len(clean_text(ground_truth)):]: - return 1 - else: - if debug: - print('INCORRECT') - print('GROUND TRUTH', ground_truth) - print('SOLUTION', parsed_answer) - if parsed_answer != llm_answer: - print('END OF OUTPUT', llm_answer[-100:]) - - return 0 diff --git a/reasoners/reward/livebench/data_analysis/tablejoin/utils.py b/reasoners/reward/livebench/data_analysis/tablejoin/utils.py deleted file mode 100644 index 2b564086b..000000000 --- a/reasoners/reward/livebench/data_analysis/tablejoin/utils.py +++ /dev/null @@ -1,82 +0,0 @@ -import re -import numpy as np -from ast import literal_eval -import json -import ast -from verl.utils.reward_score.livebench.util import last_boxed_only_string, remove_boxed - -def clean_llm_output(s): - pattern_solution = r'(.*?)' - matches = re.findall(pattern_solution, s, re.DOTALL) - if len(matches) > 0: - return clean_llm_output(matches[-1].strip()) - try: - match_d = literal_eval(s) - except: - matches = re.findall('%s(.*?)%s' % ("```python", "```"), s.replace("\n",""),re.MULTILINE) - if len(matches) == 0: - matches = re.findall('%s(.*?)%s' % ("```json", "```"), s.replace("\n",""),re.MULTILINE) - if len(matches) == 0: - matches = re.findall('%s(.*?)%s' % ("```", "```"), s.replace("\n",""),re.MULTILINE) - if len(matches) == 0: - if '\\boxed' in s: - boxed = last_boxed_only_string(s.replace('\n', '')) - if boxed: - no_boxed = remove_boxed(boxed) - matches = [re.sub(r"\\text{[\'|\"](.*?)[\'|\"]}", r"'\1'", no_boxed).replace('\\', '')] - if len(matches) == 0: - matches = [s] - if len(matches) >= 1: - matches = matches[-1] - matches = matches.replace('null', 'None') - try: - match_d = literal_eval(matches) - except: - return {} - if not isinstance(match_d, dict): - return {} - else: - keys = list(match_d.keys()) - for k in keys: - if match_d[k] is None: - del match_d[k] - return match_d - -def joinmap_process_results(ground_truth, llm, debug=True): - if type(ground_truth) != dict: - ground_truth = ast.literal_eval(ground_truth) - llm_clean = clean_llm_output(llm) - if len(llm_clean) == 0: - if debug: - print('could not parse output') - print('GROUND TRUTH', ground_truth) - print('END OF OUTPUT', llm[-min(500, len(llm)):]) - print('=' * 100) - return 0.0 - tp = 0 - fp = 0 - fn = 0 - for k, v in llm_clean.items(): - gt = ground_truth.get(k, None) - if not gt: - fp += 1 - elif gt == v: - tp += 1 - else: - fp += 1 - fn += 1 - if isinstance(ground_truth, str): - ground_truth = json.loads(ground_truth) - for k, v in ground_truth.items(): - llm_resp = llm_clean.get(k, None) - if not llm_resp: - fn += 1 - result = np.round(((2.0 * tp) / ((2.0 * tp) + fp + fn)), 2) - if debug and result < 1: - print('INCORRECT') - print('GROUND TRUTH', ground_truth) - print('SOLUTION', llm_clean) - print('END OF OUTPUT', llm[-min(500, len(llm)):]) - print('RESULT', result) - print("=" * 100) - return result diff --git a/reasoners/reward/livebench/data_analysis/tablereformat/utils.py b/reasoners/reward/livebench/data_analysis/tablereformat/utils.py deleted file mode 100644 index 958136420..000000000 --- a/reasoners/reward/livebench/data_analysis/tablereformat/utils.py +++ /dev/null @@ -1,240 +0,0 @@ -import json -import traceback -import pandas as pd -from io import StringIO -import re -import math -from pandas.api.types import is_string_dtype -from pandas.api.types import is_numeric_dtype - - -def read_df_func(df_type, df_str): - if df_type == "json": - try: - return pd.read_json(StringIO(df_str), orient="index", encoding="utf-8") - except: - pass - try: - return pd.read_json(StringIO(df_str), orient="records", lines=False, encoding="utf-8") - except: - pass - try: - return pd.read_json(StringIO(df_str), orient="records", lines=True, encoding="utf-8") - except: - pass - try: - return pd.read_json(StringIO(df_str), orient="table", encoding="utf-8") - except: - pass - return pd.read_json(StringIO(df_str), orient="values", encoding="utf-8") - elif df_type == "jsonl": - return pd.read_json(StringIO(df_str), orient="records", lines=True, encoding="utf-8") - elif df_type == "html": - return pd.concat(pd.read_html(StringIO(df_str), encoding="utf-8"), axis=0) - elif df_type == "csv": - return pd.read_csv(StringIO(df_str), encoding="utf-8") - elif df_type == "markdown": - return pd.read_table(StringIO(df_str), sep="|", header=0, index_col=1, skipinitialspace=True) - elif df_type == "tsv": - return pd.read_csv(StringIO(df_str), sep='\t', encoding="utf-8") - -def read_df_func_v2(df_type, df_str): - if df_type == "json": - try: - # Try table orientation first as it preserves dtypes - return pd.read_json(StringIO(df_str), orient="table", encoding="utf-8") - except: - try: - return pd.read_json(StringIO(df_str), orient="index", lines=False, encoding="utf-8") - except: - try: - return pd.read_json(StringIO(df_str), orient="records", lines=False, encoding="utf-8") - except: - print("Could not read JSON") - return None - elif df_type == "jsonl": - return pd.read_json(StringIO(df_str), orient="records", lines=True, encoding="utf-8") - elif df_type == "html": - return pd.concat(pd.read_html(StringIO(df_str), encoding="utf-8"), axis=0) - elif df_type == "csv": - return pd.read_csv(StringIO(df_str), encoding="utf-8") - elif df_type == "markdown": - # Process markdown table by removing the separator line - lines = df_str.strip().split("\n") - header = lines[0] - # Skip the separator line (typically line with |:-----|:-----|) - data_lines = lines[2:] if len(lines) > 2 else [] - processed_md = header + "\n" + "\n".join(data_lines) - df = pd.read_table(StringIO(processed_md), sep="|", header=0, skipinitialspace=True).iloc[:, 1:-1] - - # Strip whitespace from all string columns - for col in df.columns: - if df[col].dtype == 'object': - df[col] = df[col].astype(str).str.strip() - - return df - elif df_type == "tsv": - return pd.read_csv(StringIO(df_str), sep='\t', encoding="utf-8") - else: - raise ValueError(f"Unsupported file type: {df_type}") - -def clean_llm_output(s): - pattern_solution = r'(.*?)' - matches = re.findall(pattern_solution, s, re.DOTALL) - if len(matches) > 0: - return clean_llm_output(matches[-1].strip()) - pattern_json = r'```json\n(.*?)```' - matches = re.findall(pattern_json, s, re.DOTALL) - if len(matches) > 0: - return matches[-1].strip() - pattern_html = r'```html\n(.*?)```' - matches = re.findall(pattern_html, s, re.DOTALL) - if len(matches) > 0: - return matches[-1].strip() - pattern_a = r'^```.*\n' - s = re.sub(pattern_a, "", s) - # re.findall(pattern, text, re.MULTILINE) - s = s.replace("&", "&") - return s.replace("```", "").strip() - -def read_sep_table_from_text(text, header, sep=','): - text = text.strip() - # look for the first line that contains the header - header_line = 0 - while header_line < len(text.split('\n')) and text.split('\n')[header_line].strip() != header.strip(): - header_line += 1 - if header_line == len(text.split('\n')) or text.split('\n')[header_line].strip() != header.strip(): - return None - # read the table from the header index - table = text.split('\n')[header_line:] - # read the table as a csv - parsed_table = None - while parsed_table is None: - try: - parsed_table = pd.read_csv(StringIO('\n'.join(table)), sep=sep) - except: - # in case there's extra text after the table - table = table[:-1] - return parsed_table - -def read_jsonl_table_from_text(text, header): - text = text.strip().split('\n') - table = [] - for line in text: - if len(line) < 2 or line[0] != '{' or line[-1] != '}': - continue - if not all(key in line for key in header): - continue - try: - table.append(json.loads(line)) - except: - continue - if len(table) == 0: - return None - return pd.DataFrame(table) - - -def remove_initial_phrase(text): - # remove phrases like "Here is the updated table:" , "Here is the table in a new format:" - pattern = r'^\s*(Here|Input)\b.*?\b(format|table)\s*[:)]\s*' - modified_text = re.sub(pattern, '', text, flags=re.IGNORECASE) - return modified_text.strip() - -def table_process_results(input_command: str, ground_truth: str, llm_answer: str, version: str = "v1", debug=True) -> int: - if version == "v1": - input_format = input_command.split("Please convert the Input Table from ")[1].split(" format")[0].lower() - output_format = input_command.split("Please convert the Input Table from ")[1].split("format to ")[1].split(" format")[0].lower() - else: - command_lines = input_command.split('\n') - input_format_lines = [line for line in command_lines if "Source Format" in line] - input_format = input_format_lines[-1].split("Source Format: ")[1].strip().lower() - output_format_lines = [line for line in command_lines if "Target Format" in line] - output_format = output_format_lines[-1].split("Target Format: ")[1].strip().lower() - - df_read_func = read_df_func if version == "v1" else read_df_func_v2 - - gt_df = df_read_func(output_format, ground_truth) - - llm_clean = clean_llm_output(llm_answer) - # if there's an initial phrase before the table, remove it and try to score again - llm_clean = remove_initial_phrase(llm_clean) - # first check the raw LLM output - llm_df = None - try: - llm_df = df_read_func(output_format, llm_clean) - except: - if output_format == 'csv' or output_format == 'tsv': - header = (',', '\t')[output_format == 'tsv'].join(gt_df.columns) - llm_df = read_sep_table_from_text(llm_clean, header, sep=',' if output_format == 'csv' else '\t') - elif output_format == 'jsonl': - llm_df = read_jsonl_table_from_text(llm_clean, gt_df.columns) - if llm_df is None: - print('Could not read the LLM output') - print('GROUND TRUTH\n', ground_truth) - print('END OF OUTPUT\n', llm_answer[-min(3000, len(llm_answer)):]) - print('=' * 100) - return 0 - score = check_table_reformat(output_format, llm_df, gt_df, debug) - - if score == 0: - # try to read table directly from the text, in case the LLM added some text before/after the table - if output_format == 'csv': - header = ','.join(gt_df.columns) - llm_df = read_sep_table_from_text(llm_clean, header, sep=',' if output_format == 'csv' else '\t') - - elif output_format == 'jsonl': - llm_df = read_jsonl_table_from_text(llm_clean, gt_df.columns) - - if llm_df is not None: - score = check_table_reformat(output_format, llm_df, gt_df, debug) - if debug and score == 0: - print('INCORRECT') - print('GROUND TRUTH\n', 'None' if gt_df is None else gt_df.head()) - print('LLM DF\n', 'None' if llm_df is None else llm_df.head()) - print('LLM ANSWER\n', llm_clean) - print('=' * 100) - return score - -def check_table_reformat(output_format, llm_df, gt_df, debug=False): - try: - gt_df.columns = [s.strip() for s in gt_df.columns] - if 'index' in gt_df.columns: - gt_df = gt_df.drop(columns=['index']) - llm_df.columns = [s.strip() for s in llm_df.columns] - if 'index' in llm_df.columns: - llm_df = llm_df.drop(columns=['index']) - assert len(llm_df) == len(gt_df), f"DataFrame Length does not match, {len(llm_df)} (LLM) vs {len(gt_df)} (Ground Truth)" - assert list(sorted(llm_df.columns)) == list(sorted(gt_df.columns)), f"Columns do not match:\n{sorted(llm_df.columns)} (LLM)\n{sorted(gt_df.columns)} (Ground Truth)" - # for test_col in llm_df.columns: - # assert sorted(llm_df[test_col].tolist()) == sorted(gt_df[test_col].tolist()), f"Column content {test_col} does not match" - for i in range(len(llm_df)): - for key in llm_df.columns: - llm_val = llm_df.iloc[i][key] - gt_val = gt_df.iloc[i][key] - if isinstance(llm_val, str): - llm_val = llm_val.strip() - if isinstance(gt_val, str): - gt_val = gt_val.strip() - - if (isinstance(llm_val, float) or is_numeric_dtype(llm_val)) and (isinstance(gt_val, float) or is_numeric_dtype(gt_val)): - try: - llm_val = float(llm_val) - gt_val = float(gt_val) - except: - assert str(llm_val).strip() == str(gt_val).strip(), f"Mismatched types of values {llm_val} (LLM) vs {gt_val} (Ground Truth) for key {key} in row {i}" - continue - if math.isnan(llm_val) and math.isnan(gt_val): - continue - assert abs(llm_val - gt_val) < 1e-6, f"Unequal float values {llm_val} (LLM) vs {gt_val} (Ground Truth) for key {key} in row {i}" - else: - assert llm_val == gt_val, f"Value {llm_val} (LLM) vs {gt_val} (Ground Truth) for key {key} in row {i}" - except AssertionError as e: - if debug: - print(e) - return 0 - except Exception as e: - if debug: - print(e) - traceback.print_exc() - return 0 - return 1 diff --git a/reasoners/reward/livebench/reasoning/house_traversal/utils.py b/reasoners/reward/livebench/reasoning/house_traversal/utils.py deleted file mode 100644 index af21006cc..000000000 --- a/reasoners/reward/livebench/reasoning/house_traversal/utils.py +++ /dev/null @@ -1,41 +0,0 @@ - -import re - - -def house_traversal_process_results(ground_truth: str, llm_answer: str, debug=False) -> int: - # pull out words in bold - bold_words = re.findall(r'(\*{2,})(.*?)\1', llm_answer.lower()) - if not len(bold_words): - if debug: - print('INCORRECT') - print('GROUND TRUTH', ground_truth) - print('SOLUTION', llm_answer) - return 0 - - last_bold = bold_words[-1][1] - ground_truth_names = ground_truth.lower().split(" ") - - # check if all the ground truth names are in the last few bolded words, in order - if len(bold_words) >= len(ground_truth_names): - if all([name in bold_words[-1 - i] for i,name in enumerate(ground_truth_names[::-1])]): - return 1 - - score = 1 - # check if all the ground truth names are in the last bolded part, in order - last_index = -1 - for name in ground_truth_names: - index = last_bold.find(name) - if index == -1: - score = 0 - break - elif index <= last_index: - score = 0 - break - else: - last_index = index - - if debug and score == 0: - print('INCORRECT', score) - print('GROUND TRUTH', ground_truth) - print('SOLUTION', llm_answer) - return score \ No newline at end of file diff --git a/reasoners/reward/livebench/reasoning/spatial/utils.py b/reasoners/reward/livebench/reasoning/spatial/utils.py deleted file mode 100644 index f4e82358d..000000000 --- a/reasoners/reward/livebench/reasoning/spatial/utils.py +++ /dev/null @@ -1,61 +0,0 @@ -import re -from verl.utils.reward_score.livebench.util import last_boxed_only_string, remove_boxed - -def spatial_process_results(ground_truth: str, llm_answer: str, debug=False) -> int: - - if llm_answer == ground_truth: - return 1 - - word_to_number = { - 'zero': '0', 'one': '1', 'two': '2', 'three': '3', 'four': '4', 'five': '5', - 'six': '6', 'seven': '7', 'eight': '8', 'nine': '9', 'ten': '10', - 'eleven': '11', 'twelve': '12', 'thirteen': '13', 'fourteen': '14', 'fifteen': '15', - 'sixteen': '16', 'seventeen': '17', 'eighteen': '18', 'nineteen': '19', 'twenty': '20' - } - - bold_words = re.findall(r'\*\*([^\*]+)\*\*', llm_answer) - score = 0 - - # allow the answer to be within the last 3 bolded words - words_to_check = [] - for i in range(3): - if bold_words and len(bold_words) > i: - words_to_check.append(bold_words[-i-1].strip().lower()) - - for i, word in enumerate(words_to_check): - if word == ground_truth.strip().lower(): - score = 1 - - # allow the answer to be the number spelled out - if word in word_to_number and word_to_number[word] == ground_truth.strip().lower(): - score = 1 - - # allow certain cases like "two tetrahedra" == "tetrahedra" and "equilateral triangle" == "triangle" - # while still disallowing cases like "circle square triangle" == "triangle" - for answer in ["tetrahedra", "tetrahedron", "triangle", "square"]: - if ground_truth.strip().lower() == answer and answer in word and len(word) < (2 * len(answer) + 5): - score = 1 - - allow_boxed = True - if score == 0 and allow_boxed: - llm_answer = llm_answer.replace("\\\\fbox{", "\\\\boxed{") - last_boxed = last_boxed_only_string(llm_answer) - if last_boxed: - parsed_answer = remove_boxed(last_boxed) - parsed_answer = parsed_answer.replace("\\textbf{", "") - parsed_answer = parsed_answer.replace("\\mathbf{", "") - parsed_answer = parsed_answer.replace("}", "") - if parsed_answer == ground_truth: - score = 1 - - if debug and score == 0: - print("INCORRECT") - print("GROUND TRUTH", ground_truth.strip().lower()) - if bold_words: - print("BOLD WORDS:", bold_words[-1].strip().lower()) - if last_boxed: - print("LAST BOXED", last_boxed) - print("PARSED ANSWER", parsed_answer) - print("END OF OUTPUT", llm_answer[-50:]) - - return score diff --git a/reasoners/reward/livebench/reasoning/web_of_lies_v2/utils.py b/reasoners/reward/livebench/reasoning/web_of_lies_v2/utils.py deleted file mode 100644 index 07be3af68..000000000 --- a/reasoners/reward/livebench/reasoning/web_of_lies_v2/utils.py +++ /dev/null @@ -1,79 +0,0 @@ -import re -import itertools -from verl.utils.reward_score.livebench.util import last_boxed_only_string, remove_boxed - -def web_of_lies_process_results(ground_truth: str, llm_answer: str, debug=False) -> int: - - score = 0 - parsed_answer = None - - # extract text from tags - solution_matches = re.findall(r'(.*?)', llm_answer) - - if len(solution_matches) == 0: - solution_matches = re.findall(r'(.*?)', llm_answer) - - if len(solution_matches) > 0: - parsed_answer = solution_matches[-1] - - - # pull out words in bold - bold_words = re.findall(r'\*\*(.*?)\*\*', llm_answer) - - if parsed_answer is None and bold_words: - bold_words = [word.lower().strip().replace(',', '').replace('.', '')[0:max(len(word), 3)] for match in bold_words for word in match.split()] - parsed_answer = [] - i = len(bold_words) - 1 - while i >= 0 and len(parsed_answer) < 3: - if bold_words[i] in ["yes", "no", "unknown"]: - parsed_answer = [bold_words[i]] + parsed_answer - i -= 1 - if len(parsed_answer) > 0: - parsed_answer = ", ".join(parsed_answer) - else: - parsed_answer = None - - allow_latex = True - if parsed_answer is None or parsed_answer.strip() == '' and allow_latex: - llm_answer = llm_answer.replace("\\\\boxed{\\\\textbf{", "\\\\boxed{") - llm_answer = llm_answer.replace("\\\\fbox{", "\\\\boxed{") - llm_answer = llm_answer.replace("\\textbf{", "\\boxed{") - - last_boxed = last_boxed_only_string(llm_answer) - if last_boxed: - parsed_answer = remove_boxed(last_boxed) - - allow_plain = True - if allow_plain and parsed_answer is None: - combs = itertools.product(['yes', 'no', 'unknown'], repeat=3) - # find all instances of these combinations in the answer, then treat the last one as the actual answer - # to compare to the ground truth - final_comb = None - final_comb_index = -1 - for comb in combs: - index = llm_answer.lower().find(', '.join(comb)) - if index != -1 and index > final_comb_index: - final_comb = comb - final_comb_index = index - if final_comb is not None: - parsed_answer = ', '.join(final_comb) - - # Check if parsed_answer is an exact match of ground_truth - if parsed_answer and parsed_answer == ground_truth.lower(): - score = 1 - - # Check if parsed_answer contains the ground_truth - if parsed_answer and parsed_answer.count("yes") + parsed_answer.count("no") + parsed_answer.count("unknown") == 3 and ground_truth.lower() in parsed_answer: - score = 1 - - if debug and score == 0: - print("INCORRECT") - print('GROUND TRUTH', ground_truth) - if parsed_answer: - print('PARSED ANSWER', parsed_answer) - else: - print('NO PARSED ANSWER') - print('END OF OUTPUT', llm_answer[-50:]) - - return score - diff --git a/reasoners/reward/livebench/reasoning/web_of_lies_v3/utils.py b/reasoners/reward/livebench/reasoning/web_of_lies_v3/utils.py deleted file mode 100644 index ae9313e4e..000000000 --- a/reasoners/reward/livebench/reasoning/web_of_lies_v3/utils.py +++ /dev/null @@ -1,87 +0,0 @@ -import re -import itertools -from verl.utils.reward_score.livebench.util import last_boxed_only_string, remove_boxed - -def web_of_lies_v3_process_results(ground_truth: str, llm_answer: str, debug=False) -> float: - score = 0 - parsed_answer = None - - # extract text from tags - solution_matches = re.findall(r'(.*?)', llm_answer) - - if len(solution_matches) > 0: - parsed_answer = solution_matches[-1] - - # if parsed_answer is None: - # bold_words = re.findall(r'\*\*(.*?)\*\*', llm_answer) - # if bold_words: - # bold_words = [word.lower().strip().replace(',', '').replace('.', '')[0:max(len(word), 3)] for match in bold_words for word in match.split()] - - # parsed_answer = [] - # i = len(bold_words) - 1 - # while i >= 0 and len(parsed_answer) < 3: - # if bold_words[i] in ["yes", "no", "unknown"]: - # parsed_answer = [bold_words[i]] + parsed_answer - # i -= 1 - - # if len(parsed_answer) > 0: - # parsed_answer = ", ".join(parsed_answer) - # else: - # parsed_answer = None - - allow_latex = True - if parsed_answer is None or parsed_answer.strip() == '' and allow_latex: - llm_answer = llm_answer.replace("\\\\boxed{\\\\textbf{", "\\\\boxed{") - llm_answer = llm_answer.replace("\\\\fbox{", "\\\\boxed{") - llm_answer = llm_answer.replace("\\textbf{", "\\boxed{") - - last_boxed = last_boxed_only_string(llm_answer) - - if last_boxed: - parsed_answer = remove_boxed(last_boxed) - - allow_plain = True - if allow_plain and parsed_answer is None: - combs = itertools.product(['yes', 'no', 'unknown'], repeat=3) - - final_comb = None - final_comb_index = -1 - for comb in combs: - index = llm_answer.lower().find(', '.join(comb)) - if index != -1 and index > final_comb_index: - final_comb = comb - final_comb_index = index - - if final_comb is not None: - parsed_answer = ', '.join(final_comb) - - if parsed_answer is None: - score = 0 - else: - parsed_answer_list = parsed_answer.split(',') - ground_truth_list = ground_truth.split(',') - - num_correct = 0 - total = len(ground_truth_list) - for i in range(total): - if i >= len(parsed_answer_list): - break - gt_word = ground_truth_list[i] - llm_word = parsed_answer_list[i] - if llm_word == gt_word: - num_correct += 1 - - score = ((num_correct == total) + num_correct / total) / 2 - - # score = 1 if parsed_answer.strip().replace(' ', '') == ground_truth.strip().replace(' ', '') else 0 - - if debug and score < 1: - print("INCORRECT") - print('GROUND TRUTH', ground_truth) - if parsed_answer: - print('PARSED ANSWER', parsed_answer) - else: - print('PARSED ANSWER', parsed_answer) - print('END OF OUTPUT', llm_answer[-50:]) - - return score diff --git a/reasoners/reward/livebench/reasoning/zebra_puzzle/utils.py b/reasoners/reward/livebench/reasoning/zebra_puzzle/utils.py deleted file mode 100644 index 92d29fd28..000000000 --- a/reasoners/reward/livebench/reasoning/zebra_puzzle/utils.py +++ /dev/null @@ -1,128 +0,0 @@ - -import re -from verl.utils.reward_score.livebench.util import last_boxed_only_string, remove_boxed - -def zebra_puzzle_process_results_old(ground_truth: str, llm_answer: str, debug=False) -> int: - # Mapping of numbers to words for 1 to 9 - number_to_word = { - '1': 'one', '2': 'two', '3': 'three', '4': 'four', - '5': 'five', '6': 'six', '7': 'seven', '8': 'eight', '9': 'nine' - } - - # Pull out words in bold - bold_words = re.findall(r'\*\*\*(\w+)\*\*\*', llm_answer) - - score = 0 - - # Remove any trailing punctuation from the last bold word if exists - if bold_words: - if (bold_words[-1].lower() == ground_truth.lower() or - (bold_words[-1] in number_to_word and number_to_word[bold_words[-1]].lower() == ground_truth.lower()) - or bold_words[-1].lower() + ' movies' == ground_truth.lower()): - score = 1 - else: - # Split the text into words and remove punctuation. - words = re.findall(r'\b\w+\b', llm_answer) - last_word = words[-1] if words else '' - # Check if the last bold word is a number and matches the word representation of the ground_truth - if (last_word.lower() == ground_truth.lower() or - (last_word in number_to_word and number_to_word[last_word].lower() == ground_truth.lower()) - or last_word.lower() + ' movies' == ground_truth.lower()): - score = 1 - - if debug and score == 0: - print('INCORRECT') - print('GROUND TRUTH', ground_truth.lower()) - if bold_words: - print('LAST BOLD WORD', bold_words[-1].lower()) - else: - print('LAST WORD', last_word.lower()) - print('END OF OUTPUT', llm_answer[-50:]) - return score - - -def zebra_puzzle_process_results(ground_truth: str, llm_answer: str, debug=False) -> float: - # Mapping of numbers to words for 1 to 9 - word_to_num = { - 'one': '1', - 'two': '2', - 'three': '3', - 'four': '4', - 'five': '5', - 'six': '6', - 'seven': '7', - 'eight': '8', - 'nine': '9' - } - - ground_truth = ground_truth.split(',') - - # extract text from tags - solution_matches = re.findall(r'(.*?)', llm_answer) - - if len(solution_matches) == 0: - solution_matches = re.findall(r'(.*?)', llm_answer) - - - allow_boxed = True - if len(solution_matches) == 0 and allow_boxed: - llm_answer = llm_answer.replace("\\\\fbox{", "\\\\boxed{") - last_boxed = last_boxed_only_string(llm_answer) - if last_boxed: - boxed_removed = remove_boxed(last_boxed) - boxed_removed = boxed_removed.replace("\\text{", "").replace("}", "").replace('\\', '') - solution_matches.append(boxed_removed) - - if len(solution_matches) == 0: - last_line = llm_answer.strip().split('\n')[-1] - if last_line.count(',') == len(ground_truth) - 1: - solution_matches.append(last_line) - - - if len(solution_matches) == 0: - if debug: - print('No solution text found for zebra puzzle') - print('GROUND TRUTH', ground_truth) - print('END OF OUTPUT', llm_answer[-100:]) - return 0 - - if len(solution_matches) > 1: - if debug: - print('WARNING: Multiple solution texts found for zebra puzzle, combining starting from last') - all_solution_text = [] - for match in solution_matches: - all_solution_text += match.split(',') - # get last len(ground_truth) elements - solution_text = all_solution_text[-len(ground_truth):] - else: - solution_text = solution_matches[-1].split(',') - - if debug and len(solution_text) < len(ground_truth): - print(f'WARNING: LLM does not have an answer for all zebra puzzle questions (expected {len(ground_truth)}, got {len(solution_text)})') - - num_correct = 0 - total = len(ground_truth) - for i in range(total): - gt_word = ground_truth[i].strip().lower().replace('-', ' ') - if i >= len(solution_text): - continue - llm_word = solution_text[i].strip().lower().replace('-', ' ').replace('position', '') - if llm_word in word_to_num: - # llm_word = word_to_num[llm_word] - if debug: - print('WARNING: LLM answer contains numbers in word form') - - if i < len(solution_text) and (gt_word == llm_word or gt_word in llm_word): - num_correct += 1 - if debug and num_correct < total: - print('INCORRECT', num_correct / total) - print('GROUND TRUTH', ground_truth) - print('SOLUTION', solution_text) - print('END OF OUTPUT', llm_answer[-50:]) - return ((num_correct == total) + num_correct / total) / 2 - - -def get_zebra_puzzle_evaluator(release_date: str): - if release_date < '2024-11-25': - return zebra_puzzle_process_results_old - return zebra_puzzle_process_results diff --git a/reasoners/reward/livebench/util.py b/reasoners/reward/livebench/util.py deleted file mode 100644 index 084a3e17b..000000000 --- a/reasoners/reward/livebench/util.py +++ /dev/null @@ -1,46 +0,0 @@ -from typing import Optional - - -def last_boxed_only_string(string: str) -> Optional[str]: - 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].replace("$", "").replace("fbox","boxed") - - return retval - - -def remove_boxed(s: str) -> str: - 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] \ No newline at end of file diff --git a/reasoners/reward/livebench/writing/connections/utils.py b/reasoners/reward/livebench/writing/connections/utils.py deleted file mode 100644 index 7f4877ebd..000000000 --- a/reasoners/reward/livebench/writing/connections/utils.py +++ /dev/null @@ -1,108 +0,0 @@ -import re -from verl.utils.reward_score.livebench.util import last_boxed_only_string, remove_boxed - - -def group_words(words): - groups = [set()] - words = [w.strip().lower() for w in words] - for word in words: - if len(groups[-1]) == 4: - groups.append(set()) - groups[-1].add(word) - return groups - - -def connections_process_results_old(ground_truth: str, llm_answer: str, debug=False) -> int: - - # pull out words in bold - bold_words = re.findall(r'\*\*(.*?)\*\*', llm_answer.replace('\n', '')) - - if not bold_words: - if debug: - print('No bold words found for connections') - print('END OF OUTPUT', llm_answer[-50:]) - return 0 - - bold_words = [words.split(',') for words in bold_words] - - ground_truth_groups = group_words(ground_truth.split(',')) - max_score = 0 - for output_groups in list(map(group_words, bold_words)): - - correct_groups = 0 - for ground_truth_group in ground_truth_groups: - for output_group in output_groups: - if all([word in output_group for word in ground_truth_group]): - correct_groups += 1 - break - - max_score = max(max_score, correct_groups / len(ground_truth_groups)) - if debug and max_score < 1: - print('INCORRECT', max_score) - print('GROUND TRUTH', ground_truth_groups) - print('SOLUTION', output_groups) - print('END OF OUTPUT', llm_answer[-50:]) - return max_score - - -def connections_process_results(ground_truth: str, llm_answer: str, debug=False) -> int: - - # extract text from tags - solution_matches = re.findall(r'(.*?)<\/solution>', llm_answer) - if len(solution_matches) == 0: - solution_matches = re.findall(r'(.*?)<\/solution>', llm_answer.replace('\n', '')) - if len(solution_matches) == 0: - solution_matches = re.findall(r'(.*?)<\/solution>', llm_answer) - if len(solution_matches) == 0: - solution_matches = re.findall(r'(.*?)<\/solution>', llm_answer.replace('\n', '')) - - ground_truth_words = ground_truth.split(',') - - if len(solution_matches) == 0 and '\\boxed' in llm_answer: - boxed = last_boxed_only_string(llm_answer) - no_box = remove_boxed(boxed) - solution_matches = [no_box.replace('\\text{', '').replace('}', '').replace('\\', '')] - - solution_matches = [match.replace('\n', '') for match in solution_matches] - - if len(solution_matches) == 0: - if debug: - print('No solution text found for connections') - print(llm_answer[-500:]) - return 0 - elif len(solution_matches) > 1: - if debug: - print('Multiple solution texts found for connections. Combining starting from last') - all_words = [] - num_words = len(ground_truth_words) - for match in solution_matches: - all_words.extend(match.split(',')) - solution_words = all_words[-num_words:] - else: - solution_words = solution_matches[-1].split(',') - - if len(solution_words) != len(ground_truth_words): - if debug: - print(f'Number of words in solution does not match number of words in ground truth ({len(solution_words)} vs expected {len(ground_truth_words)})') - - llm_groups = group_words(solution_words) - ground_truth_groups = group_words(ground_truth_words) - - correct_groups = 0 - for llm_group in llm_groups: - if llm_group in ground_truth_groups: - correct_groups += 1 - - if debug and correct_groups / len(ground_truth_groups) < 1: - print('INCORRECT', correct_groups / len(ground_truth_groups)) - print('GROUND TRUTH', sorted([sorted(list(group)) for group in ground_truth_groups])) - print('SOLUTION', sorted([sorted(list(group)) for group in llm_groups])) - print('END OF OUTPUT', llm_answer[-500:]) - - return correct_groups / len(ground_truth_groups) - - -def get_connections_puzzle_evaluator(release_date: str): - if release_date < '2024-11-25': - return connections_process_results_old - return connections_process_results diff --git a/reasoners/reward/livebench/writing/plot_unscrambling/utils.py b/reasoners/reward/livebench/writing/plot_unscrambling/utils.py deleted file mode 100644 index 084fc26f4..000000000 --- a/reasoners/reward/livebench/writing/plot_unscrambling/utils.py +++ /dev/null @@ -1,66 +0,0 @@ -import difflib -import re - - -def levenshtein_distance(A, B): - N, M = len(A), len(B) - # Create an array of size NxM - dp = [[0 for i in range(M + 1)] for j in range(N + 1)] - - # Base Case: When N = 0 - for j in range(M + 1): - dp[0][j] = j - # Base Case: When M = 0 - for i in range(N + 1): - dp[i][0] = i - # Transitions - for i in range(1, N + 1): - for j in range(1, M + 1): - if A[i - 1] == B[j - 1]: - dp[i][j] = dp[i-1][j-1] - else: - dp[i][j] = 1 + min( - dp[i-1][j], # Insertion - dp[i][j-1], # Deletion - dp[i-1][j-1] # Replacement - ) - - return dp[N][M] - - -def extract_plot_summary(text: str) -> str: - pattern = r'(.*)' - match = re.search(pattern, text, re.DOTALL) # re.DOTALL allows . to match newline characters - if not match: - pattern = r'(.*)' - match = re.search(pattern, text, re.DOTALL) - return match.group(1) if match else text - - -def plot_unscrambling_process_results(ground_truth: str, llm_answer: str, debug=False) -> float: - # Split the ground truth and answer into sentences based on full stops - llm_answer = extract_plot_summary(llm_answer) - - gt_sentences = [s.strip() for s in ground_truth.split('.')] - ans_sentences = [s.strip() for s in llm_answer.split('.') if s.strip() != '' and s.strip() != '**End of Plot Summary**'] - - # Remove empty sentences resulting from trailing or double full stops. - gt_sentences = [s for s in gt_sentences if s] - ans_sentences = [s for s in ans_sentences if s] - - ans_ordering = [] - for x in gt_sentences: - best_match = difflib.get_close_matches(x, ans_sentences, n=1, cutoff=0.0) - if best_match: - ans_ordering.append(ans_sentences.index(best_match[0])) - - n_sentences_gt = len(gt_sentences) - raw_distance = levenshtein_distance(list(range(len(gt_sentences))), ans_ordering) - score = 1 - (raw_distance / n_sentences_gt) - - if debug and score < 1: - print('INCORRECT', score) - print('GROUND TRUTH', gt_sentences) - print('SOLUTION', ans_sentences) - print('END OF OUTPUT', llm_answer[-50:]) - return score diff --git a/reasoners/reward/livebench/writing/typos/utils.py b/reasoners/reward/livebench/writing/typos/utils.py deleted file mode 100644 index 1b7715da2..000000000 --- a/reasoners/reward/livebench/writing/typos/utils.py +++ /dev/null @@ -1,59 +0,0 @@ -import difflib -import re - -def extract_answer(llm_answer): - pattern = r'.* --- (.*?) --- .*' - match = re.search(pattern, llm_answer) - return match.group(1) if match else llm_answer - -def typos_process_results(ground_truth: str, llm_answer: str, debug=False) -> int: - - parsed_answer = None - - # extract text from tags - solution_matches = re.findall(r'(.*?)', llm_answer) - if len(solution_matches) > 0: - match = solution_matches[-1] - parsed_answer = match - else: - parsed_answer = llm_answer.replace('', '').replace('', '') - parsed_answer = extract_answer(parsed_answer) - - parsed_answer = ' '.join(list(filter(None, parsed_answer.strip().split('\n')))) - - if int(ground_truth in parsed_answer): - return 1 - - score = 0 - - if debug and score == 0: - - a = ground_truth - b = parsed_answer - m = difflib.SequenceMatcher(a=a, b=b) - pad = 10 - - for tag, i1, i2, j1, j2 in m.get_opcodes(): - length = min(len(parsed_answer), len(ground_truth)) - mi1, mi2, mj1, mj2 = max(0,i1-pad), min(length, i2+pad), max(0, j1-pad), min(length, j2+pad) - - mistake_length = 0 - if tag == 'replace': - print("", a[i1:i2], b[j1:j2], "::::", a[mi1:mi2], "-->", b[mj1:mj2]) - mistake_length = i2 - i1 - if tag == 'delete': - print("", a[i1:i2], "::::", a[mi1:mi2], "-->", b[mj1:mj2]) - mistake_length = i2 - i1 - if tag == 'insert': - print("", b[j1:j2], "::::", a[mi1:mi2], "-->", b[mj1:mj2]) - mistake_length = j2 - j1 - #score = score - (mistake_length / len(ground_truth)) * 100 - - - if debug and score == 0: - print('INCORRECT') - print('GROUND TRUTH', ground_truth) - print('SOLUTION', parsed_answer) - print('END OF OUTPUT', llm_answer[-len(parsed_answer):]) - - return score diff --git a/reasoners/reward/math_batch.py b/reasoners/reward/math_batch.py deleted file mode 100644 index 46a7cfec6..000000000 --- a/reasoners/reward/math_batch.py +++ /dev/null @@ -1,25 +0,0 @@ -# Copyright 2025 Individual Contributor: Mert Unsal - # - # 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. - - from .math import compute_score - - - def compute_score_batched(data_sources, solution_strs, ground_truths, extra_infos): - """ - This is a demonstration of how the batched reward function should look like. - Typically, you want to use batched reward to speed up the process with parallelization - """ - return [ - compute_score(solution_str, ground_truth) for solution_str, ground_truth in zip(solution_strs, ground_truths) - ] \ No newline at end of file diff --git a/reasoners/reward/math_dapo.py b/reasoners/reward/math_dapo.py deleted file mode 100644 index 5ce35811f..000000000 --- a/reasoners/reward/math_dapo.py +++ /dev/null @@ -1,294 +0,0 @@ -# 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 - - -def last_boxed_only_string(string: str) -> Optional[str]: - """Extract the last LaTeX boxed expression from a string. - - Args: - string: Input string containing LaTeX code - - Returns: - The last boxed expression or None if not found - """ - idx = string.rfind("\\boxed{") - 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 - - return string[idx:right_brace_idx + 1] if right_brace_idx is not None else None - - -def remove_boxed(s: str) -> str: - """Remove the LaTeX boxed command from a string. - - Args: - s: String with format "\\boxed{content}" - - Returns: - The content inside the boxed command - """ - left = "\\boxed{" - assert s[:len(left)] == left, f"box error: {s}" - assert s[-1] == "}", f"box error: {s}" - return s[len(left):-1] - - -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() - - -def is_correct_minerva(solution_str: str, - gt: str, - gt_need_extract: bool = False, - answer_pattern: str = r"(?i)Answer\s*:\s*([^\n]+)") -> tuple[bool, str]: - """Check if the solution is correct according to Minerva criteria. - - Args: - solution_str: The solution string to check - gt: The ground truth answer - gt_need_extract: Whether the ground truth needs extraction - answer_pattern: Regex pattern to extract the answer - - Returns: - Tuple of (is_correct, normalized_prediction) - """ - print(f"[DAPO compute_score: gt]: {gt}, type: {type(gt)}") - print(f"[DAPO compute_score: solution_str]: {solution_str}") - # Extract answer from solution - match = re.findall(answer_pattern, solution_str) - print(f"[DAPO compute_score: match]: {match}") - extracted_answer = match[-1] if match else "[INVALID]" - pred = normalize_final_answer(extracted_answer) - print(f"[DAPO compute_score: extracted_answer]: {extracted_answer}") - print(f"[DAPO compute_score: pred]: {pred}") - - # Process ground truth - if gt_need_extract: - gt = normalize_final_answer(remove_boxed(last_boxed_only_string(gt))) - else: - gt = normalize_final_answer(gt) - - return (pred == gt), pred - - -def is_correct_strict_box(pred: str, - gt: str, - pause_tokens_index: Optional[list[int]] = None) -> tuple[int, Optional[str]]: - """Check if the prediction is correct using strict boxed answer criteria. - - Args: - pred: The prediction string - gt: The ground truth answer - pause_tokens_index: Indices of pause tokens - - Returns: - Tuple of (score, extracted_prediction) - """ - # Extract the relevant part of the prediction - if pause_tokens_index is not None: - assert len(pause_tokens_index) == 4 - pred = pred[pause_tokens_index[-1] - 100:] - else: - pred = pred[-100:] - - # Extract and check the boxed answer - boxed_pred = last_boxed_only_string(pred) - extracted_pred = remove_boxed(boxed_pred) if boxed_pred is not None else None - - return 1 if (extracted_pred == gt) else -1, extracted_pred - - -def verify(solution_str: str, - answer: str, - strict_box_verify: bool = False, - pause_tokens_index: Optional[list[int]] = None) -> bool: - """Verify if the solution is correct. - - Args: - solution_str: The solution string to verify - answer: The ground truth answer - strict_box_verify: Whether to use strict box verification - pause_tokens_index: Indices of pause tokens - - Returns: - True if the solution is correct, False otherwise - """ - if strict_box_verify: - correct, pred = is_correct_strict_box(solution_str, answer, pause_tokens_index) - return correct == 1, pred - - correct, pred = is_correct_minerva(solution_str, answer) - return correct, pred - - -def compute_score(solution_str: str, - ground_truth: str, - strict_box_verify: bool = False, - pause_tokens_index: Optional[list[int]] = None) -> float: - """Compute the reward score for a solution. - - Args: - solution_str: The solution string - ground_truth: The ground truth answer - config: Configuration object containing reward model settings - pause_tokens_index: Indices of pause tokens - - Returns: - Reward score (1.0 for correct, -1.0 for incorrect) - """ - # Limit solution length for efficiency - solution_str = solution_str[-300:] # The longest answer in MATH-500 has 159 characters - - # Verify the solution - correct, pred = verify(solution_str, ground_truth, strict_box_verify, pause_tokens_index) - - reward = 1.0 if correct else -1.0 - acc = correct - - return { - "score": reward, - "acc": acc, - } \ No newline at end of file diff --git a/reasoners/reward/math_verify.py b/reasoners/reward/math_verify.py deleted file mode 100644 index 66669ef94..000000000 --- a/reasoners/reward/math_verify.py +++ /dev/null @@ -1,33 +0,0 @@ -# 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. - -from math_verify.metric import math_metric -from math_verify.parser import LatexExtractionConfig, ExprExtractionConfig - - -def compute_score(model_output: str, ground_truth: str) -> bool: - verify_func = math_metric( - gold_extraction_target=(LatexExtractionConfig(),), - pred_extraction_target=(ExprExtractionConfig(), LatexExtractionConfig()), - ) - ret_score = 0. - - # Wrap the ground truth in \boxed{} format for verification - ground_truth_boxed = "\\boxed{" + ground_truth + "}" - try: - ret_score, _ = verify_func([ground_truth_boxed], [model_output]) - except Exception as e: - print(e) - - return ret_score diff --git a/reasoners/reward/orz/__init__.py b/reasoners/reward/orz/__init__.py deleted file mode 100644 index 0a513de80..000000000 --- a/reasoners/reward/orz/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -""" -Adapted from https://github.com/Open-Reasoner-Zero/Open-Reasoner-Zero -""" diff --git a/reasoners/reward/orz/math_utils.py b/reasoners/reward/orz/math_utils.py deleted file mode 100644 index 859a2d4f0..000000000 --- a/reasoners/reward/orz/math_utils.py +++ /dev/null @@ -1,448 +0,0 @@ -import asyncio -import re -from itertools import islice, zip_longest - -from sympy.parsing.latex import parse_latex - -try: - from math_verify import parse, verify -except ImportError: - print("math_verify is not installed in this environment") - parse = None - verify = None - - -def repeatness(s: str): - def ranks(l): - index = {v: i for i, v in enumerate(sorted(set(l)))} - return [index[v] for v in l] - - def suffixArray(s): - line = ranks(s) - n, k, ans, sa = len(s), 1, line, [0] * len(s) - while k < n - 1: - line = ranks(list(zip_longest(line, islice(line, k, None), fillvalue=-1))) - ans, k = line, k << 1 - for i, k in enumerate(ans): - sa[k] = i - return ans, sa - - def lcp(arr, suffixArr, inv_suff): - n, ans, k = len(arr), [0] * len(arr), 0 - - for i in range(n): - if inv_suff[i] == n - 1: - k = 0 - continue - - j = suffixArr[inv_suff[i] + 1] - while i + k < n and j + k < n and arr[i + k] == arr[j + k]: - k += 1 - - ans[inv_suff[i]] = k - if k > 0: - k -= 1 - - return ans - - arr = [ord(i) for i in s] - n = len(arr) - if n <= 1: - return 0 - c, sa = suffixArray(arr) - cnt = sum(lcp(arr, sa, c)) - - return (cnt * 2 / (n * (n + 1))) > 0.2 - - -SUBSTITUTIONS = [ - ("an ", ""), - ("a ", ""), - (".$", "$"), - ("\\$", ""), - (r"\ ", ""), - (" ", ""), - ("mbox", "text"), - (",\\text{and}", ","), - ("\\text{and}", ","), - ("\\text{m}", "\\text{}"), -] - - -REMOVED_EXPRESSIONS = [ - "square", - "ways", - "integers", - "dollars", - "mph", - "inches", - "ft", - "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. - This code comes from https://arxiv.org/pdf/2206.14858.pdf, page18. - """ - # final_answer = final_answer.split("=")[-1] - - for before, after in SUBSTITUTIONS: - final_answer = final_answer.replace(before, after) - for expr in REMOVED_EXPRESSIONS: - final_answer = final_answer.replace(expr, "") - - # Extract answer that is in LaTeX math, is bold, - # is surrounded by a box, etc. - 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 100,000 -> 100000 - if final_answer.replace(",", "").isdigit(): - final_answer = final_answer.replace(",", "") - - return final_answer - - -def latex_eval(latex): - sym = parse_latex(latex) - val = sym.evalf() - return sym, val - - -def _is_latex_equal(str1, str2): - try: - sym1, val1 = latex_eval(str1) - sym2, val2 = latex_eval(str2) - if sym1 == sym2 or val1 == val2: - return True - else: - raise ValueError - except Exception: # noqa - try: - norm1, norm2 = normalize_final_answer(str1), normalize_final_answer(str2) - sym1, val1 = latex_eval(norm1) - sym2, val2 = latex_eval(norm2) - if sym1 == sym2 or val1 == val2: - return True - except Exception: # noqa - return norm1 == norm2 - return False - - -async def is_latex_equal(str1, str2, executor, math_mode="legacy"): - if math_mode == "legacy": - if (len(str1) > 128 and repeatness(str1)) or ( - len(str2) > 128 and repeatness(str2) - ): - return False - - try: - loop = asyncio.get_event_loop() - task = loop.run_in_executor(executor, _is_latex_equal, str1, str2) - result = await asyncio.wait_for(task, timeout=1.0) - return result - except asyncio.exceptions.TimeoutError: - return False - elif math_mode == "math_verify": - try: - loop = asyncio.get_event_loop() - task = loop.run_in_executor(executor, verify, parse(str1), parse(str2)) - result = await asyncio.wait_for(task, timeout=1.0) - return result - except asyncio.exceptions.TimeoutError: - return False - else: - raise NotImplementedError(f"Math mode {math_mode} is not implemented") - - -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 Exception: # noqa - 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 Exception: # noqa - 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", "") - # print(string) - - # remove inverse spaces - string = string.replace("\\!", "") - # print(string) - - # replace \\ with \ - string = string.replace("\\\\", "\\") - # print(string) - - # replace tfrac and dfrac with frac - string = string.replace("tfrac", "frac") - string = string.replace("dfrac", "frac") - # print(string) - - # remove \left and \right - string = string.replace("\\left", "") - string = string.replace("\\right", "") - # print(string) - - # Remove circ (degrees) - string = string.replace("^{\\circ}", "") - string = string.replace("^\\circ", "") - - # remove dollar signs - string = string.replace("\\$", "") - string = string.replace("$", "") - 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 - - -def is_equiv(str1, str2, verbose=False) -> bool: - 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) - try: - return float(ss1) == (float(ss2)) - except Exception: # noqa - return ss1 == ss2 - except Exception: # noqa - return str1 == str2 - - -def last_boxed_only_string(string): - idx = string.rfind("\\boxed") - 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 remove_boxed(s): - left = "\\boxed{" - try: - assert s[: len(left)] == left - assert s[-1] == "}" - return s[len(left) : -1] - except Exception: - return None - - -def get_answer_str(s: str) -> str: - res = remove_boxed(last_boxed_only_string(s)) - if res is not None: - return res - return s - - -async def is_equal(str1, str2, executor, math_mode="legacy"): - first_equal = is_equiv(str1, str2) - if first_equal: - return True - return await is_latex_equal(str1, str2, executor, math_mode) - - -def solution2answer(solution: str, math_mode="eval_peeking") -> str: - answer = solution - if math_mode == "eval_peeking": - answer = get_answer_str(solution) - else: - raise ValueError(f"Invalid math_mode: {math_mode}") - return answer - - -def get_final_answer(output: str) -> str: - output = output.replace("is:", "is").replace("answer:", "answer is").strip() - if output.endswith("."): - output = output[:-1] - if ".$" in output: - output = output.replace(".$", "$") - pattern_list = [ - r"answer is (-?\d+\.?\d*)$", - r"answer is (.+?)$", - ] - matches = [] - for pat in pattern_list: - matches = re.findall(pat, output, re.S) - if matches: - return get_answer_str(matches[0]) - - return get_answer_str(output) diff --git a/reasoners/reward/orz/math_utils_sync.py b/reasoners/reward/orz/math_utils_sync.py deleted file mode 100644 index 5d453f930..000000000 --- a/reasoners/reward/orz/math_utils_sync.py +++ /dev/null @@ -1,442 +0,0 @@ -import asyncio -import re -from itertools import islice, zip_longest - -from sympy.parsing.latex import parse_latex - -try: - from math_verify import parse, verify -except ImportError: - print("math_verify is not installed in this environment") - parse = None - verify = None - - -def repeatness(s: str): - def ranks(l): - index = {v: i for i, v in enumerate(sorted(set(l)))} - return [index[v] for v in l] - - def suffixArray(s): - line = ranks(s) - n, k, ans, sa = len(s), 1, line, [0] * len(s) - while k < n - 1: - line = ranks(list(zip_longest(line, islice(line, k, None), fillvalue=-1))) - ans, k = line, k << 1 - for i, k in enumerate(ans): - sa[k] = i - return ans, sa - - def lcp(arr, suffixArr, inv_suff): - n, ans, k = len(arr), [0] * len(arr), 0 - - for i in range(n): - if inv_suff[i] == n - 1: - k = 0 - continue - - j = suffixArr[inv_suff[i] + 1] - while i + k < n and j + k < n and arr[i + k] == arr[j + k]: - k += 1 - - ans[inv_suff[i]] = k - if k > 0: - k -= 1 - - return ans - - arr = [ord(i) for i in s] - n = len(arr) - if n <= 1: - return 0 - c, sa = suffixArray(arr) - cnt = sum(lcp(arr, sa, c)) - - return (cnt * 2 / (n * (n + 1))) > 0.2 - - -SUBSTITUTIONS = [ - ("an ", ""), - ("a ", ""), - (".$", "$"), - ("\\$", ""), - (r"\ ", ""), - (" ", ""), - ("mbox", "text"), - (",\\text{and}", ","), - ("\\text{and}", ","), - ("\\text{m}", "\\text{}"), -] - - -REMOVED_EXPRESSIONS = [ - "square", - "ways", - "integers", - "dollars", - "mph", - "inches", - "ft", - "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. - This code comes from https://arxiv.org/pdf/2206.14858.pdf, page18. - """ - # final_answer = final_answer.split("=")[-1] - - for before, after in SUBSTITUTIONS: - final_answer = final_answer.replace(before, after) - for expr in REMOVED_EXPRESSIONS: - final_answer = final_answer.replace(expr, "") - - # Extract answer that is in LaTeX math, is bold, - # is surrounded by a box, etc. - 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 100,000 -> 100000 - if final_answer.replace(",", "").isdigit(): - final_answer = final_answer.replace(",", "") - - return final_answer - - -def latex_eval(latex): - sym = parse_latex(latex) - val = sym.evalf() - return sym, val - - -def _is_latex_equal(str1, str2): - try: - sym1, val1 = latex_eval(str1) - sym2, val2 = latex_eval(str2) - if sym1 == sym2 or val1 == val2: - return True - else: - raise ValueError - except Exception: # noqa - try: - norm1, norm2 = normalize_final_answer(str1), normalize_final_answer(str2) - sym1, val1 = latex_eval(norm1) - sym2, val2 = latex_eval(norm2) - if sym1 == sym2 or val1 == val2: - return True - except Exception: # noqa - return norm1 == norm2 - return False - - -def is_latex_equal(str1, str2, math_mode="legacy"): - if math_mode == "legacy": - if (len(str1) > 128 and repeatness(str1)) or ( - len(str2) > 128 and repeatness(str2) - ): - return False - - try: - return _is_latex_equal(str1, str2) - except Exception: # noqa - return False - elif math_mode == "math_verify": - try: - return verify(parse(str1), parse(str2)) - except Exception: # noqa - return False - else: - raise NotImplementedError(f"Math mode {math_mode} is not implemented") - - -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 Exception: # noqa - 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 Exception: # noqa - 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", "") - # print(string) - - # remove inverse spaces - string = string.replace("\\!", "") - # print(string) - - # replace \\ with \ - string = string.replace("\\\\", "\\") - # print(string) - - # replace tfrac and dfrac with frac - string = string.replace("tfrac", "frac") - string = string.replace("dfrac", "frac") - # print(string) - - # remove \left and \right - string = string.replace("\\left", "") - string = string.replace("\\right", "") - # print(string) - - # Remove circ (degrees) - string = string.replace("^{\\circ}", "") - string = string.replace("^\\circ", "") - - # remove dollar signs - string = string.replace("\\$", "") - string = string.replace("$", "") - 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 - - -def is_equiv(str1, str2, verbose=False) -> bool: - 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) - try: - return float(ss1) == (float(ss2)) - except Exception: # noqa - return ss1 == ss2 - except Exception: # noqa - return str1 == str2 - - -def last_boxed_only_string(string): - idx = string.rfind("\\boxed") - 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 remove_boxed(s): - left = "\\boxed{" - try: - assert s[: len(left)] == left - assert s[-1] == "}" - return s[len(left) : -1] - except Exception: - return None - - -def get_answer_str(s: str) -> str: - res = remove_boxed(last_boxed_only_string(s)) - if res is not None: - return res - return s - - -def is_equal(str1, str2, math_mode="legacy"): - first_equal = is_equiv(str1, str2) - if first_equal: - return True - return is_latex_equal(str1, str2, math_mode) - - -def solution2answer(solution: str, math_mode="eval_peeking") -> str: - answer = solution - if math_mode == "eval_peeking": - answer = get_answer_str(solution) - else: - raise ValueError(f"Invalid math_mode: {math_mode}") - return answer - - -def get_final_answer(output: str) -> str: - output = output.replace("is:", "is").replace("answer:", "answer is").strip() - if output.endswith("."): - output = output[:-1] - if ".$" in output: - output = output.replace(".$", "$") - pattern_list = [ - r"answer is (-?\d+\.?\d*)$", - r"answer is (.+?)$", - ] - matches = [] - for pat in pattern_list: - matches = re.findall(pat, output, re.S) - if matches: - return get_answer_str(matches[0]) - - return get_answer_str(output) diff --git a/reasoners/reward/prime_code/__init__.py b/reasoners/reward/prime_code/__init__.py deleted file mode 100644 index 49b0d30f5..000000000 --- a/reasoners/reward/prime_code/__init__.py +++ /dev/null @@ -1,73 +0,0 @@ -# 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. - -from .utils import check_correctness as apps_check_correctness -import json -import re -import traceback - - -def compute_score(completion, test_cases, continuous=False): - # try to get code solution from completion. if the completion is pure code, this will not take effect. - solution = completion.split('```python')[-1].split('```')[0] - try: - try: - if not isinstance(test_cases, dict): - test_cases = json.loads(test_cases) - except Exception as e: - print(f"Error:{e}") - - # Complete check on all in-out pairs first. If there is no failure, per-sample test can be skipped. - try: - res, metadata = apps_check_correctness(in_outs=test_cases, generation=solution, timeout=5, debug=False) - metadata = dict(enumerate(metadata))[0] - success = all(map(lambda x: x == True, res)) - if success: - return success, metadata - except Exception as e: - pass - - test_cases_list = [] - inputs = test_cases["inputs"] - outputs = test_cases["outputs"] - for i in range(len(inputs)): - test_cases_list.append({"inputs": [inputs[i]], "outputs": [outputs[i]]}) - - if continuous: - # per sample test: if continuous score is needed, test first 10 samples regardless of failures - # do not test all samples cuz some problems have enormous test cases - metadata_list = [] - res_list = [] - for test_case_id, test_case in enumerate(test_cases_list): - res, metadata = apps_check_correctness(in_outs=test_case, generation=solution, timeout=5, debug=False) - try: - metadata = dict(enumerate(metadata))[0] # metadata can be empty occasionally - except Exception as e: - metadata = {} - metadata["test_case"] = {} - metadata["test_case"]["input"] = str(test_case["inputs"][0]) - metadata["test_case"]["output"] = str(test_case["outputs"][0]) - metadata["test_case"]["res"] = str(res) - metadata_list.append(metadata) - res_list.extend(res) - - if test_case_id >= 9: - break - res_count = len(res_list) if len(res_list) > 0 else 1 - success = sum(map(lambda x: x == True, res_list)) / res_count - except Exception as e: - traceback.print_exc(10) - success = False - metadata_list = None - return success, metadata_list diff --git a/reasoners/reward/prime_code/testing_util.py b/reasoners/reward/prime_code/testing_util.py deleted file mode 100644 index 34845d07e..000000000 --- a/reasoners/reward/prime_code/testing_util.py +++ /dev/null @@ -1,721 +0,0 @@ -# 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. - -import ast -import json -import sys -import faulthandler -import platform - -# used for debugging to time steps -from datetime import datetime - -# to run the solution files we're using a timing based approach -import signal - -import numpy as np - -# for capturing the stdout -from io import StringIO - -# used for testing the code that reads from input -from unittest.mock import patch, mock_open - -from pyext import RuntimeModule - -from enum import Enum - -import traceback - - -def truncatefn(s, length=300): - assert isinstance(s, str) - if len(s) <= length: - return s - - return s[:length // 2] + "...(truncated) ..." + s[-length // 2:] - - -class CODE_TYPE(Enum): - call_based = 0 - standard_input = 1 - - -# stuff for setting up signal timer -class TimeoutException(Exception): - pass - - -def timeout_handler(signum, frame): - print("alarm went off") - return - # raise TimeoutException # this is an unhandled exception. just return None is OK - - -signal.signal(signal.SIGALRM, timeout_handler) - -# timeout = 6 # seconds - - -# used to capture stdout as a list -# from https://stackoverflow.com/a/16571630/6416660 -# alternative use redirect_stdout() from contextlib -class Capturing(list): - - def __enter__(self): - self._stdout = sys.stdout - sys.stdout = self._stringio = StringIO() - # Make closing the StringIO a no-op - self._stringio.close = lambda x: 1 - return self - - def __exit__(self, *args): - self.append(self._stringio.getvalue()) - del self._stringio # free up some memory - sys.stdout = self._stdout - - -def only_int_check(val): - return isinstance(val, int) - - -def string_int_check(val): - return isinstance(val, str) and val.isdigit() - - -def combined_int_check(val): - return only_int_check(val) or string_int_check(val) - - -def clean_traceback(error_traceback): - file_start = error_traceback.find('File \"\"') - # print(file_start) - error_traceback = "Traceback (most recent call last):\n " + error_traceback[file_start:] - return error_traceback - - -def run_test(in_outs, test=None, debug=False, timeout=15): - """ - if test(generated_code) is not None it'll try to run the code. - otherwise it'll just return an input and output pair. - """ - # Disable functionalities that can make destructive changes to the test. - reliability_guard() - - if debug: - print(f"start = {datetime.now().time()}") - - if in_outs: - if in_outs.get("fn_name") is None: - which_type = CODE_TYPE.standard_input # Standard input - method_name = None - else: - which_type = CODE_TYPE.call_based # Call-based - method_name = in_outs["fn_name"] - - if debug: - print(f"loaded input_output = {datetime.now().time()}") - - if test is None: - assert False, "should not happen: test code is none" - return in_outs, {"error": "no test code provided"} - elif test is not None: - results = [] - sol = "from string import *\nfrom re import *\nfrom datetime import *\nfrom collections import *\nfrom heapq import *\nfrom bisect import *\nfrom copy import *\nfrom math import *\nfrom random import *\nfrom statistics import *\nfrom itertools import *\nfrom functools import *\nfrom operator import *\nfrom io import *\nfrom sys import *\nfrom json import *\nfrom builtins import *\nfrom typing import *\nimport string\nimport re\nimport datetime\nimport collections\nimport heapq\nimport bisect\nimport copy\nimport math\nimport random\nimport statistics\nimport itertools\nimport functools\nimport operator\nimport io\nimport sys\nimport json\nsys.setrecursionlimit(6*10**5)\n" - if debug: - print(f"loading test code = {datetime.now().time()}") - - if which_type == CODE_TYPE.call_based: - - sol += test - if debug: - print(f"sol = {sol}") - signal.alarm(timeout) - try: - tmp_sol = RuntimeModule.from_string("tmp_sol", "", sol) - if "class Solution" not in test: - tmp = tmp_sol - else: - tmp = tmp_sol.Solution() - signal.alarm(0) - except Exception as e: - signal.alarm(0) - error_traceback = traceback.format_exc() - if debug: - print(f"type 0 compilation error = {e}") - results.append(-2) - return results, { - "error": repr(e), - # "error_code": -1, - # "error_message": "Compilation Error", - "traceback": clean_traceback(error_traceback), - } - signal.alarm(0) - - elif which_type == CODE_TYPE.standard_input: - # sol - # if code has if __name__ == "__main__": then remove it - try: - astree = ast.parse(test) - last_block = astree.body[-1] - if isinstance(last_block, ast.If): - condition = last_block.test - if ast.unparse(condition).strip() == "__name__ == '__main__'": - test = (ast.unparse(astree.body[:-1]) + "\n" + ast.unparse(last_block.body)) - except: - pass - - tmp_test = test.split("\n") - - new_test = [] - for x in tmp_test: - if (not x.startswith("from ")) and (not x.startswith("import ")): - new_test.append("\t" + x + "\n") - else: - new_test.append(x + "\n") - tmp_test = new_test - - new_test = "" - started = False - for i in tmp_test: - if i.startswith("\t") and not started: - new_test += "stdin = sys.stdin\nstdout = sys.stdout\n" - new_test += "def code():\n" - new_test += i - started = True - elif started and ((i.startswith("from ")) or (i.startswith("import "))): - new_test += "\t" + i - else: - new_test += i - tmp_test = new_test - - sol += tmp_test - if debug: - print(f"sol = {sol}") - method_name = "code" - signal.alarm(timeout) - try: - tmp_sol = RuntimeModule.from_string("tmp_sol", "", sol) - tmp = tmp_sol - signal.alarm(0) - except Exception as e: - signal.alarm(0) - error_traceback = traceback.format_exc() - if debug: - print(f"type 1 compilation error = {e}") - results.append(-2) - return results, { - "error": repr(e), - # "error_code": -1, - # "error_message": "Compilation Error", - "traceback": clean_traceback(error_traceback), - } - signal.alarm(0) - if debug: - print(f"get method = {datetime.now().time()}") - - try: - method = getattr(tmp, method_name) # get_attr second arg must be str - except: - signal.alarm(0) - error_traceback = traceback.format_exc() - e = sys.exc_info() - print(f"unable to get function error = {e}") - results.append(-2) - return results, { - "error": repr(e), - # "error_code": -1, - # "error_message": "Unable to extract code", - "traceback": clean_traceback(error_traceback), - } - - for index, inputs in enumerate(in_outs["inputs"]): - raw_inputs = inputs - raw_outputs = in_outs["outputs"][index] - if which_type == CODE_TYPE.call_based: - inputs = [json.loads(line) for line in inputs.split("\n")] - in_outs["outputs"][index] = json.loads(in_outs["outputs"][index]) - - truncate_line_size = 300 // (raw_inputs.count("\n") + 1) - raw_inputs = "\n".join( - [truncatefn(line, truncate_line_size) for line in raw_inputs.strip().split("\n")]) - raw_outputs = truncatefn(raw_outputs, 200) - else: - raw_inputs = truncatefn(raw_inputs) - raw_outputs = truncatefn(raw_outputs, 200) - # JSON forces dictionaries to have string keys; this undoes this (assuming a singleton list) - try: - if isinstance(inputs[0], dict): - inputs = [{int(k): v for k, v in inputs[0].items()}] - except: - True - try: - if isinstance(in_outs["outputs"][index], dict): - in_outs["outputs"][index] = [{int(k): v for k, v in in_outs["outputs"][index].items()}] - except: - True - try: - if isinstance(in_outs["outputs"][index][0], dict): - in_outs["outputs"][index] = [{int(k): v for k, v in in_outs["outputs"][index][0].items()}] - except: - True - - if debug: - print( - f"time: {datetime.now().time()} testing index = {index} inputs = {inputs}, {type(inputs)}. type = {which_type}" - ) - if which_type == CODE_TYPE.call_based: # Call-based - signal.alarm(timeout) - faulthandler.enable() - try: - output = method(*inputs) - raw_true_output = output - - raw_true_output_copy = json.dumps(output) - raw_true_output_copy = truncatefn(raw_true_output_copy, 200) - - # ground truth sequences are not tuples - if isinstance(output, tuple): - output = list(output) - - tmp_result = output == in_outs["outputs"][index] - if (isinstance(in_outs["outputs"][index], list) and in_outs["outputs"][index]): - tmp_result = tmp_result or (output == in_outs["outputs"][index][0]) - - # ground truth sequences are not tuples - try: - if isinstance(output[0], tuple): - tmp_result = tmp_result or ([list(x) for x in output] == in_outs["outputs"][index][0]) - except: - True - results.append(tmp_result) - if tmp_result != True: - return results, { - "output": raw_true_output_copy, - "expected": raw_outputs, - "inputs": raw_inputs, - # "error_code": -2, - "error_message": "Wrong Answer", - } - # reset the alarm - signal.alarm(0) - except Exception as e: - signal.alarm(0) - error_traceback = traceback.format_exc() - faulthandler.disable() - if debug: - print(f"Standard input runtime error or time limit exceeded error = {e}") - results.append(-1) - if "timeoutexception" in repr(e).lower(): - return results, { - "error": repr(e), - # "error_code": -3, - # "error_message": "Time Limit Exceeded", - # "inputs": raw_inputs, - # "expected": raw_outputs, - "traceback": clean_traceback(error_traceback), - } - else: - return results, { - "error": repr(e), - # "error_code": -4, - # "error_message": "Runtime Error", - # "inputs": raw_inputs, - # "expected": raw_outputs, - "traceback": clean_traceback(error_traceback), - } - faulthandler.disable() - signal.alarm(0) - if debug: - print( - f"outputs = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs}, {type(inputs)}, {output == [in_outs['outputs'][index]]}" - ) - elif which_type == CODE_TYPE.standard_input: # Standard input - faulthandler.enable() - passed = False - - if isinstance(inputs, list): - inputs = "\n".join(inputs) - if isinstance(in_outs["outputs"][index], list): - in_outs["outputs"][index] = "\n".join(in_outs["outputs"][index]) - - signal.alarm(timeout) - with Capturing() as output: - try: - call_method(method, inputs) - # reset the alarm - signal.alarm(0) - passed = True - except Exception as e: - # runtime error or took too long - signal.alarm(0) - error_traceback = traceback.format_exc() - print(f"Call-based runtime error or time limit exceeded error = {repr(e)}{e}") - results.append(-1) - if "timeoutexception" in repr(e).lower(): - return results, { - "error": repr(e), - # "error_code": -3, - # "error_message": "Time Limit Exceeded", - # "inputs": raw_inputs, - # "expected": raw_outputs, - "traceback": clean_traceback(error_traceback), - } - else: - return results, { - "error": repr(e), - # "error_code": -4, - # "error_message": "Runtime Error", - # "inputs": raw_inputs, - # "expected": raw_outputs, - "traceback": clean_traceback(error_traceback), - } - signal.alarm(0) - raw_true_output = output[0] - raw_true_output_copy = truncatefn(raw_true_output, 200) - output = raw_true_output.splitlines() - if not passed: - if debug: - nl = "\n" - if not isinstance(inputs, list): - print( - f"not passed output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs.replace(nl,' new-line ')}, {type(inputs)}, {output == [in_outs['outputs'][index]]}" - ) - else: - print( - f"not passed output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs}, {type(inputs)}, {output == [in_outs['outputs'][index]]}" - ) - continue - - if passed and debug: - print(f"==> output = {output}, test outputs = {in_outs['outputs'][index]}") - - if custom_compare_(output, in_outs["outputs"][index]): - tmp_result = True - results.append(tmp_result) - continue - - # ground truth sequences are expressed as lists not tuples - if isinstance(output, tuple): - output = list(output) - - tmp_result = False - try: - tmp_result = output == [in_outs["outputs"][index]] - if isinstance(in_outs["outputs"][index], list): - tmp_result = tmp_result or (output == in_outs["outputs"][index]) - if isinstance(output[0], str): - tmp_result = tmp_result or ([e.strip() for e in output] == in_outs["outputs"][index]) - except Exception as e: - if debug: - print(f"Failed check1 exception = {e}") - pass - - if tmp_result == True: - results.append(tmp_result) - continue - - # try one more time without \n - if isinstance(in_outs["outputs"][index], list): - for tmp_index, i in enumerate(in_outs["outputs"][index]): - in_outs["outputs"][index][tmp_index] = i.split("\n") - in_outs["outputs"][index][tmp_index] = [ - x.strip() for x in in_outs["outputs"][index][tmp_index] if x - ] - else: - in_outs["outputs"][index] = in_outs["outputs"][index].split("\n") - in_outs["outputs"][index] = list(filter(len, in_outs["outputs"][index])) - in_outs["outputs"][index] = list(map(lambda x: x.strip(), in_outs["outputs"][index])) - - try: - tmp_result = output == [in_outs["outputs"][index]] - if isinstance(in_outs["outputs"][index], list): - tmp_result = tmp_result or (output == in_outs["outputs"][index]) - except Exception as e: - if debug: - print(f"Failed check2 exception = {e}") - pass - - if tmp_result == True: - results.append(tmp_result) - continue - - # try by converting the output into a split up list too - if isinstance(output, list): - output = list(filter(len, output)) - - if debug: - nl = "\n" - if not isinstance(inputs, list): - print( - f"@1 output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs.replace(nl,' new-line ')}, {type(inputs)}, {output == [in_outs['outputs'][index]]} {tmp_result=}" - ) - else: - print( - f"@1 output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs}, {type(inputs)}, {output == [in_outs['outputs'][index]]} {tmp_result=}" - ) - - if tmp_result == True: - results.append(tmp_result) - continue - - if debug: - print(f"{tmp_result=} @a") - - try: - tmp_result = output == [in_outs["outputs"][index]] - if isinstance(in_outs["outputs"][index], list): - tmp_result = tmp_result or (output == in_outs["outputs"][index]) - except Exception as e: - if debug: - print(f"Failed check3 exception = {e}") - pass - - if debug: - print(f"{tmp_result=} @b") - - try: - all_ints = all( - combined_int_check(e1) and combined_int_check(e2) - for e1, e2 in zip(output, in_outs["outputs"][index])) - if not all_ints: - if debug: - print([ - combined_int_check(e1) and combined_int_check(e2) - for e1, e2 in zip(output, in_outs["outputs"][index]) - ]) - output_float = [float(e) for e in output] - gt_float = [float(e) for e in in_outs["outputs"][index]] - tmp_result = tmp_result or ((len(output_float) == len(gt_float)) and - np.allclose(output_float, gt_float)) - except Exception as e: - pass - - if debug: - print(f"{tmp_result=} @c") - - try: - if isinstance(output[0], list): - all_ints = all( - combined_int_check(e1) and combined_int_check(e2) - for e1, e2 in zip(output[0], in_outs["outputs"][index])) - if not all_ints: - output_float = [float(e) for e in output[0]] - gt_float = [float(e) for e in in_outs["outputs"][index][0]] - tmp_result = tmp_result or ((len(output_float) == len(gt_float)) and - np.allclose(output_float, gt_float)) - except Exception as e: - pass - - if tmp_result == True: - results.append(tmp_result) - continue - - if debug: - print(f"{tmp_result=} @d") - # try by converting the stuff into split up list - if isinstance(in_outs["outputs"][index], list): - for tmp_index, i in enumerate(in_outs["outputs"][index]): - in_outs["outputs"][index][tmp_index] = set(i.split()) - else: - in_outs["outputs"][index] = set(in_outs["outputs"][index].split()) - - if debug: - print(f"{tmp_result=} @e") - - try: - tmp_result = output == in_outs["outputs"][index] - except Exception as e: - if debug: - print(f"Failed check4 exception = {e}") - continue - - if tmp_result == True: - results.append(tmp_result) - continue - - if debug: - print(f"{tmp_result=} @f") - - # try by converting the output into a split up list too - if isinstance(output, list): - for tmp_index, i in enumerate(output): - output[tmp_index] = i.split() - output = list(filter(len, output)) - for tmp_index, i in enumerate(output): - output[tmp_index] = set(i) - else: - output = output.split() - output = list(filter(len, output)) - output = set(output) - - if debug: - print(f"{tmp_result=} @g") - - if tmp_result == True and debug: - print("PASSED") - - results.append(tmp_result) - if tmp_result != True: - return results, { - "output": raw_true_output_copy, - "expected": raw_outputs, - "inputs": raw_inputs, - # "error_code": -2, - "error_message": "Wrong Answer", - } - - if debug: - nl = "\n" - if not isinstance(inputs, list): - print( - f"@2 output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs.replace(nl,' new-line ')}, {type(inputs)}, {output == [in_outs['outputs'][index]]}" - ) - else: - print( - f"@2 output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs}, {type(inputs)}, {output == [in_outs['outputs'][index]]}" - ) - - print(f"results = {results}") - - return results, {} - - -def custom_compare_(output, ground_truth): - - if isinstance(output, list): - output_1 = "\n".join(output) - if stripped_string_compare(output_1, ground_truth): - return True - - if isinstance(output, list): - output_2 = [o.lstrip().rstrip() for o in output] - output_2 = "\n".join(output_2) - if stripped_string_compare(output_2, ground_truth): - return True - - return False - - -def stripped_string_compare(s1, s2): - s1 = s1.lstrip().rstrip() - s2 = s2.lstrip().rstrip() - return s1 == s2 - - -def call_method(method, inputs): - - if isinstance(inputs, list): - inputs = "\n".join(inputs) - - inputs_line_iterator = iter(inputs.split("\n")) - - # sys.setrecursionlimit(10000) - - # @patch('builtins.input', side_effect=inputs.split("\n")) - @patch("builtins.open", mock_open(read_data=inputs)) - @patch("sys.stdin", StringIO(inputs)) - @patch("sys.stdin.readline", lambda *args: next(inputs_line_iterator)) - @patch("sys.stdin.readlines", lambda *args: inputs.split("\n")) - @patch("sys.stdin.read", lambda *args: inputs) - # @patch('sys.stdout.write', print) - def _inner_call_method(_method): - try: - return _method() - except SystemExit as e: - pass - finally: - pass - - return _inner_call_method(method) - - -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 # 防止干扰repl评测 - 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 diff --git a/reasoners/reward/prime_code/utils.py b/reasoners/reward/prime_code/utils.py deleted file mode 100644 index 07f3db88e..000000000 --- a/reasoners/reward/prime_code/utils.py +++ /dev/null @@ -1,59 +0,0 @@ -# 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. - -# Borrowed from: https://huggingface.co/spaces/codeparrot/apps_metric/blob/main/utils.py - -import multiprocessing -from typing import Dict, Optional -from datasets import load_dataset -from .testing_util import run_test -import traceback -import os, sys - - -def _temp_run(sample, generation, debug, result, metadata_list, timeout): - with open(os.devnull, 'w') as devnull: - sys.stdout = devnull - sys.stderr = devnull - try: - res, metadata = run_test(in_outs=sample, test=generation, debug=debug, timeout=timeout) - result.append(res) - metadata_list.append(metadata) - except Exception as e: - # print(e) # some tracebacks are extremely long. - traceback.print_exc(10) - result.append([-1 for i in range(len(sample['inputs']))]) - metadata_list.append({}) - - -def check_correctness(in_outs: Optional[dict], generation, timeout=10, debug=True): - """Check correctness of code generation with a global timeout. - The global timeout is to catch some extreme/rare cases not handled by the timeouts - inside `run_test`""" - - manager = multiprocessing.Manager() - result = manager.list() - metadata_list = manager.list() - p = multiprocessing.Process(target=_temp_run, args=(in_outs, generation, debug, result, metadata_list, timeout)) - p.start() - p.join(timeout=timeout + 1) - if p.is_alive(): - p.kill() - # p.terminate() - if not result: - # consider that all tests failed - result = [[-1 for i in range(len(in_outs["inputs"]))]] - if debug: - print(f"global timeout") - return result[0], metadata_list diff --git a/reasoners/reward/tablereason.py b/reasoners/reward/tablereason.py index 187d68f0a..55fcd74d8 100644 --- a/reasoners/reward/tablereason.py +++ b/reasoners/reward/tablereason.py @@ -1,5 +1,5 @@ -from verl.utils.reward_score.prime_math.grader import math_equal -from verl.utils.reward_score import math +from reasoners.reward.prime_math.grader import math_equal +from reasoners.reward import math import re diff --git a/reasoners/tools/model_filtering/README.md b/reasoners/tools/model_filtering/README.md deleted file mode 100644 index 040a15edc..000000000 --- a/reasoners/tools/model_filtering/README.md +++ /dev/null @@ -1,194 +0,0 @@ -# Difficulty Filter Pipeline - -A pipeline for difficulty filtering using reward functions, supporting data-parallel and tensor-parallel processing. The pipeline now operates in two stages: -1. **Inference stage**: Runs the model to generate responses (GPU-intensive) -2. **Reward stage**: Evaluates responses with reward functions (CPU-intensive) - -## Resource Requirements - -- 8x H200 (140GB) GPUs - -## Example Usage - -### Stage 1: Inference - -> [!IMPORTANT] -> To run the model, you have to install verl and upgrade vllm to 0.8.5, please first follow the instructions in the main README.md, and then run the following command: - -```bash -pip uninstall vllm -pip install vllm==0.8.5 -``` - -#### Qwen2.5-7B-Instruct (~0.09s 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 -``` - -#### Qwen3-30B-A3B (~12s 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 -``` -### Stage 2: Reward Calculation - -After the inference stage is complete, run the 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 -``` - -## Checkpoint & Resumption - -The pipeline automatically saves batch results to JSON files in the output directory as it processes data. If a run is interrupted, the pipeline will automatically resume from where it left off by: - -1. Scanning the output directory for existing batch files -2. Loading results from these files -3. Continuing from the next unprocessed batch - -### Advanced Checkpoint Options - -The pipeline provides flags to control checkpoint behavior: - -- **Inference stage**: - - `--force_regenerate`: Ignores all existing checkpoint files and starts processing from the beginning, overwriting previous results - -- **Reward stage**: - - `--recalculate_rewards`: Recalculates all reward scores using the previously generated model responses (useful when implementing new evaluation metrics without re-running the expensive model inference) - -## Utility Functions - -After running models with data-parallel processing, you can use the utility functions to concatenate and analyze results. - -### Generating Pass Rate Mappings - -Generate idx to pass_rate mapping for a specific model and dataset: - -```bash -python -m model_filtering.utils map \ - --output_dir "./diff_filter_output" \ - --dataset "codegen__leetcode2k_2.4k" \ - --model "Qwen2.5-7B-Instruct" -``` - -This command will create an `idx_to_passrate.json` file in the model's directory that can be used for analysis. - -### Analyzing Results - -Analyze difficulty distributions of one or more datasets using a specific model: - -```bash -python -m model_filtering.utils analyze \ - --output_dir "./diff_filter_output" \ - --datasets "codegen__leetcode2k_2.4k" \ - --model "Qwen2.5-7B-Instruct" -``` - -Analyze multiple dataset chunks with the 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" -``` - -Some advanced options: -- `--regenerate`: Force regeneration of idx_to_passrate mapping -- `--save_combined`: Save the combined mapping when analyzing multiple datasets - -The analysis output categorizes problems into seven difficulty levels: -- Impossible (pass rate exactly 0.0) -- Very Hard (pass rate 0.0-0.2, exclusive) -- Hard (pass rate 0.2-0.4) -- Medium (pass rate 0.4-0.6) -- Easy (pass rate 0.6-0.8) -- Very Easy (pass rate 0.8-1.0, exclusive) -- Perfect (pass rate exactly 1.0) - - -# Annotate Pass Rates Script *run_add_pr.py* - -This script adds two new columns to each row of a Parquet dataset, reflecting pass-rate values loaded from two JSON files: - -* **`low_pass_rate`**: pass rate from the `--low_json` mapping (default is 0.0 if missing) -* **`high_pass_rate`**: pass rate from the `--high_json` mapping (default is 0.0 if missing) - -## Usage - -```bash -python run_add_pr.py \ - --parquet_in .parquet \ - --low_json low_pass_rates.json \ - --high_json high_pass_rates.json \ - --parquet_out_dir directory for saving processed parquet -``` - -## Command-Line Arguments - -| Argument | Description | Required | -| --------------- | -------------------------------------------------------- | -------- | -| `--parquet_in` | Path to the input Parquet file to annotate. | Yes | -| `--low_json` | JSON file mapping `idx → pass_rate` for low thresholds. | Yes | -| `--high_json` | JSON file mapping `idx → pass_rate` for high thresholds. | Yes | -| `--parquet_out_dir` | Directory where the annotated Parquet file will be saved. | Yes | - -# Filter by Pass-Rate Thresholds *run_filter.py* - -This script filters a Parquet dataset based on pass-rate thresholds defined in two JSON files. It uses **Rich** for a colored, tabular console summary. - -## Usage - -```bash -python run_filter.py \ - --parquet_in .parquet \ - --parquet_out_dir \ - --low_json low_pass_rates.json \ - --high_json high_pass_rates.json \ - --low_thresh 0.1 \ - --high_thresh 0.9 -``` - -## Arguments - -| Flag | Description | -| ------------------- | ----------------------------------------------------------- | -| `--parquet_in` | Path to the source Parquet file | -| `--parquet_out_dir` | Directory where filtered Parquet files will be saved | -| `--low_json` | JSON mapping `idx → pass_rate` for low-threshold filtering | -| `--high_json` | JSON mapping `idx → pass_rate` for high-threshold filtering | -| `--low_thresh` | Drop rows with pass rate **≤** this value in low json | -| `--high_thresh` | Drop rows with pass rate **≥** this value in high json | diff --git a/reasoners/tools/model_filtering/filter.py b/reasoners/tools/model_filtering/filter.py new file mode 100644 index 000000000..a8b406479 --- /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/test_parquet.py b/reasoners/tools/model_filtering/test_parquet.py deleted file mode 100644 index 2f7a0c194..000000000 --- a/reasoners/tools/model_filtering/test_parquet.py +++ /dev/null @@ -1,113 +0,0 @@ -#!/usr/bin/env python -# coding: utf-8 - -# In[1]: - - -import polars as pl -import pandas as pd -import glob -import os - -paths_and_dfs = {} -for file in glob.glob("/lustrefs/users/zhuojun.cheng/Reasoning360/data/train_guru15k/code*.parquet"): - if "livecodebench" in file.lower(): - df = pl.read_parquet(file).to_pandas() - else: - df = pd.read_parquet(file) - paths_and_dfs[file] = df - - -# In[3]: - - -import json -import matplotlib.pyplot as plt - -for key, df in paths_and_dfs.items(): - print(key) - print(f"Number of rows: {len(df)}") - unit_test_counts, all_unit_test_counts = [], [] - for i in range(len(df)): - ground_truth = json.loads(df['reward_model'][i]['ground_truth']) - if 'inputs' in ground_truth: - assert len(ground_truth['inputs']) == len(ground_truth['outputs']) - unit_test_counts.append(len(ground_truth['inputs'])) - all_unit_test_counts.append(len(ground_truth['inputs'])) - - # distribution of the unit test counts - plt.figure(figsize=(10, 6)) - plt.hist(unit_test_counts, bins=30, edgecolor='black') - plt.title('Distribution of Unit Test Counts') - plt.xlabel('Number of Unit Tests') - plt.ylabel('Frequency') - plt.grid(True, alpha=0.3) - plt.show() - - -# In[4]: - - -import numpy as np - -for key, df in paths_and_dfs.items(): - print(f"\nProcessing: {os.path.basename(key)}") - - # Set seed for reproducibility - np.random.seed(42) - - for idx in range(len(df)): - ground_truth = json.loads(df['reward_model'][idx]['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'] - - # Get 8 random indices from this test case's inputs/outputs - n_tests = len(complete_inputs) - 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 - ground_truth['complete_outputs'] = complete_outputs - ground_truth['inputs'] = [complete_inputs[i] for i in sample_indices] - ground_truth['outputs'] = [complete_outputs[i] for i in sample_indices] - - df.at[idx, 'reward_model']['ground_truth'] = json.dumps(ground_truth) - - # Update the dictionary with modified dataframe - paths_and_dfs[key] = df - - -# In[ ]: - - -import json -import matplotlib.pyplot as plt - -for key, df in paths_and_dfs.items(): - print(key) - print(f"Number of rows: {len(df)}") - unit_test_counts = [] - for i in range(len(df)): - ground_truth = json.loads(df['reward_model'][i]['ground_truth']) - if 'inputs' in ground_truth: - unit_test_counts.append(len(ground_truth['inputs'])) - - # distribution of the unit test counts - plt.figure(figsize=(10, 6)) - plt.hist(unit_test_counts, bins=30, edgecolor='black') - plt.title('Distribution of Unit Test Counts') - plt.xlabel('Number of Unit Tests') - plt.ylabel('Frequency') - plt.grid(True, alpha=0.3) - plt.show() - - -# In[ ]: - - -# save the df back to parquet -for key, df in paths_and_dfs.items(): - df.to_parquet(key) - diff --git a/reasoners/tools/model_filtering/utils.py b/reasoners/tools/model_filtering/utils.py index 0179d6b93..b91f0f978 100644 --- a/reasoners/tools/model_filtering/utils.py +++ b/reasoners/tools/model_filtering/utils.py @@ -13,6 +13,61 @@ 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 #