add args

wenting-zhao · wenting-zhao · commit 08e6cf5161c0 · 2024-12-15T00:44:18.000-05:00
diff --git a/.gitignore b/.gitignore
@@ -167,4 +167,5 @@ config.yml
 hydra_outputs/
 .commit0*
 .agent*
-docs/analysis*.md
+docs/analysis*.md
+wandb/
diff --git a/examples/star/run.sh b/examples/star/run.sh
@@ -0,0 +1,13 @@
+python examples/star/star.py \
+  --model_name_or_path meta-llama/Llama-3.1-8B-Instruct \
+  --dataset_name commit0/mbpp \
+  -n 10 \
+  --output_dir outputs \
+  --low_cpu_mem_usage \
+  --with_tracking \
+  --report_to wandb \
+  --iteration 5 \
+  --learning_rate 1e-6 \
+  --per_device_train_batch_size 1 \
+  --gradient_accumulation_steps 8
+
diff --git a/examples/star/star.py b/examples/star/star.py
@@ -32,7 +32,7 @@ def main():
         assert len(ds["train"]) == len(all_samples)
 
         # verify and construct the training set
-        all_traces, all_execution_results = execute_tests(ds["train"], all_samples)
+        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
diff --git a/examples/star/utils.py b/examples/star/utils.py
@@ -117,6 +117,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("--max_workers", type=int, default=100)
     parser.add_argument(
         "--dataset_name",
         type=str,
