itp-interface 1.0.0__py3-none-any.whl

itp_interface/main/config/repo/coq_repos.yaml

@@ -0,0 +1,191 @@
+repos:
+  - UniMath:
+      url: https://github.com/UniMath/UniMath
+      commit: 7432feea2113a460eb5a69fbbba5fda02e2bf234
+      branch: master
+      build_cmds: # Verified on 2024-05-04
+        - opam switch create UniMath 4.14.1
+        - eval $(opam env --switch=UniMath --set-switch)
+        - opam pin add -y coq 8.18.0
+        - eval $(opam env)
+        - opam install coq-lsp --yes # 0.1.8+8.18 version
+        - make -j `nproc`
+      setup_cmds:
+        - opam env --switch=UniMath --set-switch
+  - math-comp:
+      url: https://github.com/math-comp/math-comp
+      commit: 936a3fd89c621caec7bae631147234f2e4d389f7
+      branch: master
+      build_cmds: # Verified on 2024-05-04
+        - cd mathcomp
+        - opam switch create MathComp 4.14.1
+        - opam repo add coq-released https://coq.inria.fr/opam/released
+        - eval $(opam env --switch=MathComp --set-switch)
+        - opam pin add -y coq 8.18.0
+        - opam pin add -y coq-lsp 0.1.8+8.18
+        - opam pin add -n -y -k path coq-mathcomp-ssreflect .
+        - opam install -y coq-mathcomp-ssreflect --deps-only
+        - make -j `nproc`
+      setup_cmds:
+        - opam env --switch=MathComp --set-switch
+  - GeoCoq:
+      url: https://github.com/GeoCoq/GeoCoq
+      commit: 2f038905da84b8e63f07e2bf5b1b1bf348e1e0a1
+      branch: master
+      build_cmds: # Verified on 2024-05-04
+        - opam switch create GeoCoq 4.14.1
+        - eval $(opam env --switch=GeoCoq --set-switch)
+        - opam repo add coq-released https://coq.inria.fr/opam/released
+        - opam pin add -y coq 8.18.0
+        - opam pin add -y coq-lsp 0.1.8+8.18
+        - opam pin -n . --yes
+        - opam install coq-geocoq-coinc --yes
+        - opam install coq-geocoq-axioms --yes
+        - opam install coq-geocoq-elements --yes
+        - opam install coq-geocoq-main --yes
+        - opam install coq-geocoq-algebraic --yes
+        - opam install ./coq-geocoq.opam --deps-only --yes
+        - ./configure.sh
+        - make -j `nproc`
+      setup_cmds:
+        - opam env --switch=GeoCoq --set-switch
+  - category-theory:
+      url: https://github.com/jwiegley/category-theory
+      commit: f8295f0d77ab0dd9f989e8e45d43670a69f424df
+      branch: master
+      build_cmds: # Verified on 2024-05-05
+        - pushd ..
+        - opam switch create category-theory 4.14.1
+        - eval $(opam env --switch=category-theory --set-switch)
+        - opam pin add -y coq 8.18.0
+        - opam pin add -y coq-lsp 0.1.8+8.18
+        - opam repo add coq-released https://coq.inria.fr/opam/released
+        - opam install ./category-theory/coq-category-theory.opam --deps-only --yes
+        - make -C category-theory -j `nproc`
+        - popd
+      setup_cmds:
+        - opam env --switch=category-theory --set-switch
+  - CompCert:
+      url: https://github.com/AbsInt/CompCert
+      commit: 76a4ff8f5b37429a614a2a97f628d9d862c93f46
+      build_cmds: # Verified on 2024-05-04
+        - pushd ..
+        - pwd
+        - opam switch create CompCert 4.07.1
+        - eval $(opam env --switch=CompCert --set-switch)
+        - eval `opam config env`
+        - opam pin -yn add coq 8.10.2
+        - opam pin -yn add menhir 20190626
+        - opam install -y coq-serapi coq menhir
+        - if [[ ! -f "CompCert/Makefile.config" ]]; then
+        -     (cd CompCert && ./configure x86_64-linux)
+        - fi
+        - make -C CompCert -j `nproc`
+        - popd
+    setup_cmds:
+      - opam env --switch=CompCert --set-switch
+
+
+
+# Will do these in future
+  # - Coq-HoTT:
+  #     url: https://github.com/HoTT/Coq-HoTT
+  #     commit: 72ced09bc29602791755a025838bb4d48805ae86
+  #     branch: master
+  #     build_cmds:
+  #       - opam switch create Coq-HoTT 4.09.0
+  #       - eval $(opam env --switch=Coq-HoTT --set-switch)
+  #       - opam pin add -y coq 8.10.2
+  #       - make -j `nproc`
+  #     setup_cmds:
+  #       - opam env --switch=Coq-HoTT --set-switch
+  # - verdi:
+  #     url: https://github.com/uwplse/verdi
+  #     commit: 8f7fe47ae25c9ba0faeb1e9f62417ba238ef15bf
+  #     branch: master
+  #     build_cmds:
+  #       - opam switch create verdi 4.09.0
+  #       - eval $(opam env --switch=verdi --set-switch)
+  #       - opam pin add -y coq 8.10.2
+  #       - make -j `nproc`
+  #     setup_cmds:
+  #       - opam env --switch=verdi --set-switch
+  # - fourcolor:
+  #     url: https://github.com/coq-community/fourcolor
+  #     commit: 91ff6b8b846c8ad683260a5e6ce400e186f43c6e
+  #     branch: master
+  #     build_cmds:
+  #       - opam switch create fourcolor 4.
+  #       - eval $(opam env --switch=fourcolor --set-switch)
+  #       - opam pin add -y coq 8.10.2
+  #       - make -j `nproc`
+  #     setup_cmds:
+  #       - opam env --switch=fourcolor --set-switch
+  # - coq-tricks:
+  #     url: https://github.com/tchajed/coq-tricks
+  #     commit: 8636f839bebd101286aaa9f41350de40f1f911c1
+  #     branch: master
+  #     build_cmds:
+  #       - opam switch create coq-tricks 4.09.0
+  #       - eval $(opam env --switch=coq-tricks --set-switch)
+  #       - opam pin add -y coq 8.10.2
+  #       - make -j `nproc`
+  #     setup_cmds:
+  #       - opam env --switch=coq-tricks --set-switch
+  # - VST:
+  #     url: https://github.com/PrincetonUniversity/VST
+  #     commit: 56e6886288918f1b78ad53dd3dcdb1a58a9916cb
+  #     branch: master
+  #     build_cmds:
+  #       - opam switch create VST 4.09.0
+  #       - eval $(opam env --switch=VST --set-switch)
+  #       - opam pin add -y coq 8.10.2
+  #       - make -j `nproc`
+  #     setup_cmds:
+  #       - opam env --switch=VST --set-switch
+  # - Set-Theory:
+  #     url: https://github.com/choukh/Set-Theory
+  #     commit: 5677d0d9cc3814adfb9bc1286a826f9d620fcc2e
+  #     branch: master
+  #     build_cmds:
+  #       - opam switch create Set-Theory 4
+  #       - eval $(opam env --switch=Set-Theory --set-switch)
+  #       - opam pin add -y coq 8.10.2
+  #       - make -j `nproc`
+  #     setup_cmds:
+  #       - opam env --switch=Set-Theory --set-switch
+  # - math-classes:
+  #     url: https://github.com/coq-community/math-classes
+  #     commit: 2a8e12360cceee510f39e3ef4d0a7472d70fa684
+  #     branch: master
+  #     build_cmds:
+  #       - opam switch create math-classes
+  #       - eval $(opam env --switch=math-classes --set-switch)
+  #       - opam pin add -y coq 8.10.2
+  #       - make -j `nproc`
+  #     setup_cmds:
+  #       - opam env --switch=math-classes --set-switch
+  # - coq-100-theorems:
+  #     url: https://github.com/coq-community/coq-100-theorems
+  #     commit: 36ba6d3abae801fb92bd74a26cd27ebf5de0c234
+  #     branch: master
+  #     build_cmds:
+  #       - opam switch create coq-100-theore
+  #       - eval $(opam env --switch=coq-100-the --set-switch)
+  #       - opam pin add -y coq 8.10.2
+  #       - make -j `nproc`
+  #     setup_cmds:
+  #       - opam env --switch=coq-100-the --set-switch
+  # - coq:
+  #     url: https://github.com/coq/coq
+  #     commit: 6fdccaf0346754dda384c6847f1f2aded2b88c41
+  #     build_cmds: # Verified on 2024-05-04
+  #       - opam switch create coq --packages="ocaml-variants.4.14.1+options,ocaml-option-flambda"
+  #       - eval $(opam env --switch=coq --set-switch)
+  #       - opam pin add -y coq 8.18.0
+  #       - opam install dune ocamlfind zarith lablgtk3-sourceview3 --yes
+  #       - ./configure
+  #       - make dunestrap
+  #       - dune build -p coq-core,coq-stdlib,coq,coqide-server
+  #     setup_cmds:
+  #       - opam env --switch=coq --set-switch

