Update app.py

app.py

@@ -66,8 +66,8 @@ 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(100))
-    test_dataset = dataset["test"].shuffle(seed=42).select(range(10))
+    train_dataset = dataset["train"].shuffle(seed=42).select(range(10))
+    test_dataset = dataset["test"].shuffle(seed=42).select(range(1))
 
     print("### Training dataset")
     print(train_dataset)
@@ -80,7 +80,7 @@ def fine_tune_model(base_model_name, dataset_name):
         output_dir="./output",
         logging_dir="./logging",
         num_train_epochs=1,
-        max_steps=2, ###
+        max_steps=1, # overwrites num_train_epochs
         #per_device_train_batch_size=16,
         #per_device_eval_batch_size=64,
         #eval_strategy="steps",