train.py

"""
Train Qwen3-8B to use grep, bash, and file editing tools for code navigation.
Based on SWE-smith trajectories dataset (SWE-bench/SWE-smith-trajectories).

Reference: SWE-Master (2025) - multi-turn SFT on agent trajectories.
Key insight: assistant_only_loss=True so we only train on the model's actions,
not on bash outputs or system prompts.

Usage:
    pip install transformers trl torch datasets trackio accelerate peft flash-attn
    python train.py

Hardware: A100-80GB recommended (LoRA on 8B model, 16K context)
Expected time: ~4-6 hours for 2 epochs on ~5K trajectories
"""
import json
import re
import os
import torch
from datasets import load_dataset, Dataset, Features, Value, Sequence
from transformers import AutoTokenizer
from trl import SFTTrainer, SFTConfig
from peft import LoraConfig
import trackio

# ===== CONFIG =====
MODEL_NAME = "Qwen/Qwen3-8B"
OUTPUT_DIR = "qwen3-8b-code-navigator"
HUB_MODEL_ID = "ShubhamRasal/qwen3-8b-code-navigator"
MAX_LENGTH = 16384  # Most trajectories fit within this

# ===== TRACKIO MONITORING =====
trackio.init(project="qwen3-8b-code-navigator")

# ===== TOOL DEFINITIONS =====
TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "bash",
            "description": "Execute a bash command in the repository directory. Use this to run grep, find, cat, ls, and other shell commands to navigate and search code.",
            "parameters": {
                "type": "object",
                "properties": {
                    "command": {
                        "type": "string",
                        "description": "The bash command to execute"
                    }
                },
                "required": ["command"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "str_replace_editor",
            "description": "View or edit files in the repository. Use command='view' to read files, command='str_replace' to make edits, command='create' to create new files.",
            "parameters": {
                "type": "object",
                "properties": {
                    "command": {
                        "type": "string",
                        "description": "The command: 'view', 'str_replace', 'create', 'insert'",
                        "enum": ["view", "str_replace", "create", "insert"]
                    },
                    "path": {
                        "type": "string",
                        "description": "Absolute path to the file"
                    },
                    "old_str": {
                        "type": "string",
                        "description": "String to replace (for str_replace command)"
                    },
                    "new_str": {
                        "type": "string",
                        "description": "Replacement string (for str_replace command)"
                    },
                    "file_text": {
                        "type": "string",
                        "description": "Full file content (for create command)"
                    },
                    "insert_line": {
                        "type": "integer",
                        "description": "Line number to insert at (for insert command)"
                    },
                    "view_range": {
                        "type": "array",
                        "items": {"type": "integer"},
                        "description": "Line range to view [start, end] (for view command)"
                    }
                },
                "required": ["command", "path"]
            }
        }
    }
]

# ===== PREPROCESSING =====
# Pattern to match function calls in the XML format
FUNC_PATTERN = re.compile(r'<function=(\w+)>(.*?)</function>', re.DOTALL)
PARAM_PATTERN = re.compile(r'<parameter=(\w+)>(.*?)</parameter>', re.DOTALL)


def parse_function_call(func_name, body):
    """Parse parameters from the function call body."""
    params = {}
    for match in PARAM_PATTERN.finditer(body):
        pname = match.group(1)
        pval = match.group(2).strip()
        # Try to parse as JSON for structured values
        if pname == "view_range":
            try:
                pval = json.loads(pval)
            except:
                pass
        elif pname == "insert_line":
            try:
                pval = int(pval)
            except:
                pass
        params[pname] = pval
    return params


def convert_trajectory(example):
    """Convert a trajectory from XML function format to proper tool_calls format."""
    msgs = example["messages"]
    if isinstance(msgs, str):
        msgs = json.loads(msgs)
    
    converted = []
    tool_call_id = 0
    pending_tool_call_ids = []  # Track IDs of pending tool calls
    
    for msg in msgs:
        role = msg.get("role", "")
        content = msg.get("content", "") or ""
        
        if not isinstance(content, str):
            # Skip non-string content (shouldn't happen but be safe)
            continue
        
        if role == "system":
            # Replace the system prompt with a cleaner version
            converted.append({
                "role": "system",
                "content": (
                    "You are an expert software engineer that navigates and modifies code repositories "
                    "using bash commands (grep, find, cat, etc.) and a file editor. You systematically "
                    "explore the codebase to understand the structure, find relevant files, and make "
                    "precise changes. Think step by step about what to search for and why."
                )
            })
        
        elif role == "assistant":
            # Check for function calls
            func_matches = list(FUNC_PATTERN.finditer(content))
            
            if func_matches:
                # Extract reasoning text before first function call
                reasoning = content[:func_matches[0].start()].strip()
                
                tool_calls = []
                pending_tool_call_ids = []
                for fmatch in func_matches:
                    func_name = fmatch.group(1)
                    body = fmatch.group(2)
                    params = parse_function_call(func_name, body)
                    
                    call_id = f"call_{tool_call_id}"
                    tool_calls.append({
                        "id": call_id,
                        "type": "function",
                        "function": {
                            "name": func_name,
                            "arguments": json.dumps(params)
                        }
                    })
                    pending_tool_call_ids.append(call_id)
                    tool_call_id += 1
                
                new_msg = {"role": "assistant"}
                if reasoning:
                    new_msg["content"] = reasoning
                else:
                    new_msg["content"] = None
                new_msg["tool_calls"] = tool_calls
                converted.append(new_msg)
            else:
                # Pure text assistant message
                converted.append({"role": "assistant", "content": content})
        
        elif role == "user":
            # Check if this is an OBSERVATION (tool response)
            if content.startswith("OBSERVATION:"):
                obs_content = content[len("OBSERVATION:"):].strip()
                # Truncate very long observations (bash output can be huge)
                if len(obs_content) > 4000:
                    obs_content = obs_content[:2000] + "\n... [truncated] ...\n" + obs_content[-1000:]
                
                if pending_tool_call_ids:
                    call_id = pending_tool_call_ids.pop(0)
                else:
                    call_id = f"call_{tool_call_id - 1}"
                
                converted.append({
                    "role": "tool",
                    "tool_call_id": call_id,
                    "content": obs_content
                })
            else:
                # Regular user message
                converted.append({"role": "user", "content": content})
    
    return {"messages": converted, "tools": TOOLS}