itp_interface/main/config/run_settings/default_coq_data_generation_transforms.yaml

@@ -0,0 +1,24 @@
+name: default_coq_data_generation_transforms
+use_human_readable: True
+save_intermidiate_transforms: True
+buffer_size: 10000
+pool_size: 20
+transform_type: LOCAL
+dep_depth: 0
+max_search_results:
+output_dir: .log/run_data_generation_transforms/data/train/coq
+setting_type: Agent
+timeout_in_secs: 60 # coq tactic execution timeout
+proof_retries: 1
+max_theorems_in_prompt: 7
+max_number_of_episodes: 1
+max_steps_per_episode: 60
+render: False
+checkpoint_dir: .log/checkpoints
+should_checkpoint: False
+random_seed: 42
+random_split: False
+train_eval_test_split:
+  - 1.0
+  - 0.0
+  - 0.0

itp_interface/main/config/run_settings/default_isabelle_data_generation_transforms.yaml

@@ -0,0 +1,24 @@
+name: default_isabelle_data_generation_transforms
+use_human_readable: True
+save_intermidiate_transforms: True
+buffer_size: 10000
+pool_size: 7
+transform_type: LOCAL
+dep_depth: 0
+max_search_results:
+output_dir: .log/run_data_generation_transforms/data/train/isabelle
+setting_type: Agent
+timeout_in_secs: 60 # tactic execution timeout
+proof_retries: 1
+max_theorems_in_prompt: 7
+max_number_of_episodes: 1
+max_steps_per_episode: 60
+render: False
+checkpoint_dir: .log/checkpoints
+should_checkpoint: False
+random_seed: 42
+random_split: False
+train_eval_test_split:
+  - 1.0
+  - 0.0
+  - 0.0

