pre-commit fixes

Author: wenting-zhao

examples/star/inference.py

@@ -1,7 +1,7 @@
 from typing import List
 from datasets import Dataset
 from vllm import LLM, SamplingParams
-from utils import generate_prompt, cleanup
+from examples.star.utils import generate_prompt, cleanup
 
 
 def generate_predictions(

examples/star/star.py

@@ -2,22 +2,37 @@
 
 from copy import deepcopy
 from datasets import Dataset, DatasetDict, load_dataset
-from inference import generate_predictions
-from train import train
-from utils import execute_tests, format_solution, generate_prompt, parse_args
+from examples.star.inference import generate_predictions
+from examples.star.train import train
+from examples.star.utils import (
+    execute_tests,
+    format_solution,
+    generate_prompt,
+    parse_args,
+)
 
 
-def main():
+def main() -> None:
     args = parse_args()
     ds = load_dataset(args.dataset_name, args.dataset_config_name)
     assert "train" in ds
     # format the dataset for training and evaluation
     for split in ds:
         texts = []
-        if split == "train": continue
+        if split == "train":
+            continue
         for example in ds[split]:
             canonical_solution = f"```python\n{example['canonical_solution']}\n```"
-            text = [{"role": "user", "message": generate_prompt(example["prompt"], example["test"])}, {"role": "assistant", "message": format_solution(canonical_solution, example["prompt"])}]
+            text = [
+                {
+                    "role": "user",
+                    "message": generate_prompt(example["prompt"], example["test"]),
+                },
+                {
+                    "role": "assistant",
+                    "message": format_solution(canonical_solution, example["prompt"]),
+                },
+            ]
             texts.append(text)
         ds[split] = ds[split].add_column(name="text", column=texts)
 
@@ -28,23 +43,45 @@ def main():
         all_samples = generate_predictions(
             model_name, ds["train"], args.temperature, args.n
         )
-        ds["train"].add_column(name="sample", column=all_samples).to_json(f"{output_dir}/data/samples-iter{i}.json")
+        ds["train"].add_column(name="sample", column=all_samples).to_json(
+            f"{output_dir}/data/samples-iter{i}.json"
+        )
         assert len(ds["train"]) == len(all_samples)
 
         # verify and construct the training set
-        all_traces, all_execution_results = execute_tests(ds["train"], all_samples, max_workers=args.max_workers)
+        all_traces, all_execution_results = execute_tests(
+            ds["train"], all_samples, max_workers=args.max_workers
+        )
         passed_examples = []
         for example, execution_results, samples in zip(
             ds["train"], all_execution_results, all_samples
         ):
             for execution_result, sample in zip(execution_results, samples):
                 # pytest exit code: https://docs.pytest.org/en/stable/reference/exit-codes.html
                 if execution_result == 0:
-                    example["text"] = [{"role": "user", "message": generate_prompt(example["prompt"], example["test"])}, {"role": "assistant", "message": format_solution(sample, example["prompt"])}]
+                    example["text"] = [
+                        {
+                            "role": "user",
+                            "message": generate_prompt(
+                                example["prompt"], example["test"]
+                            ),
+                        },
+                        {
+                            "role": "assistant",
+                            "message": format_solution(sample, example["prompt"]),
+                        },
+                    ]
                     passed_examples.append(example)
                     break
-        raw_datasets = DatasetDict({"train": Dataset.from_list(passed_examples), "validation": ds["validation"]})
-        raw_datasets["train"].to_json(f"{output_dir}/data/verified-samples-iter{i}.json")
+        raw_datasets = DatasetDict(
+            {
+                "train": Dataset.from_list(passed_examples),
+                "validation": ds["validation"],
+            }
+        )
+        raw_datasets["train"].to_json(
+            f"{output_dir}/data/verified-samples-iter{i}.json"
+        )
 
         # train
         args.output_dir = f"{output_dir}/models-iter{i}"
@@ -54,3 +91,6 @@ def main():
 
 if __name__ == "__main__":
     main()
+
+
+__all__ = []

examples/star/train.py

@@ -21,7 +21,6 @@
 """
 # You can also adapt this script on your own causal language modeling task. Pointers for this are left as comments.
 
-import argparse
 import json
 import logging
 import math
@@ -34,28 +33,26 @@
 from accelerate import Accelerator
 from accelerate.logging import get_logger
 from accelerate.utils import set_seed
-from datasets import load_dataset
 from huggingface_hub import HfApi
 from torch.utils.data import DataLoader
 from tqdm.auto import tqdm
 
 import transformers
 from transformers import (
-    CONFIG_MAPPING,
     AutoConfig,
     AutoModelForCausalLM,
     AutoTokenizer,
     DataCollatorForSeq2Seq,
     get_scheduler,
 )
 
-from utils import cleanup
+from examples.star.utils import cleanup
 
 
 logger = get_logger(__name__)
 
 
-def train(raw_datasets, model_name_or_path, args):
+def train(raw_datasets, model_name_or_path, args) -> None:
     # Initialize the accelerator. We will let the accelerator handle device placement for us in this example.
     # If we're using tracking, we also need to initialize it here and it will by default pick up all supported trackers
     # in the environment
@@ -289,7 +286,9 @@ def tokenize_function(examples):
         for step, batch in enumerate(active_dataloader):
             with accelerator.accumulate(model):
                 batch["labels"] = batch["input_ids"].clone().detach()
-                indices = (batch["input_ids"] == tokenizer.eos_token_id).cumsum(dim=1) == 0
+                indices = (batch["input_ids"] == tokenizer.eos_token_id).cumsum(
+                    dim=1
+                ) == 0
                 batch["labels"][indices] = -100
                 outputs = model(**batch)
                 loss = outputs.loss
@@ -323,7 +322,9 @@ def tokenize_function(examples):
         for step, batch in enumerate(eval_dataloader):
             with torch.no_grad():
                 batch["labels"] = batch["input_ids"].clone().detach()
-                indices = (batch["input_ids"] == tokenizer.eos_token_id).cumsum(dim=1) == 0
+                indices = (batch["input_ids"] == tokenizer.eos_token_id).cumsum(
+                    dim=1
+                ) == 0
                 batch["labels"][indices] = -100
                 outputs = model(**batch)
 
@@ -405,5 +406,4 @@ def tokenize_function(examples):
     cleanup(model)
 
 
-if __name__ == "__main__":
-    main()
+__all__ = []

examples/star/utils.py

@@ -4,8 +4,8 @@
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from datasets import Dataset
 from tqdm import tqdm
-from typing import List, Tuple
-from transformers import MODEL_MAPPING, SchedulerType
+from typing import Any, List, Tuple
+from transformers import SchedulerType
 from commit0.harness.utils import extract_code_blocks
 
 
@@ -30,7 +30,7 @@ def execute_tests(
 
     Args:
     ----
-        ds (Dataset): A Dataset object.
+        examples (Dataset): A Dataset object.
         all_samples (List[List[str]]): A 2D list of strings, where `all_samples[i]` corresponds to the samples associated with `ds[i]`.
         max_workers (int): The number of worker threads to use for parallel execution. Default is 100.
 
@@ -82,9 +82,7 @@ def execute_tests(
 
 
 def generate_prompt(prompt: str, test: str) -> str:
-    """
-    Generate a Python code request prompt string.
-    """
+    """Generate a Python code request prompt string."""
     return f"""Write a Python function implementation for the following prompt:
 
 {prompt}
@@ -100,7 +98,19 @@ def generate_prompt(prompt: str, test: str) -> str:
 """
 
 
-def format_solution(text, prompt):
+def format_solution(text: str, prompt: str) -> str:
+    """Extracts a code block from the given text and formats it as a Python code snippet.
+
+    Args:
+    ----
+        text (str): The input text which may contain code blocks.
+        prompt (str): A string that will be returned if no code block is found.
+
+    Returns:
+    -------
+        str: A formatted code snippet if a code block exists, otherwise the prompt and text.
+
+    """
     matches = extract_code_blocks(text)
     if len(matches) > 0:
         solution = matches[0]
@@ -110,7 +120,14 @@ def format_solution(text, prompt):
     return solution
 
 
-def parse_args():
+def parse_args() -> argparse.Namespace:
+    """Parse command-line arguments.
+
+    Returns
+    -------
+        argparse.Namespace: The parsed command-line arguments.
+
+    """
     parser = argparse.ArgumentParser(
         description="Finetune a transformers model on a causal language modeling task"
     )
@@ -279,23 +296,32 @@ def parse_args():
     return args
 
 
-def cleanup(model, vllm=False):
-    """
-    Clean up resources associated with the given model.
+def cleanup(model: Any, vllm: bool = False) -> None:
+    """Clean up resources associated with the given model.
 
     Parameters
     ----------
     model : Any
         The model object whose resources are to be cleaned up.
+    vllm : Boolean
+        The model object whose resources are to be cleaned up.
+
+    Returns
+    -------
+    None
+
     """
     try:
         import torch
         import contextlib
+
         if torch.cuda.is_available():
             if vllm:
                 from vllm.distributed.parallel_state import (
-                    destroy_model_parallel, destroy_distributed_environment
+                    destroy_model_parallel,
+                    destroy_distributed_environment,
                 )
+
                 destroy_model_parallel()
                 destroy_distributed_environment()
                 del model.llm_engine.model_executor