README.md

# 🛠️ Code LLM Toolkit: Fine-tune + RAG + Tool-Calling for Internal Codebases

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.

## Architecture

```
┌─────────────────────────────────────────────────────────────────┐
│                     Code LLM System                             │
│                                                                 │
│  ┌──────────────┐   ┌──────────────┐   ┌──────────────────────┐│
│  │  Fine-tuned   │   │  RAG Pipeline │   │  Tool Executor       ││
│  │  Qwen2.5-     │◄──│  (AST-aware   │   │  - search_codebase  ││
│  │  Coder-7B     │   │   chunking +  │   │  - execute_python   ││
│  │  + LoRA       │   │   embeddings) │   │  - read_file        ││
│  └──────┬───────┘   └──────────────┘   │  - run_tests        ││
│         │                               └──────────┬───────────┘│
│         │            ReAct Agent Loop              │            │
│         └──────────────────┬───────────────────────┘            │
│                            │                                     │
│                     ┌──────▼──────┐                              │
│                     │   Response   │                              │
│                     └─────────────┘                              │
└─────────────────────────────────────────────────────────────────┘
```

## Research Foundation

Every component is grounded in published research with verified results:

| Component | Paper | Key Result |
|-----------|-------|------------|
| **Base Model** | [Qwen2.5-Coder](https://arxiv.org/abs/2409.12186) | HumanEval 88.4% (7B), SOTA open-source |
| **Tool-Calling Data** | [ToolACE](https://arxiv.org/abs/2409.00920) | Beats GPT-4-turbo on BFCL benchmark |
| **Multi-turn Agent Data** | [APIGen-MT](https://arxiv.org/abs/2504.03601) | 78.19% BFCL v3 (#1, beats o1/GPT-4o) |
| **Code SFT Data** | [Magicoder](https://arxiv.org/abs/2312.02120) | HumanEval 70.7% from 7B with 185K samples |
| **RAG Chunking** | [cAST](https://arxiv.org/abs/2506.15655) | +5.6pp over fixed-size on RepoEval |
| **RAG Strategy** | [AllianceCoder](https://arxiv.org/abs/2503.20589) | API signatures > similar code (+20%) |
| **Retriever-Aware Training** | [Gorilla](https://arxiv.org/abs/2305.15334) | Outperforms GPT-4 on API accuracy |
| **Code Embeddings** | [CodeSage-v2](https://huggingface.co/codesage/codesage-large-v2) | Best open code embedding model |
| **LoRA for Code** | [Astraios](https://arxiv.org/abs/2401.00788) | LoRA matches FFT at ≥16B scale |

## Quick Start

### Step 1: Prepare Training Data

Merges 4 verified datasets (ToolACE + APIGen-MT + Magicoder + CodeAct) into a unified ChatML format:

```bash
pip install datasets

# Test with small sample first
python prepare_data.py --max_per_source 100 --dry_run

# Full run — pushes merged dataset to Hub
python prepare_data.py --output_repo your-username/code-toolcall-sft-data
```

**Dataset composition (~110K examples):**
| Source | Examples | Purpose |
|--------|----------|---------|
| [Team-ACE/ToolACE](https://huggingface.co/datasets/Team-ACE/ToolACE) | 26K | Tool-calling (single-turn) |
| [Salesforce/APIGen-MT-5k](https://huggingface.co/datasets/Salesforce/APIGen-MT-5k) | 5K | Multi-turn agentic tool use |
| [Magicoder-OSS-Instruct-75K](https://huggingface.co/datasets/ise-uiuc/Magicoder-OSS-Instruct-75K) | ~25K (Python) | Python code generation |
| [xingyaoww/code-act](https://huggingface.co/datasets/xingyaoww/code-act) | 7K | Code-as-action (tools via Python) |

### Step 2: Add Your Internal Codebase Data

**This is the most impactful step.** Use the Gorilla/Magicoder pattern:

1. **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
2. **Retriever-aware examples:** Include retrieved code context in training prompts so the model learns to use RAG at inference time
3. **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
# Edit train_sft.py to set your dataset and model repo IDs, then:

# Option A: Run on HF Jobs (recommended for A100/H100 hardware)
# Use the hf_jobs API or CLI

# Option B: Run locally with GPU
pip install trl peft transformers datasets trackio accelerate torch
python train_sft.py
```

**Training configuration (from literature):**
- **Base:** [Qwen/Qwen2.5-Coder-7B-Instruct](https://huggingface.co/Qwen/Qwen2.5-Coder-7B-Instruct) (Apache 2.0)
- **Method:** LoRA (r=32, alpha=64) on all linear layers
- **LR:** 1e-4 with cosine schedule, 10% warmup
- **Epochs:** 2
- **Context:** 8192 tokens
- **Loss:** Assistant-only (masks user/system/tool tokens)
- **Hardware:** 1x A100-80GB (or 2x A10G-24GB)
- **Time:** ~4-6 hours for 110K examples

### Step 4: Index Your Codebase (RAG)

```python
from rag_pipeline import CodebaseIndexer

# Index your internal Python codebase
indexer = CodebaseIndexer(
    "/path/to/your/repo",
    embedding_model="jinaai/jina-embeddings-v2-base-code"  # or codesage-large-v2
)
retriever = indexer.index()

# Save for reuse
retriever.save_index("./my_index")

# Search!
results = retriever.search("authentication token validation", top_k=5)
for chunk, score in results:
    print(f"[{score:.3f}] {chunk.file_path}/{chunk.name}: {chunk.signature}")
```

**RAG pipeline features:**
- **AST-aware chunking** (cAST): Functions, methods, classes stay intact — no mid-function cuts
- **Dual embeddings**: Code content + metadata strings (NL descriptions) for hybrid search
- **AllianceCoder context assembly**: API signatures prioritized over full code bodies
- **In-context dependencies**: Automatically extracts imports and class signatures from the current file
- **Embedding models**: Jina-Code-v2 (8K context, 161M) or CodeSage-v2 (best quality, 1.3B)

### Step 5: Run the Agent

```bash
# Interactive mode
python inference.py \
  --model your-username/qwen25-coder-7b-code-toolcall \
  --repo /path/to/your/codebase \
  --index-dir ./my_index

# Single query
python inference.py \
  --model your-username/qwen25-coder-7b-code-toolcall \
  --repo /path/to/your/codebase \
  --query "Add pagination to the product search endpoint"
```

The agent uses a **ReAct loop**:
1. Pre-fetches relevant code via RAG
2. Sends query + context to the LLM
3. If the LLM calls tools → executes them → feeds results back
4. Repeats until the LLM gives a final answer (max 10 turns)

## Recommended Embedding Models

| Model | Size | Context | Best For | HF Link |
|-------|------|---------|----------|---------|
| `codesage/codesage-large-v2` | 1.3B | 2048 tok | Best quality (NL→Code 69.4) | [Link](https://huggingface.co/codesage/codesage-large-v2) |
| `jinaai/jina-embeddings-v2-base-code` | 161M | **8192 tok** | Long files, 30 languages | [Link](https://huggingface.co/jinaai/jina-embeddings-v2-base-code) |
| `codesage/codesage-small-v2` | 130M | 2048 tok | Fast, lightweight | [Link](https://huggingface.co/codesage/codesage-small-v2) |

## Advanced: Full Fine-Tuning (FFT)

For maximum performance, skip LoRA and do full fine-tuning:

```python
# In train_sft.py, remove peft_config and adjust:
LEARNING_RATE = 2e-5    # 10x lower than LoRA
BATCH_SIZE = 1          # Lower to fit in memory
GRAD_ACCUM = 16         # Keep effective batch = 16
# Hardware: 2x A100-80GB minimum for 7B FFT
```

Per [Astraios](https://arxiv.org/abs/2401.00788): 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
from trl import GRPOConfig, GRPOTrainer

# Reward function: does the generated code pass unit tests?
def reward_fn(completions, prompts):
    rewards = []
    for code in completions:
        try:
            exec(code, {})  # Sandbox this properly!
            rewards.append(1.0)
        except:
            rewards.append(0.0)
    return rewards

# Train with GRPO
config = GRPOConfig(
    learning_rate=1e-6,
    num_train_epochs=1,
    per_device_train_batch_size=4,
)
```

## File Structure

```
├── prepare_data.py     # Dataset merging & formatting
├── train_sft.py        # SFT training script (TRL + LoRA)
├── rag_pipeline.py     # AST-aware indexing & retrieval
├── inference.py        # ReAct agent with tool calling
└── README.md           # This file
```

## Requirements

```
transformers>=4.45.0
trl>=1.0.0
peft>=0.12.0
datasets>=3.0.0
accelerate>=1.0.0
trackio>=0.2.0
torch>=2.0.0
sentence-transformers>=3.0.0  # For RAG embeddings
numpy
scikit-learn  # TF-IDF fallback
```

## Citation

If you use this toolkit, please cite the underlying research:

```bibtex
@article{qwen2.5coder,
  title={Qwen2.5-Coder Technical Report},
  author={Hui, Binyuan and others},
  journal={arXiv:2409.12186},
  year={2024}
}
@article{toolace,
  title={ToolACE: Winning the Points of LLM Function Calling},
  author={Liu, Weiwen and others},
  journal={arXiv:2409.00920},
  year={2024}
}
@article{gorilla,
  title={Gorilla: Large Language Model Connected with Massive APIs},
  author={Patil, Shishir G. and others},
  journal={arXiv:2305.15334},
  year={2023}
}
```