# How to Fine-Tune the RAG Model for Specific Code Domains in Code-Graph-RAG

> Learn to fine-tune the RAG model for your code domain. Adapt UnixCoder embeddings and re-index your codebase for improved retrieval accuracy with Code-Graph-RAG.

- Repository: [Vitali Avagyan/code-graph-rag](https://github.com/vitali87/code-graph-rag)
- Tags: how-to-guide
- Published: 2026-08-14

---

**Fine-tune Code-Graph-RAG for custom code domains by adapting the UnixCoder embedding model (RoBERTa architecture) to your domain corpus, then re-index your codebase to generate specialized embeddings that improve retrieval accuracy.**

Code-Graph-RAG is an open-source retrieval-augmented generation framework that combines graph-based code representation with dense vector search. To specialize it for proprietary languages, legacy frameworks, or domain-specific patterns, you must **fine-tune the embedding model** that converts code snippets into dense vectors. This guide walks through the complete workflow using the actual source paths and classes from the `vitali87/code-graph-rag` repository.

---

## Understanding the RAG Architecture

Before modifying anything, you need to know how the pipeline fits together. Code-Graph-RAG relies on two core components:

1. **Embedding Model** – [`codebase_rag/unixcoder.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/unixcoder.py) wraps a **RoBERTa-style** transformer that encodes code fragments into 768-dimensional vectors.
2. **Graph Store & Query Engine** – Memgraph stores structural relationships (call graphs, imports), while Qdrant handles approximate nearest neighbor search.

The `codebase_rag/` package coordinates these pieces. When you run `cgr start`, the ingestor ([`codebase_rag/ingestor.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/ingestor.py)) parses files, the `UnixCoder` class generates embeddings via [`codebase_rag/vector_store.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/vector_store.py), and both stores are populated.

---

## Step 1: Prepare Your Domain-Specific Corpus

Collect representative source files from your target domain. The quality of fine-tuning depends entirely on corpus relevance.

**Recommended extraction paths by language:**

- **C/C++, Go, Python, TypeScript, Java** – Use the Tree-sitter parsers in `codebase_rag/parsers/`. The C++ utilities at [`codebase_rag/parsers/cpp/utils.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/parsers/cpp/utils.py) demonstrate AST-based extraction.
- **Other languages or raw files** – Fall back to [`codebase_rag/parsers/io_access/descriptor.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/parsers/io_access/descriptor.py) for generic file reading.

Aim for 10,000–100,000 lines of code minimum. Include realistic patterns: error handling, API usage, testing idioms, and architectural conventions unique to your domain.

---

## Step 2: Convert to Hugging Face Dataset Format

The fine-tuning script expects a `datasets.Dataset` with a `"text"` column. Each row should contain a contiguous code snippet.

```python
from pathlib import Path
from datasets import Dataset

def load_domain_corpus(root: Path) -> Dataset:
    """Load all source files under root as a Hugging Face Dataset."""
    snippets = []
    for file in root.rglob("*.*"):
        # Skip binaries and very large files

        if file.stat().st_size > 100_000:
            continue
        try:
            text = file.read_text(encoding="utf-8")
            snippets.append({"text": text})
        except UnicodeDecodeError:
            pass
    return Dataset.from_dict({"text": [s["text"] for s in snippets]})

train_ds = load_domain_corpus(Path("/path/to/domain/code"))
train_ds.save_to_disk("./domain_corpus")

```

This mirrors the tensor preparation logic used internally by [`codebase_rag/unixcoder.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/unixcoder.py), which expects `source_ids` tensors derived from tokenized text.

---

## Step 3: Fine-Tune the Embedding Model

Create a training script that runs **Masked Language Modeling (MLM)** on your corpus. You must use the same architecture as the base UnixCoder to ensure compatibility with the rest of the pipeline.