def filter_and_preprocess(dataset):
    """Filter for resolved trajectories and convert format."""
    # Filter for resolved
    dataset = dataset.filter(lambda x: x["resolved"] == True)
    
    # Convert format
    dataset = dataset.map(
        convert_trajectory,
        remove_columns=[c for c in dataset.column_names if c not in ["messages", "tools"]],
        num_proc=4,
    )
    
    return dataset


# ===== MAIN =====
def main():
    print("=" * 60)
    print("Training Qwen3-8B Code Navigator")
    print("=" * 60)
    
    # Load tokenizer to validate format
    print("\nLoading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    
    # Load dataset
    print("\nLoading SWE-smith trajectories (tool split)...")
    dataset = load_dataset("SWE-bench/SWE-smith-trajectories", split="tool")
    print(f"Total trajectories: {len(dataset)}")
    
    # Filter and preprocess
    print("\nFiltering for resolved trajectories and converting format...")
    dataset = filter_and_preprocess(dataset)
    print(f"Resolved trajectories: {len(dataset)}")
    
    # Validate a few examples with the tokenizer
    print("\nValidating format with Qwen3 tokenizer...")
    for i in range(min(3, len(dataset))):
        try:
            text = tokenizer.apply_chat_template(
                dataset[i]["messages"],
                tools=dataset[i]["tools"],
                tokenize=False
            )
            tokens = tokenizer.encode(text)
            print(f"  Example {i}: {len(dataset[i]['messages'])} messages -> {len(tokens)} tokens")
        except Exception as e:
            print(f"  Example {i}: ERROR - {e}")
    
    # LoRA config
    peft_config = LoraConfig(
        r=64,
        lora_alpha=128,
        target_modules=[
            "q_proj", "k_proj", "v_proj", "o_proj",
            "gate_proj", "up_proj", "down_proj"
        ],
        lora_dropout=0.05,
        task_type="CAUSAL_LM",
    )
    
    # Training config
    training_args = SFTConfig(
        output_dir=OUTPUT_DIR,
        hub_model_id=HUB_MODEL_ID,
        push_to_hub=True,
        
        # Training hyperparams (from SWE-Master recipe)
        num_train_epochs=2,
        per_device_train_batch_size=1,
        gradient_accumulation_steps=8,
        gradient_checkpointing=True,
        gradient_checkpointing_kwargs={"use_reentrant": False},
        learning_rate=1e-4,
        warmup_ratio=0.1,
        lr_scheduler_type="cosine",
        weight_decay=0.01,
        bf16=True,
        
        # Max sequence length
        max_length=MAX_LENGTH,
        
        # Train only on assistant outputs (SWE-Master: exclude env outputs from loss)
        assistant_only_loss=True,
        
        # Logging
        logging_steps=5,
        logging_strategy="steps",
        logging_first_step=True,
        disable_tqdm=True,
        report_to=["trackio"],
        
        # Save
        save_strategy="steps",
        save_steps=200,
        save_total_limit=3,
        
        # Model loading kwargs
        model_init_kwargs={
            "torch_dtype": "bfloat16",
            "attn_implementation": "flash_attention_2",
        },
        
        # Other
        dataloader_num_workers=4,
        seed=42,
    )
    
    # Create trainer
    print("\nInitializing SFTTrainer...")
    trainer = SFTTrainer(
        model=MODEL_NAME,
        args=training_args,
        train_dataset=dataset,
        peft_config=peft_config,
    )
    
    # Train
    print("\nStarting training...")
    trainer.train()
    
    # Save and push
    print("\nSaving final model...")
    trainer.save_model()
    trainer.push_to_hub()
    
    print("\n" + "=" * 60)
    print("Training complete!")
    print(f"Model pushed to: https://huggingface.co/{HUB_MODEL_ID}")
    print("=" * 60)


if __name__ == "__main__":
    main()