itp_interface/main/config/run_settings/default_lean4_data_generation_transforms.yaml

@@ -0,0 +1,24 @@
+name: default_lean4_data_generation_transforms
+use_human_readable: True
+save_intermidiate_transforms: True
+buffer_size: 10000
+pool_size: 8
+transform_type: LOCAL
+dep_depth: 0
+max_search_results:
+output_dir: .log/run_data_generation_transforms/data/train/lean4
+setting_type: Agent
+timeout_in_secs: 200 # coq tactic execution timeout
+proof_retries: 1
+max_theorems_in_prompt: 7
+max_number_of_episodes: 1
+max_steps_per_episode: 60
+render: False
+checkpoint_dir: .log/checkpoints
+should_checkpoint: False
+random_seed: 42
+random_split: False
+train_eval_test_split:
+  - 1.0
+  - 0.0
+  - 0.0

itp_interface/main/config/run_settings/default_lean_data_generation_transforms.yaml

@@ -0,0 +1,24 @@
+name: default_lean_data_generation_transforms
+use_human_readable: True
+save_intermidiate_transforms: True
+buffer_size: 10000
+pool_size: 20
+transform_type: LOCAL
+dep_depth: 0
+max_search_results:
+output_dir: .log/run_data_generation_transforms/data/train/lean
+setting_type: Agent
+timeout_in_secs: 60 # coq tactic execution timeout
+proof_retries: 1
+max_theorems_in_prompt: 7
+max_number_of_episodes: 1
+max_steps_per_episode: 60
+render: False
+checkpoint_dir: .log/checkpoints
+should_checkpoint: False
+random_seed: 42
+random_split: False
+train_eval_test_split:
+  - 1.0
+  - 0.0
+  - 0.0

