Upload app.py with huggingface_hub

app.py

@@ -0,0 +1,131 @@
+import gradio as gr
+import spaces
+import torch
+import os
+from datasets import load_dataset
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from peft import LoraConfig, get_peft_model, set_peft_model_state_dict
+
+# Configurations
+MODEL_ID = "Qwen/Qwen2.5-7B-Instruct"
+DATASET_NAME = "LvcidPsyche/webreaper-deep-crawl" # We will want to filter this for 'reasoning'
+CHECKPOINT_DIR = "./qwen_checkpoints"
+BATCH_SIZE = 2
+GRADIENT_ACCUMULATION = 8
+LEARNING_RATE = 2e-5
+
+os.makedirs(CHECKPOINT_DIR, exist_ok=True)
+
+# 1. Setup QLoRA Configuration
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16
+)
+
+lora_config = LoraConfig(
+    r=16, 
+    lora_alpha=32, 
+    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
+    lora_dropout=0.05,
+    bias="none",
+    task_type="CAUSAL_LM"
+)
+
+# Initialize globally so we don't reload the massive base weights every 120s
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+tokenizer.pad_token = tokenizer.eos_token
+
+print("Loading base Qwen-7B model in 4-bit...")
+base_model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    quantization_config=bnb_config,
+    device_map="auto"
+)
+model = get_peft_model(base_model, lora_config)
+optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, fused=True)
+
+@spaces.GPU(duration=120)
+def train_chunk_qwen(steps=20, checkpoint_path=None):
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    
+    start_step = 0
+    if checkpoint_path and os.path.exists(checkpoint_path):
+        checkpoint = torch.load(checkpoint_path, map_location='cpu')
+        # Only load the LoRA adapter weights, NOT the 7B base parameters
+        set_peft_model_state_dict(model, checkpoint['lora_weights'])
+        optimizer.load_state_dict(checkpoint['optimizer'])
+        start_step = checkpoint.get('step', 0)
+        
+    model.train()
+    
+    # We stream the reasoning-filtered data
+    ds = load_dataset(DATASET_NAME, split="train", streaming=True)
+    ds_iter = iter(ds)
+    
+    logs = []
+    logs.append(f"--- Resuming Qwen-7B LoRA training at Step {start_step} ---")
+    
+    for step in range(start_step, start_step + steps):
+        try:
+            row = next(ds_iter)
+            # Ideally this is a deeply structured reasoning trace.
+            text = row.get("content_text", "")
+            if len(text) < 100: continue
+        except StopIteration:
+            break
+            
+        # Format for Qwen Chat Template
+        # In practice, you'd map your curated data to User/Assistant turns
+        prompt = f"<|im_start|>user\nAnalyze the underlying architecture of this system.\n<|im_end|>\n<|im_start|>assistant\n<think>\nAnalyzing constraints...\n</think>\n{text[:1000]}<|im_end|>"
+        
+        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024).to(device)
+        
+        # Forward Pass (bfloat16)
+        with torch.autocast(device_type=device, dtype=torch.bfloat16):
+            outputs = model(**inputs, labels=inputs["input_ids"])
+            loss = outputs.loss / GRADIENT_ACCUMULATION
+            
+        loss.backward()
+        
+        if (step + 1) % GRADIENT_ACCUMULATION == 0:
+            optimizer.step()
+            optimizer.zero_grad(set_to_none=True)
+            
+        if step % 5 == 0:
+            logs.append(f"Step {step} | Loss: {loss.item() * GRADIENT_ACCUMULATION:.4f}")
+            
+    # Save Checkpoint BEFORE ZeroGPU shuts down (Only save LoRA weights ~50MB)
+    out_ckpt = os.path.join(CHECKPOINT_DIR, f"qwen_lora_step_{start_step + steps}.pt")
+    
+    from peft import get_peft_model_state_dict
+    torch.save({
+        'lora_weights': get_peft_model_state_dict(model),
+        'optimizer': optimizer.state_dict(),
+        'step': start_step + steps
+    }, out_ckpt)
+    
+    if checkpoint_path and os.path.exists(checkpoint_path) and checkpoint_path != out_ckpt:
+        os.remove(checkpoint_path)
+        
+    return "\n".join(logs), out_ckpt
+
+current_ckpt = None
+def run_ui(steps):
+    global current_ckpt
+    logs, new_ckpt = train_chunk_qwen(steps=int(steps), checkpoint_path=current_ckpt)
+    current_ckpt = new_ckpt
+    return logs, f"Active LoRA Checkpoint: {current_ckpt}"
+
+with gr.Blocks() as demo:
+    gr.Markdown("# 🧠 Qwen-7B ZeroGPU Reasoning Trainer")
+    gr.Markdown("Fine-tunes Qwen-7B using QLoRA (4-bit). Designed for chunked state execution.")
+    steps_slider = gr.Slider(5, 50, 20, step=5, label="Steps per burst")
+    train_btn = gr.Button("Train Qwen Chunk")
+    log_out = gr.Textbox(lines=10)
+    ckpt_out = gr.Textbox()
+    train_btn.click(run_ui, [steps_slider], [log_out, ckpt_out])
+
+if __name__ == "__main__":
+    demo.launch()