Upload folder using huggingface_hub

all_code.py

@@ -0,0 +1,66 @@
+import tqdm
+import  json
+def filtered():
+    import json
+    with open("data/lean4_basic/1k_test.jsonl", "r") as f:
+        test_data = json.load(f)
+        test_data = [item['statement_poof'] for item in test_data]
+
+    # Function to filter items based on existence in test data
+
+
+    with open("data/lean4_random/5k_second.json", "r") as f:
+        second_5k = json.load(f)
+
+    def filter_items(data, test_data):
+        filtered_data = [item for item in tqdm.tqdm(data) if item['statement_poof'][:-2] not in test_data]
+        return filtered_data
+
+    # Filter and save filtered data
+
+    filtered_second_5k = filter_items(second_5k, test_data)
+    with open("data/lean4_random/5k_second_filtered.json", "w") as f:
+        json.dump(filtered_second_5k, f, ensure_ascii=False, indent=2)
+    print("Filtered second file length:", len(filtered_second_5k))
+
+
+
+
+def insert_label_for_autoformalization():
+    input_lists = ["data/lean4_statement_translate/15k_state_problem_translation.json"]
+    for input_file in input_lists:
+        with open(input_file, "r") as f:
+            test_data = json.load(f)
+            for item in test_data:
+                item['task']='statement_form'
+
+            with open("data/lean4_statement_translate/15k_state_problem_translation_statement_form.json", "w") as f:
+                json.dump(test_data, f, ensure_ascii=False, indent=2)
+
+
+    input_lists = ["data/lean4_statement_translate/15k_state_problem_translation.json"]
+    for input_file in input_lists:
+        with open(input_file, "r") as f:
+            test_data = json.load(f)
+            for item in test_data:
+                item['task']='statementproof_inform'
+            with open("data/lean4_statement_translate/15k_state_problem_translation_statementproof_inform.json", "w") as f:
+                json.dump(test_data, f, ensure_ascii=False, indent=2)
+
+
+    input_lists = ["all_theorem.jsonl"]
+    for input_file in input_lists:
+        with open(input_file, "r") as f:
+            test_data = json.load(f)
+            for item in test_data:
+                item['task']='solver'
+
+            with open("data/all_theorem_solver.jsonl", "w") as f:
+                json.dump(test_data, f, ensure_ascii=False, indent=2)
+
+
+
+
+if __name__ == '__main__':
+    insert_label_for_autoformalization()
+    # filtered()

autobackform_train.py

@@ -0,0 +1,305 @@
+#    Copyright 2023 Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li
+#
+#    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.
+#    Modified by Zheng Yuan and Hongyi Yuan
+
+import os
+import copy
+import logging
+from dataclasses import dataclass, field
+from typing import Optional, Dict, Sequence
+import io
+import torch
+import transformers
+from torch.utils.data import Dataset
+from transformers import Trainer
+import argparse
+import json
+import random;random.seed(42)
+
+def _make_r_io_base(f, mode: str):
+    if not isinstance(f, io.IOBase):
+        f = open(f, mode=mode)
+    return f
+
+def jload(f, mode="r"):
+    """Load a .json file into a dictionary."""
+    f = _make_r_io_base(f, mode)
+    jdict = json.load(f)
+    f.close()
+    return jdict
+
+IGNORE_INDEX = -100
+DEFAULT_PAD_TOKEN = "[PAD]"
+DEFAULT_EOS_TOKEN = "</s>"
+DEFAULT_BOS_TOKEN = "<s>"
+DEFAULT_UNK_TOKEN = "<unk>"
+PROMPT_DICT = {
+    "lean4": (
+        "Statement and proof in natural language:\n\n"
+        "{statement_text}\n\n"
+        "Translate the statement and proof in natural language to lean4:"
+    ),
+    "backform": (
+        "Statement and proof in lean4:\n\n"
+        "{statement_text}\n\n"
+        "Translate the statement and proof in lean4 to natural language:"
+    ),
+    "prompt_no_input": (
+    "Below is an instruction that describes a task. "
+    "Write a response that appropriately completes the request.\n\n"
+    "### Instruction:\n{instruction}\n\n### Response:"
+    ),
+}
+#### 28
+@dataclass
+class ModelArguments:
+    model_name_or_path: Optional[str] = field(default="facebook/opt-125m")
+
+
+@dataclass
+class DataArguments:
+    data_path: str = field(default=None, metadata={"help": "Path to the training data."})
+
+
+@dataclass
+class TrainingArguments(transformers.TrainingArguments):
+    cache_dir: Optional[str] = field(default=None)
+    optim: str = field(default="adamw_torch")
+    model_max_length: int = field(
+        default=2048,
+        metadata={"help": "Maximum sequence length. Sequences will be right padded (and possibly truncated)."},
+    )
+    overwrite_output_dir: bool = field(default=True)
+
+
+def safe_save_model_for_hf_trainer(trainer: transformers.Trainer, output_dir: str):
+    """Collects the state dict and dump to disk."""
+    state_dict = trainer.model.state_dict()
+    if trainer.args.should_save:
+        cpu_state_dict = {key: value.cpu() for key, value in state_dict.items()}
+        del state_dict
+        trainer._save(output_dir, state_dict=cpu_state_dict)  # noqa
+
+
+def smart_tokenizer_and_embedding_resize(
+    special_tokens_dict: Dict,
+    tokenizer: transformers.PreTrainedTokenizer,
+    model: transformers.PreTrainedModel,
+):
+    """Resize tokenizer and embedding.
+    Note: This is the unoptimized version that may make your embedding size not be divisible by 64.
+    """
+    num_new_tokens = tokenizer.add_special_tokens(special_tokens_dict)
+    model.resize_token_embeddings(len(tokenizer))
+
+    if num_new_tokens > 0:
+        input_embeddings = model.get_input_embeddings().weight.data
+        output_embeddings = model.get_output_embeddings().weight.data
+
+        input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(dim=0, keepdim=True)
+        output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(dim=0, keepdim=True)
+
+        input_embeddings[-num_new_tokens:] = input_embeddings_avg
+        output_embeddings[-num_new_tokens:] = output_embeddings_avg
+
+
+def _tokenize_fn(strings: Sequence[str], tokenizer: transformers.PreTrainedTokenizer) -> Dict:
+    """Tokenize a list of strings."""
+    tokenized_list = [
+        tokenizer(
+            text,
+            return_tensors="pt",
+            padding="longest",
+            max_length=tokenizer.model_max_length,
+            truncation=True,
+        )
+        for text in strings
+    ]
+    input_ids = labels = [tokenized.input_ids[0] for tokenized in tokenized_list]
+    input_ids_lens = labels_lens = [
+        tokenized.input_ids.ne(tokenizer.pad_token_id).sum().item() for tokenized in tokenized_list
+    ]
+    return dict(
+        input_ids=input_ids,
+        labels=labels,
+        input_ids_lens=input_ids_lens,
+        labels_lens=labels_lens,
+    )
+
+
+def preprocess(
+    sources: Sequence[str],
+    targets: Sequence[str],
+    tokenizer: transformers.PreTrainedTokenizer,
+) -> Dict:
+    """Preprocess the data by tokenizing."""
+    examples = [s + t for s, t in zip(sources, targets)]
+    examples_tokenized, sources_tokenized = [_tokenize_fn(strings, tokenizer) for strings in (examples, sources)]
+    input_ids = examples_tokenized["input_ids"]
+    labels = copy.deepcopy(input_ids)
+    for label, source_len in zip(labels, sources_tokenized["input_ids_lens"]):
+        label[:source_len] = IGNORE_INDEX
+    return dict(input_ids=input_ids, labels=labels)
+
+class SupervisedDataset(Dataset):
+    """Dataset for supervised fine-tuning."""
+    def __init__(self, data_args, tokenizer: transformers.PreTrainedTokenizer):
+        super(SupervisedDataset, self).__init__()
+        logging.warning("Loading data...")
+        data_path = data_args.data_path
+        try:
+            data_path = data_path_map[data_path]
+        except:
+            data_path = data_path
+        list_data_dict = []
+        for item in data_path.split(','):
+            try:
+                list_data_dict += jload(item)
+
+            except BaseException:
+                with open(item, 'r') as f:
+                    lines = f.readlines()
+                list_data_dict += [json.loads(line.strip()) for line in lines]
+
+        list_data_dict = random.sample(list_data_dict,  len(list_data_dict))
+        list_data_dict = list_data_dict[:data_args.data_length]
+
+        logging.warning("Formatting inputs...")
+        prompt_lean4  = PROMPT_DICT["backform"]
+
+        # list_data_dict = [{'instruction':data['items'][0]['value'], 'input':'', 'output':data['items'][1]['value']} for data in list_data_dict]
+
+        list_data_dict = [{'instruction':prompt_lean4.format(statement_text = data['statement_poof']), 'input':'', 'output':data['model_response']} for data in list_data_dict]
+        print(f"len of {len(list_data_dict)}")
+        sources = [example['instruction'] for example in list_data_dict]
+        targets = [f"{example['output']}{tokenizer.eos_token}" for example in list_data_dict]
+        # targets = [example['output'] for example in list_data_dict]
+
+        self.sources = sources
+        self.targets = targets
+
+    def __len__(self):
+        return len(self.sources)
+
+    def naive__getitem__(self, i) -> Dict[str, torch.Tensor]:
+        return dict(input_ids=self.input_ids[i], labels=self.labels[i])
+
+    def __getitem__(self, i):
+        return dict(input_ids=self.sources[i], labels=self.targets[i])
+
+@dataclass
+class DataCollatorForSupervisedDataset(object):
+    """Collate examples for supervised fine-tuning."""
+
+    tokenizer: transformers.PreTrainedTokenizer
+
+    def naive__call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
+        input_ids, labels = tuple([instance[key] for instance in instances] for key in ("input_ids", "labels"))
+        input_ids = torch.nn.utils.rnn.pad_sequence(
+            input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id
+        )
+        labels = torch.nn.utils.rnn.pad_sequence(labels, batch_first=True, padding_value=IGNORE_INDEX)
+        return dict(
+            input_ids=input_ids,
+            labels=labels,
+            attention_mask=input_ids.ne(self.tokenizer.pad_token_id),
+        )
+
+    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
+        sources = []
+        targets = []
+        for instance in instances:
+            source = instance['input_ids']
+            target = instance['labels']
+            sources.append(source)
+            targets.append(target)
+
+        data_dict = preprocess(sources, targets, self.tokenizer)
+        input_ids, labels = data_dict['input_ids'], data_dict['labels']
+        # input_ids, labels = tuple([instance[key] for instance in instances] for key in ("input_ids", "labels"))
+        input_ids = torch.nn.utils.rnn.pad_sequence(
+            input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id
+        )
+        labels = torch.nn.utils.rnn.pad_sequence(labels, batch_first=True, padding_value=IGNORE_INDEX)
+        return dict(
+            input_ids=input_ids,
+            labels=labels,
+            attention_mask=input_ids.ne(self.tokenizer.pad_token_id),
+        )
+
+def make_supervised_data_module(tokenizer: transformers.PreTrainedTokenizer, data_args) -> Dict:
+    """Make dataset and collator for supervised fine-tuning."""
+    train_dataset = SupervisedDataset(tokenizer=tokenizer, data_args=data_args)
+    data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)
+    return dict(train_dataset=train_dataset, eval_dataset=None, data_collator=data_collator)
+
+
+os.environ["WANDB_PROJECT"] = "auto_backform"
+
+def train():
+
+    parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
+    model_args, data_args, training_args, remaining_args = parser.parse_args_into_dataclasses(return_remaining_strings=True)
+    data_args.data_length = int(remaining_args[1])
+
+    model = transformers.AutoModelForCausalLM.from_pretrained(
+        model_args.model_name_or_path,
+        cache_dir=training_args.cache_dir,
+        trust_remote_code=True,
+        torch_dtype=torch.bfloat16,
+        attn_implementation="flash_attention_2",
+    )
+
+    model.config.use_cache = False
+    model.gradient_checkpointing_enable()
+
+    tokenizer = transformers.AutoTokenizer.from_pretrained(
+        model_args.model_name_or_path,
+        cache_dir=training_args.cache_dir,
+        model_max_length=training_args.model_max_length,
+        padding_side="right",
+        use_fast=False,
+    )
+    if tokenizer.pad_token is None:
+        smart_tokenizer_and_embedding_resize(
+            special_tokens_dict=dict(pad_token=DEFAULT_PAD_TOKEN),
+            tokenizer=tokenizer,
+            model=model,
+        )
+    if "llama" in model_args.model_name_or_path:
+        tokenizer.add_special_tokens(
+            {
+                "eos_token": DEFAULT_EOS_TOKEN,
+                "bos_token": DEFAULT_BOS_TOKEN,
+                "unk_token": DEFAULT_UNK_TOKEN,
+            }
+        )
+        try:
+            tokenizer.pad_token = tokenizer.unk_token
+        except:
+            pass
+
+    data_module = make_supervised_data_module(tokenizer=tokenizer, data_args=data_args)
+    trainer = Trainer(model=model, tokenizer=tokenizer, args=training_args, **data_module)
+    trainer.train()
+    model.config.use_cache = True
+    # trainer.save_state()
+    # if os.environ.get('LOCAL_RANK') == '0':
+    safe_save_model_for_hf_trainer(trainer=trainer, output_dir=training_args.output_dir)
+
+
+
+
+if __name__ == "__main__":
+    train()

autoform_train.py

@@ -0,0 +1,300 @@
+#    Copyright 2023 Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li
+#
+#    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.
+#    Modified by Zheng Yuan and Hongyi Yuan
+
+import os
+import copy
+import logging
+from dataclasses import dataclass, field
+from typing import Optional, Dict, Sequence
+import io
+import torch
+import transformers
+from torch.utils.data import Dataset
+from transformers import Trainer
+import argparse
+import json
+import random;random.seed(42)
+
+def _make_r_io_base(f, mode: str):
+    if not isinstance(f, io.IOBase):
+        f = open(f, mode=mode)
+    return f
+
+def jload(f, mode="r"):
+    """Load a .json file into a dictionary."""
+    f = _make_r_io_base(f, mode)
+    jdict = json.load(f)
+    f.close()
+    return jdict
+
+IGNORE_INDEX = -100
+DEFAULT_PAD_TOKEN = "[PAD]"
+DEFAULT_EOS_TOKEN = "</s>"
+DEFAULT_BOS_TOKEN = "<s>"
+DEFAULT_UNK_TOKEN = "<unk>"
+PROMPT_DICT = {
+    "lean4": (
+        "Statement and proof in natural language:\n\n"
+        "{statement_text}\n\n"
+        "Translate the statement and proof in natural language to lean4:"
+    ),
+    "prompt_no_input": (
+    "Below is an instruction that describes a task. "
+    "Write a response that appropriately completes the request.\n\n"
+    "### Instruction:\n{instruction}\n\n### Response:"
+    ),
+}
+#### 28
+@dataclass
+class ModelArguments:
+    model_name_or_path: Optional[str] = field(default="facebook/opt-125m")
+
+
+@dataclass
+class DataArguments:
+    data_path: str = field(default=None, metadata={"help": "Path to the training data."})
+
+
+@dataclass
+class TrainingArguments(transformers.TrainingArguments):
+    cache_dir: Optional[str] = field(default=None)
+    optim: str = field(default="adamw_torch")
+    model_max_length: int = field(
+        default=2048,
+        metadata={"help": "Maximum sequence length. Sequences will be right padded (and possibly truncated)."},
+    )
+    overwrite_output_dir: bool = field(default=True)
+
+
+def safe_save_model_for_hf_trainer(trainer: transformers.Trainer, output_dir: str):
+    """Collects the state dict and dump to disk."""
+    state_dict = trainer.model.state_dict()
+    if trainer.args.should_save:
+        cpu_state_dict = {key: value.cpu() for key, value in state_dict.items()}
+        del state_dict
+        trainer._save(output_dir, state_dict=cpu_state_dict)  # noqa
+
+
+def smart_tokenizer_and_embedding_resize(
+    special_tokens_dict: Dict,
+    tokenizer: transformers.PreTrainedTokenizer,
+    model: transformers.PreTrainedModel,
+):
+    """Resize tokenizer and embedding.
+    Note: This is the unoptimized version that may make your embedding size not be divisible by 64.
+    """
+    num_new_tokens = tokenizer.add_special_tokens(special_tokens_dict)
+    model.resize_token_embeddings(len(tokenizer))
+
+    if num_new_tokens > 0:
+        input_embeddings = model.get_input_embeddings().weight.data
+        output_embeddings = model.get_output_embeddings().weight.data
+
+        input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(dim=0, keepdim=True)
+        output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(dim=0, keepdim=True)
+
+        input_embeddings[-num_new_tokens:] = input_embeddings_avg
+        output_embeddings[-num_new_tokens:] = output_embeddings_avg
+
+
+def _tokenize_fn(strings: Sequence[str], tokenizer: transformers.PreTrainedTokenizer) -> Dict:
+    """Tokenize a list of strings."""
+    tokenized_list = [
+        tokenizer(
+            text,
+            return_tensors="pt",
+            padding="longest",
+            max_length=tokenizer.model_max_length,
+            truncation=True,
+        )
+        for text in strings
+    ]
+    input_ids = labels = [tokenized.input_ids[0] for tokenized in tokenized_list]
+    input_ids_lens = labels_lens = [
+        tokenized.input_ids.ne(tokenizer.pad_token_id).sum().item() for tokenized in tokenized_list
+    ]
+    return dict(
+        input_ids=input_ids,
+        labels=labels,
+        input_ids_lens=input_ids_lens,
+        labels_lens=labels_lens,
+    )
+
+
+def preprocess(
+    sources: Sequence[str],
+    targets: Sequence[str],
+    tokenizer: transformers.PreTrainedTokenizer,
+) -> Dict:
+    """Preprocess the data by tokenizing."""
+    examples = [s + t for s, t in zip(sources, targets)]
+    examples_tokenized, sources_tokenized = [_tokenize_fn(strings, tokenizer) for strings in (examples, sources)]
+    input_ids = examples_tokenized["input_ids"]
+    labels = copy.deepcopy(input_ids)
+    for label, source_len in zip(labels, sources_tokenized["input_ids_lens"]):
+        label[:source_len] = IGNORE_INDEX
+    return dict(input_ids=input_ids, labels=labels)
+
+class SupervisedDataset(Dataset):
+    """Dataset for supervised fine-tuning."""
+    def __init__(self, data_args, tokenizer: transformers.PreTrainedTokenizer):
+        super(SupervisedDataset, self).__init__()
+        logging.warning("Loading data...")
+        data_path = data_args.data_path
+        try:
+            data_path = data_path_map[data_path]
+        except:
+            data_path = data_path
+        list_data_dict = []
+        for item in data_path.split(','):
+            try:
+                list_data_dict += jload(item)
+
+            except BaseException:
+                with open(item, 'r') as f:
+                    lines = f.readlines()
+                list_data_dict += [json.loads(line.strip()) for line in lines]
+
+        list_data_dict = random.sample(list_data_dict,  len(list_data_dict))
+        list_data_dict = list_data_dict[:data_args.data_length]
+
+        logging.warning("Formatting inputs...")
+        prompt_lean4  = PROMPT_DICT["lean4"]
+
+        # list_data_dict = [{'instruction':data['items'][0]['value'], 'input':'', 'output':data['items'][1]['value']} for data in list_data_dict]
+
+        list_data_dict = [{'instruction':prompt_lean4.format(statement_text = data['model_response']), 'input':'', 'output':data['statement_poof']} for data in list_data_dict]
+        print(f"len of {len(list_data_dict)}")
+        sources = [example['instruction'] for example in list_data_dict]
+        targets = [f"{example['output']}{tokenizer.eos_token}" for example in list_data_dict]
+        # targets = [example['output'] for example in list_data_dict]
+
+        self.sources = sources
+        self.targets = targets
+
+    def __len__(self):
+        return len(self.sources)
+
+    def naive__getitem__(self, i) -> Dict[str, torch.Tensor]:
+        return dict(input_ids=self.input_ids[i], labels=self.labels[i])
+
+    def __getitem__(self, i):
+        return dict(input_ids=self.sources[i], labels=self.targets[i])
+
+@dataclass
+class DataCollatorForSupervisedDataset(object):
+    """Collate examples for supervised fine-tuning."""
+
+    tokenizer: transformers.PreTrainedTokenizer
+
+    def naive__call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
+        input_ids, labels = tuple([instance[key] for instance in instances] for key in ("input_ids", "labels"))
+        input_ids = torch.nn.utils.rnn.pad_sequence(
+            input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id
+        )
+        labels = torch.nn.utils.rnn.pad_sequence(labels, batch_first=True, padding_value=IGNORE_INDEX)
+        return dict(
+            input_ids=input_ids,
+            labels=labels,
+            attention_mask=input_ids.ne(self.tokenizer.pad_token_id),
+        )
+
+    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
+        sources = []
+        targets = []
+        for instance in instances:
+            source = instance['input_ids']
+            target = instance['labels']
+            sources.append(source)
+            targets.append(target)
+
+        data_dict = preprocess(sources, targets, self.tokenizer)
+        input_ids, labels = data_dict['input_ids'], data_dict['labels']
+        # input_ids, labels = tuple([instance[key] for instance in instances] for key in ("input_ids", "labels"))
+        input_ids = torch.nn.utils.rnn.pad_sequence(
+            input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id
+        )
+        labels = torch.nn.utils.rnn.pad_sequence(labels, batch_first=True, padding_value=IGNORE_INDEX)
+        return dict(
+            input_ids=input_ids,
+            labels=labels,
+            attention_mask=input_ids.ne(self.tokenizer.pad_token_id),
+        )
+
+def make_supervised_data_module(tokenizer: transformers.PreTrainedTokenizer, data_args) -> Dict:
+    """Make dataset and collator for supervised fine-tuning."""
+    train_dataset = SupervisedDataset(tokenizer=tokenizer, data_args=data_args)
+    data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)
+    return dict(train_dataset=train_dataset, eval_dataset=None, data_collator=data_collator)
+
+
+os.environ["WANDB_PROJECT"] = "auto_form"
+
+def train():
+
+    parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
+    model_args, data_args, training_args, remaining_args = parser.parse_args_into_dataclasses(return_remaining_strings=True)
+    data_args.data_length = int(remaining_args[1])
+
+    model = transformers.AutoModelForCausalLM.from_pretrained(
+        model_args.model_name_or_path,
+        cache_dir=training_args.cache_dir,
+        trust_remote_code=True,
+        torch_dtype=torch.bfloat16,
+        attn_implementation="flash_attention_2",
+    )
+
+    model.config.use_cache = False
+    model.gradient_checkpointing_enable()
+
+    tokenizer = transformers.AutoTokenizer.from_pretrained(
+        model_args.model_name_or_path,
+        cache_dir=training_args.cache_dir,
+        model_max_length=training_args.model_max_length,
+        padding_side="right",
+        use_fast=False,
+    )
+    if tokenizer.pad_token is None:
+        smart_tokenizer_and_embedding_resize(
+            special_tokens_dict=dict(pad_token=DEFAULT_PAD_TOKEN),
+            tokenizer=tokenizer,
+            model=model,
+        )
+    if "llama" in model_args.model_name_or_path:
+        tokenizer.add_special_tokens(
+            {
+                "eos_token": DEFAULT_EOS_TOKEN,
+                "bos_token": DEFAULT_BOS_TOKEN,
+                "unk_token": DEFAULT_UNK_TOKEN,
+            }
+        )
+        try:
+            tokenizer.pad_token = tokenizer.unk_token
+        except:
+            pass
+
+    data_module = make_supervised_data_module(tokenizer=tokenizer, data_args=data_args)
+    trainer = Trainer(model=model, tokenizer=tokenizer, args=training_args, **data_module)
+    trainer.train()
+    model.config.use_cache = True
+    # trainer.save_state()
+    # if os.environ.get('LOCAL_RANK') == '0':
+    safe_save_model_for_hf_trainer(trainer=trainer, output_dir=training_args.output_dir)
+
+
+
+
+if __name__ == "__main__":
+    train()