itp_interface/main/config/simple_coq_data_gen.yaml

@@ -0,0 +1,12 @@
+defaults:
+  # - benchmark: simple_benchmark_lean_training_data
+  # - run_settings: default_lean_data_generation_transforms
+  # - benchmark: simple_benchmark_1
+  # - run_settings: default_lean4_data_generation_transforms
+  - benchmark: simple_benchmark_1
+  - run_settings: default_coq_data_generation_transforms
+  - env_settings: no_retrieval
+  - override hydra/job_logging: 'disabled'
+
+run_settings:
+ output_dir: .log/data_generation/benchmark/simple_benchmark_1

itp_interface/main/config/simple_coq_data_gen_random.yaml

@@ -0,0 +1,17 @@
+defaults:
+  # - benchmark: simple_benchmark_lean_training_data
+  # - run_settings: default_lean_data_generation_transforms
+  # - benchmark: simple_benchmark_1
+  # - run_settings: default_lean4_data_generation_transforms
+  - benchmark: simple_benchmark_1
+  - run_settings: default_coq_data_generation_transforms
+  - env_settings: no_retrieval
+  - override hydra/job_logging: 'disabled'
+
+run_settings:
+  output_dir: .log/data_generation/benchmark/simple_benchmark_1
+  random_split: True
+  train_eval_test_split:
+    - 0.6
+    - 0.2
+    - 0.2

itp_interface/main/config/simple_lean_data_gen.yaml

@@ -0,0 +1,12 @@
+defaults:
+  # - benchmark: simple_benchmark_lean_training_data
+  # - run_settings: default_lean_data_generation_transforms
+  # - benchmark: simple_benchmark_1
+  # - run_settings: default_lean4_data_generation_transforms
+  - benchmark: simple_benchmark_lean
+  - run_settings: default_lean4_data_generation_transforms
+  - env_settings: no_retrieval
+  - override hydra/job_logging: 'disabled'
+
+run_settings:
+ output_dir: .log/data_generation/benchmark/simple_benchmark_lean

itp_interface/main/config/simple_rl_lean_data_gen.yaml

@@ -0,0 +1,12 @@
+defaults:
+  # - benchmark: simple_benchmark_lean_training_data
+  # - run_settings: default_lean_data_generation_transforms
+  # - benchmark: simple_benchmark_1
+  # - run_settings: default_lean4_data_generation_transforms
+  - benchmark: simple_rl_benchmark_lean
+  - run_settings: default_lean4_data_generation_transforms
+  - env_settings: no_retrieval
+  - override hydra/job_logging: 'disabled'
+
+run_settings:
+ output_dir: .log/data_generation/benchmark/simple_rl_benchmark_lean

itp_interface/main/config/uni_math_data_gen.yaml

@@ -0,0 +1,14 @@
+defaults:
+  # - benchmark: simple_benchmark_lean_training_data
+  # - run_settings: default_lean_data_generation_transforms
+  - benchmark: UniMath
+  - run_settings: default_coq_data_generation_transforms
+  - env_settings: no_retrieval
+  - override hydra/job_logging: 'disabled'
+
+run_settings:
+ output_dir: <root>/data/proofsteps/UniMath
+ train_eval_test_split:
+   - 0.925
+   - 0.0375
+   - 0.0375

itp_interface/main/config.py

