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

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 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) parses files, the UnixCoder class generates embeddings via 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 demonstrate AST-based extraction.
• Other languages or raw files – Fall back to 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.

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, 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.

# 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.
• Keep max_length=512 to match the production embedding behavior.
• The final checkpoint must contain pytorch_model.bin, config.json, and tokenizer.json.

---

Step 4: Register the Fine-Tuned Model

Code-Graph-RAG loads the model via 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 respects this override:

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 directly:

# 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.

# 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 to parse source files.
2. Parsed nodes are passed to UnixCoder.encode() for vectorization.
3. 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:

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	Embedding model wrapper; modify UNIXCODER_MODEL or __init__ to load custom checkpoints
codebase_rag/utils/dependencies.py	Lazy loading utilities; checks UNIXCODER_MODEL environment variable
codebase_rag/vector_store.py	Qdrant client; consumes embeddings from UnixCoder
codebase_rag/parsers/cpp/utils.py	Example Tree-sitter extraction for domain corpus preparation
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.
• 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. 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.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}