autosolver_train.py

@@ -0,0 +1,303 @@
+#    Copyright 2023 Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li
+#
+#    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.
+#    Modified by Zheng Yuan and Hongyi Yuan
+
+import os
+import copy
+import logging
+from dataclasses import dataclass, field
+from typing import Optional, Dict, Sequence
+import io
+import torch
+import transformers
+from torch.utils.data import Dataset
+from transformers import Trainer
+import argparse
+import json
+import random;random.seed(42)
+
+def _make_r_io_base(f, mode: str):
+    if not isinstance(f, io.IOBase):
+        f = open(f, mode=mode)
+    return f
+
+def jload(f, mode="r"):
+    """Load a .json file into a dictionary."""
+    f = _make_r_io_base(f, mode)
+    jdict = json.load(f)
+    f.close()
+    return jdict
+
+IGNORE_INDEX = -100
+DEFAULT_PAD_TOKEN = "[PAD]"
+DEFAULT_EOS_TOKEN = "</s>"
+DEFAULT_BOS_TOKEN = "<s>"
+DEFAULT_UNK_TOKEN = "<unk>"
+PROMPT_DICT = {
+    "lean4": (
+        "Statement and proof in natural language:\n\n"
+        "{statement_text}\n\n"
+        "Translate the statement and proof in natural language to lean4:"
+    ),
+    "plain": (
+        "{statement_text}"
+    ),
+    "prompt_no_input": (
+    "Below is an instruction that describes a task. "
+    "Write a response that appropriately completes the request.\n\n"
+    "### Instruction:\n{instruction}\n\n### Response:"
+    ),
+}
+#### 28
+@dataclass
+class ModelArguments:
+    model_name_or_path: Optional[str] = field(default="facebook/opt-125m")
+
+
+@dataclass
+class DataArguments:
+    data_path: str = field(default=None, metadata={"help": "Path to the training data."})
+
+
+@dataclass
+class TrainingArguments(transformers.TrainingArguments):
+    cache_dir: Optional[str] = field(default=None)
+    optim: str = field(default="adamw_torch")
+    model_max_length: int = field(
+        default=2048,
+        metadata={"help": "Maximum sequence length. Sequences will be right padded (and possibly truncated)."},
+    )
+    overwrite_output_dir: bool = field(default=True)
+
+
+def safe_save_model_for_hf_trainer(trainer: transformers.Trainer, output_dir: str):
+    """Collects the state dict and dump to disk."""
+    state_dict = trainer.model.state_dict()
+    if trainer.args.should_save:
+        cpu_state_dict = {key: value.cpu() for key, value in state_dict.items()}
+        del state_dict
+        trainer._save(output_dir, state_dict=cpu_state_dict)  # noqa
+
+
+def smart_tokenizer_and_embedding_resize(
+    special_tokens_dict: Dict,
+    tokenizer: transformers.PreTrainedTokenizer,
+    model: transformers.PreTrainedModel,
+):
+    """Resize tokenizer and embedding.
+    Note: This is the unoptimized version that may make your embedding size not be divisible by 64.
+    """
+    num_new_tokens = tokenizer.add_special_tokens(special_tokens_dict)
+    model.resize_token_embeddings(len(tokenizer))
+
+    if num_new_tokens > 0:
+        input_embeddings = model.get_input_embeddings().weight.data
+        output_embeddings = model.get_output_embeddings().weight.data
+
+        input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(dim=0, keepdim=True)
+        output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(dim=0, keepdim=True)
+
+        input_embeddings[-num_new_tokens:] = input_embeddings_avg
+        output_embeddings[-num_new_tokens:] = output_embeddings_avg
+
+
+def _tokenize_fn(strings: Sequence[str], tokenizer: transformers.PreTrainedTokenizer) -> Dict:
+    """Tokenize a list of strings."""
+    tokenized_list = [
+        tokenizer(
+            text,
+            return_tensors="pt",
+            padding="longest",
+            max_length=tokenizer.model_max_length,
+            truncation=True,
+        )
+        for text in strings
+    ]
+    input_ids = labels = [tokenized.input_ids[0] for tokenized in tokenized_list]
+    input_ids_lens = labels_lens = [
+        tokenized.input_ids.ne(tokenizer.pad_token_id).sum().item() for tokenized in tokenized_list
+    ]
+    return dict(
+        input_ids=input_ids,
+        labels=labels,
+        input_ids_lens=input_ids_lens,
+        labels_lens=labels_lens,
+    )
+
+
+def preprocess(
+    sources: Sequence[str],
+    targets: Sequence[str],
+    tokenizer: transformers.PreTrainedTokenizer,
+) -> Dict:
+    """Preprocess the data by tokenizing."""
+    examples = [s + t for s, t in zip(sources, targets)]
+    examples_tokenized, sources_tokenized = [_tokenize_fn(strings, tokenizer) for strings in (examples, sources)]
+    input_ids = examples_tokenized["input_ids"]
+    labels = copy.deepcopy(input_ids)
+    for label, source_len in zip(labels, sources_tokenized["input_ids_lens"]):
+        label[:source_len] = IGNORE_INDEX
+    return dict(input_ids=input_ids, labels=labels)
+
+class SupervisedDataset(Dataset):
+    """Dataset for supervised fine-tuning."""
+    def __init__(self, data_args, tokenizer: transformers.PreTrainedTokenizer):
+        super(SupervisedDataset, self).__init__()
+        logging.warning("Loading data...")
+        data_path = data_args.data_path
+        try:
+            data_path = data_path_map[data_path]
+        except:
+            data_path = data_path
+        list_data_dict = []
+        for item in data_path.split(','):
+            try:
+                list_data_dict += jload(item)
+
+            except BaseException:
+                with open(item, 'r') as f:
+                    lines = f.readlines()
+                list_data_dict += [json.loads(line.strip()) for line in lines]
+
+        list_data_dict = random.sample(list_data_dict,  len(list_data_dict))
+        list_data_dict = list_data_dict[:data_args.data_length]
+
+        logging.warning("Formatting inputs...")
+        prompt_lean4  = PROMPT_DICT["plain"]
+
+        # list_data_dict = [{'instruction':data['items'][0]['value'], 'input':'', 'output':data['items'][1]['value']} for data in list_data_dict]
+
+        list_data_dict = [{'instruction':prompt_lean4.format(statement_text = data['statement']), 'input':'', 'output':data['proof']} for data in list_data_dict]
+        print(f"len of {len(list_data_dict)}")
+        sources = [example['instruction'] for example in list_data_dict]
+        targets = [f"{example['output']}{tokenizer.eos_token}" for example in list_data_dict]
+        # targets = [example['output'] for example in list_data_dict]
+
+        self.sources = sources
+        self.targets = targets
+
+    def __len__(self):
+        return len(self.sources)
+
+    def naive__getitem__(self, i) -> Dict[str, torch.Tensor]:
+        return dict(input_ids=self.input_ids[i], labels=self.labels[i])
+
+    def __getitem__(self, i):
+        return dict(input_ids=self.sources[i], labels=self.targets[i])
+
+@dataclass
+class DataCollatorForSupervisedDataset(object):
+    """Collate examples for supervised fine-tuning."""
+
+    tokenizer: transformers.PreTrainedTokenizer
+
+    def naive__call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
+        input_ids, labels = tuple([instance[key] for instance in instances] for key in ("input_ids", "labels"))
+        input_ids = torch.nn.utils.rnn.pad_sequence(
+            input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id
+        )
+        labels = torch.nn.utils.rnn.pad_sequence(labels, batch_first=True, padding_value=IGNORE_INDEX)
+        return dict(
+            input_ids=input_ids,
+            labels=labels,
+            attention_mask=input_ids.ne(self.tokenizer.pad_token_id),
+        )
+
+    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
+        sources = []
+        targets = []
+        for instance in instances:
+            source = instance['input_ids']
+            target = instance['labels']
+            sources.append(source)
+            targets.append(target)
+
+        data_dict = preprocess(sources, targets, self.tokenizer)
+        input_ids, labels = data_dict['input_ids'], data_dict['labels']
+        # input_ids, labels = tuple([instance[key] for instance in instances] for key in ("input_ids", "labels"))
+        input_ids = torch.nn.utils.rnn.pad_sequence(
+            input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id
+        )
+        labels = torch.nn.utils.rnn.pad_sequence(labels, batch_first=True, padding_value=IGNORE_INDEX)
+        return dict(
+            input_ids=input_ids,
+            labels=labels,
+            attention_mask=input_ids.ne(self.tokenizer.pad_token_id),
+        )
+
+def make_supervised_data_module(tokenizer: transformers.PreTrainedTokenizer, data_args) -> Dict:
+    """Make dataset and collator for supervised fine-tuning."""
+    train_dataset = SupervisedDataset(tokenizer=tokenizer, data_args=data_args)
+    data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)
+    return dict(train_dataset=train_dataset, eval_dataset=None, data_collator=data_collator)
+
+
+os.environ["WANDB_PROJECT"] = "auto_solver"
+
+def train():
+
+    parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
+    model_args, data_args, training_args, remaining_args = parser.parse_args_into_dataclasses(return_remaining_strings=True)
+    data_args.data_length = int(remaining_args[1])
+
+    model = transformers.AutoModelForCausalLM.from_pretrained(
+        model_args.model_name_or_path,
+        cache_dir=training_args.cache_dir,
+        trust_remote_code=True,
+        torch_dtype=torch.bfloat16,
+        attn_implementation="flash_attention_2",
+    )
+
+    model.config.use_cache = False
+    model.gradient_checkpointing_enable()
+
+    tokenizer = transformers.AutoTokenizer.from_pretrained(
+        model_args.model_name_or_path,
+        cache_dir=training_args.cache_dir,
+        model_max_length=training_args.model_max_length,
+        padding_side="right",
+        use_fast=False,
+    )
+    if tokenizer.pad_token is None:
+        smart_tokenizer_and_embedding_resize(
+            special_tokens_dict=dict(pad_token=DEFAULT_PAD_TOKEN),
+            tokenizer=tokenizer,
+            model=model,
+        )
+    if "llama" in model_args.model_name_or_path:
+        tokenizer.add_special_tokens(
+            {
+                "eos_token": DEFAULT_EOS_TOKEN,
+                "bos_token": DEFAULT_BOS_TOKEN,
+                "unk_token": DEFAULT_UNK_TOKEN,
+            }
+        )
+        try:
+            tokenizer.pad_token = tokenizer.unk_token
+        except:
+            pass
+
+    data_module = make_supervised_data_module(tokenizer=tokenizer, data_args=data_args)
+    trainer = Trainer(model=model, tokenizer=tokenizer, args=training_args, **data_module)
+    trainer.train()
+    model.config.use_cache = True
+    # trainer.save_state()
+    # if os.environ.get('LOCAL_RANK') == '0':
+    safe_save_model_for_hf_trainer(trainer=trainer, output_dir=training_args.output_dir)
+
+
+
+
+if __name__ == "__main__":
+    train()

autostatement_train.py