@@ -0,0 +1,192 @@
+#!/usr/bin/env python3
+
+import sys
+
+root_dir = f"{__file__.split('itp_interface')[0]}"
+if root_dir not in sys.path:
+    sys.path.append(root_dir)
+import typing
+from dataclasses import dataclass, field
+from dataclasses_json import dataclass_json
+from enum import Enum
+from itp_interface.rl.proof_tree import ProofSearchResult
+from itp_interface.rl.proof_action import ProofAction
+from itp_interface.rl.simple_proof_env import ProofEnvReRankStrategy
+
+class SettingType(Enum):
+    Agent = "Agent"
+    GptF = "GptF"
+
+    def __str__(self):
+        return self.value
+
+class PolicyName(Enum):
+    # WARN: Don't make enums dataclasses because deserialization has some weird bug which matches the deserialized enum to all the enum values
+    Dfs = "Dfs"
+    FewShot = "FewShot"
+
+    def __str__(self):
+        return self.value
+    
+class TransformType(Enum):
+    LOCAL = "LOCAL"
+    GLOBAL = "GLOBAL"
+    FULL = "FULL"
+
+    def __str__(self):
+        return self.value
+
+@dataclass_json
+@dataclass
+class EnvSettings(object):
+    name: str
+    retrieval_strategy: ProofEnvReRankStrategy
+
+@dataclass_json
+@dataclass
+class RunSettings(object):
+    name: str
+    use_human_readable: bool
+    save_intermidiate_transforms: bool
+    buffer_size: int
+    pool_size: int
+    transform_type: TransformType
+    dep_depth: int
+    output_dir: str
+    setting_type: SettingType
+    timeout_in_secs: int # coq tactic execution timeout
+    proof_retries: int
+    max_theorems_in_prompt: int
+    max_number_of_episodes: int
+    max_steps_per_episode: int
+    render: bool
+    checkpoint_dir: str
+    should_checkpoint: bool
+    max_search_results: typing.Optional[int] = None
+    random_seed: int = 42
+    random_split: bool = False
+    train_eval_test_split: typing.List[float] = field(default_factory=lambda: [1.0, 0.0, 0.0])
+
+@dataclass_json
+@dataclass
+class EvalFile(object):
+    path: str
+    theorems: typing.Union[str, typing.List[str]]
+
+@dataclass_json
+@dataclass
+class EvalDataset(object):
+    project: str
+    files: typing.List[EvalFile]
+
+@dataclass_json
+@dataclass
+class EvalBenchmark(object):
+    name: str
+    num_files: int
+    language: ProofAction.Language
+    datasets: typing.List[EvalDataset]
+    few_shot_data_path_for_retrieval: str = None
+    few_shot_metadata_filename_for_retrieval: str = None
+    dfs_data_path_for_retrieval: str = None
+    dfs_metadata_filename_for_retrieval: str = None
+    setup_cmds: typing.List[str] = field(default_factory=list)
+
+@dataclass_json
+@dataclass
+class Experiments(object):
+    env_settings: EnvSettings
+    run_settings: RunSettings
+    benchmark: EvalBenchmark
+
+@dataclass_json
+@dataclass
+class EvalRunCheckpointInfo(object):
+    checkpoint_file: str
+    logging_dirs: typing.List[str]
+    proof_dump_dir: str
+    theorem_maps: typing.Dict[str, typing.Dict[str, bool]]
+
+    def add_path_to_maps(self, path: str):
+        if path not in self.theorem_maps:
+            self.theorem_maps[path] = {}
+
+    def add_theorem_to_maps(self, path: str, theorem: str, success: bool):
+        self.theorem_maps[path][theorem] = success
+        with open(self.checkpoint_file, "w") as f:
+            f.write(self.to_json(indent=4))
+    
+@dataclass_json
+@dataclass
+class EvalProofResults(object):
+    path: str
+    theorem_map: typing.Dict[str, typing.Dict[str, ProofSearchResult]]
+
+    def add_path_to_maps(self, path: str):
+        if path not in self.theorem_map:
+            self.theorem_map[path] = {}
+    
+    def add_theorem_to_maps(self, path: str, theorem: str, proof_result: ProofSearchResult):
+        self.theorem_map[path][theorem] = proof_result
+        with open(self.path, "w") as f:
+            f.write(self.to_json(indent=4))
+
+
+def parse_config(cfg):
+    env_settings_cfg = cfg["env_settings"]
+    env_settings = EnvSettings(
+        name=env_settings_cfg["name"],
+        retrieval_strategy=ProofEnvReRankStrategy(env_settings_cfg["retrieval_strategy"]))
+    run_settings_cfg = cfg["run_settings"]
+    eval_settings = RunSettings(
+        name=run_settings_cfg["name"],
+        use_human_readable=run_settings_cfg["use_human_readable"],
+        save_intermidiate_transforms=run_settings_cfg["save_intermidiate_transforms"],
+        buffer_size=run_settings_cfg["buffer_size"],
+        pool_size=run_settings_cfg["pool_size"],
+        transform_type=TransformType(run_settings_cfg["transform_type"]),
+        dep_depth=run_settings_cfg["dep_depth"],
+        output_dir=run_settings_cfg["output_dir"],
+        max_search_results=run_settings_cfg["max_search_results"],
+        setting_type=SettingType(run_settings_cfg["setting_type"]),
+        timeout_in_secs=run_settings_cfg["timeout_in_secs"],
+        proof_retries=run_settings_cfg["proof_retries"],
+        max_theorems_in_prompt=run_settings_cfg["max_theorems_in_prompt"],
+        max_number_of_episodes=run_settings_cfg["max_number_of_episodes"],
+        max_steps_per_episode=run_settings_cfg["max_steps_per_episode"],
+        render=run_settings_cfg["render"],
+        checkpoint_dir=run_settings_cfg["checkpoint_dir"],
+        should_checkpoint=run_settings_cfg["should_checkpoint"],
+        random_seed=run_settings_cfg["random_seed"],
+        random_split=run_settings_cfg["random_split"],
+        train_eval_test_split=run_settings_cfg["train_eval_test_split"])
+    benchmark_cfg = cfg["benchmark"]
+    datasets_cfg = benchmark_cfg["datasets"]
+    eval_datasets = []
+    for dataset_cfg in datasets_cfg:
+        files_cfg = list(dataset_cfg["files"])
+        eval_files = []
+        for file_cfg in files_cfg:
+            theorems = None
+            if type(file_cfg["theorems"]) == str:
+                theorems = file_cfg["theorems"]
+            else:
+                theorems = list(file_cfg["theorems"])
+            eval_files.append(EvalFile(
+                path=file_cfg["path"],
+                theorems=theorems))
+        eval_datasets.append(EvalDataset(
+            project=dataset_cfg["project"],
+            files=eval_files))
+    language = ProofAction.Language(benchmark_cfg["language"])
+    benchmark = EvalBenchmark(
+        name=benchmark_cfg["name"],
+        num_files=benchmark_cfg["num_files"],
+        language=language,
+        datasets=eval_datasets,
+        few_shot_data_path_for_retrieval=benchmark_cfg["few_shot_data_path_for_retrieval"],
+        few_shot_metadata_filename_for_retrieval=benchmark_cfg["few_shot_metadata_filename_for_retrieval"],
+        dfs_data_path_for_retrieval=benchmark_cfg["dfs_data_path_for_retrieval"],
+        dfs_metadata_filename_for_retrieval=benchmark_cfg["dfs_metadata_filename_for_retrieval"],
+        setup_cmds=benchmark_cfg["setup_cmds"] if "setup_cmds" in benchmark_cfg else [])
+    return Experiments(env_settings=env_settings, run_settings=eval_settings, benchmark=benchmark)

