README.md

---
title: Fine-Tuned RAG Framework For Code Analysis
colorFrom: gray
colorTo: gray
sdk: gradio
sdk_version: 5.41.1
app_file: app.py
pinned: false
license: mit
short_description: RAG + LoRA Fine-Tuning for Code Analysis
---

# Fine-Tuned RAG Framework for Code Analysis

This project is a production-ready, Retrieval-Augmented Generation (RAG) system specifically designed for code analysis and software development queries. It integrates a fine-tuned Large Language Model (LLM) with a vector database to provide accurate, context-aware answers based on a comprehensive knowledge base of programming concepts.

A key feature of this system is its **automatic fine-tuning process**. On initialization, it automatically fine-tunes the base model (`Salesforce/codegen-350M-mono`) on a curated dataset of code-related questions and answers using Parameter-Efficient Fine-Tuning (PEFT) with LoRA. This ensures the model is specialized for the domain of software engineering, resulting in higher quality and more relevant responses.

## Core Features

* **Automatic Model Fine-Tuning**: The system automatically fine-tunes a code-specific language model (`Salesforce/codegen-350M-mono`) on startup using LoRA to specialize it for code analysis tasks.
* **Retrieval-Augmented Generation (RAG)**: Leverages a `ChromaDB` vector store and a `sentence-transformers` model to retrieve the most relevant documents from a knowledge base to ground the LLM's responses in factual information.
* **Code-Specific Knowledge Base**: The system is pre-loaded with a detailed knowledge base covering software architecture patterns, clean code principles, testing strategies, performance optimization, and API design best practices.
* **Comprehensive Evaluation Metrics**: Every generated response is evaluated in real-time for relevance, context grounding, hallucination, and technical accuracy. The system also calculates a performance improvement score based on whether the fine-tuned model is active.
* **Performance & Cost Tracking**: A built-in `PerformanceTracker` monitors system usage, including query latency, tokens processed, and estimated operational costs, providing a full overview of system efficiency.
* **Source Attribution**: To ensure transparency and trustworthiness, every answer is accompanied by a list of the source documents from the knowledge base that were used to generate the response.

## How It Works

The system follows an automated pipeline from startup to query processing:

1.  **Automatic Fine-Tuning (On Startup)**: The `ModelFineTuner` class initiates an automatic fine-tuning process. It loads the base `Salesforce/codegen-350M-mono` model, applies a LoRA configuration, and trains it on a specialized dataset of code analysis Q&A pairs. The resulting fine-tuned model is then used for generation.
2.  **Knowledge Base Indexing**: The `RAGSystem` class initializes a `ChromaDB` vector store. It processes and chunks the provided code documentation, computes embeddings using a `SentenceTransformer` model, and indexes these embeddings for efficient retrieval.
3.  **Query Processing**: A user submits a query through the Gradio interface.
4.  **Retrieval**: The system encodes the user's query into an embedding and uses it to search the `ChromaDB` vector store, retrieving the top-k most relevant text chunks from the knowledge base.
5.  **Generation**: The retrieved chunks are formatted as context and prepended to the user's query in a prompt. This combined prompt is then passed to the fine-tuned LLM, which generates a context-aware and accurate answer.
6.  **Evaluation & Display**: The final response is evaluated for quality, and all relevant information—the answer, sources, metrics, and performance data—is presented to the user in the interactive dashboard.

## Technical Stack

* **AI & Machine Learning**: `transformers`, `peft`, `bitsandbytes`, `torch`, `accelerate`
* **Retrieval & Vector Search**: `chromadb`, `sentence-transformers`, `langchain`
* **Data Processing**: `pandas`, `numpy`, `datasets`
* **Web Interface & Dashboard**: `gradio`

## How to Use the Demo

The interface is designed for simplicity and provides a wealth of information with each query.

1.  Enter a question related to software development, architecture, or best practices into the text box. You can use the provided sample queries as inspiration.
2.  Click the **Analyze Query** button.
3.  Review the output in the panels:
    * **Analysis Result**: The generated answer from the fine-tuned RAG system.
    * **Referenced Sources**: The documents from the knowledge base used to formulate the answer.
    * **Response Metrics**: A detailed breakdown of the response quality, including relevance, grounding, and technical accuracy scores.
    * **Performance Data**: Information on the query processing time, tokens used, and estimated cost.
    * **System Statistics**: An overview of cumulative usage and performance.

## Disclaimer

This project is a demonstration of a production-ready RAG system for a specialized domain. The automatic fine-tuning process is computationally intensive and may be slow on CPU. For optimal performance, running this demo on a GPU-accelerated environment is recommended. All generated responses are for informational purposes and should be validated by a qualified professional.