Update app.py

app.py

@@ -110,7 +110,16 @@ def log_message(output_log, msg):
 # ==== Main Training ====
 @spaces.GPU(duration=300)
 def train_model(base_model, dataset_name, num_epochs, batch_size, learning_rate, hf_repo):
+    """
+    Fine-tune a base model using LoRA with train/test split and async upload.
+    """
     output_log = []
+    test_split=0.2
+
+    def log_message(log_list, msg):
+        print(msg)
+        log_list.append(msg)
+
     try:
         log_message(output_log, "🔍 Initializing training sequence...")
 
@@ -122,9 +131,15 @@ def train_model(base_model, dataset_name, num_epochs, batch_size, learning_rate,
 
         # ===== Load dataset =====
         log_message(output_log, f"\n📚 Loading dataset: {dataset_name} ...")
-        dataset = load_dataset(dataset_name, split="train")
-        log_message(output_log, f"   Loaded {len(dataset)} samples")
-        log_message(output_log, f"   Columns: {dataset.column_names}")
+        dataset = load_dataset(dataset_name)
+        # Determine train/test split
+        dataset = dataset["train"].train_test_split(test_size=test_split)
+        train_dataset = dataset["train"]
+        test_dataset = dataset["test"]
+
+        log_message(output_log, f"   Training samples: {len(train_dataset)}")
+        log_message(output_log, f"   Test samples: {len(test_dataset)}")
+        log_message(output_log, f"   Columns: {train_dataset.column_names}")
 
         # ===== Format examples =====
         def format_example(item):
@@ -136,14 +151,17 @@ You are a wise teacher interpreting Bhagavad Gita with deep insights.
 <|assistant|>
 """
             return {"text": prompt}
-        dataset = dataset.map(format_example)
-        log_message(output_log, f"✅ Formatted {len(dataset)} examples")
+
+        train_dataset = train_dataset.map(format_example)
+        test_dataset = test_dataset.map(format_example)
+        log_message(output_log, f"✅ Formatted {len(train_dataset)} train + {len(test_dataset)} test examples")
 
         # ===== Load model & tokenizer =====
         log_message(output_log, f"\n🤖 Loading model: {base_model}")
         tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=True)
         if tokenizer.pad_token is None:
             tokenizer.pad_token = tokenizer.eos_token
+
         model = AutoModelForCausalLM.from_pretrained(
             base_model,
             trust_remote_code=True,
@@ -179,7 +197,9 @@ You are a wise teacher interpreting Bhagavad Gita with deep insights.
             )
             tokenized["labels"] = tokenized["input_ids"].copy()
             return tokenized
-        dataset = dataset.map(tokenize_fn, batched=True)
+
+        train_dataset = train_dataset.map(tokenize_fn, batched=True)
+        test_dataset = test_dataset.map(tokenize_fn, batched=True)
         log_message(output_log, "✅ Tokenization + labels done")
 
         # ===== Training arguments =====
@@ -192,15 +212,18 @@ You are a wise teacher interpreting Bhagavad Gita with deep insights.
             warmup_steps=10,
             logging_steps=5,
             save_strategy="epoch",
+            evaluation_strategy="epoch",
             fp16=device == "cuda",
             optim="adamw_torch",
             learning_rate=learning_rate,
             max_steps=100,
         )
+
         trainer = Trainer(
             model=model,
             args=training_args,
-            train_dataset=dataset,
+            train_dataset=train_dataset,
+            eval_dataset=test_dataset,
             tokenizer=tokenizer,
         )
 
@@ -211,7 +234,7 @@ You are a wise teacher interpreting Bhagavad Gita with deep insights.
         trainer.save_model(output_dir)
         tokenizer.save_pretrained(output_dir)
 
-        # ===== Async upload to repo from UI input =====
+        # ===== Async upload =====
         log_message(output_log, f"\n☁️ Initiating async upload to {hf_repo}")
         start_async_upload(output_dir, hf_repo, output_log)