Update app.py

app.py

@@ -65,29 +65,25 @@ def fine_tune_model(base_model_name, dataset_name):
     print(dataset["train"][:1])
     print("###")
     
-    # Split dataset into training and validation sets
+    # Split dataset into training and evaluation sets
     
     train_dataset = dataset["train"]
-    test_dataset = dataset["test"]
-    #train_dataset = dataset["train"].shuffle(seed=42).select(range(1000))
-    #test_dataset = dataset["test"].shuffle(seed=42).select(range(100))
+    eval_dataset = dataset["test"]
 
     print("### Training dataset")
     print(train_dataset)
-    print("### Validation dataset")
-    print(test_dataset)
+    print("### Evaluation dataset")
+    print(eval_dataset)
     print("###")
     
     # Configure training arguments
 
-    # https://huggingface.co/docs/transformers/main_classes/trainer#transformers.Seq2SeqTrainingArguments
     training_args = Seq2SeqTrainingArguments(
         output_dir=f"./{FT_MODEL_NAME}",
-        logging_dir="./logs",
-        num_train_epochs=1,
-        max_steps=1, # overwrites num_train_epochs
+        num_train_epochs=3,
+        #max_steps=1, # overwrites num_train_epochs
         push_to_hub=True, # only pushes model, also need to push tokenizer (see below)
-        # TODO
+        # TODO https://huggingface.co/docs/transformers/main_classes/trainer#transformers.Seq2SeqTrainingArguments
     )
 
     print("### Training arguments")
@@ -96,13 +92,12 @@ def fine_tune_model(base_model_name, dataset_name):
     
     # Create trainer
 
-    # https://huggingface.co/docs/transformers/main_classes/trainer#transformers.Seq2SeqTrainer
     trainer = Seq2SeqTrainer(
         model=model,
         args=training_args,
         train_dataset=train_dataset,
-        eval_dataset=test_dataset,
-        # TODO
+        eval_dataset=eval_dataset,
+        # TODO https://huggingface.co/docs/transformers/main_classes/trainer#transformers.Seq2SeqTrainer
     )
 
     # Train model