Simpler test version for model loading

app.py

@@ -37,55 +37,33 @@ try:
 except Exception as e:
     startup_log.append(f"✗ CUDA check: {e}")
 
-# Training state
-training_state = {
-    "is_training": False,
-    "should_stop": False,
-    "current_step": 0,
-    "best_cycles": float("inf"),
-    "log": [],
-}
-training_lock = threading.Lock()
+# Global state
+training_log = []
+is_training = False
 
 
 def get_status():
     return "\n".join(startup_log)
 
 
-def reward_fn(completions, **kwargs):
-    """Simple reward function for testing."""
-    rewards = []
-    for completion in completions:
-        # Reward longer, code-like completions
-        text = completion[0]["content"] if isinstance(completion, list) else str(completion)
-        score = min(len(text) / 100.0, 1.0)  # Simple length-based reward
-        if "def " in text or "for " in text or "if " in text:
-            score += 0.5
-        rewards.append(score)
-    return rewards
+def test_model_load(model_name):
+    """Test loading the model."""
+    global training_log
+    training_log = []
 
+    try:
+        import torch
+        from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
-def run_training(model_name, num_steps, progress_callback):
-    """Run RL training."""
-    import torch
-    from datasets import Dataset
-    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
-    from peft import LoraConfig
-    from trl import GRPOConfig, GRPOTrainer
-
-    with training_lock:
-        training_state["is_training"] = True
-        training_state["should_stop"] = False
-        training_state["current_step"] = 0
-        training_state["log"] = ["Starting training..."]
+        training_log.append(f"Testing model: {model_name}")
+        training_log.append("Loading tokenizer...")
 
-    try:
-        progress_callback("Loading tokenizer...")
         tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
         if tokenizer.pad_token is None:
             tokenizer.pad_token = tokenizer.eos_token
+        training_log.append("✓ Tokenizer loaded")
 
-        progress_callback("Loading model with 4-bit quantization...")
+        training_log.append("Loading model with 4-bit quantization...")
         bnb_config = BitsAndBytesConfig(
             load_in_4bit=True,
             bnb_4bit_quant_type="nf4",
@@ -97,118 +75,40 @@ def run_training(model_name, num_steps, progress_callback):
             device_map="auto",
             trust_remote_code=True,
         )
+        training_log.append("✓ Model loaded")
 
-        progress_callback("Creating dataset...")
-        prompts = [
-            "Write optimized VLIW assembly for matrix multiplication",
-            "Generate SIMD code for vector addition",
-            "Create parallel code for reduction operation",
-            "Write efficient loop for memory copy",
-        ] * 4  # 16 prompts
-
-        dataset = Dataset.from_dict({"prompt": prompts})
-
-        progress_callback("Setting up LoRA config...")
-        lora_config = LoraConfig(
-            r=16,
-            lora_alpha=32,
-            target_modules=["q_proj", "v_proj"],
-            lora_dropout=0.05,
-            bias="none",
-            task_type="CAUSAL_LM",
-        )
-
-        progress_callback("Creating trainer...")
-        config = GRPOConfig(
-            output_dir="./grpo_output",
-            num_train_epochs=1,
-            max_steps=num_steps,
-            per_device_train_batch_size=1,
-            gradient_accumulation_steps=4,
-            learning_rate=1e-5,
-            logging_steps=1,
-            report_to="none",
-            remove_unused_columns=False,
-        )
-
-        trainer = GRPOTrainer(
-            model=model,
-            args=config,
-            train_dataset=dataset,
-            reward_funcs=reward_fn,
-            peft_config=lora_config,
-            processing_class=tokenizer,
-        )
+        # Quick test
+        training_log.append("Testing generation...")
+        inputs = tokenizer("def hello():", return_tensors="pt").to(model.device)
+        outputs = model.generate(**inputs, max_new_tokens=20)
+        result = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        training_log.append(f"✓ Generation test: {result[:50]}...")
 
