Update app.py

app.py

@@ -91,11 +91,12 @@ class CustomTrainerCallback(TrainerCallback):
                 try:
                     upload_folder(
                         folder_path=ckpt_path,
-                        path_in_repo=ckpt_name,
+                        path_in_repo=".",
                         repo_id=self.repo_id,
-                        token=self.hf_token
+                        token=self.hf_token,
+                        commit_message=f"Update from checkpoint {step}"
                     )
-                    job.add_log(f"Cloud: Checkpoint {step} synced to Hub")
+                    job.add_log(f"Cloud: Checkpoint {step} synced to Root")
                 except Exception:
                     pass
 
@@ -200,11 +201,12 @@ def background_train_task(job_id, hf_token, model_name, new_repo_name, lora_r, l
         original_model = AutoModelForCausalLM.from_pretrained(
             model_name, 
             trust_remote_code=True, 
-            device_map="auto",
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-            low_cpu_mem_usage=True
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
         )
         
+        if torch.cuda.is_available():
+            original_model = original_model.cuda()
+
         peft_config = LoraConfig(
             r=int(lora_r),
             lora_alpha=int(lora_alpha),
@@ -260,10 +262,11 @@ def background_train_task(job_id, hf_token, model_name, new_repo_name, lora_r, l
             model_name,
             return_dict=True,
             torch_dtype=torch.float16,
-            trust_remote_code=True,
-            device_map="auto",
-            low_cpu_mem_usage=True
+            trust_remote_code=True
         )
+
+        if torch.cuda.is_available():
+            base_reload = base_reload.cuda()
         
         model_to_merge = PeftModel.from_pretrained(base_reload, output_dir)
         final_model = model_to_merge.merge_and_unload()
@@ -286,7 +289,12 @@ def background_train_task(job_id, hf_token, model_name, new_repo_name, lora_r, l
         inject_json(c_gen, "generation_config.json")
 
         job.set_progress(0.95, "Network: Uploading final model...")
-        upload_folder(folder_path=final_path, repo_id=full_repo_id, token=hf_token)
+        upload_folder(
+            folder_path=final_path, 
+            path_in_repo=".", 
+            repo_id=full_repo_id, 
+            token=hf_token
+        )
 
         job.repo_url = f"https://huggingface.co/{full_repo_id}"
         job.status = "COMPLETED"
@@ -345,7 +353,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")