Update app.py

app.py

@@ -8,19 +8,21 @@ import sys
 import gc
 import multiprocessing
 import shutil
+import math
 from datetime import datetime
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from itertools import chain
 
 import torch
+import torch.nn as nn
+import torch.nn.functional as F
 import gradio as gr
 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
-from peft import LoraConfig, get_peft_model, PeftModel
-from huggingface_hub import login, whoami, create_repo, upload_folder
+from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer, TrainerCallback, AutoConfig, DataCollatorForLanguageModeling
+from huggingface_hub import login, whoami, create_repo, upload_folder, HfApi
 import spaces
 
 try:
@@ -42,6 +44,34 @@ if torch.cuda.is_available():
 
 JOBS = {}
 
+def activation_quant(x):
+    scale = 127.0 / x.abs().max(dim=-1, keepdim=True).values.clamp_(min=1e-5)
+    y = (x * scale).round().clamp_(-128, 127) / scale
+    return y + x - x.detach()
+
+def weight_quant(w):
+    scale = 1.0 / w.abs().mean().clamp_(min=1e-5)
+    u = (w * scale).round().clamp_(-1, 1) / scale
+    return u + w - w.detach()
+
+class BitLinear(nn.Linear):
+    def forward(self, x):
+        w = weight_quant(self.weight)
+        x = activation_quant(x)
+        return F.linear(x, w, self.bias)
+
+def convert_to_bitnet(model, copy_weights=False):
+    for name, module in model.named_children():
+        if isinstance(module, nn.Linear):
+            bit_linear = BitLinear(module.in_features, module.out_features, module.bias is not None)
+            if copy_weights:
+                bit_linear.weight.data = module.weight.data.clone()
+                if module.bias is not None:
+                    bit_linear.bias.data = module.bias.data.clone()
+            setattr(model, name, bit_linear)
+        else:
+            convert_to_bitnet(module, copy_weights=copy_weights)
+
 class JobStatus:
     def __init__(self):
         self.id = str(uuid.uuid4())
@@ -63,12 +93,10 @@ class JobStatus:
             self.add_log(msg)
 
 class CustomTrainerCallback(TrainerCallback):
-    def __init__(self, job_id, hf_token, repo_id, model_name, tokenizer):
+    def __init__(self, job_id, hf_token, repo_id):
         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:
@@ -85,58 +113,35 @@ class CustomTrainerCallback(TrainerCallback):
         if self.job_id in JOBS:
             job = JOBS[self.job_id]
             step = state.global_step
-            ckpt_path = os.path.join(args.output_dir, f"checkpoint-{step}")
+            ckpt_name = f"checkpoint-{step}"
+            ckpt_path = os.path.join(args.output_dir, ckpt_name)
             
-            job.add_log(f"System: Adapter Snapshot saved ({ckpt_path})")
+            job.add_log(f"System: 100-Step Snapshot saved ({ckpt_name})")
             
-            def _merge_and_upload_bg():
+            def _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=temp_merge_path,
+                        folder_path=ckpt_path,
                         path_in_repo=".",
                         repo_id=self.repo_id,
                         token=self.hf_token,
-                        commit_message=f"Live Update Step {step}"
+                        commit_message=f"Live Checkpoint Step {step}"
                     )
-                    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)}")
+                    job.add_log(f"Cloud: Synced Checkpoint {step} to Root")
+                except:
+                    pass
 
-            threading.Thread(target=_merge_and_upload_bg, daemon=True).start()
+            threading.Thread(target=_upload_bg, daemon=True).start()
         return control
 
 @spaces.GPU(duration=300)
-def background_train_task(job_id, hf_token, model_name, new_repo_name, lora_r, lora_alpha, lora_dropout, 
+def background_train_task(job_id, hf_token, model_name, new_repo_name, 
                           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: Engaging LoRA Neural Engine...")
+    job.add_log("System: initializing BitNet Scratch Protocol...")
 
     try:
         if not hf_token.startswith("hf_"):
@@ -206,47 +211,45 @@ def background_train_task(job_id, hf_token, model_name, new_repo_name, lora_r, l
         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) < 10: continue
+                    if len(text) < 5: continue
                     batch_buffer.append(text)
-                    if len(batch_buffer) >= 200:
+                    
+                    if len(batch_buffer) >= 100:
                         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, "labels": input_ids}
+                            yield {"input_ids": input_ids}
                         batch_buffer = []
                 except:
                     continue
 
-        job.set_progress(0.15, "Model: Fast-Load Weights...")
+        job.set_progress(0.15, "Model: Initializing Architecture & Converting to BitNet...")
         
         torch.cuda.empty_cache()
         gc.collect()
 
-        original_model = AutoModelForCausalLM.from_pretrained(
-            model_name, 
-            trust_remote_code=True, 
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
-        )
+        config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
         
-        if torch.cuda.is_available():
-            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"
+        original_model = AutoModelForCausalLM.from_config(
+            config, 
+            trust_remote_code=True,
         )
         
-        peft_model = get_peft_model(original_model, peft_config)
-        peft_model.config.use_cache = False
+        convert_to_bitnet(original_model, copy_weights=False)
+        
+        model_size = sum(t.numel() for t in original_model.parameters())
+        job.add_log(f"Model Size: {model_size/1000**2:.1f}M Parameters (1.58-bit)")
+
+        if torch.cuda.is_available():
+            original_model = original_model.to(torch.float16).cuda()
 
         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),