itp_interface/main/extract_benchmark_dataset.py

@@ -0,0 +1,106 @@
+#!/usr/bin/env python3
+
+import sys
+root_dir = f"{__file__.split('itp_interface')[0]}"
+if root_dir not in sys.path:
+    sys.path.append(root_dir)
+import argparse
+import logging
+import time
+import os
+import typing
+import copy
+import yaml
+from itp_interface.tools.log_utils import setup_logger
+from itp_interface.tools.training_data import TrainingData
+
+
+def create_yaml(project_to_theorems, name, language, output_file):
+    data = {
+        "name": name,
+        "num_files": 0,
+        "language": str(language),
+        "few_shot_data_path_for_retrieval": None,
+        "few_shot_metadata_filename_for_retrieval": None,
+        "dfs_data_path_for_retrieval": None,
+        "dfs_metadata_filename_for_retrieval": "local.meta.json",
+        "theorem_cnt": 0,
+        "datasets": []
+    }
+    for project_root, file_dict in project_to_theorems.items():
+        dataset = {"project": project_root, "files": []}
+        for file_path, theorems in file_dict.items():
+            data["num_files"] += 1
+            data["theorem_cnt"] += len(theorems)
+            dataset["files"].append({"path": file_path, "theorems": theorems}) 
+        data["datasets"].append(dataset)
+
+    with open(output_file, 'w') as yaml_file:
+        yaml.dump(data, yaml_file, sort_keys=False)
+
+def extract_project_to_theorems(training_data : TrainingData, project_to_theorems: typing.Dict[str, typing.Dict[str, typing.List[str]]]):
+    for tdf in training_data:
+        project_path = tdf.project_id
+        file_path = tdf.file_path
+        # Check if both project and file_path are the same absolute path
+        both_abs_path = os.path.isabs(project_path) and os.path.isabs(file_path)
+        if both_abs_path:
+            file_path = os.path.relpath(file_path, project_path)
+        if project_path not in project_to_theorems:
+            project_to_theorems[project_path] = {}
+        if file_path not in project_to_theorems[project_path]:
+            project_to_theorems[project_path][file_path] = []
+        if tdf.theorem_name not in project_to_theorems[project_path][file_path]:
+            project_to_theorems[project_path][file_path].append(tdf.theorem_name)
+
+def extract_benchmarks_from_datasets(datasets, metafilenames, name, language, output, max_parallelism=8, logger=None):
+    """
+    extract datasets
+    """
+    assert len(datasets) == len(metafilenames), "Length of datasets and metafilenames must be the same"
+    assert max_parallelism > 0, "Max parallelism must be greater than 0"
+    logger = logger or logging.getLogger("dataset_merge")
+    tds : typing.List[TrainingData] = []
+    for dataset, metafilename in zip(datasets, metafilenames):
+        training_data = TrainingData(
+            folder=dataset,
+            training_meta_filename=metafilename,
+            max_parallelism=max_parallelism
+        )
+        training_data.load_meta()
+        tds.append(training_data)
+        logger.info(f"Inited training data for {dataset}")
+    project_to_theorems_map = {}
+    for td in tds:
+        logger.info(f"Start loading {td.folder} ...")
+        td.load()
+        logger.info(f"Loaded {td.folder}.")
+        logger.info(f"Start extraction from {td.folder} ...")
+        extract_project_to_theorems(td, project_to_theorems_map)
+        logger.info(f"Finished extraction from {td.folder}.")
+    logger.info(f"Saving the extracted theorems to {output} ...")
+    create_yaml(project_to_theorems_map, name, language, output)
+    logger.info(f"Saved the extracted theorems to {output}.")
+    pass
+
+
+if __name__ == "__main__":
+    args = argparse.ArgumentParser()
+    args.add_argument("--datasets", type=str, nargs="+", help="List of datasets to merge")
+    args.add_argument("--output", type=str, help="Output file")
+    args.add_argument("--metafilenames", type=str, nargs="+", help="List of metafilenames", default=None)
+    args.add_argument("--name", type=str, help="Name of the merged dataset")
+    args.add_argument("--language", type=str, help="Language of the merged dataset")
+    args.add_argument("--max_parallelism", type=int, help="Max parallelism", default=8)
+    args = args.parse_args()
+    # Add root dir to python path
+    os.environ["PYTHONPATH"] = f"{root_dir}:{os.environ.get('PYTHONPATH', '')}"
+    time_str = time.strftime("%Y-%m-%d_%H-%M-%S")
+    log_folder = os.path.join(".log", "extract_benchmark", time_str)
+    os.makedirs(log_folder, exist_ok=True)
+    logger = setup_logger("extract_benchmark", os.path.join(log_folder, "extract.log"))
+    if args.metafilenames is None:
+        args.metafilenames = ['local.meta.json' for _ in range(len(args.datasets))]
+    extract_benchmarks_from_datasets(
+        args.datasets, args.metafilenames, args.name, args.language, args.output, args.max_parallelism, logger
+    )