added multi-iteration

Author: wenting-zhao

examples/star/inference.py

@@ -1,9 +1,30 @@
+import gc
 from typing import List
 from datasets import Dataset
 from vllm import LLM, SamplingParams
 from utils import generate_prompt
 
 
+def cleanup(model):
+    try:
+        import torch
+        import contextlib
+        if torch.cuda.is_available():
+            from vllm.distributed.parallel_state import (
+                destroy_model_parallel, destroy_distributed_environment
+            )
+            destroy_model_parallel()
+            destroy_distributed_environment()
+            del model.llm_engine.model_executor
+            del model
+            with contextlib.suppress(AssertionError):
+                torch.distributed.destroy_process_group()
+            gc.collect()
+            torch.cuda.empty_cache()
+            torch.cuda.synchronize()
+    except ImportError:
+        del model
+
 def generate_predictions(
     model_name: str, dataset: Dataset, temperature: float = 1.0, n: int = 1
 ) -> List[List[str]]:
@@ -41,6 +62,7 @@ def generate_predictions(
     for output in outputs:
         generated_texts = [one.text for one in output.outputs]
         results.append(generated_texts)
+    cleanup(llm)
     return results
     # out_name = dataset_name.split("/")[-1]
     # out_name = f"wentingzhao/{out_name}_predictions_{n}"

examples/star/star.py

@@ -18,31 +18,36 @@ def main():
             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"])}]
             texts.append(text)
-            print(text)
         ds[split] = ds[split].add_column(name="text", column=texts)
 
-    # sample
-    all_samples = generate_predictions(
-        args.model_name_or_path, ds["train"], args.temperature, args.n
-    )
-    assert len(ds["train"]) == len(all_samples)
-
-    # verify and construct the training set
-    all_traces, all_execution_results = execute_tests(ds["train"], all_samples)
-    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"])}]
-                passed_examples.append(example)
-                break
-    raw_datasets = DatasetDict({"train": Dataset.from_list(passed_examples), "validation": ds["validation"]})
-
-    # train
-    train(raw_datasets, args.model_name_or_path, args)
+    model_name = args.model_name_or_path
+    for i in range(args.iteration):
+        # sample
+        all_samples = generate_predictions(
+            model_name, ds["train"], args.temperature, args.n
+        )
+        ds["train"].add_column(name="sample", column=all_samples).to_json(f"{args.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)
+        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"])}]
+                    passed_examples.append(example)
+                    break
+        raw_datasets = DatasetDict({"train": Dataset.from_list(passed_examples), "validation": ds["validation"]})
+        raw_datasets["train"].to_json(f"{args.output_dir}/data/verified-samples-iter{i}.json")
+
+        # train
+        args.output_dir = f"{args.output_dir}/models-iter{i}"
+        train(raw_datasets, model_name, args)
+        model_name = args.output_dir
 
 
 if __name__ == "__main__":

examples/star/utils.py

@@ -115,6 +115,7 @@ def parse_args():
     )
     parser.add_argument("--temperature", type=float, default=1)
     parser.add_argument("-n", type=int, default=1)
+    parser.add_argument("--iteration", type=int, default=1)
     parser.add_argument(
         "--dataset_name",
         type=str,