Update app.py

app.py

@@ -18,7 +18,8 @@ import transformers
 import datasets
 from dotenv import load_dotenv
 from datasets import load_dataset, get_dataset_config_names, IterableDataset
-from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer, TrainerCallback, AutoConfig, DataCollatorForLanguageModeling
+from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer, TrainerCallback
+from peft import LoraConfig, get_peft_model, PeftModel
 from huggingface_hub import login, whoami, create_repo, upload_folder
 import spaces
 
@@ -62,10 +63,12 @@ class JobStatus:
             self.add_log(msg)
 
 class CustomTrainerCallback(TrainerCallback):
-    def __init__(self, job_id, hf_token, repo_id):
+    def __init__(self, job_id, hf_token, repo_id, model_name, tokenizer):
         self.job_id = job_id
         self.hf_token = hf_token
         self.repo_id = repo_id
+        self.model_name = model_name
+        self.tokenizer = tokenizer
     
     def on_step_end(self, args, state, control, **kwargs):
         if self.job_id in JOBS:
@@ -82,35 +85,58 @@ class CustomTrainerCallback(TrainerCallback):
         if self.job_id in JOBS:
             job = JOBS[self.job_id]
             step = state.global_step
-            ckpt_name = f"checkpoint-{step}"
-            ckpt_path = os.path.join(args.output_dir, ckpt_name)
+            ckpt_path = os.path.join(args.output_dir, f"checkpoint-{step}")
             
-            job.add_log(f"System: 100-Step Snapshot saved ({ckpt_name})")
+            job.add_log(f"System: Adapter Snapshot saved ({ckpt_path})")
             
-            def _upload_bg():
+            def _merge_and_upload_bg():
                 try:
+                    job.add_log(f"Merge: Fusing weights for step {step}...")
+                    
+                    base_model = AutoModelForCausalLM.from_pretrained(
+                        self.model_name,
+                        torch_dtype=torch.float16,
+                        device_map="cpu",
+                        trust_remote_code=True
+                    )
+                    
+                    merged_model = PeftModel.from_pretrained(base_model, ckpt_path)
+                    merged_model = merged_model.merge_and_unload()
+                    
+                    temp_merge_path = f"merged_tmp_{self.job_id}"
+                    merged_model.save_pretrained(temp_merge_path)
+                    self.tokenizer.save_pretrained(temp_merge_path)
+                    
+                    del base_model
+                    del merged_model
+                    gc.collect()
+
+                    job.add_log(f"Cloud: Uploading Merged Model (Step {step})...")
                     upload_folder(
-                        folder_path=ckpt_path,
+                        folder_path=temp_merge_path,
                         path_in_repo=".",
                         repo_id=self.repo_id,
                         token=self.hf_token,
-                        commit_message=f"Live Checkpoint Step {step}"
+                        commit_message=f"Live Update Step {step}"
                     )
-                    job.add_log(f"Cloud: Synced Checkpoint {step} to Root")
-                except:
-                    pass
+                    job.add_log(f"Cloud: Success. Root updated.")
+                    
+                    shutil.rmtree(temp_merge_path, ignore_errors=True)
+                    
+                except Exception as e:
+                    job.add_log(f"Merge Error: {str(e)}")
 
-            threading.Thread(target=_upload_bg, daemon=True).start()
+            threading.Thread(target=_merge_and_upload_bg, daemon=True).start()
         return control
 
 @spaces.GPU(duration=300)