```python

# fine_tune_unixcoder.py

import torch
from transformers import (
    RobertaTokenizer,
    RobertaForMaskedLM,
    Trainer,
    TrainingArguments,
)
from datasets import load_from_disk

MODEL_NAME = "microsoft/unixcoder-base"  # Same base as Code-Graph-RAG

def main():
    # Load tokenizer—identical class to that imported in codebase_rag/unixcoder.py

    tokenizer = RobertaTokenizer.from_pretrained(MODEL_NAME)
    
    # Load and tokenize domain corpus

    dataset = load_from_disk("./domain_corpus")
    
    def tokenize_function(examples):
        return tokenizer(
            examples["text"],
            truncation=True,
            max_length=512,
            padding="max_length",
            return_special_tokens_mask=True,
        )
    
    tokenized = dataset.map(
        tokenize_function,
        batched=True,
        remove_columns=["text"],
    )
    
    # Initialize RoBERTa model for MLM

    model = RobertaForMaskedLM.from_pretrained(MODEL_NAME)
    
    training_args = TrainingArguments(
        output_dir="./fine_tuned_unixcoder",
        per_device_train_batch_size=8,
        num_train_epochs=3,
        learning_rate=5e-5,
        warmup_ratio=0.1,
        weight_decay=0.01,
        save_strategy="epoch",
        logging_steps=100,
        fp16=torch.cuda.is_available(),
    )
    
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized,
    )
    
    trainer.train()
    trainer.save_model("./fine_tuned_unixcoder_final")

if __name__ == "__main__":
    main()

```

**Critical compatibility notes:**

- **`RobertaTokenizer`** and **`RobertaModel`** are the exact classes used in [`codebase_rag/unixcoder.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/unixcoder.py).
- Keep `max_length=512` to match the production embedding behavior.
- The final checkpoint must contain `pytorch_model.bin`, [`config.json`](https://github.com/vitali87/code-graph-rag/blob/main/config.json), and [`tokenizer.json`](https://github.com/vitali87/code-graph-rag/blob/main/tokenizer.json).

---

## Step 4: Register the Fine-Tuned Model

Code-Graph-RAG loads the model via [`codebase_rag/unixcoder.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/unixcoder.py). You have two integration options.

### Option A: Environment Variable Override (Recommended)

Set `UNIXCODER_MODEL` before running any CLI command. The loading logic in [`codebase_rag/utils/dependencies.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/utils/dependencies.py) respects this override:

```bash
export UNIXCODER_MODEL="/absolute/path/to/fine_tuned_unixcoder_final"
cgr start --repo-path /my/project --clean

```

Use `local_files_only=True` implicitly by ensuring the path is local and valid.

### Option B: Hardcoded Path in Source

For persistent configuration, modify [`codebase_rag/unixcoder.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/unixcoder.py) directly:

```python

# codebase_rag/unixcoder.py (excerpt)

from transformers import RobertaTokenizer, RobertaConfig, RobertaModel

# Override the default model identifier

UNIXCODER_MODEL = "fine_tuned_unixcoder_final"  # ← your checkpoint path

class UnixCoder:
    def __init__(self, model_name: str = UNIXCODER_MODEL) -> None:
        self.tokenizer = RobertaTokenizer.from_pretrained(
            model_name, local_files_only=True
        )
        self.config = RobertaConfig.from_pretrained(
            model_name, local_files_only=True
        )
        self.model = RobertaModel.from_pretrained(
            model_name, config=self.config, local_files_only=True
        )
        self.model.eval()

```

The `local_files_only=True` flag ensures the library never falls back to downloading the base model from Hugging Face.

---

## Step 5: Re-Index Your Codebase

Fine-tuned models produce different embedding distributions. You must **fully re-index** to avoid mixing incompatible vectors in Qdrant.

```bash

# Ensure infrastructure is running

cgr daemon up

# Clean and re-index with new embeddings

cgr start --repo-path /my/project --clean

# Verify: check that embedding dimension matches (should be 768)

```

