Upload README.md

README.md

@@ -0,0 +1,246 @@
+# 🛠️ 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}
+}
+```