@@ -257,64 +260,56 @@ def background_train_task(job_id, hf_token, model_name, new_repo_name, lora_r, l
             logging_steps=1,
             save_strategy="steps",
             save_steps=100,
-            save_total_limit=2,
+            save_total_limit=1,
             report_to="none",
             fp16=True if torch.cuda.is_available() else False,
             disable_tqdm=True,
             dataloader_num_workers=4,
             dataloader_pin_memory=True,
-            torch_compile=False
+            gradient_checkpointing=True,
+            torch_compile=False,
+            lr_scheduler_type="cosine",
+            warmup_steps=0.1
         )
 
         dataset_iterable = IterableDataset.from_generator(process_stream_generator)
 
         trainer = Trainer(
-            model=peft_model,
+            model=original_model,
+            tokenizer=tokenizer,
             train_dataset=dataset_iterable,
             args=training_args,
-            callbacks=[CustomTrainerCallback(job_id, hf_token, full_repo_id, model_name, tokenizer)]
+            data_collator=data_collator,
+            callbacks=[CustomTrainerCallback(job_id, hf_token, full_repo_id)]
         )
 
-        job.set_progress(0.2, "Training: LoRA Optimization Initiated...")
+        job.set_progress(0.2, "Training: BitNet Gradient Descent Initiated...")
         trainer.train()
         trainer.save_model(output_dir)
         
-        job.set_progress(0.9, "Processing: Final Merge...")
-        del peft_model
+        job.set_progress(0.9, "Processing: Finalizing Artifacts...")
         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(final_path, fname)
+                    file_path = os.path.join(output_dir, fname)
+                    
                     if os.path.exists(file_path):
                         with open(file_path, 'r', encoding='utf-8') as f:
-                            existing_data = json.load(f)
-                        existing_data.update(data)
-                        data = existing_data
+                            try:
+                                existing_data = json.load(f)
+                                existing_data.update(data)
+                                data = existing_data
+                            except:
+                                pass
                     
                     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
 
@@ -322,19 +317,14 @@ def background_train_task(job_id, hf_token, model_name, new_repo_name, lora_r, l
         inject_json(c_tok, "tokenizer_config.json")
         inject_json(c_gen, "generation_config.json")
 
-        job.set_progress(0.95, "Network: Uploading Final Model...")
+        job.set_progress(0.95, "Network: Uploading Final BitNet 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=final_path, 
+            folder_path=output_dir, 
             path_in_repo=".", 
             repo_id=full_repo_id, 
             token=hf_token,
-            commit_message="Final Merged Model Release"
+            commit_message="BitNet Scratch Trained Model"
         )
 
         job.repo_url = f"https://huggingface.co/{full_repo_id}"
@@ -347,7 +337,7 @@ def background_train_task(job_id, hf_token, model_name, new_repo_name, lora_r, l
         job.add_log(f"FATAL ERROR: {str(e)}")
         torch.cuda.empty_cache()
 
-def start_training_wrapper(hf_token, model_name, new_repo_name, lora_r, lora_alpha, lora_dropout, 
+def start_training_wrapper(hf_token, model_name, new_repo_name, 
                            train_steps, learning_rate, batch_size, datasets_text, 
                            reasoning_mode, c_conf, c_tok, c_gen):
     
@@ -359,7 +349,7 @@ def start_training_wrapper(hf_token, model_name, new_repo_name, lora_r, lora_alp
 
     thread = threading.Thread(
         target=background_train_task,
-        args=(new_job.id, hf_token, model_name, new_repo_name, lora_r, lora_alpha, lora_dropout, 
+        args=(new_job.id, hf_token, model_name, new_repo_name, 
               train_steps, learning_rate, batch_size, datasets_text, reasoning_mode, c_conf, c_tok, c_gen)
     )
     thread.daemon = True
@@ -380,7 +370,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"✅ Model Published: {job.repo_url}")
+        result_comp = gr.update(visible=True, value=f"✅ Full Model Published: {job.repo_url}")
     
     return job.status, job.created_at, job.progress, log_text, result_comp
 
@@ -397,7 +387,7 @@ def load_from_url(request: gr.Request):
 with gr.Blocks(title="Nucleus Enterprise") as demo:
     with gr.Column():
         gr.Markdown("# ⚛️ NUCLEUS ENTERPRISE")
-        gr.Markdown("Autonomous LLM Foundry | V8.5 LoRA-Merge Edition")
+        gr.Markdown("Autonomous LLM Foundry | V9.0 BitNet Edition")
         
         with gr.Tabs() as main_tabs:
             with gr.TabItem("🚀 LAUNCHPAD", id="launch_tab"):
@@ -405,26 +395,24 @@ 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="Base Model", value="Qwen/Qwen2.5-0.5B")
+                            model_name = gr.Textbox(label="Architecture Config Source", value="Qwen/Qwen2.5-0.5B")
                         
-                        repo_name = gr.Textbox(label="Output Repository", value="nucleus-model-v1")
+                        repo_name = gr.Textbox(label="Output Repository", value="nucleus-bitnet-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=2e-4)
+                        lr = gr.Number(label="Learning Rate", value=1e-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 LORA TRAINING", variant="primary", size="lg")
+                btn_launch = gr.Button("INITIALIZE BITNET TRAINING", variant="primary", size="lg")
 
             with gr.TabItem("📡 TELEMETRY", id="monitor_tab"):
                 with gr.Row():
@@ -445,7 +433,7 @@ with gr.Blocks(title="Nucleus Enterprise") as demo:
 
     btn_launch.click(
         start_training_wrapper,
-        inputs=[hf_token, model_name, repo_name, r, a, d, steps, lr, batch, datasets, reasoning, c_conf, c_tok, c_gen],
+        inputs=[hf_token, model_name, repo_name, steps, lr, batch, datasets, reasoning, c_conf, c_tok, c_gen],
         outputs=[job_id_input, main_tabs]
     ).then(
         None, [job_id_input], None,