""" 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'(.*?)', re.DOTALL) PARAM_PATTERN = re.compile(r'(.*?)', 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()