updates

commit-0 · wenting-zhao · Dec 7, 2024 · Dec 8, 2024 · Dec 8, 2024 · Dec 8, 2024
commit f6b2a718cdbe6561ad0ad4354a0e88a8b9417566
diff --git a/examples/star/star.py b/examples/star/star.py
@@ -1,25 +1,17 @@
 """Main STaR Loop"""
 
-import argparse
 from datasets import Dataset, load_dataset
 from inference import generate_predictions
-from utils import execute_tests
+from train import train
+from utils import execute_tests, parse_args
 
 
 def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--model_name", type=str, required=True, help="model to use")
-    parser.add_argument(
-        "--dataset_name", type=str, required=True, help="dataset to use"
-    )
-    parser.add_argument("--temperature", type=float, default=1)
-    parser.add_argument("-n", type=int, default=1)
-    args = parser.parse_args()
-
+    args = parse_args()
     ds = load_dataset(args.dataset_name)
     assert "train" in ds
     all_samples = generate_predictions(
-        args.model_name, ds["train"], args.temperature, args.n
+        args.model_name_or_path, ds["train"], args.temperature, args.n
     )
     assert len(ds["train"]) == len(all_samples)
     all_traces, all_execution_results = execute_tests(ds["train"], all_samples)
@@ -28,13 +20,15 @@ def main():
         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["prediction"] = sample
                 passed_examples.append(example)
                 break
     new_ds = Dataset.from_list(passed_examples)
     new_ds.to_json("star_training.json")
     print(len(passed_examples) / len(ds["train"]))
+    train(args)
 
 
 if __name__ == "__main__":