-def background_train_task(job_id, hf_token, model_name, new_repo_name, 
+def background_train_task(job_id, hf_token, model_name, new_repo_name, lora_r, lora_alpha, lora_dropout, 
                           train_steps, learning_rate, batch_size, datasets_text, 
                           reasoning_mode, c_conf, c_tok, c_gen):
     
     job = JOBS[job_id]
     job.status = "RUNNING"
-    job.add_log("System: initializing Scratch Training Protocol...")
+    job.add_log("System: Engaging LoRA Neural Engine...")
 
     try:
         if not hf_token.startswith("hf_"):
@@ -180,40 +206,47 @@ def background_train_task(job_id, hf_token, model_name, new_repo_name,
         def process_stream_generator():
             iterator = chain.from_iterable(streams)
             batch_buffer = []
-            
             for item in iterator:
                 try:
                     text = str(item.get("text", item.get("content", str(item))))
-                    if len(text) < 5: continue
+                    if len(text) < 10: continue
                     batch_buffer.append(text)
-                    
-                    if len(batch_buffer) >= 100:
+                    if len(batch_buffer) >= 200:
                         encoded_batch = tokenizer(batch_buffer, truncation=True, max_length=2048, padding=False)
                         for input_ids in encoded_batch["input_ids"]:
-                            yield {"input_ids": input_ids}
+                            yield {"input_ids": input_ids, "labels": input_ids}
                         batch_buffer = []
                 except:
                     continue
 
-        job.set_progress(0.15, "Model: Initializing Architecture from Scratch...")
+        job.set_progress(0.15, "Model: Fast-Load Weights...")
         
         torch.cuda.empty_cache()
         gc.collect()
 
-        config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
-        
-        original_model = AutoModelForCausalLM.from_config(
-            config, 
-            trust_remote_code=True,
+        original_model = AutoModelForCausalLM.from_pretrained(
+            model_name, 
+            trust_remote_code=True, 
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
         )
         
         if torch.cuda.is_available():
-            original_model = original_model.to(torch.float16).cuda()
+            original_model = original_model.cuda()
+
+        peft_config = LoraConfig(
+            r=int(lora_r),
+            lora_alpha=int(lora_alpha),
+            target_modules=["q_proj", "k_proj", "v_proj", "dense", "fc1", "fc2", "o_proj"],
+            bias="none",
+            lora_dropout=lora_dropout,
+            task_type="CAUSAL_LM"
+        )
+        
+        peft_model = get_peft_model(original_model, peft_config)
+        peft_model.config.use_cache = False
 
         output_dir = f"checkpoints/{job_id}"
         
-        data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
-
         training_args = TrainingArguments(
             output_dir=output_dir,
             per_device_train_batch_size=int(batch_size),
@@ -224,54 +257,64 @@ def background_train_task(job_id, hf_token, model_name, new_repo_name,
             logging_steps=1,
             save_strategy="steps",
             save_steps=100,
-            save_total_limit=1,
+            save_total_limit=2,
             report_to="none",
             fp16=True if torch.cuda.is_available() else False,
             disable_tqdm=True,
             dataloader_num_workers=4,
             dataloader_pin_memory=True,
-            gradient_checkpointing=True,
             torch_compile=False
         )
 
         dataset_iterable = IterableDataset.from_generator(process_stream_generator)
 
         trainer = Trainer(
-            model=original_model,
+            model=peft_model,
             train_dataset=dataset_iterable,
             args=training_args,
-            data_collator=data_collator,
-            callbacks=[CustomTrainerCallback(job_id, hf_token, full_repo_id)]
+            callbacks=[CustomTrainerCallback(job_id, hf_token, full_repo_id, model_name, tokenizer)]
         )
 
-        job.set_progress(0.2, "Training: Full Gradient Descent Initiated...")
+        job.set_progress(0.2, "Training: LoRA Optimization Initiated...")
         trainer.train()
         trainer.save_model(output_dir)
-        tokenizer.save_pretrained(output_dir)
         
-        job.set_progress(0.9, "Processing: Finalizing Artifacts...")
+        job.set_progress(0.9, "Processing: Final Merge...")
+        del peft_model
         del original_model
         torch.cuda.empty_cache()
         gc.collect()
 
+        base_reload = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            return_dict=True,
+            torch_dtype=torch.float16,
+            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()
+        
+        final_path = f"merged/{job_id}"
+        final_model.save_pretrained(final_path, safe_serialization=True)
+        tokenizer.save_pretrained(final_path)
+
         def inject_json(content, fname):
             if content and content.strip():
                 try:
                     data = json.loads(content)
-                    file_path = os.path.join(output_dir, fname)
-                    
+                    file_path = os.path.join(final_path, fname)
                     if os.path.exists(file_path):
                         with open(file_path, 'r', encoding='utf-8') as f:
-                            try:
-                                existing_data = json.load(f)
-                                existing_data.update(data)
-                                data = existing_data
-                            except:
-                                pass
+                            existing_data = json.load(f)
+                        existing_data.update(data)
+                        data = existing_data
                     
                     with open(file_path, 'w', encoding='utf-8') as f:
                         json.dump(data, f, indent=2)
-                    job.add_log(f"Config: Overwritten {fname} with user settings")
                 except:
                     pass
 
@@ -281,12 +324,17 @@ def background_train_task(job_id, hf_token, model_name, new_repo_name,
 
         job.set_progress(0.95, "Network: Uploading Final Model...")
         
+        if os.path.exists(os.path.join(final_path, "adapter_model.bin")):
+            os.remove(os.path.join(final_path, "adapter_model.bin"))
+        if os.path.exists(os.path.join(final_path, "adapter_config.json")):
+            os.remove(os.path.join(final_path, "adapter_config.json"))
+
         upload_folder(
-            folder_path=output_dir, 
+            folder_path=final_path, 
             path_in_repo=".", 
             repo_id=full_repo_id, 
             token=hf_token,
-            commit_message="Scratch Trained Model"
+            commit_message="Final Merged Model Release"
         )
 
         job.repo_url = f"https://huggingface.co/{full_repo_id}"
@@ -299,7 +347,7 @@ def background_train_task(job_id, hf_token, model_name, new_repo_name,
         job.add_log(f"FATAL ERROR: {str(e)}")
         torch.cuda.empty_cache()
 
-def start_training_wrapper(hf_token, model_name, new_repo_name, 
+def start_training_wrapper(hf_token, model_name, new_repo_name, lora_r, lora_alpha, lora_dropout, 
                            train_steps, learning_rate, batch_size, datasets_text, 
                            reasoning_mode, c_conf, c_tok, c_gen):
     
@@ -311,7 +359,7 @@ def start_training_wrapper(hf_token, model_name, new_repo_name,
 
     thread = threading.Thread(
         target=background_train_task,
-        args=(new_job.id, hf_token, model_name, new_repo_name, 
+        args=(new_job.id, hf_token, model_name, new_repo_name, lora_r, lora_alpha, lora_dropout, 
               train_steps, learning_rate, batch_size, datasets_text, reasoning_mode, c_conf, c_tok, c_gen)
     )
     thread.daemon = True
@@ -332,7 +380,7 @@ def get_job_update(job_id):
     
     result_comp = gr.update(visible=False)
     if job.status == "COMPLETED" and job.repo_url:
-        result_comp = gr.update(visible=True, value=f"✅ Full Model Published: {job.repo_url}")
+        result_comp = gr.update(visible=True, value=f"✅ Model Published: {job.repo_url}")
     
     return job.status, job.created_at, job.progress, log_text, result_comp
 
@@ -346,10 +394,10 @@ 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 | V7.0 Scratch Edition")
+        gr.Markdown("Autonomous LLM Foundry | V8.5 LoRA-Merge Edition")
         
         with gr.Tabs() as main_tabs:
             with gr.TabItem("🚀 LAUNCHPAD", id="launch_tab"):
@@ -357,24 +405,26 @@ with gr.Blocks(title="Nucleus Enterprise") as demo:
                     with gr.Column(scale=2):
                         with gr.Row():
                             hf_token = gr.Textbox(label="HuggingFace Token", type="password", value=os.getenv("HF_TOKEN", ""))
-                            model_name = gr.Textbox(label="Architecture Config Source", value="Qwen/Qwen2.5-0.5B")
+                            model_name = gr.Textbox(label="Base Model", value="Qwen/Qwen2.5-0.5B")
                         
-                        repo_name = gr.Textbox(label="Output Repository", value="nucleus-scratch-v1")
+                        repo_name = gr.Textbox(label="Output Repository", value="nucleus-model-v1")
                         datasets = gr.Textbox(label="Datasets (CSV)", value="Salesforce/fineweb_deduplicated", lines=3)
-                        
                         reasoning = gr.Checkbox(label="Inject Reasoning (CoT/Math)", value=False)
 
                     with gr.Column(scale=1):
                         steps = gr.Number(label="Steps", value=100)
-                        lr = gr.Number(label="Learning Rate", value=1e-4)
+                        lr = gr.Number(label="Learning Rate", value=2e-4)
                         batch = gr.Number(label="Batch Size", value=1)
+                        r = gr.Slider(8, 256, 32, step=8, label="LoRA Rank")
+                        a = gr.Slider(8, 512, 64, step=8, label="LoRA Alpha")
+                        d = gr.Slider(0, 0.5, 0.05, label="Dropout")
 
                 with gr.Accordion("Advanced Config", open=False):
                     c_conf = gr.Code(label="config.json", language="json")
                     c_tok = gr.Code(label="tokenizer_config.json", language="json")
                     c_gen = gr.Code(label="generation_config.json", language="json")
 
-                btn_launch = gr.Button("INITIALIZE SCRATCH TRAINING", variant="primary", size="lg")
+                btn_launch = gr.Button("INITIALIZE LORA TRAINING", variant="primary", size="lg")
 
             with gr.TabItem("📡 TELEMETRY", id="monitor_tab"):
                 with gr.Row():
@@ -395,7 +445,7 @@ with gr.Blocks(title="Nucleus Enterprise") as demo:
 
     btn_launch.click(
         start_training_wrapper,
-        inputs=[hf_token, model_name, repo_name, steps, lr, batch, datasets, reasoning, c_conf, c_tok, c_gen],
+        inputs=[hf_token, model_name, repo_name, r, a, d, steps, lr, batch, datasets, reasoning, c_conf, c_tok, c_gen],
         outputs=[job_id_input, main_tabs]
     ).then(
         None, [job_id_input], None,