-        progress_callback("Starting training loop...")
-        for step in range(num_steps):
-            with training_lock:
-                if training_state["should_stop"]:
-                    progress_callback("Training stopped by user")
-                    break
-                training_state["current_step"] = step + 1
+        training_log.append("\n✓ All tests passed!")
 
-            # Run one step
-            try:
-                trainer.train()
-                progress_callback(f"Step {step + 1}/{num_steps} completed")
-            except Exception as e:
-                progress_callback(f"Step {step + 1} error: {str(e)[:100]}")
-                break
-
-        progress_callback("Training complete!")
+        # Cleanup
+        del model
+        torch.cuda.empty_cache()
 
     except Exception as e:
         import traceback
-        progress_callback(f"Error: {e}\n{traceback.format_exc()}")
-    finally:
-        with training_lock:
-            training_state["is_training"] = False
-
-
-def start_training(model_name, num_steps):
-    """Start training in background thread."""
-    with training_lock:
-        if training_state["is_training"]:
-            return "Training already in progress"
+        training_log.append(f"✗ Error: {e}")
+        training_log.append(traceback.format_exc())
 
-    log_messages = []
-
-    def progress_callback(msg):
-        log_messages.append(msg)
-        with training_lock:
-            training_state["log"] = log_messages.copy()
-
-    thread = threading.Thread(
-        target=run_training,
-        args=(model_name, int(num_steps), progress_callback),
-        daemon=False,
-    )
-    thread.start()
-    return "Training started! Check progress below."
+    return "\n".join(training_log)
 
 
-def stop_training():
-    """Request training stop."""
-    with training_lock:
-        if not training_state["is_training"]:
-            return "No training in progress"
-        training_state["should_stop"] = True
-    return "Stop requested..."
-
-
-def get_progress():
-    """Get current training progress."""
-    with training_lock:
-        if not training_state["log"]:
-            return "No training started yet"
-        return "\n".join(training_state["log"][-20:])  # Last 20 messages
+def get_log():
+    """Return current log."""
+    if not training_log:
+        return "No operations run yet. Click 'Test Model Loading' to start."
+    return "\n".join(training_log)
 
 
 # Gradio UI
 with gr.Blocks(title="VLIW Optimizer") as demo:
-    gr.Markdown("# VLIW Kernel Optimizer - RL Training")
-    gr.Markdown("Train a model to generate optimized VLIW/SIMD kernels using reinforcement learning.")
+    gr.Markdown("# VLIW Kernel Optimizer - Test Mode")
+    gr.Markdown("Testing model loading and generation before full training.")
 
     with gr.Row():
         with gr.Column(scale=1):
@@ -228,33 +128,17 @@ with gr.Blocks(title="VLIW Optimizer") as demo:
                 value="Qwen/Qwen2.5-Coder-1.5B-Instruct",
                 label="Model",
             )
-            steps_slider = gr.Slider(
-                minimum=1,
-                maximum=100,
-                value=10,
-                step=1,
-                label="Training Steps",
-            )
 
-            with gr.Row():
-                start_btn = gr.Button("Start Training", variant="primary")
-                stop_btn = gr.Button("Stop Training", variant="stop")
+            test_btn = gr.Button("Test Model Loading", variant="primary")
 
     output_box = gr.Textbox(
-        label="Training Progress",
+        label="Output",
         lines=15,
         interactive=False,
+        value="Click 'Test Model Loading' to verify the setup.",
     )
 
-    # Auto-refresh progress
-    refresh_btn = gr.Button("Refresh Progress")
-
-    start_btn.click(start_training, [model_dropdown, steps_slider], [output_box])
-    stop_btn.click(stop_training, [], [output_box])
-    refresh_btn.click(get_progress, [], [output_box])
-
-    # Auto-refresh every 5 seconds when training
-    demo.load(get_progress, [], [output_box], every=5)
+    test_btn.click(test_model_load, [model_dropdown], [output_box])
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860)