@@ -0,0 +1,308 @@
+#    Copyright 2023 Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li
+#
+#    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.
+#    Modified by Zheng Yuan and Hongyi Yuan
+
+import os
+import copy
+import logging
+from dataclasses import dataclass, field
+from typing import Optional, Dict, Sequence
+import io
+import torch
+import transformers
+from torch.utils.data import Dataset
+from transformers import Trainer
+import argparse
+import json
+import random;random.seed(42)
+
+def _make_r_io_base(f, mode: str):
+    if not isinstance(f, io.IOBase):
+        f = open(f, mode=mode)
+    return f
+
+def jload(f, mode="r"):
+    """Load a .json file into a dictionary."""
+    f = _make_r_io_base(f, mode)
+    jdict = json.load(f)
+    f.close()
+    return jdict
+
+IGNORE_INDEX = -100
+DEFAULT_PAD_TOKEN = "[PAD]"
+DEFAULT_EOS_TOKEN = "</s>"
+DEFAULT_BOS_TOKEN = "<s>"
+DEFAULT_UNK_TOKEN = "<unk>"
+PROMPT_DICT = {
+    "lean4": (
+        "Statement and proof in natural language:\n\n"
+        "{statement_text}\n\n"
+        "Translate the statement and proof in natural language to lean4:"
+    ),
+    "plain": (
+        "{statement_text}"
+    ),
+    "statement": (
+        "Statement in natural language:\n"
+        "{problem}\n"
+        "Translate the statement in natural language to Lean4:"
+    ),
+    "prompt_no_input": (
+    "Below is an instruction that describes a task. "
+    "Write a response that appropriately completes the request.\n\n"
+    "### Instruction:\n{instruction}\n\n### Response:"
+    ),
+}
+#### 28
+@dataclass
+class ModelArguments:
+    model_name_or_path: Optional[str] = field(default="facebook/opt-125m")
+
+
+@dataclass
+class DataArguments:
+    data_path: str = field(default=None, metadata={"help": "Path to the training data."})
+
+
+@dataclass
+class TrainingArguments(transformers.TrainingArguments):
+    cache_dir: Optional[str] = field(default=None)
+    optim: str = field(default="adamw_torch")
+    model_max_length: int = field(
+        default=2048,
+        metadata={"help": "Maximum sequence length. Sequences will be right padded (and possibly truncated)."},
+    )
+    overwrite_output_dir: bool = field(default=True)
+
+
+def safe_save_model_for_hf_trainer(trainer: transformers.Trainer, output_dir: str):
+    """Collects the state dict and dump to disk."""
+    state_dict = trainer.model.state_dict()
+    if trainer.args.should_save:
+        cpu_state_dict = {key: value.cpu() for key, value in state_dict.items()}
+        del state_dict
+        trainer._save(output_dir, state_dict=cpu_state_dict)  # noqa
+
+
+def smart_tokenizer_and_embedding_resize(
+    special_tokens_dict: Dict,
+    tokenizer: transformers.PreTrainedTokenizer,
+    model: transformers.PreTrainedModel,
+):
+    """Resize tokenizer and embedding.
+    Note: This is the unoptimized version that may make your embedding size not be divisible by 64.
+    """
+    num_new_tokens = tokenizer.add_special_tokens(special_tokens_dict)
+    model.resize_token_embeddings(len(tokenizer))
+
+    if num_new_tokens > 0:
+        input_embeddings = model.get_input_embeddings().weight.data
+        output_embeddings = model.get_output_embeddings().weight.data
+
+        input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(dim=0, keepdim=True)
+        output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(dim=0, keepdim=True)
+
+        input_embeddings[-num_new_tokens:] = input_embeddings_avg
+        output_embeddings[-num_new_tokens:] = output_embeddings_avg
+
+
+def _tokenize_fn(strings: Sequence[str], tokenizer: transformers.PreTrainedTokenizer) -> Dict:
+    """Tokenize a list of strings."""
+    tokenized_list = [
+        tokenizer(
+            text,
+            return_tensors="pt",
+            padding="longest",
+            max_length=tokenizer.model_max_length,
+            truncation=True,
+        )
+        for text in strings
+    ]
+    input_ids = labels = [tokenized.input_ids[0] for tokenized in tokenized_list]
+    input_ids_lens = labels_lens = [
+        tokenized.input_ids.ne(tokenizer.pad_token_id).sum().item() for tokenized in tokenized_list
+    ]
+    return dict(
+        input_ids=input_ids,
+        labels=labels,
+        input_ids_lens=input_ids_lens,
+        labels_lens=labels_lens,
+    )
+
+
+def preprocess(
+    sources: Sequence[str],
+    targets: Sequence[str],
+    tokenizer: transformers.PreTrainedTokenizer,
+) -> Dict:
+    """Preprocess the data by tokenizing."""
+    examples = [s + t for s, t in zip(sources, targets)]
+    examples_tokenized, sources_tokenized = [_tokenize_fn(strings, tokenizer) for strings in (examples, sources)]
+    input_ids = examples_tokenized["input_ids"]
+    labels = copy.deepcopy(input_ids)
+    for label, source_len in zip(labels, sources_tokenized["input_ids_lens"]):
+        label[:source_len] = IGNORE_INDEX
+    return dict(input_ids=input_ids, labels=labels)
+
+class SupervisedDataset(Dataset):
+    """Dataset for supervised fine-tuning."""
+    def __init__(self, data_args, tokenizer: transformers.PreTrainedTokenizer):
+        super(SupervisedDataset, self).__init__()
+        logging.warning("Loading data...")
+        data_path = data_args.data_path
+        try:
+            data_path = data_path_map[data_path]
+        except:
+            data_path = data_path
+        list_data_dict = []
+        for item in data_path.split(','):
+            try:
+                list_data_dict += jload(item)
+
+            except BaseException:
+                with open(item, 'r') as f:
+                    lines = f.readlines()
+                list_data_dict += [json.loads(line.strip()) for line in lines]
+
+        list_data_dict = random.sample(list_data_dict,  len(list_data_dict))
+        list_data_dict = list_data_dict[:data_args.data_length]
+
+        logging.warning("Formatting inputs...")
+        prompt_lean4  = PROMPT_DICT["statement"]
+
+        # list_data_dict = [{'instruction':data['items'][0]['value'], 'input':'', 'output':data['items'][1]['value']} for data in list_data_dict]
+
+        list_data_dict = [{'instruction':prompt_lean4.format(problem= data['problem']), 'input':'', 'output':data['statement']} for data in list_data_dict]
+        print(f"len of {len(list_data_dict)}")
+        sources = [example['instruction'] for example in list_data_dict]
+        targets = [f"{example['output']}{tokenizer.eos_token}" for example in list_data_dict]
+        # targets = [example['output'] for example in list_data_dict]
+
+        self.sources = sources
+        self.targets = targets
+
+    def __len__(self):
+        return len(self.sources)
+
+    def naive__getitem__(self, i) -> Dict[str, torch.Tensor]:
+        return dict(input_ids=self.input_ids[i], labels=self.labels[i])
+
+    def __getitem__(self, i):
+        return dict(input_ids=self.sources[i], labels=self.targets[i])
+
+@dataclass
+class DataCollatorForSupervisedDataset(object):
+    """Collate examples for supervised fine-tuning."""
+
+    tokenizer: transformers.PreTrainedTokenizer
+
+    def naive__call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
+        input_ids, labels = tuple([instance[key] for instance in instances] for key in ("input_ids", "labels"))
+        input_ids = torch.nn.utils.rnn.pad_sequence(
+            input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id
+        )
+        labels = torch.nn.utils.rnn.pad_sequence(labels, batch_first=True, padding_value=IGNORE_INDEX)
+        return dict(
+            input_ids=input_ids,
+            labels=labels,
+            attention_mask=input_ids.ne(self.tokenizer.pad_token_id),
+        )
+
+    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
+        sources = []
+        targets = []
+        for instance in instances:
+            source = instance['input_ids']
+            target = instance['labels']
+            sources.append(source)
+            targets.append(target)
+
+        data_dict = preprocess(sources, targets, self.tokenizer)
+        input_ids, labels = data_dict['input_ids'], data_dict['labels']
+        # input_ids, labels = tuple([instance[key] for instance in instances] for key in ("input_ids", "labels"))
+        input_ids = torch.nn.utils.rnn.pad_sequence(
+            input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id
+        )
+        labels = torch.nn.utils.rnn.pad_sequence(labels, batch_first=True, padding_value=IGNORE_INDEX)
+        return dict(
+            input_ids=input_ids,
+            labels=labels,
+            attention_mask=input_ids.ne(self.tokenizer.pad_token_id),
+        )
+
+def make_supervised_data_module(tokenizer: transformers.PreTrainedTokenizer, data_args) -> Dict:
+    """Make dataset and collator for supervised fine-tuning."""
+    train_dataset = SupervisedDataset(tokenizer=tokenizer, data_args=data_args)
+    data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)
+    return dict(train_dataset=train_dataset, eval_dataset=None, data_collator=data_collator)
+
+
+os.environ["WANDB_PROJECT"] = "auto_statement"
+
+def train():
+
+    parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
+    model_args, data_args, training_args, remaining_args = parser.parse_args_into_dataclasses(return_remaining_strings=True)
+    data_args.data_length = int(remaining_args[1])
+
+    model = transformers.AutoModelForCausalLM.from_pretrained(
+        model_args.model_name_or_path,
+        cache_dir=training_args.cache_dir,
+        trust_remote_code=True,
+        torch_dtype=torch.bfloat16,
+        attn_implementation="flash_attention_2",
+    )
+
+    model.config.use_cache = False
+    model.gradient_checkpointing_enable()
+
+    tokenizer = transformers.AutoTokenizer.from_pretrained(
+        model_args.model_name_or_path,
+        cache_dir=training_args.cache_dir,
+        model_max_length=training_args.model_max_length,
+        padding_side="right",
+        use_fast=False,
+    )
+    if tokenizer.pad_token is None:
+        smart_tokenizer_and_embedding_resize(
+            special_tokens_dict=dict(pad_token=DEFAULT_PAD_TOKEN),
+            tokenizer=tokenizer,
+            model=model,
+        )
+    if "llama" in model_args.model_name_or_path:
+        tokenizer.add_special_tokens(
+            {
+                "eos_token": DEFAULT_EOS_TOKEN,
+                "bos_token": DEFAULT_BOS_TOKEN,
+                "unk_token": DEFAULT_UNK_TOKEN,
+            }
+        )
+        try:
+            tokenizer.pad_token = tokenizer.unk_token
+        except:
+            pass
+
+    data_module = make_supervised_data_module(tokenizer=tokenizer, data_args=data_args)
+    trainer = Trainer(model=model, tokenizer=tokenizer, args=training_args, **data_module)
+    trainer.train()
+    model.config.use_cache = True
+    # trainer.save_state()
+    # if os.environ.get('LOCAL_RANK') == '0':
+    safe_save_model_for_hf_trainer(trainer=trainer, output_dir=training_args.output_dir)
+
+
+
+
+if __name__ == "__main__":
+    train()

generation_method.py

