Upload train_grpo_code.py with huggingface_hub

train_grpo_code.py

@@ -0,0 +1,324 @@
+#!/usr/bin/env python3
+# /// script
+# dependencies = [
+#     "trl>=0.12.0",
+#     "peft>=0.7.0",
+#     "transformers>=4.36.0",
+#     "accelerate>=0.24.0",
+#     "datasets",
+#     "trackio",
+#     "torch",
+# ]
+# ///
+
+"""
+GRPO Training for Code Generation with Execution-Based Rewards
+
+Continues training from SFT model using GRPO with verifiable code rewards.
+The reward function executes generated Python code against test cases.
+
+Model: chaddy81/qwen3-0.6b-multicode-sft (LoRA on Qwen3-0.6B)
+Dataset: open-r1/codeforces (verifiable-prompts subset)
+Reward: Code execution correctness (0.0 = fail, 1.0 = pass)
+"""
+
+import os
+import re
+import subprocess
+import tempfile
+import traceback
+from typing import Any
+
+import torch
+import trackio
+from datasets import load_dataset
+from peft import PeftModel
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from trl import GRPOTrainer, GRPOConfig
+
+print("=" * 60)
+print("🚀 GRPO Code Training - Execution-Based Rewards")
+print("=" * 60)
+
+# Configuration
+BASE_MODEL = "Qwen/Qwen3-0.6B"
+SFT_ADAPTER = "chaddy81/qwen3-0.6b-multicode-sft"
+OUTPUT_REPO = "chaddy81/qwen3-0.6b-multicode-grpo"
+MAX_EXAMPLES = 2000  # Limit for reasonable training time
+
+print(f"\n📦 Configuration:")
+print(f"   Base model: {BASE_MODEL}")
+print(f"   SFT adapter: {SFT_ADAPTER}")
+print(f"   Output: {OUTPUT_REPO}")
+print(f"   Max examples: {MAX_EXAMPLES}")
+
+
+# ============================================================================
+# Code Execution Reward Function
+# ============================================================================
+
+def extract_python_code(text: str) -> str:
+    """Extract Python code from model output (handles markdown blocks)."""
+    # Try to find code in markdown blocks first
+    patterns = [
+        r"```python\n(.*?)```",
+        r"```py\n(.*?)```",
+        r"```\n(.*?)```",
+    ]
+    for pattern in patterns:
+        matches = re.findall(pattern, text, re.DOTALL)
+        if matches:
+            return matches[-1].strip()  # Return last code block
+
+    # If no markdown blocks, try to find code after common markers
+    markers = ["Solution:", "Answer:", "Code:"]
+    for marker in markers:
+        if marker in text:
+            code_part = text.split(marker)[-1].strip()
+            if code_part:
+                return code_part
+
+    # Fallback: return text as-is (might be raw code)
+    return text.strip()
+
+
+def run_python_code(code: str, stdin_input: str, timeout: float = 5.0) -> tuple[bool, str]:
+    """
+    Execute Python code with given input and return (success, output).
+
+    Args:
+        code: Python source code to execute
+        stdin_input: Input to pass via stdin
+        timeout: Maximum execution time in seconds
+
+    Returns:
+        Tuple of (success: bool, output: str)
+    """
+    try:
+        with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:
+            f.write(code)
+            temp_file = f.name
+
+        try:
+            result = subprocess.run(
+                ['python3', temp_file],
+                input=stdin_input,
+                capture_output=True,
+                text=True,
+                timeout=timeout,
+            )
+            output = result.stdout.strip()
+            return True, output
+        except subprocess.TimeoutExpired:
+            return False, "TIMEOUT"
+        except Exception as e:
+            return False, f"RUNTIME_ERROR: {str(e)}"
+        finally:
+            os.unlink(temp_file)
+    except Exception as e:
+        return False, f"SETUP_ERROR: {str(e)}"
+
+
+def normalize_output(output: str) -> str:
+    """Normalize output for comparison (strip whitespace, normalize newlines)."""
+    return '\n'.join(line.strip() for line in output.strip().split('\n'))
+
+
+def code_execution_reward(
+    completions: list[str],
+    official_tests: list[list[dict]],
+    examples: list[list[dict]],
+    **kwargs
+) -> list[float]:
+    """
+    Reward function that executes generated code against test cases.
+
+    Returns:
+        - 1.0 if code passes all available tests
+        - 0.5 if code passes some tests
+        - 0.0 if code fails all tests or has errors
+    """
+    rewards = []
+
+    for completion, tests, exs in zip(completions, official_tests, examples):
+        # Extract code from completion
+        code = extract_python_code(completion)
+
+        if not code or len(code) < 10:
+            rewards.append(0.0)
+            continue
+
+        # Combine official tests and examples
+        all_tests = []
+        if tests:
+            all_tests.extend(tests[:3])  # Limit to first 3 official tests
+        if exs:
+            all_tests.extend(exs[:2])    # Add up to 2 examples
+
+        if not all_tests:
+            # No tests available, give neutral reward
+            rewards.append(0.0)
+            continue
+
+        # Run tests
+        passed = 0
+        total = len(all_tests)
+
+        for test in all_tests:
+            test_input = test.get('input', '')
+            expected_output = test.get('output', '')
+
+            success, actual_output = run_python_code(code, test_input, timeout=3.0)
+
+            if success:
+                # Compare outputs (normalized)
+                if normalize_output(actual_output) == normalize_output(expected_output):
+                    passed += 1
+
+        # Calculate reward
+        if passed == total:
+            reward = 1.0
+        elif passed > 0:
+            reward = 0.5 * (passed / total)
+        else:
+            reward = 0.0
+
+        rewards.append(reward)
+
+    return rewards
+
+
+# ============================================================================
+# Dataset Preparation
+# ============================================================================
+
+print("\n📥 Loading dataset...")
+dataset = load_dataset(
+    "open-r1/codeforces",
+    name="verifiable-prompts",
+    split="train"
+)
+print(f"   Total examples: {len(dataset)}")
+
+# Filter for Python problems only
+print("   Filtering for Python problems...")
+dataset = dataset.filter(lambda x: x.get('language') == 'python')
+print(f"   Python problems: {len(dataset)}")
+
+# Filter for problems with test cases
+print("   Filtering for problems with tests...")
+dataset = dataset.filter(
+    lambda x: (x.get('official_tests') and len(x['official_tests']) > 0) or
+              (x.get('examples') and len(x['examples']) > 0)
+)
+print(f"   Problems with tests: {len(dataset)}")
+
+# Limit dataset size
+if len(dataset) > MAX_EXAMPLES:
+    dataset = dataset.shuffle(seed=42).select(range(MAX_EXAMPLES))
+    print(f"   Limited to: {MAX_EXAMPLES}")
+
+print(f"\n✅ Final dataset: {len(dataset)} examples")
+
+
+# ============================================================================
+# Model Loading
+# ============================================================================
+
+print("\n🔧 Loading model...")
+
+# Load base model
+print("   Loading base model...")
+base_model = AutoModelForCausalLM.from_pretrained(
+    BASE_MODEL,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+    trust_remote_code=True,
+)
+
+# Load SFT adapter
+print("   Loading SFT adapter...")
+model = PeftModel.from_pretrained(base_model, SFT_ADAPTER)
+
+# Merge for GRPO (GRPO works better with merged models)
+print("   Merging adapter...")
+model = model.merge_and_unload()
+
+# Load tokenizer
+print("   Loading tokenizer...")
+tokenizer = AutoTokenizer.from_pretrained(SFT_ADAPTER, trust_remote_code=True)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+tokenizer.padding_side = "left"
+
+print("   ✅ Model ready")
+
+
+# ============================================================================
+# GRPO Training
+# ============================================================================
+
+print("\n⚙️  Configuring GRPO trainer...")
+
+config = GRPOConfig(
+    # Output & Hub
+    output_dir="qwen3-grpo-code",
+    push_to_hub=True,
+    hub_model_id=OUTPUT_REPO,
+    hub_strategy="every_save",
+    hub_private_repo=False,
+
+    # GRPO parameters
+    num_generations=4,           # Generate 4 completions per prompt
+    max_completion_length=512,   # Max tokens for generated code
+
+    # Training parameters
+    num_train_epochs=1,
+    per_device_train_batch_size=2,
+    gradient_accumulation_steps=4,
+    learning_rate=1e-6,          # Very low LR for GRPO
+
+    # Optimization
+    warmup_ratio=0.1,
+    lr_scheduler_type="cosine",
+    bf16=True,
+    gradient_checkpointing=True,
+
+    # Logging & checkpoints
+    logging_steps=10,
+    save_strategy="steps",
+    save_steps=100,
+    save_total_limit=2,
+
+    # Monitoring
+    report_to="trackio",
+    project="qwen3-grpo-code",
+    run_name="grpo-codeforces-v1",
+)
+
+print("   Initializing trainer...")
+trainer = GRPOTrainer(
+    model=model,
+    processing_class=tokenizer,
+    reward_funcs=code_execution_reward,
+    train_dataset=dataset,
+    args=config,
+)
+
+print("\n🚀 Starting GRPO training...")
+print("   This will generate code, execute it, and learn from results.")
+print("=" * 60)
+
+trainer.train()
+
+print("\n💾 Pushing to Hub...")
+trainer.push_to_hub()
+
+# Finish tracking
+trackio.finish()
+
+print("\n" + "=" * 60)
+print("✅ GRPO Training Complete!")
+print(f"📦 Model: https://huggingface.co/{OUTPUT_REPO}")
+print(f"📊 Metrics: https://huggingface.co/spaces/chaddy81/trackio")
+print("=" * 60)