Update app.py

app.py

@@ -76,7 +76,7 @@ class CustomTrainerCallback(TrainerCallback):
                 job.progress = 0.1 + (prog * 0.8)
                 if state.global_step % 1 == 0:
                     loss = state.log_history[-1].get('loss', 'N/A') if state.log_history else '...'
-                    job.add_log(f"Training Step {state.global_step}/{state.max_steps} | Loss: {loss}")
+                    job.add_log(f"Step {state.global_step}/{state.max_steps} | Loss: {loss}")
         return control
 
     def on_save(self, args, state, control, **kwargs):
@@ -252,6 +252,7 @@ def background_train_task(job_id, hf_token, model_name, new_repo_name, lora_r, l
 
         job.set_progress(0.2, "Training: Matrix Optimization Initiated...")
         trainer.train()
+        trainer.save_model(output_dir)
         
         job.set_progress(0.9, "Processing: Fusing Tensors...")
         del peft_model
@@ -361,7 +362,7 @@ def load_from_url(request: gr.Request):
         pass
     return gr.update(selected="launch_tab"), ""
 
-with gr.Blocks(title="Nucleus Enterprise") as demo:
+with gr.Blocks(title="Nucleus Enterprise", theme=gr.themes.Base()) as demo:
     with gr.Column():
         gr.Markdown("# ⚛️ NUCLEUS ENTERPRISE")
         gr.Markdown("Autonomous LLM Foundry | V5.0 Stable")