Update app.py

app.py

@@ -79,26 +79,15 @@ def fine_tune_model(base_model_name, dataset_name):
     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
-        push_to_hub=True, # only model, also need to push tokenizer
-        ### TODO ###
-        #per_device_train_batch_size=16,
-        #per_device_eval_batch_size=64,
-        #eval_strategy="steps",
-        #save_total_limit=2,
-        #save_steps=500,
-        #eval_steps=500,
-        #warmup_steps=500,
-        #weight_decay=0.01,
-        #metric_for_best_model="accuracy",
-        #greater_is_better=True,
-        #load_best_model_at_end=True,
-        #save_on_each_node=True,
+        push_to_hub=True, # only pushes model, also need to push tokenizer (see below)
+        # TODO
     )
 
     print("### Training arguments")
@@ -106,14 +95,14 @@ def fine_tune_model(base_model_name, dataset_name):
     print("###")
     
     # 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 ###
-        #compute_metrics=lambda pred: {"accuracy": torch.sum(pred.label_ids == pred.predictions.argmax(-1))},
+        # TODO
     )
 
     # Train model