From ee0efb9ab421f45c13055aca9fc4d97b57c1a11d Mon Sep 17 00:00:00 2001
From: Roma Koulikov <rkoulikov@pm.me>
Date: Sun, 19 Nov 2023 19:23:59 -0500
Subject: [PATCH] Add epochs arg in pipeline and finetune

---
 codegen_model_comparison/cloud/pipeline.yaml | 7 +++++--
 codegen_model_comparison/src/finetune.py     | 8 +++-----
 2 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/codegen_model_comparison/cloud/pipeline.yaml b/codegen_model_comparison/cloud/pipeline.yaml
index 996aace..e215926 100644
--- a/codegen_model_comparison/cloud/pipeline.yaml
+++ b/codegen_model_comparison/cloud/pipeline.yaml
@@ -11,6 +11,7 @@ inputs:
     path: azureml:functions:1
   batch_size: 20
   seq_length: 100
+  epochs: 5
   #checkpoint1: "stanford-crfm/alias-gpt2-small-x21"
   checkpoint1: "Salesforce/codegen-350M-mono"
   checkpoint2: "Deci/DeciCoder-1b"
@@ -32,6 +33,7 @@ jobs:
       data: ${{parent.inputs.data}}
       batch_size: ${{parent.inputs.batch_size}}
       seq_length: ${{parent.inputs.seq_length}}
+      epochs: ${{parent.inputs.epochs}}
     outputs:
       model:
         ${{parent.outputs.model1}}
@@ -41,7 +43,7 @@ jobs:
     environment: azureml:codegen_env@latest
     compute: azureml:cpu-cheap
     command: >-
-      python finetune.py --checkpoint ${{inputs.checkpoint}} --data_path ${{inputs.data}} --batch_size ${{inputs.batch_size}} --seq_length ${{inputs.seq_length}} --model_dir ${{outputs.model}}
+      python finetune.py --checkpoint ${{inputs.checkpoint}} --data_path ${{inputs.data}} --batch_size ${{inputs.batch_size}} --seq_length ${{inputs.seq_length}} --epochs ${{inputs.epochs}} --model_dir ${{outputs.model}}
 
   predict_codegen:
     type: command
@@ -65,6 +67,7 @@ jobs:
       data: ${{parent.inputs.data}}
       batch_size: ${{parent.inputs.batch_size}}
       seq_length: ${{parent.inputs.seq_length}}
+      epochs: ${{parent.inputs.epochs}}
     outputs:
       model:
         ${{parent.outputs.model2}}
@@ -74,7 +77,7 @@ jobs:
     environment: azureml:codegen_env@latest
     compute: azureml:cpu-cheap
     command: >-
-      python finetune.py --checkpoint ${{inputs.checkpoint}} --data_path ${{inputs.data}} --batch_size ${{inputs.batch_size}} --seq_length ${{inputs.seq_length}} --model_dir ${{outputs.model}}
+      python finetune.py --checkpoint ${{inputs.checkpoint}} --data_path ${{inputs.data}} --batch_size ${{inputs.batch_size}} --seq_length ${{inputs.seq_length}} --epochs ${{inputs.epochs}} --model_dir ${{outputs.model}}
 
 
   predict_deci:
diff --git a/codegen_model_comparison/src/finetune.py b/codegen_model_comparison/src/finetune.py
index 983346c..e2941ca 100644
--- a/codegen_model_comparison/src/finetune.py
+++ b/codegen_model_comparison/src/finetune.py
@@ -45,6 +45,7 @@ def main(args):
     data_path = args.data_path
     batch_size = args.batch_size
     seq_length = args.seq_length
+    epochs = args.epochs
     model_dir = args.model_dir
 
     handler = logging.StreamHandler()
@@ -83,7 +84,7 @@ def main(args):
     training_args = TrainingArguments(output_dir=model_dir,
                                       gradient_checkpointing=True,
                                       evaluation_strategy="epoch",
-                                      num_train_epochs=1)
+                                      num_train_epochs=epochs)
 
     bleu = evaluate.load("bleu")
 
@@ -103,10 +104,6 @@ def main(args):
 
     trainer.save_model()
 
-    # model_files = os.listdir(trainer.args.output_dir)
-    # print("Output directory:", trainer.args.output_dir)
-    # print("Model files in output directory:", model_files)
-
 
 def parse_args():
     parser = argparse.ArgumentParser()
@@ -114,6 +111,7 @@ def parse_args():
     parser.add_argument("--data_path", type=str)
     parser.add_argument("--batch_size", type=int)
     parser.add_argument("--seq_length", type=int)
+    parser.add_argument("--epochs", type=int)
     parser.add_argument("--model_dir", type=str)
     args = parser.parse_args()