@@ -0,0 +1,472 @@
+import argparse
+import random
+import glob
+
+from tqdm import tqdm
+import re
+import sys
+import os
+import numpy as np
+
+
+def generate_few_shot(prompt):
+    base_gsm8k_list = [
+        {
+            'question': "John and his best friend Steve bought 12 cupcakes together.  Each cupcake cost $1.50. If they split the costs evenly, how much did each person pay?",
+            'answer': "The total cost of cupcakes was 1.5*12=$<<1.5*12=18>>18\\nSo they each paid 18/2=$<<18/2=9>>9.",
+            'direct_answer': "9"
+        },
+        {
+            'question': "Lizzy has to ship 540 pounds of fish that are packed into 30-pound crates. If the shipping cost of each crate is $1.5, how much will Lizzy pay for the shipment?",
+            'answer': "There are 540 pounds / 30 pounds/crate = <<540/30=18>>18 crates of fish needed.\\nHence, the total cost for the shipment is $1.5/crate x 18 crates = $<<1.5*18=27>>27.",
+            'direct_answer': "27"
+        },
+        {
+            'question': "Tom, Tim, and Paul are collecting photos of cars. Paul has 10 photos more than Tim. Tim has one hundred photos less than the total amount of photos which is 152. How many photos does Tom have?",
+            'answer': "Tim has 152 photos - 100 photos = <<152-100=52>>52 photos.\\nWhen Tim has 52 photos, then Paul has 52 + 10 photos = <<52+10=62>>62 photos.\\nTim and Paul have together 52 photos + 62 photos = <<52+62=114>>114 photos.\\nThat leaves Tom with 152 photos - 114 photos = <<152-114=38>>38 photos.",
+            'direct_answer': "38"
+        },
+
+    ]
+    index_list = list(range(len(base_gsm8k_list)))
+    random.shuffle(index_list)
+    few_shot_example = ""
+    for i in index_list:
+        item = base_gsm8k_list[i]
+        few_shot_example += "Q: " + item['question'] + "\n" + "A: "+ item['answer']  + "\nThe answer is " + item['direct_answer'] + "\n"
+
+    few_shot_example += "Q: " + prompt + "A: "
+    return few_shot_example
+
+
+
+def generate_prompt_generation(args, question):
+    if args.evaluation_mode == 'generation':
+        if args.method == 'zero_shot_cot':
+            content = question + " Let's think step by step."
+        elif args.method == 'zero_shot':
+            content = question
+        elif args.method == 'few_shot':
+            content = generate_few_shot(question)
+        else:
+            raise ValueError("we do not method for such model type yet")
+
+    if "generator" not in args.model_type:
+        MODEL_DICT = {
+            "llama": (
+                "[INST] \n{content}\n [/INST]"
+            ),
+            "mistral": (
+                "<s>[INST] {content} [/INST]"
+            ),
+            "chatglm": (
+                "<|user|> \n{content}\n <|assistant|>"
+            ),
+            "qianwen": (
+                "<|im_start|>user\n{content}<|im_end|>\n<|im_start|>assistant\n"
+            ),
+            "baichuan": (
+                "<reserved_106>{content}<reserved_107>"
+            )
+        }
+
+        if args.model_type in ["qianwen", "qianwen-13b", "qianwen-70b"]:
+            content = MODEL_DICT['qianwen'].format_map(
+                {'content': content}
+            )
+
+        elif args.model_type in ["chatglm"]:
+            pass
+
+
+        elif args.model_type in ['llama2-7b-chat']:
+            content = MODEL_DICT['llama'].format_map(
+                {'content': content}
+            )
+
+        elif args.model_type in ["mistral", 'mixtral']:
+            content = MODEL_DICT['mistral'].format_map(
+                {'content': content}
+            )
+
+
+    return content
+
+
+
+
+few_shot_list = [
+    {
+        'question': "There are 15 trees in the grove. Grove workers will plant trees in the grove today. After they are done, there will be 21 trees. How many trees did the grove workers plant today?",
+        'answer': "There are 15 trees originally. Then there were 21 trees after some more were planted. So there must have been 21 - 15 = 6.",
+        'direct_answer': "6"
+    },
+    {
+        'question': "If there are 3 cars in the parking lot and 2 more cars arrive, how many cars are in the parking lot?",
+        'answer': "There are originally 3 cars. 2 more cars arrive. 3 + 2 = 5.",
+        'direct_answer': "5",
+    },
+    {
+        'question': "Leah had 32 chocolates and her sister had 42. If they ate 35, how many pieces do they have left in total?",
+        'answer': "Originally, Leah had 32 chocolates. Her sister had 42. So in total they had 32 + 42 = 74. After eating 35, they had 74 - 35 = 39.",
+        'direct_answer': "39",
+    },
+    {
+        'question': "Jason had 20 lollipops. He gave Denny some lollipops. Now Jason has 12 lollipops. How many lollipops did Jason give to Denny?",
+        'answer': "Jason started with 20 lollipops. Then he had 12 after giving some to Denny. So he gave Denny 20 - 12 = 8.",
+        'direct_answer': "8",
+    },
+    {
+        'question': "Shawn has five toys. For Christmas, he got two toys each from his mom and dad. How many toys does he have now?",
+        'answer': "Shawn started with 5 toys. If he got 2 toys each from his mom and dad, then that is 4 more toys. 5 + 4 = 9.",
+        'direct_answer': "9",
+    },
+    {
+        'question': "There were nine computers in the server room. Five more computers were installed each day, from monday to thursday. How many computers are now in the server room?",
+        'answer': "There were originally 9 computers. For each of 4 days, 5 more computers were added. So 5 * 4 = 20 computers were added. 9 + 20 is 29.",
+        'direct_answer': "29",
+    },
+    {
+        'question': "Michael had 58 golf balls. On tuesday, he lost 23 golf balls. On wednesday, he lost 2 more. How many golf balls did he have at the end of wednesday?",
+        'answer': "Michael started with 58 golf balls. After losing 23 on tuesday, he had 58 - 23 = 35. After losing 2 more, he had 35 - 2 = 33 golf balls.",
+        'direct_answer': "33",
+    },
+    {
+        'question': "Olivia has $23. She bought five bagels for $3 each. How much money does she have left?",
+        'answer': "Olivia had 23 dollars. 5 bagels for 3 dollars each will be 5 x 3 = 15 dollars. So she has 23 - 15 dollars left. 23 - 15 is 8.",
+        'direct_answer': "8",
+    },
+]
+import json
+
+from collections import Counter
+
+
+def self_consistency(pairs):
+    val_counts = Counter(value for key, value in pairs)
+    most = val_counts.most_common(1)[0][0]
+    for key, value in pairs:
+        if value == most:
+            return key
+
+
+#
+def find_feedback(content):
+    match = re.search(r'Judgement: (.+)', content)
+    if match:
+        judgement = match.group(1)
+    else:
+        judgement = "None"
+    return judgement
+
+
+def str2bool(s):
+    s = s.lower()
+    if s == 'true':
+        return True
+    elif s == 'false':
+        return False
+    else:
+        raise ValueError('invalid value: {}, must be true or false'.format(s))
+
+
+def parse_arguments():
+    parser = argparse.ArgumentParser(description="Zero-shot-CoT")
+
+    # parser.add_argument(
+    #     "--dataset", type=str, default="plan",
+    #     choices=["plan", 'tool_use_awareness', 'tool_selection', 'tool_selection_harder', 'tool_creation_awareness',
+    #              'tool_creation_awareness_harder', 'tool_creation',
+    #              'arguments_filling'], help="dataset used for experiment")
+    parser.add_argument(
+        "--cot_trigger_no", type=int, default=1,
+        help="A trigger sentence that elicits a model to execute chain of thought"
+    )
+    parser.add_argument("--dataset", type=str, default="")
+    parser.add_argument("--data_path", type=str, default="")
+    parser.add_argument("--evaluation_mode", type=str, default="")
+    parser.add_argument("--batch_size", type=int, default=1)
+    parser.add_argument("--eval_method", type=str, default="")
+
+    parser.add_argument("--model_path", type=str, default="")
+
+    parser.add_argument("--model_type", type=str, default="chatglm")
+
+    parser.add_argument("--output_dir", type=str, default="generation_test")
+
+    parser.add_argument("--lora_path", type=str, default="")
+
+    parser.add_argument("--iter_num", type=int, default=1)
+    parser.add_argument("--method", type=str, default="few_shot_cot")
+    parser.add_argument("--data_question_key", type=str, default="question")
+    parser.add_argument("--data_answer_key", type=str, default="answer")
+
+    parser.add_argument("--sample_num", type=int, default=1)
+
+    parser.add_argument("--cuda_ind", type=int, default=0)
+    parser.add_argument("--tensor_parallel", type=int, default=1)
+    parser.add_argument("--cuda_start", type=int, default=0)
+    parser.add_argument("--cuda_num", type=int, default=8)
+
+    parser.add_argument("--load_in_8bit", type=str2bool, default=False)
+    parser.add_argument("--rewrite", type=str2bool, default=True)
+    parser.add_argument("--notlean", type=str2bool, default=True)
+
+    parser.add_argument("--use_typewriter", type=int, default=0)
+
+    parser.add_argument("--temperature", type=float, default=0.0)
+    parser.add_argument("--top_p", type=float, default=1)
+    parser.add_argument("--iter_max_new_tokens", type=int, default=512)
+    parser.add_argument("--init_max_new_tokens", type=int, default=2048)
+    parser.add_argument("--min_new_tokens", type=int, default=1)
+    parser.add_argument("--correct_response_format", type=str, default="The correct response is:")
+
+    args = parser.parse_args()
+    if args.evaluation_mode == 'generation':
+        if "lean" in args.dataset:
+            args.data_question_key = 'model_response'
+            args.data_answer_key = 'statement_poof'
+
+        if args.dataset == "lean4_5k_test":
+            args.data_path = "data/lean4_gpt_5k/test/data.jsonl"
+
+        elif args.dataset == "math_train":
+            args.data_path = "data/test/math/train.jsonl"
+
+        elif args.dataset == "gsm8k_train":
+            args.data_path = "data/test/gsm8k/train.jsonl"
+            
+        elif args.dataset == "wild_test":
+            args.data_path = "/hpc2hdd/home/zyang398/data_2/wild_sample1k.jsonl"
+
+        elif args.dataset == "lean4_basic_test":
+            args.data_path = "data/lean4_basic/1k_test.jsonl"
+        elif args.dataset == "lean4_random_test":
+            args.data_path = "data/lean4_random/1k_test.json"
+        elif args.dataset == "lean4_random_first_train":
+            args.data_path = "data/lean4_random/5k_first.json"
+        elif args.dataset == "lean4_random_second_train":
+            args.data_path = "data/lean4_random/5k_second.json"
+        elif args.dataset == "lean4_random_third_train":
+            args.data_path = "data/lean4_random/5k_third.json"
+
+    if args.model_type == 'mistral_generator':
+        args.model_path = 'models/gsm8k/generators/mistral-ep2/'
+    elif args.model_type == 'mistral_generator_original':
+        args.model_path = '/data/OVM-Mistral-7b/mistral7b-ep2/'
+    elif args.model_type == 'gemma_generator':
+        args.model_path = 'models/gsm8k/generators/gemma2b2-ep2/'
+    elif args.model_type == 'phi2_generator':
+        args.model_path = 'models/gsm8k/generators/phi2b-ep2/'
+
+    elif args.model_type == 'mixtral':
+        args.model_path = '/data/Mixtral-8x7B-Instruct-v0.1'
+
+    elif args.model_type == 'mistral':
+        args.model_path = '/data/mistral-instruct'
+
+    elif args.model_type == 'qianwen-70b':
+        args.model_path = '/data/Qwen-72B-Chat'
+
+
+    elif args.model_type == 'llama2-7b-chat':
+        args.model_path = '/data/Llama-2-7b-chat/'
+
+    if args.cot_trigger_no == 1:
+        args.cot_trigger = "Let's think step by step."
+
+    return args
+
+
+def create_demo_text(args, cot_flag, index_list):
+    # Concatenate demonstration examples ...
+    demo_text = ""
+    for i in index_list:
+        item = few_shot_list[i]
+        if cot_flag:
+            demo_text += "Q: " + item['question'] + "\nA: " + item['answer'] + " " + \
+                         args.direct_answer_trigger_for_fewshot + " " + item['direct_answer'] + ".\n\n"
+        else:
+            demo_text += "Q: " + item['question'] + "\nA: " + \
+                         args.direct_answer_trigger_for_fewshot + " " + item['direct_answer'] + ".\n\n"
+
+    return demo_text
+
+
+def str2bool(s):
+    s = s.lower()
+    if s == 'true':
+        return True
+    elif s == 'false':
+        return False
+    else:
+        raise ValueError('invalid value: {}, must be true or false'.format(s))
+
+
+def batchify(pairs, batch_size):
+
+    """将列表分成指定大小的批次"""
+    for i in range(0, len(pairs), batch_size):
+        yield pairs[i:i + batch_size]
+
+
+def generate_prompts(questions, args):
+    """为每个问题生成提示"""
+    prompts = [generate_prompt_generation(args, question) for question in questions]
+    return prompts
+
+PROMPT_DICT = {
+    "wild": (
+        "Statement and proof in natural language:\n\n"
+        "# Problem:\n{question}\n\n"
+        "# Proof:\n{answer}\n\n"
+        "Translate the statement and proof in natural language to lean4:"
+    ),
+    "lean4": (
+        "Statement and proof in natural language:\n\n"
+        "{statement_text}\n\n"
+        "Translate the statement and proof in natural language to lean4:"
+    ),
+    "prompt_no_input": (
+    "Below is an instruction that describes a task. "
+    "Write a response that appropriately completes the request.\n\n"
+    "### Instruction:\n{instruction}\n\n### Response:"
+    ),
+}
+
+def get_question_answer(args):
+    allfilepath = args.data_path
+    questions = []
+    answers = []
+
+    # Attempt to read the file as a regular JSON file
+    for filepath in allfilepath.split(','):
+        try:
+            with open(filepath, 'r') as file:
+                data = json.load(file)
+                # If the data is a list, assume it's an array of objects
+                if isinstance(data, list):
+                    for json_item in data:
+                        questions.append(json_item[args.data_question_key])
+                        answers.append(json_item)
+                # If the data is a dict, assume it's a single object (or adjust logic as needed)
+                elif isinstance(data, dict):
+                    questions.append(data[args.data_question_key])
+                    answers.append(json_item)
+
+        except ValueError:
+            # If it fails, assume the file is in JSON Lines format
+            with open(filepath, 'r') as file:
+                for line in file:
+                    json_item = json.loads(line)
+                    questions.append(json_item[args.data_question_key])
+                    answers.append(json_item)
+
+    if args.notlean :
+        questions  = [ PROMPT_DICT['wild'].format(question= questions[id], answer =answers[id][args.data_answer_key] )  for id in range(len(questions))]
+
+    else:
+        questions  = [ PROMPT_DICT['lean4'].format(statement_text = item)  for item in questions]
+
+
+    return questions, answers
+
+
+def main3(args):
+    from vllm import LLM, SamplingParams
+    import torch
+
+    model = LLM(model=args.model_path, dtype="bfloat16", trust_remote_code=True,
+                tensor_parallel_size=args.tensor_parallel, gpu_memory_utilization = 0.95)
+    print(args.model_path)
+
+    if "qianwen" in args.model_type:
+        model.llm_engine.tokenizer.eos_token_id = 151645
+        # model.llm_engine.tokenizer.pad_token_id = 151645
+        model.llm_engine.tokenizer.pad_token_id = None
+        # model.llm_engine.tokenizer.eos_token_id = None
+
+
+    print("load data")
+
+
+    questions, answers = get_question_answer(args)
+
+
+
+    question_exist_list = []
+    write_pattern = 'w' if args.rewrite else "a+"
+    if os.path.exists(args.output_dir) and not args.rewrite :
+        # 如果文件存在，从文件中读取数据加载到response_list
+        # Loop through each file that matches the pattern
+        file_pattern = os.path.join(args.output_dir, '[0-9]*.json')
+        for file_path in glob.glob(file_pattern):
+            # Open and read the JSON file
+            with open(file_path, 'r') as fp:
+                # Extract the 'question' field from each line and add it to the list
+                for line in fp.readlines():
+                    question_exist_list.append(json.loads(line)['question'])
+    else:
+        try:
+            os.mkdir(args.output_dir)
+        except:
+            pass
+    qa_pairs = [(questions[idx], answers[idx]) for idx in range(len(questions)) if questions[idx] not in question_exist_list ]
+    cuda_pieces = np.array_split(range(len(qa_pairs)), args.cuda_num // args.tensor_parallel)
+    print(f"fitered {len(questions) - len(qa_pairs)} already")
+
+    with open(f"{args.output_dir}/{args.cuda_ind // args.tensor_parallel + args.cuda_start}.json", write_pattern,
+              encoding='utf-8') as wf:
+        start = cuda_pieces[args.cuda_start + args.cuda_ind // args.tensor_parallel][0]
+        end = cuda_pieces[args.cuda_start + args.cuda_ind // args.tensor_parallel][-1] + 1
+        subset_length = end - start
+        total_batches = (subset_length + args.batch_size - 1) // args.batch_size  # Calculate the total number of batches
+        for batch in tqdm(batchify(qa_pairs[start:end], args.batch_size), total=total_batches):
+            questions, answers = zip(*batch)  # 解压问题和答案
+            prompts = generate_prompts(questions, args)
+
+            with torch.no_grad():
+                output_all = []
+                try:
+                    for i in range(args.sample_num):
+                        sample_list = []
+                        sampling_params = SamplingParams(temperature=args.temperature, top_p=args.top_p,
+                                                         max_tokens=args.init_max_new_tokens)
+                        generations = model.generate(prompts, sampling_params, use_tqdm=False)
+                        for generation_output in generations:
+                            output = generation_output.outputs[0].text
+                            sample_list.append(output)
+                        output_all.append(sample_list)
+
+                    output_all = list(map(list, zip(*output_all)))
+                except Exception as e:
+                    print(str(e))
+                    exit
+                dicts = []
+                for question, answer, output, prompt in zip(questions, answers, output_all, prompts):
+                    dicts.append({
+                        "question": question,
+                        "prompt": prompt,
+                        "content": answer,
+                        "total output": output,
+                    })
+
+                for dict in dicts:
+                    wf.writelines(json.dumps(dict, ensure_ascii=False) + '\n')
+
+                wf.flush()
+
+
+def main(argv=None):
+    args = parse_arguments()
+    print('*****************************')
+    print(args)
+    print('*****************************')
+    if args.evaluation_mode == 'generation':
+        main3(args)
+    else:
+        raise ValueError("we do not yet inplement")
+
+
+if __name__ == "__main__":
+    main()

model_train.py

@@ -0,0 +1,318 @@
+#    Copyright 2023 Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li
+#
+#    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.
+#    Modified by Zheng Yuan and Hongyi Yuan
+
+import os
+import copy
+import logging
+from dataclasses import dataclass, field
+from typing import Optional, Dict, Sequence
+import io
+import torch
+import transformers
+from torch.utils.data import Dataset
+from transformers import Trainer
+import argparse
+import json
+import random;
+
+random.seed(42)
+
+
+def _make_r_io_base(f, mode: str):
+    if not isinstance(f, io.IOBase):
+        f = open(f, mode=mode)
+    return f
+
+
+def jload(f, mode="r"):
+    """Load a .json file into a dictionary."""
+    f = _make_r_io_base(f, mode)
+    jdict = json.load(f)
+    f.close()
+    return jdict
+
+
+IGNORE_INDEX = -100
+DEFAULT_PAD_TOKEN = "[PAD]"
+DEFAULT_EOS_TOKEN = "</s>"
+DEFAULT_BOS_TOKEN = "<s>"
+DEFAULT_UNK_TOKEN = "<unk>"
+PROMPT_DICT = {
+    "statement_form": (
+        "Statement in natural language:\n"
+        "{problem}\n"
+        "Translate the statement in natural language to Lean4:"
+    ),
+    "solver": (
+        "{statement_text}"
+    ),
+    "statementproof_inform": (
+        "Statement and proof in lean4:\n\n"
+        "{statement_text}\n\n"
+        "Translate the statement and proof in lean4 to natural language:"
+    ),
+}
+KEY_DICT = {
+    "statement_form" : ["problem", "statement"],
+    "solver": ["statement", "proof"],
+    "statementproof_inform": ["statement_poof", "model_response"]
+}
+
+
+#### 28
+@dataclass
+class ModelArguments:
+    model_name_or_path: Optional[str] = field(default="facebook/opt-125m")
+
+
+@dataclass
+class DataArguments:
+    data_path: str = field(default=None, metadata={"help": "Path to the training data."})
+
+
+@dataclass
+class TrainingArguments(transformers.TrainingArguments):
+    cache_dir: Optional[str] = field(default=None)
+    optim: str = field(default="adamw_torch")
+    model_max_length: int = field(
+        default=2048,
+        metadata={"help": "Maximum sequence length. Sequences will be right padded (and possibly truncated)."},
+    )
+    overwrite_output_dir: bool = field(default=True)
+
+
+def safe_save_model_for_hf_trainer(trainer: transformers.Trainer, output_dir: str):
+    """Collects the state dict and dump to disk."""
+    state_dict = trainer.model.state_dict()
+    if trainer.args.should_save:
+        cpu_state_dict = {key: value.cpu() for key, value in state_dict.items()}
+        del state_dict
+        trainer._save(output_dir, state_dict=cpu_state_dict)  # noqa
+
+
+def smart_tokenizer_and_embedding_resize(
+        special_tokens_dict: Dict,
+        tokenizer: transformers.PreTrainedTokenizer,
+        model: transformers.PreTrainedModel,
+):
+    """Resize tokenizer and embedding.
+    Note: This is the unoptimized version that may make your embedding size not be divisible by 64.
+    """
+    num_new_tokens = tokenizer.add_special_tokens(special_tokens_dict)
+    model.resize_token_embeddings(len(tokenizer))
+
+    if num_new_tokens > 0:
+        input_embeddings = model.get_input_embeddings().weight.data
+        output_embeddings = model.get_output_embeddings().weight.data
+
+        input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(dim=0, keepdim=True)
+        output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(dim=0, keepdim=True)
+
+        input_embeddings[-num_new_tokens:] = input_embeddings_avg
+        output_embeddings[-num_new_tokens:] = output_embeddings_avg
+
+
+def _tokenize_fn(strings: Sequence[str], tokenizer: transformers.PreTrainedTokenizer) -> Dict:
+    """Tokenize a list of strings."""
+    tokenized_list = [
+        tokenizer(
+            text,
+            return_tensors="pt",
+            padding="longest",
+            max_length=tokenizer.model_max_length,
+            truncation=True,
+        )
+        for text in strings
+    ]
+    input_ids = labels = [tokenized.input_ids[0] for tokenized in tokenized_list]
+    input_ids_lens = labels_lens = [
+        tokenized.input_ids.ne(tokenizer.pad_token_id).sum().item() for tokenized in tokenized_list
+    ]
+    return dict(
+        input_ids=input_ids,
+        labels=labels,
+        input_ids_lens=input_ids_lens,
+        labels_lens=labels_lens,
+    )
+
+
+def preprocess(
+        sources: Sequence[str],
+        targets: Sequence[str],
+        tokenizer: transformers.PreTrainedTokenizer,
+) -> Dict:
+    """Preprocess the data by tokenizing."""
+    examples = [s + t for s, t in zip(sources, targets)]
+    examples_tokenized, sources_tokenized = [_tokenize_fn(strings, tokenizer) for strings in (examples, sources)]
+    input_ids = examples_tokenized["input_ids"]
+    labels = copy.deepcopy(input_ids)
+    for label, source_len in zip(labels, sources_tokenized["input_ids_lens"]):
+        label[:source_len] = IGNORE_INDEX
+    return dict(input_ids=input_ids, labels=labels)
+
+
+class SupervisedDataset(Dataset):
+    """Dataset for supervised fine-tuning."""
+
+    def __init__(self, data_args, tokenizer: transformers.PreTrainedTokenizer):
+        super(SupervisedDataset, self).__init__()
+        logging.warning("Loading data...")
+        data_path = data_args.data_path
+        try:
+            data_path = data_path_map[data_path]
+        except:
+            data_path = data_path
+        list_data_dict = []
+        for item in data_path.split(','):
+            try:
+                list_data_dict += jload(item)
+
+            except BaseException:
+                with open(item, 'r') as f:
+                    lines = f.readlines()
+                list_data_dict += [json.loads(line.strip()) for line in lines]
+
+        list_data_dict = random.sample(list_data_dict, len(list_data_dict))
+        list_data_dict = list_data_dict[:data_args.data_length]
+
+        logging.warning("Formatting inputs...")
+
+        # list_data_dict = [{'instruction':data['items'][0]['value'], 'input':'', 'output':data['items'][1]['value']} for data in list_data_dict]
+
+        list_data_dict = [{'instruction': PROMPT_DICT[data['task']].format(statement_text=data[KEY_DICT[data['task']][0]]), 'input': '',
+                           'output':data[KEY_DICT[data['task']][1]] } for data in list_data_dict]
+        print(f"len of {len(list_data_dict)}")
+        sources = [example['instruction'] for example in list_data_dict]
+        targets = [f"{example['output']}{tokenizer.eos_token}" for example in list_data_dict]
+        # targets = [example['output'] for example in list_data_dict]
+
+        self.sources = sources
+        self.targets = targets
+
+    def __len__(self):
+        return len(self.sources)
+
+    def naive__getitem__(self, i) -> Dict[str, torch.Tensor]:
+        return dict(input_ids=self.input_ids[i], labels=self.labels[i])
+
+    def __getitem__(self, i):
+        return dict(input_ids=self.sources[i], labels=self.targets[i])
+
+
+@dataclass
+class DataCollatorForSupervisedDataset(object):
+    """Collate examples for supervised fine-tuning."""
+
+    tokenizer: transformers.PreTrainedTokenizer
+
+    def naive__call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
+        input_ids, labels = tuple([instance[key] for instance in instances] for key in ("input_ids", "labels"))
+        input_ids = torch.nn.utils.rnn.pad_sequence(
+            input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id
+        )
+        labels = torch.nn.utils.rnn.pad_sequence(labels, batch_first=True, padding_value=IGNORE_INDEX)
+        return dict(
+            input_ids=input_ids,
+            labels=labels,
+            attention_mask=input_ids.ne(self.tokenizer.pad_token_id),
+        )
+
+    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
+        sources = []
+        targets = []
+        for instance in instances:
+            source = instance['input_ids']
+            target = instance['labels']
+            sources.append(source)
+            targets.append(target)
+
+        data_dict = preprocess(sources, targets, self.tokenizer)
+        input_ids, labels = data_dict['input_ids'], data_dict['labels']
+        # input_ids, labels = tuple([instance[key] for instance in instances] for key in ("input_ids", "labels"))
+        input_ids = torch.nn.utils.rnn.pad_sequence(
+            input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id
+        )
+        labels = torch.nn.utils.rnn.pad_sequence(labels, batch_first=True, padding_value=IGNORE_INDEX)
+        return dict(
+            input_ids=input_ids,
+            labels=labels,
+            attention_mask=input_ids.ne(self.tokenizer.pad_token_id),
+        )
+
+
+def make_supervised_data_module(tokenizer: transformers.PreTrainedTokenizer, data_args) -> Dict:
+    """Make dataset and collator for supervised fine-tuning."""
+    train_dataset = SupervisedDataset(tokenizer=tokenizer, data_args=data_args)
+    data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)
+    return dict(train_dataset=train_dataset, eval_dataset=None, data_collator=data_collator)
+
+
+os.environ["WANDB_PROJECT"] = "train_in_one_model"
+
+
+def train():
+    parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
+    model_args, data_args, training_args, remaining_args = parser.parse_args_into_dataclasses(
+        return_remaining_strings=True)
+    data_args.data_length = int(remaining_args[1])
+
+    model = transformers.AutoModelForCausalLM.from_pretrained(
+        model_args.model_name_or_path,
+        cache_dir=training_args.cache_dir,
+        trust_remote_code=True,
+        torch_dtype=torch.bfloat16,
+        attn_implementation="flash_attention_2",
+    )
+
+    model.config.use_cache = False
+    model.gradient_checkpointing_enable()
+
+    tokenizer = transformers.AutoTokenizer.from_pretrained(
+        model_args.model_name_or_path,
+        cache_dir=training_args.cache_dir,
+        model_max_length=training_args.model_max_length,
+        padding_side="right",
+        use_fast=False,
+    )
+    if tokenizer.pad_token is None:
+        smart_tokenizer_and_embedding_resize(
+            special_tokens_dict=dict(pad_token=DEFAULT_PAD_TOKEN),
+            tokenizer=tokenizer,
+            model=model,
+        )
+    if "llama" in model_args.model_name_or_path:
+        tokenizer.add_special_tokens(
+            {
+                "eos_token": DEFAULT_EOS_TOKEN,
+                "bos_token": DEFAULT_BOS_TOKEN,
+                "unk_token": DEFAULT_UNK_TOKEN,
+            }
+        )
+        try:
+            tokenizer.pad_token = tokenizer.unk_token
+        except:
+            pass
+
+    data_module = make_supervised_data_module(tokenizer=tokenizer, data_args=data_args)
+    trainer = Trainer(model=model, tokenizer=tokenizer, args=training_args, **data_module)
+    trainer.train()
+    model.config.use_cache = True
+    # trainer.save_state()
+    # if os.environ.get('LOCAL_RANK') == '0':
+    safe_save_model_for_hf_trainer(trainer=trainer, output_dir=training_args.output_dir)
+
+
+if __name__ == "__main__":
+    train()

repl/.lake/packages/mathlib/scripts/align-import.py

@@ -0,0 +1,61 @@
+#!/usr/bin/env python3
+
+# This script was written by ChatGPT.
+# https://chat.openai.com/share/e0363ebf-ed6f-4fd8-9b76-ebf422ed9f62
+
+import re
+import sys
+
+def update_file_header(file_path):
+    with open(file_path, 'r') as f:
+        lines = f.readlines()
+
+    # Initialize variables
+    end_of_header_index = 0
+    source_module = ""
+    repo_ref = ""
+    commit_id = ""
+
+    # Lines to delete
+    delete_indices = []
+
+    for i, line in enumerate(lines):
+        # Check for the end of the "import" lines
+        if line.startswith('import'):
+            end_of_header_index = i
+        elif end_of_header_index != 0 and not line.startswith('import'):
+            break
+
+        # Extract the necessary info for the align import line and mark lines for deletion
+        if line.startswith('! This file was ported from'):
+            source_module = line.split()[-1]
+            delete_indices.append(i)
+        elif line.startswith('!') and 'commit' in line and commit_id == "":
+            split_line = line.split()
+            repo_ref = split_line[1]
+            commit_id = split_line[-1]
+            delete_indices.append(i)
+        elif line.startswith('!'):
+            delete_indices.append(i)
+        elif line == "\n" and lines[i+1].startswith("!"):
+            delete_indices.append(i)
+
+    # Only proceed if we have found the necessary info for the align import line
+    if source_module and repo_ref and commit_id:
+        # Generate the new line
+        new_line = f'\n#align_import {source_module} from "{repo_ref}"@"{commit_id}"\n'
+
+        # Delete the marked lines
+        for index in sorted(delete_indices, reverse=True):
+            del lines[index]
+
+        # Insert the new line after the "import" lines
+        lines.insert(end_of_header_index - len(delete_indices) + 1, new_line)
+
+        # Write the updated lines back to the file
+        with open(file_path, 'w') as f:
+            f.writelines(lines)
+
+# The first command line argument is the file path
+file_path = sys.argv[1]
+update_file_header(file_path)

repl/.lake/packages/mathlib/scripts/align.py

@@ -0,0 +1,64 @@
+#!/usr/bin/env python3
+
+# Tool to add source headers to ported theory files,
+# archived for historical purposes.
+
+from pathlib import Path
+import re
+import yaml
+
+excepts = {
+    'categorytheory.category.rel': 'categorytheory.category.relcat',
+    'categorytheory.isomorphism': 'categorytheory.iso',
+    'categorytheory.naturalisomorphism': 'categorytheory.natiso',
+    'categorytheory.naturaltransformation': 'categorytheory.nattrans',
+    'leancore.data.vector': 'data.vector',
+    'order.monovary': 'order.monotone.monovary'
+    }
+
+def condense(s):
+    if s.startswith('Mathlib/'):
+        s = s[len('Mathlib/'):]
+    if s.endswith('.lean'):
+        s = s[:-5]
+    s = s.lower()
+    s = s.replace('/', '.')
+    s = s.replace('_', '')
+    if s in excepts:
+        s = excepts[s]
+    return s
+
+port_status = yaml.load(open("mathlib4-port-status.yaml").read())
+
+# map from condensed names to mathlib4 paths
+map = {}
+for path in Path('Mathlib').glob('**/*.lean'):
+    path = str(path)
+    map[condense(path)] = path
+
+count = 0
+for key, val in port_status.items():
+    if val.startswith('Yes'):
+        sha = val.split()[2]
+        mathlib3 = key
+        mathlib4 = map[condense(key)]
+
+        place = '(\n-/\n\n?import )'
+        blob = "\n\n! This file was ported from Lean 3 source module " + mathlib3 + "\n" + \
+            "! leanprover-community/mathlib commit " + sha + "\n" + \
+            "! Please do not edit these lines, except to modify the commit id\n" + \
+            "! if you have ported upstream changes."
+        old = open(mathlib4).read()
+
+        if blob[1:] in old:     # match even without leading newline
+            print(f'{mathlib4} already has header')
+        elif "! leanprover-community/mathlib commit " in old:
+            m = re.search("^! leanprover-community/mathlib commit (.*)$", old, flags=re.MULTILINE)
+            print(f'file says {m.groups()[0]} but we want {sha}')
+            assert(False)
+        else:
+            new = re.sub(place, blob + '\\1', old, flags=re.MULTILINE)
+            open(mathlib4, 'w').write(new)
+            count += 1
+
+print(count)

repl/.lake/packages/mathlib/scripts/bench/accumulate_profile.py

@@ -0,0 +1,21 @@
+#!/usr/bin/env python3
+# sum up times of lines a la `elaboration 100ms`
+
+import collections
+import re
+import sys
+
+cats = collections.defaultdict(lambda: 0)
+for line in sys.stdin:
+    sys.stderr.write(line)
+    if m := re.match("(.+?) ([\d.]+)(m?)s$", line):
+        cats[m[1].strip()] += float(m[2]) * (1e-3 if m[3] else 1)
+
+for cat in sorted(cats.keys()):
+    cat2 = cat
+    if len(sys.argv) > 1:
+        cat2 = f"{sys.argv[1]} {cat}"
+    # default unit to `s`
+    if "|" not in cat2:
+        cat2 += "|s"
+    print(f"{cat2!r}: {cats[cat]:f}")

repl/.lake/packages/mathlib/scripts/detect_sha_changes.py

@@ -0,0 +1,104 @@
+"""
+This script is called by a github action to verify that commit SHAs used in porting are valid.
+It also produces links to the port-status webpage.
+
+Note that only the first 10 annotations created with this action are guaranteed to appear, so we
+produce the errors first.
+"""
+
+import dataclasses
+import re
+import sys
+from typing import Optional
+
+import git
+
+# upstream bug
+git.Git.CatFileContentStream.__next__ = git.Git.CatFileContentStream.next
+
+align_import_re = re.compile(
+    r'^#align_import ([^ ]*) from "(leanprover-community/[a-z]*)" ?@ ?"([0-9a-f]*)"')
+
+@dataclasses.dataclass(eq=True, frozen=True)
+class VersionInfo:
+    module: str
+    repo: Optional[str]
+    commit: Optional[str]
+    commit_line_no: Optional[int] = dataclasses.field(compare=False)
+
+    def to_commit(self):
+        try:
+            repo = git.Repo('port-repos/' + self.repo)
+        except git.exc.NoSuchPathError:
+            raise ValueError(f"Repo {self.repo} not recognized")
+        try:
+            repo.remotes.origin.fetch(self.commit, depth=1)
+        except Exception:
+            pass
+        return repo.commit(self.commit)
+
+def get_mathlib4_module_commit_infos(contents):
+    for i, line in enumerate(contents, 1):
+        m = align_import_re.match(line)
+        if m:
+            module = m.group(1)
+            repo = m.group(2)
+            commit = m.group(3)
+            yield VersionInfo(module, repo, commit, i)
+
+def get_mathlib4_module_commit_info_from_blob(blob: Optional[git.Blob]):
+    if blob is None:
+        return
+    yield from get_mathlib4_module_commit_infos(
+            l.decode('utf8') for l in blob.data_stream.stream)
+
+def encode_msg_text_for_github(msg):
+    # even though this is probably url quoting, we match the implementation at
+    # https://github.com/actions/toolkit/blob/af821474235d3c5e1f49cee7c6cf636abb0874c4/packages/core/src/command.ts#L36-L94
+    return msg.replace('%', '%25').replace('\r', '%0D').replace('\n', '%0A')
+
+if __name__ == '__main__':
+    repo = git.Repo('.')
+    base = repo.commit(sys.argv[1])
+    head = repo.commit(sys.argv[2])
+    any_errors = False
+
+    diff_infos = []
+    for diff in base.diff(head, paths=['Mathlib']):
+        a_info = set(get_mathlib4_module_commit_info_from_blob(diff.a_blob))
+        b_info = set(get_mathlib4_module_commit_info_from_blob(diff.b_blob))
+        if b_info <= a_info:
+            continue
+        diff_infos.append((diff, a_info, b_info))
+
+    all_refs = {}
+
+    # produce errors first
+    for diff, a_infos, b_infos in diff_infos:
+        for b_info in b_infos:
+            try:
+                b_info.to_commit()
+            except Exception as e:
+                print(f"::error file={diff.b_blob.path},line={b_info.commit_line_no},title=Invalid header::{encode_msg_text_for_github(str(e))}")
+                any_errors = True
+                continue
+
+    for diff, a_info, b_info in diff_infos:
+        same = a_info.intersection(b_info)
+        a_info -= same
+        b_info -= same
+        if a_info != {} and b_info != {}:
+            a_info_by_mod = {a.module: a for a in a_info}
+            b_info_by_mod = {b.module: b for b in b_info}
+            for k in set(a_info_by_mod.keys()) | set(b_info_by_mod.keys()):
+                a_info = a_info_by_mod.get(k, None)
+                b_info = b_info_by_mod.get(k, None)
+                if a_info is None or b_info is None:
+                    pass
+                elif a_info.module == b_info.module:
+                    mod_path = a_info.module.replace('.', '/')
+                    msg = f"See review instructions and diff at\nhttps://leanprover-community.github.io/mathlib-port-status/file/{mod_path}?range={a_info.commit}..{b_info.commit}"
+                    print(f"::notice file={diff.b_blob.path},line={b_info.commit_line_no},title=Synchronization::{encode_msg_text_for_github(msg)}")
+
+    if any_errors:
+        raise SystemExit("Setting a failure due to errors above")

repl/.lake/packages/mathlib/scripts/fix-comments.py

@@ -0,0 +1,162 @@
+#!/usr/bin/env python3
+
+import sys
+import os
+from pathlib import Path
+import subprocess
+import re
+
+if len(sys.argv) != 2 or not sys.argv[1].endswith('.lean'):
+    print("usage: fix-comments.py X.lean")
+    sys.exit(1)
+
+leanfile = sys.argv[1]
+
+is_clean = subprocess.run(
+    ['git', 'status', '--untracked-files=no', '--porcelain'],
+    capture_output=True,
+    check=True,
+    encoding='utf-8').stdout.rstrip()
+
+if is_clean != "":
+    print("Certain files tracked by git have uncommitted changes.\n")
+    os.system("git status --untracked-files=no")
+    print("\n")
+    s = input("Type y to continue. ")
+    if s != 'y':
+        sys.exit(1)
+
+root_dir = subprocess.run(
+    ['git', 'rev-parse', '--show-toplevel'],
+    capture_output=True,
+    check=True,
+    encoding='utf-8').stdout.rstrip()
+
+align_files = subprocess.run(
+    ['git', 'grep', '-l', '^#align'],
+    cwd=root_dir,
+    capture_output=True,
+    check=True,
+    encoding='utf-8')
+
+name_map = dict()
+for f in align_files.stdout.splitlines():
+    with open(os.path.join(root_dir, f), encoding="utf-8") as fh:
+        contents = fh.read()
+        for p in contents.split(sep='\n#align')[1:]:
+            n3, n4, *_ = p.split(maxsplit=2)
+            name_map[n3] = n4
+
+def replace_names(s):
+    # Terrible hack to treat `.` as a word character
+    # (to match qualified names)
+    s = s.replace('.', 'Ᾰ')
+    # re.DOTALL means that `.` can also match a newline.
+    # `\A` and `\Z` match only at the start/end of the string respectively.
+    w = re.findall(r'(?:\b|\A).+?(?:\b|\Z)', s, flags=re.DOTALL)
+    for i in range(len(w)):
+        w[i] = w[i].replace('Ᾰ', '.')
+        w[i] = name_map.get(w[i], w[i])
+    return ''.join(w)
+
+def process_backticked_names(s):
+    w = s.split(sep='`')
+    for i in range(len(w)):
+        if i % 2 == 1:
+            w[i] = replace_names(w[i])
+    return '`'.join(w)
+
+rewritten_contents = ''
+
+in_block_comment = False
+in_line_comment = False
+prev_char = None
+comment_so_far = None           # contains end marker but not begin marker
+
+def finish_comment():
+    global rewritten_contents
+    global in_block_comment
+    global in_line_comment
+    global comment_so_far
+    if comment_so_far is not None:
+        rewritten_contents += process_backticked_names(comment_so_far)
+        in_block_comment = False
+        in_line_comment = False
+        comment_so_far = None
+
+with open(leanfile, encoding="utf-8") as F:
+    while 1:
+        char = F.read(1)
+        if not char:
+            finish_comment()
+            break
+
+        if in_block_comment or in_line_comment:
+            comment_so_far = comment_so_far + char
+        else:
+            rewritten_contents += char
+
+        if in_block_comment and prev_char == '-' and char == '/':
+            finish_comment()
+
+        if in_line_comment and char == '\n':
+            finish_comment()
+
+        if comment_so_far is None and prev_char == '/' and char == '-':
+            in_block_comment = True
+            comment_so_far = ''
+
+        if comment_so_far is None and prev_char == '-' and char == '-':
+            in_line_comment = True
+            comment_so_far = ''
+
+        prev_char = char
+
+def mktree(path, sha, tree=True):
+    if path == Path('.'):
+        return sha
+    if tree:
+        inp = f"040000 tree {sha}\t{path.name}"
+    else:
+        inp = f"100644 blob {sha}\t{path.name}"
+    tree_sha = subprocess.run(
+        ['git', 'mktree'],
+        cwd=root_dir,
+        input=inp,
+        capture_output=True,
+        check=True,
+        encoding='utf8').stdout.rstrip()
+    return mktree(path.parent, tree_sha)
+
+path = Path(subprocess.run(
+    ['git', 'ls-files', '--full-name', leanfile],
+    capture_output=True,
+    check=True,
+    encoding='utf-8').stdout.rstrip())
+
+blob_sha = subprocess.run(
+    ['git', 'hash-object', '-w', '--stdin'],
+    input=rewritten_contents,
+    cwd=root_dir,
+    capture_output=True,
+    check=True,
+    encoding='utf-8').stdout.rstrip()
+
+tree_sha = mktree(path, blob_sha, tree=False)
+
+print(f"The script will now interactively suggest changes to {leanfile}.\n")
+s = input("Type y to continue. ")
+if s != 'y':
+    sys.exit(1)
+
+subprocess.run(['git', 'restore', '--patch', '--source=' + tree_sha, '--', leanfile], check=True)
+
+r = subprocess.run(['git', 'diff', '--quiet', leanfile])
+if r.returncode == 0:
+    pass
+elif r.returncode == 1:           # file was changed
+    print("\nPerhaps you would now like to run:")
+    print(f"git add {leanfile} && git commit -m 'auto: naming'")
+else:
+    # something went wrong
+    r.check_returncode()

repl/.lake/packages/mathlib/scripts/fix-line-breaks.py

@@ -0,0 +1,23 @@
+#!/usr/bin/env python3
+import sys
+from collections import deque
+
+lns = deque([], 2)
+with open(sys.argv[1], "r", encoding="utf-8", newline="\n") as f, \
+      open(sys.argv[2], "w", encoding="utf-8", newline="\n") as g:
+    for ln_raw in f:
+        ln = ln_raw.strip("\n")
+        lns.append(ln)
+        if len(lns) <= 1:
+            continue
+        if lns[1].lstrip() == "by" and len(lns[0]) < 98 and not lns[0].lstrip().startswith("--"):
+            lns.pop()
+            lns[0] += " by"
+        elif lns[1].lstrip() == "where" and len(lns[0]) < 95 and not lns[0].lstrip().startswith("--"):
+            lns.pop()
+            lns[0] += " where"
+        else:
+            print(lns[0], file=g)
+    lns.popleft()
+    for ln in lns:
+        print(ln, file=g)

repl/.lake/packages/mathlib/scripts/fix-lints.py

@@ -0,0 +1,48 @@
+#!/usr/bin/env python3
+
+import sys
+import os
+import subprocess
+import shutil
+
+def getpos(line):
+    _, line, col, _ = line.split(sep=':', maxsplit=3)
+    return int(line), int(col)
+
+if len(sys.argv) != 2 or not sys.argv[1].endswith('.lean'):
+    print("usage: fix-lints.py Mathlib/A/B/C.lean")
+    sys.exit(1)
+
+leanfile = sys.argv[1]
+leanmodule = leanfile[:-5].replace('/', '.')
+
+# try to build
+log = subprocess.run(
+    ['lake', 'build', leanmodule],
+    capture_output=True, encoding='utf8')
+if log.returncode == 0:
+    print("no errors 🎉")
+    exit(0)
+
+shutil.copyfile(leanfile, leanfile + '.bak')
+
+with open(leanfile + '.bak', encoding='utf8') as fp:
+    f = list(fp)
+count = 0
+for l in reversed(log.stderr.splitlines()):
+    if 'linter.unusedVariables' in l:
+        line, col = getpos(l)
+        f[line-1] = f[line-1][0:col] + '_' + f[line-1][col:]
+        count += 1
+    elif 'linter.unnecessarySeqFocus' in l:
+        line, col = getpos(l)
+        f[line-1] = f[line-1][0:col].rstrip() + ';' + f[line-1][col+3:]
+        count += 1
+    else:
+        print(l, file=sys.stderr)
+
+print(f'Fixed {count} warnings', file=sys.stderr)
+
+with open(leanfile, 'w', encoding='utf8') as fp:
+    fp.write(''.join(f))
+os.remove(leanfile + '.bak')

repl/.lake/packages/mathlib/scripts/lint-style.py

@@ -0,0 +1,451 @@
+#!/usr/bin/env python3
+"""
+Lint a file or files from mathlib for style.
+
+Sample usage:
+
+    $ ./scripts/lint-style.py $(find Mathlib -name '*.lean')
+
+which will lint all of the Lean files in the specified directories.
+
+The resulting error output will contain one line for each style error
+encountered that isn't in the list of allowed / ignored style exceptions.
+
+Paths with no errors will not appear in the output, and the script will
+exit with successful return code if there are no errors encountered in
+any provided paths.
+
+Paths emitted in the output will match the paths provided on the
+command line for any files containing errors -- in particular, linting
+a relative path (like ``Mathlib/Foo/Bar.lean``) will produce errors
+that contain the relative path, whilst linting absolute paths (like
+``/root/mathlib4/Mathlib/Foo/Bar.lean``) will produce errors with the
+absolute path.
+
+This script can also be used to regenerate the list of allowed / ignored style
+exceptions by redirecting the output to ``style-exceptions.txt``. Use:
+
+    $ ./scripts/update-style-exceptions.sh
+
+to perform this update.
+"""
+
+# TODO: This is adapted from the linter for mathlib3. It should be rewritten in Lean.
+
+from pathlib import Path
+import sys
+import re
+import shutil
+
+ERR_COP = 0 # copyright header
+ERR_MOD = 2 # module docstring
+ERR_LIN = 3 # line length
+ERR_OPT = 6 # set_option
+ERR_AUT = 7 # malformed authors list
+ERR_TAC = 9 # imported Mathlib.Tactic
+ERR_IBY = 11 # isolated by
+ERR_DOT = 12 # isolated or low focusing dot
+ERR_SEM = 13 # the substring " ;"
+ERR_WIN = 14 # Windows line endings "\r\n"
+ERR_TWS = 15 # trailing whitespace
+ERR_CLN = 16 # line starts with a colon
+ERR_IND = 17 # second line not correctly indented
+ERR_ARR = 18 # space after "←"
+ERR_NUM_LIN = 19 # file is too large
+ERR_NSP = 20 # non-terminal simp
+
+exceptions = []
+
+SCRIPTS_DIR = Path(__file__).parent.resolve()
+ROOT_DIR = SCRIPTS_DIR.parent
+
+
+with SCRIPTS_DIR.joinpath("style-exceptions.txt").open(encoding="utf-8") as f:
+    for exline in f:
+        filename, _, _, _, _, errno, *extra = exline.split()
+        path = ROOT_DIR / filename
+        if errno == "ERR_COP":
+            exceptions += [(ERR_COP, path, None)]
+        elif errno == "ERR_MOD":
+            exceptions += [(ERR_MOD, path, None)]
+        elif errno == "ERR_LIN":
+            exceptions += [(ERR_LIN, path, None)]
+        elif errno == "ERR_OPT":
+            exceptions += [(ERR_OPT, path, None)]
+        elif errno == "ERR_AUT":
+            exceptions += [(ERR_AUT, path, None)]
+        elif errno == "ERR_TAC":
+            exceptions += [(ERR_TAC, path, None)]
+        elif errno == "ERR_NUM_LIN":
+            exceptions += [(ERR_NUM_LIN, path, extra[1])]
+        else:
+            print(f"Error: unexpected errno in style-exceptions.txt: {errno}")
+            sys.exit(1)
+
+new_exceptions = False
+
+def annotate_comments(enumerate_lines):
+    """
+    Take a list of tuples of enumerated lines of the form
+    (line_number, line, ...)
+    and return a list of
+    (line_number, line, ..., True/False)
+    where lines have True attached when they are in comments.
+    """
+    nesting_depth = 0 # We're in a comment when `nesting_depth > 0`.
+    starts_in_comment = False # Whether we're in a comment when starting the line.
+    for line_nr, line, *rem in enumerate_lines:
+        # We assume multiline comments do not begin or end within single-line comments.
+        if line == "\n" or line.lstrip().startswith("--"):
+            yield line_nr, line, *rem, True
+            continue
+        # We assume that "/-/" and "-/-" never occur outside of "--" comments.
+        # We assume that we do not encounter "... -/ <term> /- ...".
+        # We also don't account for "/-" and "-/" appearing in strings.
+        starts_in_comment = (nesting_depth > 0)
+        nesting_depth = nesting_depth + line.count("/-") - line.count("-/")
+        in_comment = (starts_in_comment or line.lstrip().startswith("/-")) and \
+            (nesting_depth > 0 or line.rstrip().endswith("-/"))
+        yield line_nr, line, *rem, in_comment
+
+def annotate_strings(enumerate_lines):
+    """
+    Take a list of tuples of enumerated lines of the form
+    (line_number, line, ...)
+    and return a list of
+    (line_number, line, ..., True/False)
+    where lines have True attached when they are in strings.
+    """
+    in_string = False
+    in_comment = False
+    for line_nr, line, *rem in enumerate_lines:
+        # ignore comment markers inside string literals
+        if not in_string:
+            if "/-" in line:
+                in_comment = True
+            if "-/" in line:
+                in_comment = False
+        # ignore quotes inside comments
+        if not in_comment:
+            # crude heuristic: if the number of non-escaped quote signs is odd,
+            # we're starting / ending a string literal
+            if line.count("\"") - line.count("\\\"") % 2 == 1:
+                in_string = not in_string
+            # if there are quote signs in this line,
+            # a string literal probably begins and / or ends here,
+            # so we skip this line
+            if line.count("\"") > 0:
+                yield line_nr, line, *rem, True
+                continue
+            if in_string:
+                yield line_nr, line, *rem, True
+                continue
+        yield line_nr, line, *rem, False
+
+def set_option_check(lines, path):
+    errors = []
+    newlines = []
+    for line_nr, line, in_comment, in_string in annotate_strings(annotate_comments(lines)):
+        if line.strip().startswith('set_option') and not in_comment and not in_string:
+            option_prefix = line.strip().split(' ', 2)[1].split('.', 1)[0]
+            # forbidden options: pp, profiler, trace
+            if option_prefix in {'pp', 'profiler', 'trace'}:
+                errors += [(ERR_OPT, line_nr, path)]
+                # skip adding this line to newlines so that we suggest removal
+                continue
+        newlines.append((line_nr, line))
+    return errors, newlines
+
+def line_endings_check(lines, path):
+    errors = []
+    newlines = []
+    for line_nr, line in lines:
+        if "\r\n" in line:
+            errors += [(ERR_WIN, line_nr, path)]
+            line = line.replace("\r\n", "\n")
+        if line.endswith(" \n"):
+            errors += [(ERR_TWS, line_nr, path)]
+            line = line.rstrip() + "\n"
+        newlines.append((line_nr, line))
+    return errors, newlines
+
+def four_spaces_in_second_line(lines, path):
+    # TODO: also fix the space for all lines before ":=", right now we only fix the line after
+    # the first line break
+    errors = []
+    # We never alter the first line, as it does not occur as next_line in the iteration over the
+    # zipped lines below, hence we add it here
+    newlines = [lines[0]]
+    annotated_lines = list(annotate_comments(lines))
+    for (_, line, is_comment), (next_line_nr, next_line, _) in zip(annotated_lines,
+                                                                   annotated_lines[1:]):
+        # Check if the current line matches "(lemma|theorem) .* :"
+        new_next_line = next_line
+        if (not is_comment) and re.search(r"^(protected )?(def|lemma|theorem) (?!.*:=).*(where)?$",
+                                          line):
+            # Calculate the number of spaces before the first non-space character in the next line
+            stripped_next_line = next_line.lstrip()
+            if not (next_line == '\n' or next_line.startswith("#") or stripped_next_line.startswith("--")):
+                num_spaces = len(next_line) - len(stripped_next_line)
+                # The match with "| " could potentially match with a different usage of the same
+                # symbol, e.g. some sort of norm. In that case a space is not necessary, so
+                # looking for "| " should be enough.
+                if stripped_next_line.startswith("| ") or line.endswith("where\n"):
+                    # Check and fix if the number of leading space is not 2
+                    if num_spaces != 2:
+                        errors += [(ERR_IND, next_line_nr, path)]
+                        new_next_line = ' ' * 2 + stripped_next_line
+                # Check and fix if the number of leading spaces is not 4
+                else:
+                    if num_spaces != 4:
+                        errors += [(ERR_IND, next_line_nr, path)]
+                        new_next_line = ' ' * 4 + stripped_next_line
+        newlines.append((next_line_nr, new_next_line))
+    return errors, newlines
+
+def nonterminal_simp_check(lines, path):
+    errors = []
+    newlines = []
+    annotated_lines = list(annotate_comments(lines))
+    for (line_nr, line, is_comment), (_, next_line, _) in zip(annotated_lines,
+                                                              annotated_lines[1:]):
+        # Check if the current line matches whitespace followed by "simp"
+        new_line = line
+        # TODO it would be better to use a regex like r"^\s*simp( \[.*\])?( at .*)?$" and thereby
+        # catch all possible simp invocations. Adding this will require more initial cleanup or
+        # nolint.
+        if (not is_comment) and re.search(r"^\s*simp$", line):
+            # Calculate the number of spaces before the first non-space character in the line
+            num_spaces = len(line) - len(line.lstrip())
+            # Calculate the number of spaces before the first non-space character in the next line
+            stripped_next_line = next_line.lstrip()
+            if not (next_line == '\n' or next_line.startswith("#") or stripped_next_line.startswith("--") or "rfl" in next_line):
+                num_next_spaces = len(next_line) - len(stripped_next_line)
+                # Check if the number of leading spaces is the same
+                if num_spaces == num_next_spaces:
+                    # If so, the simp is nonterminal
+                    errors += [(ERR_NSP, line_nr, path)]
+                    new_line = line.replace("simp", "simp?")
+        newlines.append((line_nr, new_line))
+    newlines.append(lines[-1])
+    return errors, newlines
+
+def long_lines_check(lines, path):
+    errors = []
+    # TODO: find a good way to break long lines
+    # TODO: some string literals (in e.g. tactic output messages) can be excepted from this rule
+    for line_nr, line in lines:
+        if "http" in line or "#align" in line:
+            continue
+        if len(line) > 101:
+            errors += [(ERR_LIN, line_nr, path)]
+    return errors, lines
+
+def import_only_check(lines, path):
+    for _, line, is_comment in annotate_comments(lines):
+        if is_comment:
+            continue
+        imports = line.split()
+        if imports[0] == "#align_import":
+            continue
+        if imports[0] != "import":
+            return False
+    return True
+
+def regular_check(lines, path):
+    errors = []
+    copy_started = False
+    copy_done = False
+    copy_start_line_nr = 1
+    copy_lines = ""
+    for line_nr, line in lines:
+        if not copy_started and line == "\n":
+            errors += [(ERR_COP, copy_start_line_nr, path)]
+            continue
+        if not copy_started and line == "/-\n":
+            copy_started = True
+            copy_start_line_nr = line_nr
+            continue
+        if not copy_started:
+            errors += [(ERR_COP, line_nr, path)]
+        if copy_started and not copy_done:
+            copy_lines += line
+            if "Author" in line:
+                # Validating names is not a reasonable thing to do,
+                # so we just look for the two common variations:
+                # using ' and ' between names, and a '.' at the end of line.
+                if ((not line.startswith("Authors: ")) or
+                    ("  " in line) or
+                    (" and " in line) or
+                    (line[-2] == '.')):
+                    errors += [(ERR_AUT, line_nr, path)]
+            if line == "-/\n":
+                if ((not "Copyright" in copy_lines) or
+                    (not "Apache" in copy_lines) or
+                    (not "Authors: " in copy_lines)):
+                    errors += [(ERR_COP, copy_start_line_nr, path)]
+                copy_done = True
+            continue
+        if copy_done and line == "\n":
+            continue
+        words = line.split()
+        if words[0] != "import" and words[0] != "--" and words[0] != "/-!" and words[0] != "#align_import":
+            errors += [(ERR_MOD, line_nr, path)]
+            break
+        if words[0] == "/-!":
+            break
+    return errors, lines
+
+def banned_import_check(lines, path):
+    errors = []
+    for line_nr, line, is_comment in annotate_comments(lines):
+        if is_comment:
+            continue
+        imports = line.split()
+        if imports[0] != "import":
+            break
+        if imports[1] in ["Mathlib.Tactic"]:
+            errors += [(ERR_TAC, line_nr, path)]
+    return errors, lines
+
+def isolated_by_dot_semicolon_check(lines, path):
+    errors = []
+    newlines = []
+    for line_nr, line in lines:
+        if line.strip() == "by":
+            # We excuse those "by"s following a comma or ", fun ... =>", since generally hanging "by"s
+            # should not be used in the second or later arguments of a tuple/anonymous constructor
+            # See https://github.com/leanprover-community/mathlib4/pull/3825#discussion_r1186702599
+            prev_line = lines[line_nr - 2][1].rstrip()
+            if not prev_line.endswith(",") and not re.search(", fun [^,]* (=>|↦)$", prev_line):
+                errors += [(ERR_IBY, line_nr, path)]
+        if line.lstrip().startswith(". "):
+            errors += [(ERR_DOT, line_nr, path)]
+            line = line.replace(". ", "· ", 1)
+        if line.strip() in (".", "·"):
+            errors += [(ERR_DOT, line_nr, path)]
+        if " ;" in line:
+            errors += [(ERR_SEM, line_nr, path)]
+            line = line.replace(" ;", ";")
+        if line.lstrip().startswith(":"):
+            errors += [(ERR_CLN, line_nr, path)]
+        newlines.append((line_nr, line))
+    return errors, newlines
+
+def left_arrow_check(lines, path):
+    errors = []
+    newlines = []
+    for line_nr, line, is_comment, in_string in annotate_strings(annotate_comments(lines)):
+        if is_comment or in_string:
+            newlines.append((line_nr, line))
+            continue
+        # Allow "←" to be followed by "%" or "`", but not by "`(" or "``(" (since "`()" and "``()"
+        # are used for syntax quotations). Otherwise, insert a space after "←".
+        new_line = re.sub(r'←(?:(?=``?\()|(?![%`]))(\S)', r'← \1', line)
+        if new_line != line:
+            errors += [(ERR_ARR, line_nr, path)]
+        newlines.append((line_nr, new_line))
+    return errors, newlines
+
+def output_message(path, line_nr, code, msg):
+    if len(exceptions) == 0:
+        # we are generating a new exceptions file
+        # filename first, then line so that we can call "sort" on the output
+        print(f"{path} : line {line_nr} : {code} : {msg}")
+    else:
+        if code.startswith("ERR"):
+            msg_type = "error"
+        if code.startswith("WRN"):
+            msg_type = "warning"
+        # We are outputting for github. We duplicate path, line_nr and code,
+        # so that they are also visible in the plaintext output.
+        print(f"::{msg_type} file={path},line={line_nr},code={code}::{path}#L{line_nr}: {code}: {msg}")
+
+def format_errors(errors):
+    global new_exceptions
+    for errno, line_nr, path in errors:
+        if (errno, path.resolve(), None) in exceptions:
+            continue
+        new_exceptions = True
+        if errno == ERR_COP:
+            output_message(path, line_nr, "ERR_COP", "Malformed or missing copyright header")
+        if errno == ERR_MOD:
+            output_message(path, line_nr, "ERR_MOD", "Module docstring missing, or too late")
+        if errno == ERR_LIN:
+            output_message(path, line_nr, "ERR_LIN", "Line has more than 100 characters")
+        if errno == ERR_OPT:
+            output_message(path, line_nr, "ERR_OPT", "Forbidden set_option command")
+        if errno == ERR_AUT:
+            output_message(path, line_nr, "ERR_AUT", "Authors line should look like: 'Authors: Jean Dupont, Иван Иванович Иванов'")
+        if errno == ERR_TAC:
+            output_message(path, line_nr, "ERR_TAC", "Files in mathlib cannot import the whole tactic folder")
+        if errno == ERR_IBY:
+            output_message(path, line_nr, "ERR_IBY", "Line is an isolated 'by'")
+        if errno == ERR_DOT:
+            output_message(path, line_nr, "ERR_DOT", "Line is an isolated focusing dot or uses . instead of ·")
+        if errno == ERR_SEM:
+            output_message(path, line_nr, "ERR_SEM", "Line contains a space before a semicolon")
+        if errno == ERR_WIN:
+            output_message(path, line_nr, "ERR_WIN", "Windows line endings (\\r\\n) detected")
+        if errno == ERR_TWS:
+            output_message(path, line_nr, "ERR_TWS", "Trailing whitespace detected on line")
+        if errno == ERR_CLN:
+            output_message(path, line_nr, "ERR_CLN", "Put : and := before line breaks, not after")
+        if errno == ERR_IND:
+            output_message(path, line_nr, "ERR_IND", "If the theorem/def statement requires multiple lines, indent it correctly (4 spaces or 2 for `|`)")
+        if errno == ERR_ARR:
+            output_message(path, line_nr, "ERR_ARR", "Missing space after '←'.")
+        if errno == ERR_NSP:
+            output_message(path, line_nr, "ERR_NSP", "Non-terminal simp. Replace with `simp?` and use the suggested output")
+
+def lint(path, fix=False):
+    global new_exceptions
+    with path.open(encoding="utf-8", newline="") as f:
+        # We enumerate the lines so that we can report line numbers in the error messages correctly
+        # we will modify lines as we go, so we need to keep track of the original line numbers
+        lines = f.readlines()
+        enum_lines = enumerate(lines, 1)
+        newlines = enum_lines
+        for error_check in [line_endings_check,
+                            four_spaces_in_second_line,
+                            long_lines_check,
+                            isolated_by_dot_semicolon_check,
+                            set_option_check,
+                            left_arrow_check,
+                            nonterminal_simp_check]:
+            errs, newlines = error_check(newlines, path)
+            format_errors(errs)
+
+        if not import_only_check(newlines, path):
+            # Check for too long files: either longer than 1500 lines, or not covered by an exception.
+            # Each exception contains a "watermark". If the file is longer than that, we also complain.
+            if len(lines) > 1500:
+                ex = [e for e in exceptions if e[1] == path.resolve()]
+                if ex:
+                    (_ERR_NUM, _path, watermark) = list(ex)[0]
+                    assert int(watermark) > 500 # protect against parse error
+                    is_too_long = len(lines) > int(watermark)
+                else:
+                    is_too_long = True
+                if is_too_long:
+                    new_exceptions = True
+                    # add up to 200 lines of slack, so simple PRs don't trigger this right away
+                    watermark = len(lines) // 100 * 100 + 200
+                    output_message(path, 1, "ERR_NUM_LIN", f"{watermark} file contains {len(lines)} lines, try to split it up")
+            errs, newlines = regular_check(newlines, path)
+            format_errors(errs)
+            errs, newlines = banned_import_check(newlines, path)
+            format_errors(errs)
+    # if we haven't been asked to fix errors, or there are no errors or no fixes, we're done
+    if fix and new_exceptions and enum_lines != newlines:
+        path.with_name(path.name + '.bak').write_text("".join(l for _,l in newlines), encoding = "utf8")
+        shutil.move(path.with_name(path.name + '.bak'), path)
+
+fix = "--fix" in sys.argv
+argv = (arg for arg in sys.argv[1:] if arg != "--fix")
+
+for filename in argv:
+    lint(Path(filename), fix=fix)
+
+if new_exceptions:
+    exit(1)

repl/.lake/packages/mathlib/scripts/make_port_status.py

@@ -0,0 +1,218 @@
+#!/usr/bin/env python3
+
+import pytz
+import datetime
+import github
+import os
+import re
+import requests
+import subprocess
+import sys
+import yaml
+import networkx as nx
+from collections import defaultdict
+from pathlib import Path
+
+# Must run from root of mathlib4 directory.
+
+if not os.path.exists('port-repos/mathlib'):
+    print("Make sure you are in the root of the mathlib4 directory")
+    print("and have checked out mathlib under port-repos/mathlib.")
+    sys.exit(1)
+
+GITHUB_TOKEN_FILE = 'port-repos/github-token'
+github_token = open(GITHUB_TOKEN_FILE).read().strip()
+
+mathlib3_root = 'port-repos/mathlib'
+mathlib4_root = './'
+
+source_module_re = re.compile(r"^! .*source module (.*)$")
+commit_re = re.compile(r"^! (leanprover-community/[a-z]*) commit ([0-9a-f]*)")
+import_re = re.compile(r"^import ([^ ]*)")
+
+align_import_re = re.compile(
+    r'^#align_import ([^ ]*) from "(leanprover-community/[a-z]*)" ?@ ?"([0-9a-f]*)"')
+
+def mk_label(path: Path) -> str:
+    rel = path.relative_to(Path(mathlib3_root))
+    rel = Path(*rel.parts[1:])
+    return str(rel.with_suffix('')).replace(os.sep, '.')
+
+paths = []
+for path in Path(mathlib3_root).glob('**/*.lean'):
+    if path.relative_to(mathlib3_root).parts[0] not in ['src', 'archive', 'counterexamples']:
+        continue
+    if path.relative_to(mathlib3_root).parts[1] in ['tactic', 'meta']:
+        continue
+    paths.append(path)
+
+graph = nx.DiGraph()
+for path in paths:
+    graph.add_node(mk_label(path))
+
+for path in paths:
+    label = mk_label(path)
+    for line in path.read_text().split('\n'):
+        m = import_re.match(line)
+        if m:
+            imported = m.group(1)
+            if imported.startswith('tactic.') or imported.startswith('meta.') or imported.startswith('.'):
+                continue
+            if imported not in graph.nodes:
+                if imported + '.default' in graph.nodes:
+                    imported = imported + '.default'
+                else:
+                    imported = imported
+            graph.add_edge(imported, label)
+
+def get_mathlib4_module_commit_info(contents):
+    module = repo = commit = None
+    for line in contents.split('\n'):
+        m = align_import_re.match(line)
+        if m:
+            module = m.group(1)
+            repo = m.group(2)
+            commit = m.group(3)
+            break
+        m = source_module_re.match(line)
+        if m:
+            module = m.group(1)
+        m = commit_re.match(line)
+        if m:
+            repo = m.group(1)
+            commit = m.group(2)
+    return module, repo, commit
+
+# contains ported files
+# lean 3 module name -> { mathlib4_file, mathlib3_hash }
+data = dict()
+for path4 in Path(mathlib4_root).glob('**/*.lean'):
+    # we definitely do not want to look in `port-repos` here!
+    if path4.relative_to(mathlib4_root).parts[0] not in ('Mathlib', 'Archive', 'Counterexamples'):
+        continue
+    module, repo, commit = get_mathlib4_module_commit_info(path4.read_text())
+    if module is None:
+        continue
+
+    if commit is None:
+        print(f"Commit is None for module: {module}")
+        continue
+
+    log = subprocess.run(
+        ['git', 'log', '--oneline', str(path4)],
+        capture_output=True)
+    pr_matches = re.search(r'#([0-9]+)\)$', log.stdout.decode().splitlines()[-1])
+    if pr_matches:
+        mathlib4_pr = int(pr_matches.groups()[0])
+    else:
+        mathlib4_pr = None
+
+    data[module] = {
+        'mathlib4_file': str(path4.relative_to(mathlib4_root)),
+        'mathlib4_pr': mathlib4_pr,
+        'source': dict(repo=repo, commit=commit)
+    }
+
+    graph.add_node(module)
+
+prs = {}
+fetch_args = ['git', 'fetch', 'origin']
+nums = []
+sync_prs = defaultdict(set)
+mathlib4repo = github.Github(github_token).get_repo("leanprover-community/mathlib4")
+for pr in mathlib4repo.get_pulls(state='open'):
+    if pr.created_at < datetime.datetime(2022, 12, 1, 0, 0, 0, tzinfo=pytz.UTC):
+        continue
+    if 'no-source-header' in (l.name for l in pr.labels):
+        continue
+    if 'mathlib3-pair' in (l.name for l in pr.labels):
+        for file in (f.filename for f in pr.get_files()):
+            sync_prs[file].add(pr.number)
+    num = pr.number
+    nums.append(num)
+    prs[num] = pr
+    fetch_args.append(f'pull/{num}/head:port-status-pull/{num}')
+
+os.system("git branch -D $(git branch --list 'port-status-pull/*')")
+subprocess.run(fetch_args)
+
+prs_of_import = {}
+for num in nums:
+    p = subprocess.run(
+        ['git', 'diff', '--name-only', '--diff-filter=A',
+         f'origin/master...port-status-pull/{num}'],
+        capture_output=True)
+    for l in p.stdout.decode().splitlines():
+        f = subprocess.run(
+            ['git', 'cat-file', 'blob', f'port-status-pull/{num}:{l}'],
+            capture_output=True)
+        import_, repo, commit = get_mathlib4_module_commit_info(f.stdout.decode(encoding='utf8', errors='replace'))
+        prs_of_import.setdefault(import_, []).append({'pr': num, 'repo': repo, 'commit': commit, 'fname': l})
+
+COMMENTS_URL = "https://raw.githubusercontent.com/wiki/leanprover-community/mathlib4/port-comments.md"
+comments_dict = yaml.safe_load(requests.get(COMMENTS_URL).content.replace(b"```", b""))
+
+yaml_dict = {}
+new_yaml_dict = {}
+for node in sorted(graph.nodes):
+    if node in data:
+        new_status = dict(
+            ported=True,
+            mathlib4_file=data[node]['mathlib4_file'],
+            mathlib4_pr=data[node]['mathlib4_pr'],
+            source=data[node]['source']
+        )
+        _sync_prs = [
+            dict(
+                num=sync_pr_num,
+                labels=[dict(name=l.name, color=l.color) for l in prs[sync_pr_num].labels]
+            )
+            for sync_pr_num in sync_prs[data[node]['mathlib4_file']]
+        ]
+        if _sync_prs:
+            new_status.update(mathlib4_sync_prs=_sync_prs)
+        pr_status = f"mathlib4#{data[node]['mathlib4_pr']}" if data[node]['mathlib4_pr'] is not None else "_"
+        sha = data[node]['source']['commit'] if data[node]['source']['repo'] == 'leanprover-community/mathlib' else "_"
+
+        status = f"Yes {pr_status} {sha}"
+    else:
+        new_status = dict(ported=False)
+        status = f'No'
+        if node in prs_of_import:
+            pr_info = prs_of_import[node][0]
+            if pr_info['commit'] is None:
+                print('PR seems to be missing a source header', node, pr_info)
+                assert(False)
+            new_status.update(
+                mathlib4_pr=pr_info['pr'],
+                mathlib4_file=pr_info['fname'],
+                source=dict(repo=pr_info['repo'], commit=pr_info['commit']))
+            labels = [{'name': l.name, 'color': l.color} for l in prs[pr_info['pr']].labels]
+            if labels:
+                new_status.update(labels=labels)
+            sha = pr_info['commit'] if pr_info['repo'] == 'leanprover-community/mathlib' else "_"
+            status += f" mathlib4#{pr_info['pr']} {sha}"
+    try:
+        comment_data = comments_dict[node]
+    except KeyError:
+        pass
+    else:
+        if isinstance(comment_data, str):
+            # old comment format
+            comment_data = dict(message=comment_data)
+        # new comment format
+        status += ' ' + comment_data['message']
+        new_status.update(comment=comment_data)
+    yaml_dict[node] = status
+    new_yaml_dict[node] = new_status
+
+DO_NOT_EDIT_MESSAGE = """
+# Do not edit this file.
+# If you want to add free-form comments about files that don't have PRs yet,
+# edit https://github.com/leanprover-community/mathlib4/wiki/port-comments/_edit instead.
+""" + ("\n" * 37)
+
+with open('port_status.yaml', 'w') as f:
+    f.write(DO_NOT_EDIT_MESSAGE + "```\n" + yaml.dump(yaml_dict) + "```\n")
+with open('port_status_new.yaml', 'w') as f:
+    f.write(DO_NOT_EDIT_MESSAGE + "```\n" + yaml.dump(new_yaml_dict) + "```\n")

repl/.lake/packages/mathlib/scripts/polyrith_sage.py

@@ -0,0 +1,101 @@
+# This file is part of the `polyrith` tactic in `src/tactic/polyrith.lean`.
+# It interfaces between Lean and the Sage web interface.
+
+import requests
+import json
+import sys
+from os.path import join, dirname
+
+# These functions are used to format the output of Sage for parsing in Lean.
+# They are stored here as a string since they are passed to Sage via the web API.
+with open(join(dirname(__file__), "polyrith_sage_helper.py"), encoding='utf8') as f:
+    polynomial_formatting_functions = f.read()
+
+# future extensions may change behavior depending on the base type
+def type_str(type):
+    return "QQ"
+
+def create_query(type: str, n_vars: int, eq_list, goal_type):
+    """ Create a query to invoke Sage's `MPolynomial_libsingular.lift`. See
+    https://github.com/sagemath/sage/blob/f8df80820dc7321dc9b18c9644c3b8315999670b/src/sage/rings/polynomial/multi_polynomial_libsingular.pyx#L4472-L4518
+    for a description of this method. """
+    var_list = [f"var{i}" for i in range(n_vars)] + ['aux']
+    query = f'''
+if {n_vars!r} != 0:
+    P = PolynomialRing({type_str(type)}, {var_list})
+    [{", ".join(var_list)}] = P.gens()
+    p = P({goal_type})
+    gens = {eq_list} + [1 - p*aux]
+    I = P.ideal(gens)
+    coeffs = P(1).lift(I)
+    power = max(cf.degree(aux) for cf in coeffs)
+    coeffs = [P(cf.subs(aux = 1/p)*p^power) for cf in coeffs[:int(-1)]]
+    print(str(power)+';'+serialize_polynomials(coeffs))
+else:
+    # workaround for a Sage shortcoming with `n_vars = 0`,
+    # `TypeError: no conversion of this ring to a Singular ring defined`
+    # In this case, there is no need to look for membership in the *radical*;
+    # we just check for membership in the ideal, and return exponent 1 
+    # if coefficients are found.
+    P = PolynomialRing({type_str(type)}, 'var', 1)
+    p = P({goal_type})
+    I = P.ideal({eq_list})
+    coeffs = p.lift(I)
+    print('1;'+serialize_polynomials(coeffs))
+'''
+    return query
+
+class EvaluationError(Exception):
+    def __init__(self, ename, evalue, message='Error in Sage communication'):
+        self.ename = ename
+        self.evalue = evalue
+        self.message = message
+        super().__init__(self.message)
+
+def parse_response(resp: str) -> str:
+    exp, data = resp.split(';', 1)
+    return dict(power=int(exp), coeffs=json.loads(data))
+
+
+def evaluate_in_sage(query: str) -> str:
+    data = {'code': query}
+    headers = {'content-type': 'application/x-www-form-urlencoded'}
+    response = requests.post('https://sagecell.sagemath.org/service', data, headers=headers).json()
+    if response['success']:
+        return parse_response(response.get('stdout'))
+    elif 'execute_reply' in response and 'ename' in response['execute_reply'] and 'evalue' in response['execute_reply']:
+        raise EvaluationError(response['execute_reply']['ename'], response['execute_reply']['evalue'])
+    else:
+        raise Exception(response)
+
+def main():
+    '''The system args contain the following:
+    0 - the path to this python file
+    1 - a string containing "true" or "false" depending on whether polyrith was called with trace enabled
+    2 - a string representing the base type of the target
+    3 - the number of variables used
+    4 - a list of the polynomial hypotheses/proof terms in terms of the variables
+    5 - a single polynomial representing the target
+
+    This returns a json object with format:
+    ```
+    { success: bool,
+      data: Optional[list[str]],
+      trace: Optional[str],
+      name: Optional[str],
+      value: Optional[str] }
+    ```
+    '''
+    command = create_query(sys.argv[2], int(sys.argv[3]), sys.argv[4], sys.argv[5])
+    final_query = polynomial_formatting_functions + "\n" + command
+    if sys.argv[1] == 'true': # trace dry run enabled
+        output = dict(success=True, trace=command)
+    else:
+        try:
+            output = dict(success=True, data=evaluate_in_sage(final_query))
+        except EvaluationError as e:
+            output = dict(success=False, name=e.ename, value=e.evalue)
+    print(json.dumps(output))
+
+if __name__ == "__main__":
+    main()

repl/.lake/packages/mathlib/scripts/polyrith_sage_helper.py

@@ -0,0 +1,13 @@
+# this file will be run by the remote sage server, so should not import local files.
+from typing import Iterable
+
+def q_arr(coeff: QQ) -> str:
+    return "[" + str(coeff.numerator()) + "," + str(coeff.denominator()) + "]"
+
+def arr(args: Iterable[str]) -> str:
+    return "[" + ",".join(args) + "]"
+
+def serialize_polynomials(coeffs) -> str:
+    return arr(
+        arr(arr([arr(arr([str(t[0]), str(t[1])]) for t in etuple.sparse_iter()), q_arr(coeff)])
+        for etuple, coeff in c.dict().items()) for c in coeffs)

repl/.lake/packages/mathlib/scripts/yaml_check.py

@@ -0,0 +1,69 @@
+"""
+This file is copied from the mathlib3 file of the same name.
+It reads in the three yaml files, and translates them to simpler json files that are easier to
+process in Lean.
+"""
+from typing import Dict, Optional, Union, Tuple, List
+import yaml
+import json
+import sys
+
+TieredDict = Dict[str, Union[Optional[str], 'TieredDict']]
+
+def tiered_extract(db: TieredDict) -> List[Tuple[List[str], str]]:
+  """From a nested dictionary, return a list of (key_path, values)
+  of the deepest level."""
+  out = []
+  for name, entry in db.items():
+    if isinstance(entry, dict):
+      for subname, value in tiered_extract(entry):
+        out.append(([name] + subname, value))
+    else:
+      if entry and '/' not in entry:
+        out.append(([name], entry))
+  return out
+
+def flatten_names(data: List[Tuple[List[str], str]]) -> List[Tuple[str, str]]:
+  return [(' :: '.join(id), v) for id, v in data]
+
+def print_list(fn: str, pairs: List[Tuple[str, str]]) -> None:
+  with open(fn, 'w', encoding='utf8') as out:
+    for (id, val) in pairs:
+      out.write(f'{id}\n{val.strip()}\n\n')
+
+hundred_yaml = sys.argv[1]
+overview_yaml = sys.argv[2]
+undergrad_yaml = sys.argv[3]
+
+with open(hundred_yaml, 'r', encoding='utf8') as hy:
+  hundred = yaml.safe_load(hy)
+with open(overview_yaml, 'r', encoding='utf8') as hy:
+  overview = yaml.safe_load(hy)
+with open(undergrad_yaml, 'r', encoding='utf8') as hy:
+  undergrad = yaml.safe_load(hy)
+
+hundred_decls:List[Tuple[str, str]] = []
+
+for index, entry in hundred.items():
+  title = entry['title']
+  if 'decl' in entry:
+    hundred_decls.append((f'{index} {title}', entry['decl']))
+  elif 'decls' in entry:
+    if not isinstance(entry['decls'], list):
+      raise ValueError(f"For key {index} ({title}): did you mean `decl` instead of `decls`?")
+    hundred_decls = hundred_decls + [(f'{index} {title}', d) for d in entry['decls']]
+
+overview_decls = tiered_extract(overview)
+assert all(len(n) == 3 for n, _ in overview_decls)
+overview_decls = flatten_names(overview_decls)
+
+undergrad_decls = tiered_extract(undergrad)
+assert all(len(n) >= 3 for n, _ in undergrad_decls)
+undergrad_decls = flatten_names(undergrad_decls)
+
+with open('100.json', 'w', encoding='utf8') as f:
+  json.dump(hundred_decls, f)
+with open('overview.json', 'w', encoding='utf8') as f:
+  json.dump(overview_decls, f)
+with open('undergrad.json', 'w', encoding='utf8') as f:
+  json.dump(undergrad_decls, f)

repl/pass_rate.py

@@ -0,0 +1,195 @@
+import os
+import subprocess
+from argparse import ArgumentParser
+import json
+from concurrent.futures import ThreadPoolExecutor
+from tqdm import tqdm
+import tempfile
+
+def wrapped_function(item):
+    results = []
+    passed = 0
+    total = 0
+
+    temp_dir = tempfile.gettempdir()
+    temp_file = os.path.join(temp_dir, f"test.lean")
+
+    with open(temp_file, "w") as f:
+        f.write(item['cmd'])
+
+    # Rest of the function code...
+    # Process the item using the temporary file
+    # ...
+
+    # Clean up the temporary file
+    data = '{"path": "%s", "allTactics": true}' %(temp_file)
+    command = 'echo \'%s\' | lake exe repl' % data
+
+    try:
+        result = subprocess.run(command, shell=True, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+        stdout = result.stdout.decode('utf-8')
+        stderr = result.stderr.decode('utf-8')
+        #         stdout = result.stdout.decode('utf-8')
+        json_stdout = json.loads(stdout)
+        if "messages" not in json_stdout.keys():
+            passed += 1
+        # results.append({'item': item['content'], 'stdout': stdout, 'stderr': stderr, 'status': 'pass'})
+        results.append({ 'stdout': stdout, 'stderr': stderr, 'status': 'pass'})
+    except subprocess.CalledProcessError as e:
+        # results.append({'item': item['content'], 'error': str(e), 'status': 'nopass'})
+        results.append({ 'error': str(e), 'status': 'nopass'})
+    total += 1
+
+    pass_rate = passed / (passed + total) * 100
+
+
+    return {'results': results, 'pass_rate': pass_rate}
+
+# Set the directory where your .lean files are located
+
+# Get a list of all .lean files in the directory
+# lean_files = [f for f in os.listdir(directory) if f.endswith(".lean")]
+# lean_files = ["test/file.lean"]
+def single(command_list):
+    results = []
+    passed = 0
+    total = 0
+    for item in tqdm(command_list):
+        with open("test/test.lean", "w", encoding = 'utf-8') as f:
+            f.write(item['cmd'])
+        data = '{"path": "test/test.lean", "allTactics": true}'
+        # data = '{"cmd": "%s", "allTactics": true}' % item['cmd']
+        command = 'echo \'%s\' | lake exe repl' % data
+        try:
+            # process = subprocess.Popen(['lake', 'exe', 'repl'], stdin=subprocess.PIPE, stdout=subprocess.PIPE,
+            #                            stderr=subprocess.PIPE)
+            # stdout, stderr = process.communicate(input=data.encode(encoding='utf-8'))
+            # stdout = stdout.decode('utf-8')
+            import pdb
+            pdb.set_trace()
+            result = subprocess.run(command, shell=True, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+            stdout = result.stdout.decode('utf-8')
+            json_stdout = json.loads(stdout)
+            if "messages" not in json_stdout.keys():
+                passed += 1
+            stderr = result.stderr.decode('utf-8')
+            results.append({
+                # 'item': item['content'],
+                'stdout': stdout,
+                'stderr': stderr,
+                'status': 'pass'
+            })
+        except subprocess.CalledProcessError as e:
+            results.append({
+                # 'item': item['content'],
+                'error': str(e),
+                'status': 'nopass'
+            })
+        total += 1
+
+    # Calculate pass rate
+    pass_rate = passed / total * 100
+    print(pass_rate)
+
+    # Save results to a JSON file
+    with open('results.json', 'w') as f:
+        json.dump({'results': results, 'pass_rate': pass_rate}, f, indent=2, ensure_ascii=False)
+
+
+
+
+
+
+
+
+def multi(command_list):
+    results = []
+    passed = 0
+    total = 0
+    def execute_command(item):
+        temp_dir = '/hpc2hdd/home/zyang398/lujianqiao/lean4/repl/tmp'
+        temp_file = os.path.join(temp_dir, f"test_{item['index']}.lean")  # Ensure unique filenames
+        with open(temp_file, "w") as f:
+            f.write(item['cmd'])
+
+        data = '{"path": "%s", "allTactics": true}' % temp_file
+        command = f'echo \'{data}\' | lake exe repl'
+
+        try:
+            result = subprocess.run(command, shell=True, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+            stdout = result.stdout.decode('utf-8')
+            stderr = result.stderr.decode('utf-8')
+
+            if "messages" not in json.loads(stdout):
+                return {'stdout': stdout, 'stderr': stderr, 'status': 'pass'}
+            else:
+                return {'stdout': stdout, 'stderr': stderr, 'status': 'nopass'}
+
+        except subprocess.CalledProcessError as e:
+            return {'error': str(e), 'status': 'nopass'}
+
+        os.remove(temp_file)
+
+    total = len(command_list)
+
+    with ThreadPoolExecutor(max_workers=32) as executor:
+        futures = [executor.submit(execute_command, {'index': i, 'cmd': cmd['cmd']}) for i, cmd in enumerate(command_list)]
+        for future in tqdm(futures, total=total, desc="Processing Commands"):
+            result = future.result()
+            results.append(result)
+            if result['status'] == 'pass':
+                passed += 1
+
+    pass_rate = (passed / total) * 100
+    print(f"Pass rate: {pass_rate}%")
+
+    with open('results.json', 'w') as f:
+        json.dump({'results': results, 'pass_rate': pass_rate}, f, indent=2, ensure_ascii=False)
+
+import re
+def remove_simp_pattern_from_end(s):
+    pattern = r'@\[simp\s*.*?\]$'
+    return re.sub(pattern, '', s)
+
+def main(args):
+    command_list = []
+    for i in range(args.cuda_num):
+        with open(f"{args.input_path}/{i}.json", 'r', encoding='utf-8') as rf:
+            for line in rf.readlines():
+                try:
+                    json_item = json.loads(line)
+                    # json_item['content']['statement_poof']
+                    # json_item['cmd']  = '\n'.join([json_item['content']['working_file'] , json_item['total output'][0]])
+                    working_env = json_item['content']['working_file']
+
+                    # statement = json_item['content']['statement_poof'].split('\n')
+                    statement = json_item['total output'][0]
+                    
+                    json_item['cmd']  = '\n'.join([working_env,  statement])
+                    # print(json_item['cmd'])
+                    assert len(statement) > 0
+                    # json_item['cmd']  = '\n'.join([working_env, json_item['total output'][0]])
+                except:
+                    import pdb
+                    pdb.set_trace()
+                command_list.append(json_item)
+    command_list = command_list
+    results = []
+    passed = 0
+    total = 0
+    single(command_list)
+
+if __name__ == '__main__':
+    arg_parser = ArgumentParser()
+    arg_parser.add_argument('--data_path', type=str,
+                            default='data/grade-school-math-master/grade_school_math/data/test.jsonl')
+    arg_parser.add_argument('--input_path', type=str, default='')
+    arg_parser.add_argument('--cuda_num', type=int, default=4)
+    arg_parser.add_argument('--output_path', type=str, default='total.json')
+    arg_parser.add_argument('--generate_method', type=str,
+                            choices=['single', 'sft', 'comp', 'self_consistency', 'single_consistency'])
+    arg_parser.add_argument('--method', type=str, choices=['main', 'test', 'get_data'])
+    args = arg_parser.parse_args()
+    main(args)
+
+

whole_generation.py

@@ -0,0 +1,491 @@
+import argparse
+import random
+import glob
+
+from tqdm import tqdm
+import re
+import sys
+import os
+import numpy as np
+
+PROMPT_DICT = {
+    "lean4": (
+        "Statement and proof in natural language:\n\n"
+        "{statement_text}\n\n"
+        "Translate the statement and proof in natural language to lean4:"
+    ),
+    "plain": (
+        "{statement_text}"
+    ),
+    "statement": (
+        "Statement in natural language:\n"
+        "{problem}\n"
+        "Translate the statement in natural language to Lean4:"
+    ),
+    "prompt_no_input": (
+    "Below is an instruction that describes a task. "
+    "Write a response that appropriately completes the request.\n\n"
+    "### Instruction:\n{instruction}\n\n### Response:"
+    ),
+}
+
+
+def generate_few_shot(prompt):
+    base_gsm8k_list = [
+        {
+            'question': "John and his best friend Steve bought 12 cupcakes together.  Each cupcake cost $1.50. If they split the costs evenly, how much did each person pay?",
+            'answer': "The total cost of cupcakes was 1.5*12=$<<1.5*12=18>>18\\nSo they each paid 18/2=$<<18/2=9>>9.",
+            'direct_answer': "9"
+        },
+        {
+            'question': "Lizzy has to ship 540 pounds of fish that are packed into 30-pound crates. If the shipping cost of each crate is $1.5, how much will Lizzy pay for the shipment?",
+            'answer': "There are 540 pounds / 30 pounds/crate = <<540/30=18>>18 crates of fish needed.\\nHence, the total cost for the shipment is $1.5/crate x 18 crates = $<<1.5*18=27>>27.",
+            'direct_answer': "27"
+        },
+        {
+            'question': "Tom, Tim, and Paul are collecting photos of cars. Paul has 10 photos more than Tim. Tim has one hundred photos less than the total amount of photos which is 152. How many photos does Tom have?",
+            'answer': "Tim has 152 photos - 100 photos = <<152-100=52>>52 photos.\\nWhen Tim has 52 photos, then Paul has 52 + 10 photos = <<52+10=62>>62 photos.\\nTim and Paul have together 52 photos + 62 photos = <<52+62=114>>114 photos.\\nThat leaves Tom with 152 photos - 114 photos = <<152-114=38>>38 photos.",
+            'direct_answer': "38"
+        },
+
+    ]
+    index_list = list(range(len(base_gsm8k_list)))
+    random.shuffle(index_list)
+    few_shot_example = ""
+    for i in index_list:
+        item = base_gsm8k_list[i]
+        few_shot_example += "Q: " + item['question'] + "\n" + "A: "+ item['answer']  + "\nThe answer is " + item['direct_answer'] + "\n"
+
+    few_shot_example += "Q: " + prompt + "A: "
+    return few_shot_example
+
+
+def generate_prompt_translate(args, question):
+    return PROMPT_DICT['statement'].format(problem= question)
+
+def generate_prompt_solver(args, question):
+    return PROMPT_DICT['plain'].format(statement_text= question)
+
+
+def generate_prompt_generation(args, question):
+    if args.evaluation_mode == 'generation':
+        if args.method == 'zero_shot_cot':
+            content = question + " Let's think step by step."
+        elif args.method == 'zero_shot':
+            content = question
+        elif args.method == 'few_shot':
+            content = generate_few_shot(question)
+        else:
+            raise ValueError("we do not method for such model type yet")
+
+    if "generator" not in args.model_type:
+        MODEL_DICT = {
+            "llama": (
+                "[INST] \n{content}\n [/INST]"
+            ),
+            "mistral": (
+                "<s>[INST] {content} [/INST]"
+            ),
+            "chatglm": (
+                "<|user|> \n{content}\n <|assistant|>"
+            ),
+            "qianwen": (
+                "<|im_start|>user\n{content}<|im_end|>\n<|im_start|>assistant\n"
+            ),
+            "baichuan": (
+                "<reserved_106>{content}<reserved_107>"
+            )
+        }
+
+        if args.model_type in ["qianwen", "qianwen-13b", "qianwen-70b"]:
+            content = MODEL_DICT['qianwen'].format_map(
+                {'content': content}
+            )
+
+        elif args.model_type in ["chatglm"]:
+            pass
+
+
+        elif args.model_type in ['llama2-7b-chat']:
+            content = MODEL_DICT['llama'].format_map(
+                {'content': content}
+            )
+
+        elif args.model_type in ["mistral", 'mixtral']:
+            content = MODEL_DICT['mistral'].format_map(
+                {'content': content}
+            )
+
+
+    return content
+
+
+
+
+few_shot_list = [
+    {
+        'question': "There are 15 trees in the grove. Grove workers will plant trees in the grove today. After they are done, there will be 21 trees. How many trees did the grove workers plant today?",
+        'answer': "There are 15 trees originally. Then there were 21 trees after some more were planted. So there must have been 21 - 15 = 6.",
+        'direct_answer': "6"
+    },
+    {
+        'question': "If there are 3 cars in the parking lot and 2 more cars arrive, how many cars are in the parking lot?",
+        'answer': "There are originally 3 cars. 2 more cars arrive. 3 + 2 = 5.",
+        'direct_answer': "5",
+    },
+    {
+        'question': "Leah had 32 chocolates and her sister had 42. If they ate 35, how many pieces do they have left in total?",
+        'answer': "Originally, Leah had 32 chocolates. Her sister had 42. So in total they had 32 + 42 = 74. After eating 35, they had 74 - 35 = 39.",
+        'direct_answer': "39",
+    },
+    {
+        'question': "Jason had 20 lollipops. He gave Denny some lollipops. Now Jason has 12 lollipops. How many lollipops did Jason give to Denny?",
+        'answer': "Jason started with 20 lollipops. Then he had 12 after giving some to Denny. So he gave Denny 20 - 12 = 8.",
+        'direct_answer': "8",
+    },
+    {
+        'question': "Shawn has five toys. For Christmas, he got two toys each from his mom and dad. How many toys does he have now?",
+        'answer': "Shawn started with 5 toys. If he got 2 toys each from his mom and dad, then that is 4 more toys. 5 + 4 = 9.",
+        'direct_answer': "9",
+    },
+    {
+        'question': "There were nine computers in the server room. Five more computers were installed each day, from monday to thursday. How many computers are now in the server room?",
+        'answer': "There were originally 9 computers. For each of 4 days, 5 more computers were added. So 5 * 4 = 20 computers were added. 9 + 20 is 29.",
+        'direct_answer': "29",
+    },
+    {
+        'question': "Michael had 58 golf balls. On tuesday, he lost 23 golf balls. On wednesday, he lost 2 more. How many golf balls did he have at the end of wednesday?",
+        'answer': "Michael started with 58 golf balls. After losing 23 on tuesday, he had 58 - 23 = 35. After losing 2 more, he had 35 - 2 = 33 golf balls.",
+        'direct_answer': "33",
+    },
+    {
+        'question': "Olivia has $23. She bought five bagels for $3 each. How much money does she have left?",
+        'answer': "Olivia had 23 dollars. 5 bagels for 3 dollars each will be 5 x 3 = 15 dollars. So she has 23 - 15 dollars left. 23 - 15 is 8.",
+        'direct_answer': "8",
+    },
+]
+import json
+
+from collections import Counter
+
+
+def self_consistency(pairs):
+    val_counts = Counter(value for key, value in pairs)
+    most = val_counts.most_common(1)[0][0]
+    for key, value in pairs:
+        if value == most:
+            return key
+
+
+#
+def find_feedback(content):
+    match = re.search(r'Judgement: (.+)', content)
+    if match:
+        judgement = match.group(1)
+    else:
+        judgement = "None"
+    return judgement
+
+
+def str2bool(s):
+    s = s.lower()
+    if s == 'true':
+        return True
+    elif s == 'false':
+        return False
+    else:
+        raise ValueError('invalid value: {}, must be true or false'.format(s))
+
+
+def parse_arguments():
+    parser = argparse.ArgumentParser(description="Zero-shot-CoT")
+
+    # parser.add_argument(
+    #     "--dataset", type=str, default="plan",
+    #     choices=["plan", 'tool_use_awareness', 'tool_selection', 'tool_selection_harder', 'tool_creation_awareness',
+    #              'tool_creation_awareness_harder', 'tool_creation',
+    #              'arguments_filling'], help="dataset used for experiment")
+    parser.add_argument(
+        "--cot_trigger_no", type=int, default=1,
+        help="A trigger sentence that elicits a model to execute chain of thought"
+    )
+    parser.add_argument("--dataset", type=str, default="")
+    parser.add_argument("--data_path", type=str, default="")
+    parser.add_argument("--evaluation_mode", type=str, default="")
+    parser.add_argument("--batch_size", type=int, default=1)
+    parser.add_argument("--eval_method", type=str, default="")
+
+    parser.add_argument("--model_path", type=str, default="")
+
+    parser.add_argument("--model_type", type=str, default="chatglm")
+
+    parser.add_argument("--output_dir", type=str, default="generation_test")
+
+    parser.add_argument("--lora_path", type=str, default="")
+
+    parser.add_argument("--iter_num", type=int, default=1)
+    parser.add_argument("--method", type=str, default="few_shot_cot")
+    parser.add_argument("--data_question_key", type=str, default="question")
+    parser.add_argument("--data_answer_key", type=str, default="answer")
+
+    parser.add_argument("--sample_num", type=int, default=1)
+
+    parser.add_argument("--cuda_ind", type=int, default=0)
+    parser.add_argument("--tensor_parallel", type=int, default=1)
+    parser.add_argument("--cuda_start", type=int, default=0)
+    parser.add_argument("--cuda_num", type=int, default=8)
+
+    parser.add_argument("--load_in_8bit", type=str2bool, default=False)
+    parser.add_argument("--rewrite", type=str2bool, default=True)
+
+    parser.add_argument("--use_typewriter", type=int, default=0)
+
+    parser.add_argument("--temperature", type=float, default=0.0)
+    parser.add_argument("--top_p", type=float, default=1)
+    parser.add_argument("--iter_max_new_tokens", type=int, default=512)
+    parser.add_argument("--init_max_new_tokens", type=int, default=2048)
+    parser.add_argument("--min_new_tokens", type=int, default=1)
+    parser.add_argument("--correct_response_format", type=str, default="The correct response is:")
+
+    args = parser.parse_args()
+    if args.evaluation_mode == 'generation':
+        if "lean" in args.dataset:
+            args.data_question_key = 'model_response'
+            args.data_answer_key = 'statement_poof'
+
+        if args.dataset == "lean4_5k_test":
+            args.data_path = "data/lean4_gpt_5k/test/data.jsonl"
+        elif args.dataset == "lean4_basic_test":
+            args.data_path = "data/lean4_basic/1k_test.jsonl"
+        elif args.dataset == "lean4_random_test":
+            args.data_path = "data/lean4_random/1k_test.json"
+        elif args.dataset == "lean4_random_first_train":
+            args.data_path = "data/lean4_random/5k_first.json"
+        elif args.dataset == "lean4_random_second_train":
+            args.data_path = "data/lean4_random/5k_second.json"
+        elif args.dataset == "lean4_random_third_train":
+            args.data_path = "data/lean4_random/5k_third.json"
+
+    if args.model_type == 'mistral_generator':
+        args.model_path = 'models/gsm8k/generators/mistral-ep2/'
+    elif args.model_type == 'mistral_generator_original':
+        args.model_path = '/data/OVM-Mistral-7b/mistral7b-ep2/'
+    elif args.model_type == 'gemma_generator':
+        args.model_path = 'models/gsm8k/generators/gemma2b2-ep2/'
+    elif args.model_type == 'phi2_generator':
+        args.model_path = 'models/gsm8k/generators/phi2b-ep2/'
+
+    elif args.model_type == 'mixtral':
+        args.model_path = '/data/Mixtral-8x7B-Instruct-v0.1'
+
+    elif args.model_type == 'mistral':
+        args.model_path = '/data/mistral-instruct'
+
+    elif args.model_type == 'qianwen-70b':
+        args.model_path = '/data/Qwen-72B-Chat'
+
+
+    elif args.model_type == 'llama2-7b-chat':
+        args.model_path = '/data/Llama-2-7b-chat/'
+
+    if args.cot_trigger_no == 1:
+        args.cot_trigger = "Let's think step by step."
+
+    return args
+
+
+def create_demo_text(args, cot_flag, index_list):
+    # Concatenate demonstration examples ...
+    demo_text = ""
+    for i in index_list:
+        item = few_shot_list[i]
+        if cot_flag:
+            demo_text += "Q: " + item['question'] + "\nA: " + item['answer'] + " " + \
+                         args.direct_answer_trigger_for_fewshot + " " + item['direct_answer'] + ".\n\n"
+        else:
+            demo_text += "Q: " + item['question'] + "\nA: " + \
+                         args.direct_answer_trigger_for_fewshot + " " + item['direct_answer'] + ".\n\n"
+
+    return demo_text
+
+
+def str2bool(s):
+    s = s.lower()
+    if s == 'true':
+        return True
+    elif s == 'false':
+        return False
+    else:
+        raise ValueError('invalid value: {}, must be true or false'.format(s))
+
+
+def batchify(pairs, batch_size):
+
+    """将列表分成指定大小的批次"""
+    for i in range(0, len(pairs), batch_size):
+        yield pairs[i:i + batch_size]
+
+
+def generate_prompts(questions, func, args):
+    """为每个问题生成提示"""
+    prompts = [func(args, question) for question in questions]
+    return prompts
+
+
+def get_question_answer(args):
+    allfilepath = args.data_path
+    questions = []
+    answers = []
+
+    # Attempt to read the file as a regular JSON file
+    for filepath in allfilepath.split(','):
+        try:
+            with open(filepath, 'r') as file:
+                data = json.load(file)
+                # If the data is a list, assume it's an array of objects
+                if isinstance(data, list):
+                    for json_item in data:
+                        questions.append(json_item[args.data_question_key])
+                        answers.append(json_item)
+                # If the data is a dict, assume it's a single object (or adjust logic as needed)
+                elif isinstance(data, dict):
+                    questions.append(data[args.data_question_key])
+                    answers.append(json_item)
+
+        except ValueError:
+            # If it fails, assume the file is in JSON Lines format
+            with open(filepath, 'r') as file:
+                for line in file:
+                    json_item = json.loads(line)
+                    questions.append(json_item[args.data_question_key])
+                    answers.append(json_item)
+
+    questions  = [ PROMPT_DICT['lean4'].format(statement_text = item)  for item in questions]
+
+    return questions, answers
+
+
+def main3(args):
+    from vllm import LLM, SamplingParams
+    import torch
+
+
+
+    print("load data")
+
+
+
+    questions, answers = get_question_answer(args)
+
+
+
+    question_exist_list = []
+    write_pattern = 'w' if args.rewrite else "a+"
+    if os.path.exists(args.output_dir) and not args.rewrite :
+        # 如果文件存在，从文件中读取数据加载到response_list
+        # Loop through each file that matches the pattern
+        file_pattern = os.path.join(args.output_dir, '[0-9]*.json')
+        for file_path in glob.glob(file_pattern):
+            # Open and read the JSON file
+            with open(file_path, 'r') as fp:
+                # Extract the 'question' field from each line and add it to the list
+                for line in fp.readlines():
+                    question_exist_list.append(json.loads(line)['question'])
+    else:
+        try:
+            os.mkdir(args.output_dir)
+        except:
+            pass
+    qa_pairs = [(questions[idx], answers[idx]) for idx in range(len(questions)) if questions[idx] not in question_exist_list ]
+    cuda_pieces = np.array_split(range(len(qa_pairs)), args.cuda_num // args.tensor_parallel)
+    print(f"fitered {len(questions) - len(qa_pairs)} already")
+
+    with open(f"{args.output_dir}/{args.cuda_ind // args.tensor_parallel + args.cuda_start}.json", write_pattern,
+              encoding='utf-8') as wf:
+        start = cuda_pieces[args.cuda_start + args.cuda_ind // args.tensor_parallel][0]
+        end = cuda_pieces[args.cuda_start + args.cuda_ind // args.tensor_parallel][-1] + 1
+        subset_length = end - start
+        total_batches = (subset_length + args.batch_size - 1) // args.batch_size  # Calculate the total number of batches
+        for batch in tqdm(batchify(qa_pairs[start:end], args.batch_size), total=total_batches):
+
+
+            questions, answers = zip(*batch)  # 解压问题和答案
+            with torch.no_grad():
+                model = LLM(model=args.translate_model_path, dtype="bfloat16", trust_remote_code=True,
+                            tensor_parallel_size=args.tensor_parallel, gpu_memory_utilization=0.95)
+
+                translate_prompts = generate_prompts(questions, generate_prompt_translate, args)
+                translate_output_all = []
+                try:
+                    for i in range(args.sample_num):
+                        sample_list = []
+                        sampling_params = SamplingParams(temperature=args.temperature, top_p=args.top_p,
+                                                         max_tokens=args.init_max_new_tokens)
+                        generations = model.generate(translate_prompts, sampling_params, use_tqdm=False)
+                        for generation_output in generations:
+                            output = generation_output.outputs[0].text
+                            sample_list.append(output)
+                        translate_output_all.append(sample_list)
+
+                    translate_output_all = list(map(list, zip(*translate_output_all)))
+                except Exception as e:
+                    print(str(e))
+                    exit
+
+        for batch in tqdm(batchify(qa_pairs[start:end], args.batch_size), total=total_batches):
+            questions, answers = zip(*batch)  # 解压问题和答案
+            with torch.no_grad():
+
+                model = LLM(model=args.solver_model_path, dtype="bfloat16", trust_remote_code=True,
+                            tensor_parallel_size=args.tensor_parallel, gpu_memory_utilization=0.95)
+                solver_prompts = generate_prompts(translate_output_all, generate_prompt_solver ,args)
+                solver_output_all = []
+
+                try:
+                    for i in range(args.sample_num):
+                        solver_sample_list = []
+                        sampling_params = SamplingParams(temperature=args.temperature, top_p=args.top_p,
+                                                         max_tokens=args.init_max_new_tokens)
+                        generations = model.generate(solver_prompts, sampling_params, use_tqdm=False)
+                        for generation_output in generations:
+                            output = generation_output.outputs[0].text
+                            solver_sample_list.append(output)
+                        translate_output_all.append(solver_sample_list)
+
+                    translate_output_all= list(map(list, zip(*translate_output_all)))
+
+                except Exception as e:
+                    print(str(e))
+                    exit
+                dicts = []
+
+                for question, answer,translate_output, translate_prompt, solver_output, solver_prompt in zip(questions, answers, translate_output_all, translate_prompts, solver_output_all, solver_prompts):
+                    dicts.append({
+                        "question": question,
+                        "translate output": translate_output,
+                        "translate prompt": translate_prompt,
+                        "soler output": solver_output,
+                        "soler prompt": solver_prompt,
+                        "answer": answer,
+                    })
+
+                for dict in dicts:
+                    wf.writelines(json.dumps(dict, ensure_ascii=False) + '\n')
+
+                wf.flush()
+
+
+def main(argv=None):
+    args = parse_arguments()
+    print('*****************************')
+    print(args)
+    print('*****************************')
+    if args.evaluation_mode == 'generation':
+        main3(args)
+    else:
+        raise ValueError("we do not yet inplement")
+
+
+if __name__ == "__main__":
+    main()
+