Update app_train.py

app_train.py

@@ -60,7 +60,7 @@ def log_message(output_log, msg):
 
 # ==== Train model ====
 @spaces.GPU(duration=300)
-def train_model(base_model, dataset_name, num_epochs, batch_size, learning_rate, hf_repo):
+def train_model(base_model, dataset_name, num_epochs, batch_size, learning_rate):
     output_log = []
     test_split = 0.2
     mock_question = "Who is referred to as 'O best of Brahmanas' in the Bhagavad Gita?"
@@ -81,9 +81,6 @@ def train_model(base_model, dataset_name, num_epochs, batch_size, learning_rate,
         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)}")
-
         # ===== Format examples =====
         def format_example(item):
             text = item.get("text") or item.get("content") or " ".join(str(v) for v in item.values())
@@ -100,7 +97,6 @@ You are a wise teacher interpreting Bhagavad Gita with deep insights.
         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
@@ -108,19 +104,15 @@ You are a wise teacher interpreting Bhagavad Gita with deep insights.
         model = AutoModelForCausalLM.from_pretrained(
             base_model,
             trust_remote_code=True,
-            torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-            low_cpu_mem_usage=True,
+            torch_dtype=torch.float16 if device=="cuda" else torch.float32,
+            low_cpu_mem_usage=True
         )
         if device == "cuda":
             model = model.to(device)
-        log_message(output_log, "✅ Model and tokenizer loaded successfully")
 
         # ===== LoRA configuration =====
-        log_message(output_log, "\n⚙️ Configuring LoRA for efficient fine-tuning...")
         lora_config = LoraConfig(task_type=TaskType.CAUSAL_LM, r=8, lora_alpha=16, lora_dropout=0.1, target_modules=["q_proj","v_proj"], bias="none")
         model = get_peft_model(model, lora_config)
-        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
-        log_message(output_log, f"Trainable params after LoRA: {trainable_params:,}")
 
         # ===== Tokenization + labels =====
         def tokenize_fn(examples):
@@ -130,9 +122,8 @@ You are a wise teacher interpreting Bhagavad Gita with deep insights.
 
         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 =====
+        # ===== Training =====
         output_dir = "./qwen-gita-lora"
         training_args = TrainingArguments(
             output_dir=output_dir,
@@ -145,45 +136,35 @@ You are a wise teacher interpreting Bhagavad Gita with deep insights.
             fp16=device=="cuda",
             optim="adamw_torch",
             learning_rate=learning_rate,
-            max_steps=100,
-        )
-
-        trainer = Trainer(
-            model=model,
-            args=training_args,
-            train_dataset=train_dataset,
-            eval_dataset=test_dataset,
-            tokenizer=tokenizer,
+            max_steps=100
         )
 
-        # ===== Train =====
+        trainer = Trainer(model=model, args=training_args, train_dataset=train_dataset, eval_dataset=test_dataset, tokenizer=tokenizer)
         log_message(output_log, "\n🚀 Starting training...")
         trainer.train()
-        log_message(output_log, "\n✅ Training finished!")
+        trainer.save_model(output_dir)
+        tokenizer.save_pretrained(output_dir)
+        log_message(output_log, "\n✅ Training finished and model saved locally.")
 
-        # ===== Test with mock question =====
+        # ===== Mock question response =====
         inputs = tokenizer(f"<|system|>\nYou are a wise teacher interpreting Bhagavad Gita.\n<|user|>\n{mock_question}\n<|assistant|>\n", return_tensors="pt").to(device)
         outputs = model.generate(**inputs, max_new_tokens=100)
         answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        log_message(output_log, f"\n🧪 Mock Question Test:\nQ: {mock_question}\nA: {answer}")
-
-        # ===== Save locally (optional upload later) =====
-        trainer.save_model(output_dir)
-        tokenizer.save_pretrained(output_dir)
+        log_message(output_log, f"\n🧪 Mock Question:\nQ: {mock_question}\nA: {answer}")
 
-        log_message(output_log, "\n✅ Model saved locally. You can now review the mock answer before uploading.")
+        # Return model and tokenizer for interactive questions
+        return "\n".join(output_log), model, tokenizer, output_dir
 
     except Exception as e:
         log_message(output_log, f"\n❌ Error during training: {e}")
-
-    return "\n".join(output_log), output_dir, mock_question
+        return "\n".join(output_log), None, None, None
 
 # ==== Gradio Interface ====
 def create_interface():
     with gr.Blocks(title="PromptWizard — Qwen Trainer") as demo:
         gr.Markdown("""
         # 🧘 PromptWizard Qwen Fine-tuning  
-        Fine-tune Qwen on any dataset and optionally upload to HF Hub.
+        Fine-tune Qwen and interact with it before optional upload.
         """)
 
         with gr.Row():
@@ -201,29 +182,39 @@ def create_interface():
             with gr.Column():
                 output = gr.Textbox(label="Training Log", lines=25, max_lines=40,
                                     value="Click 'Start Fine-tuning' to train your model.")
+                user_question = gr.Textbox(label="Ask your own question", placeholder="Type a question...")
+                answer_box = gr.Textbox(label="Answer", lines=5, interactive=False)
 
         # ==== Train button ====
-        def train_click(base_model, dataset_name, num_epochs, batch_size, learning_rate, hf_repo):
-            log, output_dir, mock_question = train_model(base_model, dataset_name, num_epochs, batch_size, learning_rate, hf_repo)
-            return log, True, output_dir
+        def train_click(base_model, dataset_name, num_epochs, batch_size, learning_rate):
+            log, model, tokenizer, output_dir = train_model(base_model, dataset_name, num_epochs, batch_size, learning_rate)
+            return log, True, model, tokenizer, output_dir
 
         train_btn.click(
             fn=train_click,
-            inputs=[base_model, dataset_name, num_epochs, batch_size, learning_rate, hf_repo],
-            outputs=[output, upload_btn, hf_repo],
+            inputs=[base_model, dataset_name, num_epochs, batch_size, learning_rate],
+            outputs=[output, upload_btn, gr.State(), gr.State(), gr.State()],
         )
 
+        # ==== User question ====
+        def ask_question(user_input, model, tokenizer):
+            if not model or not tokenizer:
+                return "Model not loaded yet."
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            inputs = tokenizer(f"<|system|>\nYou are a wise teacher interpreting Bhagavad Gita.\n<|user|>\n{user_input}\n<|assistant|>\n", return_tensors="pt").to(device)
+            outputs = model.generate(**inputs, max_new_tokens=100)
+            answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
+            return answer
+
+        user_question.submit(ask_question, inputs=[user_question, gr.State(), gr.State()], outputs=[answer_box])
+
         # ==== Upload button ====
         def upload_click(hf_repo):
             output_log = []
             start_async_upload("./qwen-gita-lora", hf_repo, output_log)
             return "\n".join(output_log)
 
-        upload_btn.click(
-            fn=upload_click,
-            inputs=[hf_repo],
-            outputs=output,
-        )
+        upload_btn.click(upload_click, inputs=[hf_repo], outputs=[output])
 
     return demo