The ingestion flow works as follows:

1. **`cgr start`** invokes [`codebase_rag/ingestor.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/ingestor.py) to parse source files.
2. Parsed nodes are passed to `UnixCoder.encode()` for vectorization.
3. [`codebase_rag/vector_store.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/vector_store.py) inserts vectors into Qdrant using the configured model.

The `--clean` flag wipes existing Memgraph and Qdrant data, ensuring no stale embeddings remain.

---

## Step 6: Validate Domain-Specific Retrieval

Test retrieval quality with queries that exercise domain terminology:

```bash
cgr ask "How does the custom rate limiter handle burst traffic?"
cgr ask "Show me the protobuf serialization for the AuthRequest message"

```

**Success indicators:**

- retrieved snippets contain domain-specific identifiers not present in general codebases
- generated answers reference correct internal APIs and file paths
- cross-file relationships (e.g., "this function is called by...") are accurate

If retrieval fails, verify:
- embedding dimension in Qdrant matches the fine-tuned model (768)
- corpus during fine-tuning included files similar to your test queries
- no `local_files_only` errors in logs indicating checkpoint loading failure

---

## Key Files and Their Roles

| File | Purpose |
|------|---------|
| [`codebase_rag/unixcoder.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/unixcoder.py) | Embedding model wrapper; modify `UNIXCODER_MODEL` or `__init__` to load custom checkpoints |
| [`codebase_rag/utils/dependencies.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/utils/dependencies.py) | Lazy loading utilities; checks `UNIXCODER_MODEL` environment variable |
| [`codebase_rag/vector_store.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/vector_store.py) | Qdrant client; consumes embeddings from `UnixCoder` |
| [`codebase_rag/parsers/cpp/utils.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/parsers/cpp/utils.py) | Example Tree-sitter extraction for domain corpus preparation |
| [`codebase_rag/ingestor.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/ingestor.py) | Orchestrates parsing and embedding pipeline |

---

## Summary

- **Fine-tune the embedding model** using MLM on your domain corpus, keeping the RoBERTa architecture compatible with [`codebase_rag/unixcoder.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/unixcoder.py).
- **Point Code-Graph-RAG to your checkpoint** via `UNIXCODER_MODEL` environment variable or direct source modification.
- **Re-index completely** with `--clean` to populate Qdrant with domain-aware embeddings.
- **Validate** with targeted queries to confirm improved retrieval of specialized patterns.

The entire workflow preserves the existing graph-store infrastructure while specializing the semantic layer that drives retrieval relevance.

---

## Frequently Asked Questions

### What base model should I use for fine-tuning?

Always start from `microsoft/unixcoder-base`. This is the checkpoint hardcoded in [`codebase_rag/unixcoder.py`](https://github.com/vitali87/code-graph-rag/blob/main/codebase_rag/unixcoder.py). Using a different architecture (e.g., CodeT5, CodeBERT) will break compatibility with the tokenizer and model loading logic throughout the codebase.

### How much training data do I need for effective domain adaptation?

Minimum 5,000–10,000 code snippets; 50,000+ for substantial improvement. Quality matters more than quantity—prioritize files that represent the query distribution you expect in production. For narrow domains (a single large proprietary framework), even 2,000 well-curated examples can yield noticeable gains.

### Can I fine-tune for multiple domains simultaneously?

Yes. Create a mixed corpus containing all target domains, or fine-tune separate checkpoints and switch between them using the `UNIXCODER_MODEL` environment variable. There is no built-in multi-domain routing—you must manage checkpoint selection externally.

### Will fine-tuning affect the graph structure in Memgraph?

No. Memgraph stores structural relationships (imports, calls, inheritance) parsed by Tree-sitter, which is independent of the embedding model. Only the **vector embeddings in Qdrant** change. The combined RAG results will improve because retrieval selects more relevant nodes from an unchanged graph topology.