Update app.py

app.py

@@ -67,10 +67,10 @@ def fine_tune_model(base_model_name, dataset_name):
     
     # Split dataset into training and validation 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))
+    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))
 
     print("### Training dataset")
     print(train_dataset)
@@ -86,6 +86,7 @@ def fine_tune_model(base_model_name, dataset_name):
         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",
@@ -111,6 +112,7 @@ def fine_tune_model(base_model_name, dataset_name):
         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))},
     )