A complete toolkit for building a Python code generation LLM that can search your internal codebase via RAG, call tools, and reason through multi-step tasks.
Merges 4 verified datasets (ToolACE + APIGen-MT + Magicoder + CodeAct) into a unified ChatML format:
bash
1pip install datasets
23# Test with small sample first4python prepare_data.py --max_per_source 100 --dry_run
56# Full run — pushes merged dataset to Hub7python prepare_data.py --output_repo your-username/code-toolcall-sft-data
This is the most impactful step. Use the Gorilla/Magicoder pattern:
OSS-Instruct on your code: Sample random snippets from your internal repo → use an LLM (GPT-4o, Claude) to generate instruction-solution pairs seeded from that code
Retriever-aware examples: Include retrieved code context in training prompts so the model learns to use RAG at inference time
Internal API documentation: Convert your docstrings/README into Q&A pairs
See prepare_data.py for the format — add your examples as additional sources.
Step 3: Fine-tune
bash
1# Edit train_sft.py to set your dataset and model repo IDs, then:23# Option A: Run on HF Jobs (recommended for A100/H100 hardware)4# Use the hf_jobs API or CLI56# Option B: Run locally with GPU7pip install trl peft transformers datasets trackio accelerate torch
8python train_sft.py
For maximum performance, skip LoRA and do full fine-tuning:
python
1# In train_sft.py, remove peft_config and adjust:2LEARNING_RATE =2e-5# 10x lower than LoRA3BATCH_SIZE =1# Lower to fit in memory4GRAD_ACCUM =16# Keep effective batch = 165# Hardware: 2x A100-80GB minimum for 7B FFT
Per Astraios: FFT slightly outperforms LoRA at 7B scale, but LoRA is within 1% and 30x more parameter-efficient.
Advanced: GRPO Reinforcement Learning (Stage 2)
After SFT, you can further improve the model with GRPO using execution-based rewards:
python
1from trl import GRPOConfig, GRPOTrainer
23# Reward function: does the generated code pass unit tests?4defreward_fn(completions, prompts):5 rewards =[]6for code in completions:7try:8exec(code,{})# Sandbox this properly!9 rewards.append(1.0)10except:11 rewards.append(0.0)12return rewards
1314# Train with GRPO15config = GRPOConfig(16 learning_rate=1e-6,17 num_train_epochs=1,18 per_device_train_batch_size=4,19)