Upload 54 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+data/indices/semantic_index.faiss filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,10 +1,446 @@
 ---
-title: CodeGen Emasters IIITH
-emoji: 📉
-colorFrom: yellow
-colorTo: gray
-sdk: static
+title: Text to MongoDB Generator
+emoji: 🚀
+colorFrom: blue
+colorTo: purple
+sdk: gradio
+sdk_version: 5.9.1
+app_file: app.py
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# CodeGen: Hybrid RAG Pipeline
+
+[![Python 3.10+](https://img.shields.io/badge/Python-3.10%2B-blue)](https://www.python.org/)
+[![PyTorch](https://img.shields.io/badge/PyTorch-2.5.1-red)](https://pytorch.org/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+
+A capstone project (eMasters CS: AI & ML, Group 2) that fine-tunes a small open-source code language model — **Qwen2.5-Coder-0.5B-Instruct** — with **LoRA**, builds a **hybrid semantic + AST/structural retrieval (RAG)** layer on top of it, and evaluates the result against the un-adapted base model and an independently invoked LLM baseline, across text-to-SQL, documentation-generation, program-synthesis, and commit-message-generation tasks.
+
+---
+
+## Table of Contents
+
+- [Executive Summary](#-executive-summary)
+- [Key Capabilities](#-key-capabilities)
+- [System Architecture](#-system-architecture)
+- [Pipeline Stages](#-pipeline-stages)
+- [Benchmark Results](#-benchmark-results)
+- [Dataset](#-dataset)
+- [Installation](#-installation)
+- [Quick Start](#-quick-start)
+- [Repository Layout](#-repository-layout)
+- [Configuration](#-configuration)
+- [Requirements](#-requirements)
+- [License](#-license)
+
+---
+
+## Executive Summary
+
+CodeGen combines **Retrieval-Augmented Generation (RAG)** and **parameter-efficient LoRA fine-tuning** on top of a compact instruction-tuned code model to support three developer-facing capabilities: natural-language-to-SQL translation, automatic source-code documentation generation, and program/commit-message generation. The primary evaluation dataset is the **Spider** cross-domain text-to-SQL benchmark (1,034 dev-split examples), supplemented by small hand-curated example sets for documentation, program-synthesis, and commit-message generation.
+
+Retrieval embeddings are produced by the small code LM itself (no second embedding model needs to be downloaded), and are combined with an AST/structural fingerprint index so retrieval can be reranked on structure, not just vector similarity.
+
+---
+
+## Key Capabilities
+
+| Capability | Description |
+|---|---|
+| **Multi-Task Generation** | Program synthesis, docstring generation, Text-to-SQL translation, commit-message synthesis |
+| **Code-LM Self-Embeddings** | Retrieval embeddings mean-pooled from the base model's own hidden states — no separate embedding model required (an external `BAAI/bge-small-en-v1.5` path is also available and config-switchable) |
+| **Hybrid Retrieval** | FAISS (`IndexFlatIP`) dense semantic search fused with an AST/structural fingerprint reranker (α = 0.7 dense / 0.3 structural) |
+| **LoRA Fine-Tuning** | PEFT LoRA adapter (r=16, α=32) on the attention projections, with a custom training loop (gradient checkpointing, prompt-token label masking, OOM-safe stepping) |
+| **Small-LM vs. LLM Harness** | Pluggable large-LLM baseline generator supporting local, OpenAI-compatible, and Gemini backends |
+| **Comprehensive Evaluation** | BLEU, ROUGE-1/L, token-F1, CodeBERTScore, SQL Exact Match, SQL Execution Accuracy, RAG gain-over-baseline, top-K sweep |
+
+---
+
+## System Architecture
+
+```
+┌─────────────────────────────────────────────────────────────────┐
+│ 1. Configuration & Data                                         │
+│  • Load config.yaml (paths, model, LoRA, retrieval, eval)       │
+│  • Convert & normalize Spider (dev.json + tables.json + SQLite) │
+│  • Build retrieval corpus (1,034 question/SQL pairs)            │
+└──────────────────────────┬──────────────────────────────────────┘
+                           │
+┌──────────────────────────▼──────────────────────────────────────┐
+│ 2. Model Loading & LoRA Fine-Tuning                             │
+│  • Load Qwen2.5-Coder-0.5B-Instruct (tokenizer + base weights)  │
+│  • Attach LoRA adapter, train on illustrative task examples     │
+│  • Save / reload adapter checkpoint                             │
+└──────────────────────────┬──────────────────────────────────────┘
+                           │
+┌──────────────────────────▼──────────────────────────────────────┐
+│ 3. Generation Tasks                                             │
+│  • Text-to-SQL   • Documentation   • Program Synthesis          │
+│  • Commit Messages, evaluated for base model vs. LoRA model     │
+└──────────────────────────┬──────────────────────────────────────┘
+                           │
+┌──────────────────────────▼──────────────────────────────────────┐
+│ 4. Hybrid Indexing & RAG                                        │
+│  • CodeLMEmbedder → FAISS semantic index (IndexFlatIP)          │
+│  • ASTIndexer → structural fingerprints (Python AST / SQL regex)│
+│  • HybridRetriever → dense + structural blended top-K search    │
+│  • RAGPipeline → context injection + generation                 │
+└──────────────────────────┬──────────────────────────────────────┘
+                           │
+┌──────────────────────────▼──────────────────────────────────────┐
+│ 5. Evaluation & Comparison                                      │
+│  • Base vs. LoRA vs. RAG vs. LLM-baseline (ModelComparator)     │
+│  • RAG gain-over-baseline & top-K sweep (RAGEvaluator)          │
+│  • Charts (matplotlib) + JSON metrics persisted to output/      │
+└─────────────────────────────────────────────────────────────────┘
+```
+
+---
+
+## Pipeline Stages
+
+The pipeline is run through a set of Jupyter notebooks rather than a single CLI entry point:
+
+| Notebook | Stage |
+|---|---|
+| `baseline_model.ipynb` | Loads the **un-adapted base model**, runs text-to-SQL and documentation generation, saves metrics to `outputs/generated/baseline_model_results.json` |
+| `finetuned_model.ipynb` | Loads/trains the **LoRA adapter**, re-runs the same two tasks, saves metrics to `outputs/generated/finetuned_model_results.json` |
+| `spider_eval.ipynb` | Spider-specific SQL execution evaluation |
+| `eval.ipynb` | RAG-gain measurement, top-K sweep, small-LM-vs-LLM comparison |
+| `main.ipynb` | Complete End-to-End pipeline |
+
+Supporting production code lives in `src/` (see [Repository Layout](#-repository-layout)) and is imported by the notebooks — it is not currently exposed as a standalone CLI script.
+
+### Task-Specific Generators (`src/generators/`)
+
+| Task | Module | Benchmark used | Primary Metrics |
+|---|---|---|---|
+| **Text-to-SQL** | `text_to_sql.py`, `sql_generator.py` | Spider (100-sample subset of the 1,034-question dev split) | Exact Match, Execution Accuracy, BLEU, CodeBERTScore |
+| **Documentation Generation** | `doc_gen.py` | 3 hand-authored illustrative examples | BLEU, ROUGE, CodeBERTScore, F1 |
+| **Program Synthesis** | `program_generator.py` | 3 hand-authored illustrative examples | Qualitative + shared metric suite |
+| **Commit Message Generation** | `commit_gen.py` | 3 hand-authored illustrative examples | Qualitative + shared metric suite |
+| **LLM Baseline** | `llm_baseline.py` | Same prompts as program synthesis | Comparative scoring against the small LM |
+
+### Hybrid Indexing & Retrieval (`src/embeddings/`)
+
+| Component | Technology | Purpose |
+|---|---|---|
+| `code_lm_embedder.py` | Mean-pooled hidden states of the base model itself | Retrieval embeddings without a second model download |
+| `embedding.py` | `BAAI/bge-small-en-v1.5` (external, config-switchable) | Alternative embedding backend |
+| `indexing.py` | FAISS `IndexFlatIP` | Dense semantic vector store (cosine similarity via normalized inner product) |
+| `ast_indexing.py` | Python `ast` module (Python code) / regex-based structural extractor (SQL) | Structural fingerprints — function/class names, or referenced tables/aggregates/clauses |
+| `hybrid_retriever.py` | Blended reranker | Over-fetches `k × 4` from FAISS, reranks with `α * dense + (1-α) * structural` |
+
+### Evaluation & Persistence (`src/evaluation/`)
+
+| Component | Function |
+|---|---|
+| `metrics.py` | BLEU, ROUGE, token-F1, CodeBERTScore, execution accuracy primitives |
+| `evaluation.py` | Task-level evaluation engine |
+| `sql_eval.py` | SQL normalization and Spider-style execution comparison |
+| `comparator.py` | Base vs. LoRA vs. RAG side-by-side comparison |
+| `rag_evaluation.py` | RAG gain-over-baseline measurement and top-K sweep |
+| `visualization.py` | Matplotlib bar-chart generation, saved to `output/plots/` |
+
+---
+
+## Benchmark Results
+
+Numbers below are taken directly from `outputs/generated/baseline_model_results.json` and `finetuned_model_results.json` (100-sample Spider text-to-SQL subset, 3-sample documentation-generation set).
+
+### Baseline (un-adapted) vs. LoRA Fine-Tuned
+
+| Task | Metric | Base Model | LoRA Fine-Tuned |
+|---|---|:---:|:---:|
+| **Text-to-SQL** | BLEU | 0.2329 | **0.2467** |
+| | CodeBERTScore | 0.9740 | **0.9759** |
+| | F1 | 0.4642 | **0.4879** |
+| | ROUGE-1 | 0.4378 | **0.4602** |
+| | ROUGE-L | 0.4274 | **0.4503** |
+| | Exact Match | 7% | **8%** |
+| | Execution Accuracy | 11% | **14%** |
+| **Documentation Generation** | BLEU | 0.0241 | **0.0528** |
+| | CodeBERTScore | 0.9250 | **0.9447** |
+| | F1 | 0.1181 | **0.2022** |
+| | ROUGE-1 | 0.0831 | **0.1615** |
+| | ROUGE-L | 0.0719 | **0.1494** |
+
+LoRA fine-tuning (trained on only 9 illustrative prompt/completion examples — 2.16M trainable params, 0.44% of the model) improves every metric on both tasks. Gains are proportionally larger on documentation generation.
+
+### RAG, Top-K Sweep, and LLM Baseline (5-question SQL slice)
+
+- On the small 5-question evaluation slice used for this stage, the **RAG-augmented pipeline underperformed the no-RAG LoRA baseline** on every metric — a genuine, reproducible negative result, not a placeholder. See [Known Limitations](#-known-limitations--future-work) for why.
+- A **top-K sweep** over K ∈ {1, 2, 3, 5, 8} shows the gain is least negative around **K = 2–3** and degrades further at K = 8.
+- The **LLM-baseline comparison** fell back to the same local Qwen2.5-Coder-0.5B model in this run because no external API credential was resolved at runtime — see the security note below and Known Limitations.
+- For reference, the rule-based Spider baseline scores **12.28% exact match / 33.66% execution accuracy** over the full 1,034-question dev set.
+
+Plots for all of the above are saved under `output/plots/` (`scores_baseline_text_to_sql_vs_doc_generation.png`, `scores_finetuned_text_to_sql_vs_doc_generation.png`, `scores_model_comparison_checkpoint2.png`, `scores_model_comparison_full.png`, `topk_gain_sweep.png`).
+
+
+## Dataset
+
+This project uses the **[Spider](https://yale-lily.github.io/spider)** dataset — a large-scale, cross-domain text-to-SQL benchmark — as its primary evaluation dataset.
+
+> **Note:** The Spider dataset is **not included in this repository**. GitHub's file/repo upload limits make it impractical to commit the raw dataset (question/SQL pairs, table schemas, and per-database SQLite files), so it must be downloaded separately before running the pipeline.
+
+### Download
+
+1. Download the dataset from Google Drive: **[Spider dataset (Google Drive)](https://drive.google.com/file/d/1403EGqzIDoHMdQF4c9Bkyl7dZLZ5Wt6J/view)**
+   - Alternatively, see the [official Spider dataset page](https://yale-lily.github.io/spider) for the source release and licensing terms.
+2. Unzip the download and place its contents under `data/raw/Spider/` at the repository root, so that `dev.json`, `tables.json`, and the per-database SQLite files sit alongside each other, e.g.:
+
+   ```
+   data/raw/Spider/
+   ├── dev.json
+   ├── tables.json
+   └── database/
+       ├── concert_singer/
+       │   └── concert_singer.sqlite
+       └── ...
+   ```
+
+3. Run `src/scripts/convert_spider_to_jsonl.py` to convert the raw files into the normalized `data/processed/spider_eval.jsonl` and retrieval corpus used by the notebooks (see [Quick Start](#-quick-start)).
+
+`data/raw/`, `data/processed/`, and `data/indices/` are excluded via `.gitignore` for the same reason — they're generated/downloaded locally rather than versioned.
+
+---
+
+## Installation
+
+### Prerequisites
+
+```bash
+Python 3.10 or higher
+A CUDA-capable GPU is optional — the code auto-detects CUDA > MPS (Apple Silicon) > CPU
+Recommended: 16 GB RAM, 8 GB+ VRAM if using GPU
+```
+
+### Step 1: Clone Repository
+
+```bash
+git clone https://github.com/Rithusravya/Emasters_Group-2_CapstoneProject.git
+cd Emasters_Group-2_CapstoneProject
+```
+
+### Step 2: Create Virtual Environment
+
+```bash
+python -m venv venv
+source venv/bin/activate      # macOS/Linux
+# OR
+venv\Scripts\activate         # Windows
+```
+
+### Step 3: Install Dependencies
+
+```bash
+pip install --upgrade pip
+pip install -r requirements.txt
+```
+
+### Step 4: Verify Installation
+
+```bash
+python -c "import torch; print(f'CUDA Available: {torch.cuda.is_available()}')"
+```
+
+---
+
+## Quick Start
+
+This project is driven through notebooks, not a CLI. Typical order:
+
+1. **Download & prepare data** — download the Spider dataset as described in [Dataset](#-dataset) and place it under `data/raw/Spider/`. Then run `src/scripts/convert_spider_to_jsonl.py` to convert the raw Spider files into `data/processed/spider_eval.jsonl` and a retrieval corpus; `src/scripts/prepare_processed_data.py` builds the MongoDB-style variant.
+2. **Baseline evaluation** — open `baseline_model.ipynb` and run all cells to evaluate the un-adapted base model on text-to-SQL and documentation generation.
+3. **LoRA fine-tuning + evaluation** — open `finetuned_model.ipynb` to train (or reload) the LoRA adapter at `models/checkpoints/lora_finetuned/` and re-run the same evaluations.
+4. **RAG / top-K / LLM comparison** — open `eval.ipynb` to build the FAISS + AST indices, measure RAG gain, sweep top-K, and compare against the LLM baseline.
+5. **Spider execution evaluation** — open `spider_eval.ipynb` for SQL-execution-based scoring against the SQLite databases.
+
+Programmatic usage of the core pipeline, e.g. for text-to-SQL generation:
+
+```python
+from src.config import load_config
+from src.models.load_model import ModelLoader
+from src.generators.program_generator import GenerationPipeline
+from src.generators.text_to_sql import TextToSQLGenerator
+
+config = load_config()  # reads configs/config.yaml
+loader = ModelLoader(config)
+models, tokenizer = loader.load_models(lora_path=config.lora.output_dir)
+
+pipeline = GenerationPipeline(models["lora"] or models["base"], tokenizer, config.generation)
+sql_gen = TextToSQLGenerator(pipeline)
+
+result = sql_gen.generate_queries(
+    question="How many singers are there?",
+    schema="CREATE TABLE singer (singer_id INT, name TEXT, age INT)",
+)
+print(result["sql_query"])
+```
+
+---
+
+## Repository Layout
+
+```
+Emasters_Group-2_CapstoneProject/
+├── configs/
+│   └── config.yaml                    # Project, model, LoRA, retrieval, eval & output paths
+├── data/
+│   ├── raw/                           # Raw Spider files (dev.json, tables.json, SQLite DBs)
+│   ├── processed/                     # Normalized JSONL (spider_eval, retrieval_corpus, etc.)
+│   └── indices/                       # Saved FAISS index, AST store, FAISS metadata
+├── models/
+│   └── checkpoints/lora_finetuned/    # Saved LoRA adapter + tokenizer files
+├── src/
+│   ├── config.py                      # YAML → dataclass config loader
+│   ├── data/
+│   │   └── data_loader.py             # Generic JSON/JSONL dataset loader (Spider, etc.)
+│   ├── models/
+│   │   └── load_model.py              # Base model + LoRA loading, training loop, base-vs-LoRA generation
+│   ├── generators/
+│   │   ├── program_generator.py       # Core GenerationPipeline + prompt builder
+│   │   ├── text_to_sql.py             # Text-to-SQL generator
+│   │   ├── sql_generator.py           # SQL prompt/parsing helpers
+│   │   ├── doc_gen.py                 # Documentation generator
+│   │   ├── commit_gen.py              # Commit-message generator
+│   │   └── llm_baseline.py            # Pluggable local/OpenAI/Gemini LLM baseline
+│   ├── embeddings/
+│   │   ├── code_lm_embedder.py        # Embeddings from the small code LM's own hidden states
+│   │   ├── embedding.py               # External embedding-model backend (BGE)
+│   │   ├── indexing.py                # FAISS semantic index manager
+│   │   ├── ast_indexing.py            # AST / SQL structural fingerprint indexer
+│   │   └── hybrid_retriever.py        # Dense + structural blended retrieval
+│   ├── rag/
+│   │   └── rag_pipeline.py            # Context retrieval, prompt building, RAG generation
+│   ├── evaluation/
+│   │   ├── metrics.py                 # BLEU, ROUGE, F1, CodeBERTScore, execution accuracy
+│   │   ├── evaluation.py              # Task-level evaluation engine
+│   │   ├── sql_eval.py                # Spider SQL normalization & execution comparison
+│   │   ├── comparator.py              # Base vs. LoRA vs. RAG comparison
+│   │   ├── rag_evaluation.py          # RAG gain-over-baseline + top-K sweep
+│   │   └── visualization.py           # Chart generation
+│   └── scripts/
+│       ├── convert_spider_to_jsonl.py # Raw Spider → normalized JSONL
+│       └── prepare_processed_data.py  # Build derived/MongoDB-style datasets
+├── outputs/generated/                 # Per-task generated samples & timestamped run results
+├── output/plots/                      # Saved comparison charts
+├── backup/                            # Earlier snapshot of evaluation/rag/embeddings modules
+├── baseline_model.ipynb               # Base model evaluation
+├── finetuned_model.ipynb              # LoRA fine-tuning + evaluation
+├── spider_eval.ipynb                  # Spider execution-based SQL evaluation
+├── eval.ipynb                         # RAG gain / top-K sweep / LLM comparison
+├── main.ipynb                         # Colab entry point
+├── requirements.txt
+└── README.md
+```
+
+---
+
+## Configuration (`configs/config.yaml`)
+
+Key fields actually read by `src/config.py`:
+
+```yaml
+model_name: "Qwen/Qwen2.5-Coder-0.5B-Instruct"
+lora_adapter_path: "models/checkpoints/lora_finetuned"
+embedding_model: "BAAI/bge-small-en-v1.5"   # used only when embedding.source == "external"
+
+embedding:
+  source: "code_lm"           # "code_lm" (default, self-embeddings) or "external"
+
+retrieval:
+  hybrid: true
+  structural_weight_alpha: 0.7
+  structural_language: "sql"
+  topk_sweep: [1, 2, 3, 5, 8]
+
+llm_baseline:
+  backend: "local"            # "local" or an external provider ("openai" / "gemini")
+  local_model: "Qwen/Qwen2.5-Coder-0.5B-Instruct"
+  api_model: "gpt-5.4-mini"
+  max_new_tokens: 256
+  temperature: 0.2
+
+lora:
+  r: 16
+  alpha: 32
+  dropout: 0.05
+  target_modules: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
+  batch_size: 4
+  learning_rate: 0.0003
+  epochs: 5
+  output_dir: "models/checkpoints/lora_finetuned"
+
+indexing:
+  top_k: 3
+  vector_dim: 1024            # must match the embedding model's output dimension
+
+outputs:
+  logs_dir: "output/logs"
+  metrics_path: "output/metrics/evaluation_metrics.json"
+  plots_dir: "output/plots"
+```
+
+---
+
+## Requirements
+
+| Package | Version | Purpose |
+|---|---|---|
+| `torch` | 2.5.1 | Core deep learning framework |
+| `transformers` | 4.46.3 | Pretrained model + tokenizer loading (Qwen2.5-Coder) |
+| `peft` | 0.13.2 | LoRA fine-tuning |
+| `datasets` | 3.1.0 | Dataset utilities |
+| `faiss-cpu` | 1.9.0 | Dense vector similarity search |
+| `tree-sitter` / `tree-sitter-python` | 0.23.2 / 0.23.6 | (Available for AST-based parsing) |
+| `codebleu` | 0.7.0 | CodeBLEU metric |
+| `codebert-score` | 0.3.13 | CodeBERTScore semantic similarity metric |
+| `evaluate` | 0.4.3 | Hugging Face evaluation metrics |
+| `rouge-score` | 0.1.2 | ROUGE metrics |
+| `nltk` | 3.9.1 | Text preprocessing / BLEU support |
+| `PyYAML` | 6.0.2 | Config parsing |
+| `matplotlib` / `seaborn` | 3.10.0 / 0.13.2 | Plotting |
+| `numpy` | 2.1.3 | Numerical computing |
+
+See `requirements.txt` for the full pinned list.
+
+---
+
+## Contributing
+
+1. **Fork** the repository
+2. **Create** a feature branch (`git checkout -b feature/amazing-feature`)
+3. **Commit** your changes (`git commit -m 'Add amazing feature'`)
+4. **Push** to the branch (`git push origin feature/amazing-feature`)
+5. **Open** a Pull Request
+
+---
+
+## License
+
+This project is licensed under the **MIT License** — see the `LICENSE` file for details.
+
+---
+
+## Citation
+
+```bibtex
+@software{codegen_capstone_2026,
+  title  = {CodeGen: Small Code-LM Fine-Tuning and Hybrid RAG Pipeline},
+  author = {Emasters Group 2},
+  year   = {2026},
+  url    = {https://github.com/Rithusravya/Emasters_Group-2_CapstoneProject}
+}
+```
+
+---
+
+<div align="center">
+
+**Built by IIIT-H eMasters Group 2**
+
+[⬆ Back to Top](#-codegen-small-code-lm-fine-tuning--hybrid-rag-pipeline)
+
+</div>

configs/config.yaml

@@ -0,0 +1,105 @@
+# ============================================================
+# General & System Configurations
+# ============================================================
+project_name: "Emasters_Group-2_CapstoneProject"
+seed: 42
+device: "auto"  # Options: 'auto', 'cuda', 'mps', 'cpu'
+dtype: "float16"  # Options: 'float16', 'bfloat16', 'float32'
+
+# ============================================================
+# Model & Tokenizer Configs
+# ============================================================
+model_name: "Qwen/Qwen2.5-Coder-0.5B-Instruct"
+lora_adapter_path: "models/checkpoints/lora_finetuned"
+
+# ============================================================
+# Embedding Model Configs
+# ============================================================
+embedding_model: "BAAI/bge-small-en-v1.5"  # used only when embedding.source == "external"
+embedding:
+  source: "external"  # Options: 'code_lm' or 'external'
+  external_model: "BAAI/bge-small-en-v1.5"
+  max_length: 512
+  batch_size: 32
+
+# ============================================================
+# Retrieval Settings (semantic + AST/structural hybrid search)
+# ============================================================
+retrieval:
+  hybrid: true
+  structural_weight_alpha: 0.7
+  structural_language: "mongodb"
+  topk_sweep: [1, 2, 3, 5, 8]
+
+# ============================================================
+# Large-LLM baseline (for "small code LM vs. LLM" comparison)
+# ============================================================
+llm_baseline:
+  backend: "api"  # Options: 'local' or 'api'
+  local_model: "Qwen/Qwen2.5-Coder-0.5B-Instruct"
+  api_model: "gemini-3.6-flash"
+  api_key_env: "GEMINI_API_KEY"
+  max_new_tokens: 256
+  temperature: 0.2
+
+# ============================================================
+# Dataset Paths
+# ============================================================
+data_paths:
+  raw_dir: "data/raw"
+  spider: "spider_real_execution_mongo.json"
+
+# ============================================================
+# FAISS Vector & AST Index Paths
+# ============================================================
+indices_paths:
+  faiss_index: "data/indices/faiss_semantic.index"
+  ast_store: "data/indices/ast_store.pkl"
+
+# ============================================================
+# Text Generation Parameters
+# ============================================================
+generation:
+  max_length: 512
+  temperature: 0.2
+  top_k: 50
+  top_p: 0.95
+  num_return_sequences: 1
+
+# ============================================================
+# Evaluation Settings
+# ============================================================
+evaluation:
+  device: "cpu"  # CodeBERTScore device
+  bertscore_model: "microsoft/codebert-base"
+  bertscore_local_path: "models/codebert-base"
+  timeout_seconds: 3
+
+# ============================================================
+# LoRA Training Settings
+# ============================================================
+lora:
+  r: 16
+  alpha: 32
+  dropout: 0.05
+  batch_size: 1
+  learning_rate: 0.0002
+  epochs: 1
+  output_dir: "models/checkpoints/lora_finetuned"
+  trainable_fraction: 1.0  # Fraction of LoRA parameters to train (0.0 to 1.0)
+
+# ============================================================
+# Indexing Settings
+# ============================================================
+indexing:
+  top_k: 3
+  vector_dim: 1024  # Must match embedding model's output dimension
+
+# ============================================================
+# Outputs & Evaluation Paths
+# ============================================================
+outputs:
+  logs_dir: "output/logs"
+  metrics_path: "output/metrics/evaluation_metrics.json"
+  plots_dir: "output/plots"
+  generated_dir: "output/generated"

data/indices/ast_index.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c5f4bae7dc5a9a4ef572e9f0f64123668a95382b038132bddcd11473f779db90
+size 361374

data/indices/semantic_index.faiss

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:03db2b2dd28bcb7e5165ce11e60a976728796f48a7a13b3eaf259519c507ef16
+size 1711149

data/indices/semantic_index.metadata.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bed7ff41bd109624a645eb1e91cf612c3d4811f4eb62ba7bd787c1078659f0a9
+size 329781

models/checkpoints/lora_finetuned/README.md

@@ -0,0 +1,207 @@
+---
+base_model: Qwen/Qwen2.5-Coder-0.5B-Instruct
+library_name: peft
+pipeline_tag: text-generation
+tags:
+- base_model:adapter:Qwen/Qwen2.5-Coder-0.5B-Instruct
+- lora
+- transformers
+---
+
+# Model Card for Model ID
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+
+
+## Model Details
+
+### Model Description
+
+<!-- Provide a longer summary of what this model is. -->
+
+
+
+- **Developed by:** [More Information Needed]
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+<!-- Provide the basic links for the model. -->
+
+- **Repository:** [More Information Needed]
+- **Paper [optional]:** [More Information Needed]
+- **Demo [optional]:** [More Information Needed]
+
+## Uses
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+### Direct Use
+
+<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+
+[More Information Needed]
+
+## Bias, Risks, and Limitations
+
+<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+
+[More Information Needed]
+
+### Recommendations
+
+<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+
+[More Information Needed]
+
+### Training Procedure
+
+<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+
+#### Speeds, Sizes, Times [optional]
+
+<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+
+[More Information Needed]
+
+## Evaluation
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+<!-- This should link to a Dataset Card if possible. -->
+
+[More Information Needed]
+
+#### Factors
+
+<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
+
+[More Information Needed]
+
+#### Metrics
+
+<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+<!-- Relevant interpretability work for the model goes here -->
+
+[More Information Needed]
+
+## Environmental Impact
+
+<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]
+### Framework versions
+
+- PEFT 0.20.0

models/checkpoints/lora_finetuned/adapter_config.json

@@ -0,0 +1,50 @@
+{
+  "alora_invocation_tokens": null,
+  "alpha_pattern": {},
+  "arrow_config": null,
+  "auto_mapping": null,
+  "base_model_name_or_path": "Qwen/Qwen2.5-Coder-0.5B-Instruct",
+  "bias": "none",
+  "corda_config": null,
+  "ensure_weight_tying": false,
+  "eva_config": null,
+  "exclude_modules": null,
+  "fan_in_fan_out": false,
+  "inference_mode": true,
+  "init_lora_weights": true,
+  "layer_replication": null,
+  "layers_pattern": null,
+  "layers_to_transform": null,
+  "loftq_config": {},
+  "lora_alpha": 32,
+  "lora_bias": false,
+  "lora_dropout": 0.05,
+  "lora_ga_config": null,
+  "megatron_config": null,
+  "megatron_core": "megatron.core",
+  "modules_to_save": null,
+  "monteclora_config": null,
+  "peft_type": "LORA",
+  "peft_version": "0.20.0",
+  "qalora_group_size": 16,
+  "r": 16,
+  "rank_pattern": {},
+  "revision": null,
+  "target_modules": [
+    "gate_proj",
+    "q_proj",
+    "up_proj",
+    "k_proj",
+    "down_proj",
+    "v_proj",
+    "o_proj"
+  ],
+  "target_parameters": null,
+  "task_type": "CAUSAL_LM",
+  "trainable_token_indices": null,
+  "use_bdlora": null,
+  "use_dora": false,
+  "use_qalora": false,
+  "use_rslora": false,
+  "velora_config": null
+}

models/checkpoints/lora_finetuned/adapter_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a72b6b7afbe198b677f5df314bcfe5d5452d142d433fe9b6bc1215431fe297ae
+size 35237104

models/checkpoints/lora_finetuned/added_tokens.json

@@ -0,0 +1,24 @@
+{
+  "</tool_call>": 151658,
+  "<tool_call>": 151657,
+  "<|box_end|>": 151649,
+  "<|box_start|>": 151648,
+  "<|endoftext|>": 151643,
+  "<|file_sep|>": 151664,
+  "<|fim_middle|>": 151660,
+  "<|fim_pad|>": 151662,
+  "<|fim_prefix|>": 151659,
+  "<|fim_suffix|>": 151661,
+  "<|im_end|>": 151645,
+  "<|im_start|>": 151644,
+  "<|image_pad|>": 151655,
+  "<|object_ref_end|>": 151647,
+  "<|object_ref_start|>": 151646,
+  "<|quad_end|>": 151651,
+  "<|quad_start|>": 151650,
+  "<|repo_name|>": 151663,
+  "<|video_pad|>": 151656,
+  "<|vision_end|>": 151653,
+  "<|vision_pad|>": 151654,
+  "<|vision_start|>": 151652
+}

models/checkpoints/lora_finetuned/special_tokens_map.json

@@ -0,0 +1,31 @@
+{
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "eos_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

models/checkpoints/lora_finetuned/tokenizer_config.json

@@ -0,0 +1,207 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "bos_token": null,
+  "chat_template": "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- messages[0]['content'] }}\n    {%- else %}\n        {{- 'You are Qwen, created by Alibaba Cloud. You are a helpful assistant.' }}\n    {%- endif %}\n    {{- \"\\n\\n# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0]['content'] + '<|im_end|>\\n' }}\n    {%- else %}\n        {{- '<|im_start|>system\\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- for message in messages %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) or (message.role == \"assistant\" and not message.tool_calls) %}\n        {{- '<|im_start|>' + message.role + '\\n' + message.content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {{- '<|im_start|>' + message.role }}\n        {%- if message.content %}\n            {{- '\\n' + message.content }}\n        {%- endif %}\n        {%- for tool_call in message.tool_calls %}\n            {%- if tool_call.function is defined %}\n                {%- set tool_call = tool_call.function %}\n            {%- endif %}\n            {{- '\\n<tool_call>\\n{\"name\": \"' }}\n            {{- tool_call.name }}\n            {{- '\", \"arguments\": ' }}\n            {{- tool_call.arguments | tojson }}\n            {{- '}\\n</tool_call>' }}\n        {%- endfor %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- message.content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n{%- endif %}\n",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "model_max_length": 32768,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}

requirements.txt

@@ -0,0 +1,10 @@
+spaces
+gradio==5.9.1
+torch==2.5.1
+transformers==4.46.3
+peft==0.20.0
+faiss-cpu==1.9.0.post1
+numpy==2.1.3
+PyYAML==6.0.2
+huggingface-hub==0.26.2
+tqdm==4.67.1

src/config.py

@@ -0,0 +1,163 @@
+import logging
+from dataclasses import dataclass, field, fields
+from pathlib import Path
+from typing import List, Optional, Any, Type, TypeVar
+import yaml
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+
+def _build(dataclass_type: Type[T], data: dict) -> T:
+    valid_keys = {f.name for f in fields(dataclass_type)}
+    filtered = {k: v for k, v in data.items() if k in valid_keys}
+    return dataclass_type(**filtered)
+
+
+@dataclass
+class PathConfig:
+    data_dir: Path = Path("data/raw")
+    checkpoints_dir: Path = Path("models/checkpoints")
+    lora_save_path: Path = Path("models/checkpoints/lora_finetuned")
+    faiss_index_path: Path = Path("data/indices/faiss_semantic.index")
+    ast_store_path: Path = Path("data/indices/ast_store.pkl")
+
+@dataclass
+class ModelConfig:
+    name_or_path: str = "Qwen/Qwen2.5-Coder-0.5B-Instruct"
+    embedding_name: str = "BAAI/bge-small-en-v1.5"
+
+@dataclass
+class LoRAConfig:
+    r: int = 16
+    alpha: int = 32
+    target_modules: List[str] = field(
+        default_factory=lambda: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
+    )
+    dropout: float = 0.05
+    batch_size: int = 1
+    learning_rate: float = 2e-4
+    epochs: int = 1
+    output_dir: str = "models/checkpoints/lora_finetuned"
+    trainable_fraction: float = 1.0
+
+@dataclass
+class EmbeddingConfig:
+    source: str = "external"  # 'code_lm' or 'external'
+    external_model: str = "BAAI/bge-small-en-v1.5"
+    max_length: int = 512
+    batch_size: int = 32
+
+@dataclass
+class RetrievalConfig:
+    hybrid: bool = True
+    structural_language: str = "sql"
+    structural_weight_alpha: float = 0.7
+    topk_sweep: List[int] = field(default_factory=lambda: [1, 2, 3, 5, 8])
+
+@dataclass
+class LLMBaselineConfig:
+    backend: str = "api"  # 'local' or 'api'
+    local_model: str = "Qwen/Qwen2.5-Coder-0.5B-Instruct"
+    api_model: str = "gemini-3.6-flash"
+    api_key_env: str = "GEMINI_API_KEY"
+    max_new_tokens: int = 256
+    temperature: float = 0.2
+
+@dataclass
+class GenerationConfig:
+    max_length: int = 512
+    temperature: float = 0.2
+    top_k: int = 50
+    top_p: float = 0.95
+    num_return_sequences: int = 1
+
+@dataclass
+class EvaluationConfig:
+    device: str = "cpu"
+    bertscore_model: str = "microsoft/codebert-base"
+    bertscore_local_path: str = "models/codebert-base"
+    timeout_seconds: int = 3
+
+    def resolve_bertscore_model(self, base_dir: Optional[Path] = None) -> str:
+        candidates = []
+        local = Path(self.bertscore_local_path)
+        if base_dir is not None:
+            candidates.append(Path(base_dir) / self.bertscore_local_path)
+        candidates.append(local)
+        # Fallback: resolve relative to this config.py's own directory (the
+        # project root), since that's where `models/codebert-base` is
+        # expected to live per the project layout.
+        candidates.append(Path(__file__).resolve().parent / self.bertscore_local_path)
+
+        for candidate in candidates:
+            if candidate.is_dir() and (candidate / "config.json").exists():
+                return str(candidate.resolve())
+
+        logger.info(
+            f"Local CodeBERT checkpoint not found at '{self.bertscore_local_path}' "
+            f"(tried: {[str(c) for c in candidates]}); falling back to "
+            f"downloading '{self.bertscore_model}' from the Hugging Face Hub."
+        )
+        return self.bertscore_model
+
+@dataclass
+class OutputConfig:
+    logs_dir: str = "output/logs"
+    metrics_path: str = "output/metrics/evaluation_metrics.json"
+    plots_dir: str = "output/plots"
+    generated_dir: str = "outputs/generated"
+
+@dataclass
+class IndexingConfig:
+    top_k: int = 3
+    vector_dim: int = 1024
+
+class PipelineConfig:
+    def __init__(self, yaml_path: Optional[str] = None):
+        if yaml_path is None:
+            base_path = Path(__file__).resolve().parent
+            yaml_path = base_path / "configs" / "config.yaml"
+
+        self.yaml_path = Path(yaml_path)
+        self._load_yaml()
+
+    def _load_yaml(self):
+        if self.yaml_path.exists():
+            with open(self.yaml_path, "r", encoding="utf-8") as f:
+                data = yaml.safe_load(f) or {}
+        else:
+            logger.warning(f"Config not found at {self.yaml_path}. Using defaults.")
+            data = {}
+
+        path_data = data.get('data_paths', {})
+        idx_data = data.get('indices_paths', {})
+        self.paths = PathConfig(
+            data_dir=Path(path_data.get('raw_dir', "data/raw")),
+            faiss_index_path=Path(idx_data.get('faiss_index', "data/indices/faiss_semantic.index")),
+            ast_store_path=Path(idx_data.get('ast_store', "data/indices/ast_store.pkl"))
+        )
+
+        self.model = ModelConfig(
+            name_or_path=data.get('model_name', "Qwen/Qwen2.5-Coder-0.5B-Instruct"),
+            embedding_name=data.get('embedding_model', "BAAI/bge-small-en-v1.5")
+        )
+
+        self.lora = _build(LoRAConfig, data.get('lora', {}))
+        self.embedding = _build(EmbeddingConfig, data.get('embedding', {}))
+        self.retrieval = _build(RetrievalConfig, data.get('retrieval', {}))
+        self.llm_baseline = _build(LLMBaselineConfig, data.get('llm_baseline', {}))
+        self.generation = _build(GenerationConfig, data.get('generation', {}))
+        self.evaluation = _build(EvaluationConfig, data.get('evaluation', {}))
+        self.outputs = _build(OutputConfig, data.get('outputs', {}))
+        self.indexing = _build(IndexingConfig, data.get('indexing', {}))
+
+        # Top-level attributes for backward compatibility
+        self.model_name = self.model.name_or_path
+        self.embedding_model = self.model.embedding_name
+        self.data_paths = self.paths
+        self.indices_paths = self.paths
+
+def load_config(yaml_path: Optional[str] = None) -> PipelineConfig:
+    return PipelineConfig(yaml_path=yaml_path)

src/data/data_loader.py

@@ -0,0 +1,95 @@
+import os
+import json
+import logging
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+
+class DatasetLoader:
+    def __init__(self, target_filename="qwen_spider_mongodb_conversion.json"):
+        self.target_filename = target_filename
+        self.data = None
+        self.file_path = None
+
+    def scan_and_find_file(self, search_dirs=None):
+        """Scans directories recursively to find the target JSON file."""
+        if search_dirs is None:
+            search_dirs = [
+                Path.cwd(),
+                Path.cwd().parent,
+                Path.cwd() / "data",
+                Path.cwd() / "data" / "raw",
+                Path.cwd() / "data" / "raw" / "Spider",
+                Path.cwd() / "outputs",
+            ]
+
+        search_dirs = list({d.resolve() for d in search_dirs if d.exists()})
+
+        for directory in search_dirs:
+            for path in directory.rglob(self.target_filename):
+                if path.is_file():
+                    logger.info(f"✅ Found {self.target_filename} at: {path}")
+                    return path
+
+        logger.warning(f"File not found in predefined paths. Scanning entire {Path.cwd()}...")
+        for path in Path.cwd().rglob(self.target_filename):
+            if path.is_file():
+                logger.info(f"✅ Found {self.target_filename} at: {path}")
+                return path
+
+        raise FileNotFoundError(f"❌ Could not find {self.target_filename} in any scanned directories.")
+
+    def load_data(self, file_path=None):
+        """Loads the JSON data and parses it into a list of dictionaries."""
+        if file_path is None:
+            file_path = self.scan_and_find_file()
+        else:
+            file_path = Path(file_path)
+            if not file_path.exists():
+                raise FileNotFoundError(f"❌ Provided path does not exist: {file_path}")
+
+        self.file_path = file_path
+        logger.info(f"📂 Loading data from {file_path}...")
+
+        with open(file_path, "r", encoding="utf-8") as f:
+            raw_data = json.load(f)
+
+        if "detailed_results" in raw_data:
+            self.data = raw_data["detailed_results"]
+        elif isinstance(raw_data, list):
+            self.data = raw_data
+        else:
+            self.data = [raw_data]
+
+        # Normalize: flatten generated_mongodb_query if it's a list
+        for item in self.data:
+            if "generated_mongodb_query" in item and isinstance(item["generated_mongodb_query"], list) and len(item["generated_mongodb_query"]) > 0:
+                item["generated_mongodb_query"] = item["generated_mongodb_query"][0]
+
+        logger.info(f"✅ Successfully loaded {len(self.data)} records.")
+        return self.data
+
+    def get_training_data(self):
+        """Returns data formatted for LoRA training (Question -> MongoDB Query)."""
+        if self.data is None:
+            self.load_data()
+
+        training_data = []
+        for item in self.data:
+            question = item.get("question", "")
+            mongodb_query = item.get("generated_mongodb_query", "")
+
+            # Flatten if still a list (safety net)
+            if isinstance(mongodb_query, list) and len(mongodb_query) > 0:
+                mongodb_query = mongodb_query[0]
+
+            if question and mongodb_query:
+                training_data.append({
+                    "question": question,
+                    "generated_mongodb_query": mongodb_query,
+                    "database_id": item.get("database_id", "")
+                })
+
+        logger.info(f"✅ Prepared {len(training_data)} training examples.")
+        return training_data

src/evaluation/__init__.py

@@ -0,0 +1,11 @@
+from evaluation.metrics import EvaluationMetrics
+from evaluation.comparator import ModelComparator
+from evaluation.rag_evaluation import RAGEvaluator
+from evaluation.visualization import ResultVisualizer
+
+__all__ = [
+    "EvaluationMetrics",
+    "ModelComparator",
+    "RAGEvaluator",
+    "ResultVisualizer"
+]

src/evaluation/comparator.py

@@ -0,0 +1,74 @@
+import logging
+from pathlib import Path
+from typing import List, Dict, Any, Optional
+from evaluation.metrics import EvaluationMetrics
+
+logger = logging.getLogger(__name__)
+
+
+def _default_bertscore_model() -> str:
+    local_dir = Path("models/codebert-base")
+    if local_dir.is_dir() and (local_dir / "config.json").exists():
+        return str(local_dir.resolve())
+    return "microsoft/codebert-base"
+
+
+class ModelComparator:
+    """
+    Compares generation outputs across 4 categories:
+    1. Base_Model (Small Code LM)
+    2. LoRA_Model (Fine-tuned)
+    3. RAG_Pipeline (Retrieval Augmented)
+    4. LLM_Baseline (OpenAI/Gemini API)
+    """
+    def __init__(self, device: str = "cpu", bertscore_model: Optional[str] = None):
+        self.device = device
+        self.bertscore_model = bertscore_model or _default_bertscore_model()
+        logger.info(f"ModelComparator using BERTScore/CodeBERTScore backbone: {self.bertscore_model}")
+        self.metrics = EvaluationMetrics()
+
+    def compare(
+        self,
+        references: List[str],
+        base_preds: List[str],
+        lora_preds: List[str],
+        rag_preds: Optional[List[str]] = None,
+        llm_preds: Optional[List[str]] = None,
+        test_cases: Optional[List[str]] = None,
+        task_type: str = "sql" # "sql" or "doc"
+    ) -> Dict[str, Dict[str, float]]:
+        
+        results = {}
+        categories = {
+            "Base_Model": base_preds,
+            "LoRA_Model": lora_preds
+        }
+        if rag_preds: categories["RAG_Pipeline"] = rag_preds
+        if llm_preds: categories["LLM_Baseline"] = llm_preds
+
+        for name, preds in categories.items():
+            logger.info(f"Calculating metrics for {name}...")
+            cat_metrics = {}
+            
+            # Text-based metrics (BLEU, ROUGE, BERTScore)
+            cat_metrics["BLEU"] = self.metrics.compute_bleu(references, preds)
+            rouge_scores = self.metrics.compute_rouge(references, preds)
+            cat_metrics.update(rouge_scores)
+            cat_metrics["BERTScore"] = self.metrics.compute_bertscore(
+                references, preds, device=self.device, model_type=self.bertscore_model
+            )
+            
+            # Exact & Response Accuracy
+            # exact_matches = [self.metrics.compute_exact_match(r, p) for r, p in zip(references, preds)]
+            # cat_metrics["Exact_Match_Accuracy"] = round(sum(exact_matches) / len(exact_matches), 4) if exact_matches else 0.0
+            
+            resp_acc = [self.metrics.compute_response_accuracy(r, p) for r, p in zip(references, preds)]
+            cat_metrics["Response_Accuracy"] = round(sum(resp_acc) / len(resp_acc), 4) if resp_acc else 0.0
+
+            # SQL Execution Accuracy (if task is SQL and test cases/DBs are provided)
+            if task_type == "sql" and test_cases:
+                pass 
+
+            results[name] = cat_metrics
+
+        return results

src/evaluation/metrics.py

@@ -0,0 +1,258 @@
+import re
+import torch
+import logging
+from typing import List, Dict, Optional
+from nltk.translate.bleu_score import corpus_bleu, SmoothingFunction
+
+logger = logging.getLogger(__name__)
+
+_BERTSCORE_NUM_LAYERS_OVERRIDES = {
+    "models/codebert-base": 10,
+    "models/codebert-base-mlm": 10
+}
+
+
+def _resolve_bertscore_num_layers(model_type: str) -> Optional[int]:
+    try:
+        from bert_score.utils import model2layers
+    except Exception:
+        model2layers = {}
+
+    if model_type in model2layers:
+        return model2layers[model_type]
+
+    short_name = model_type.split("/")[-1]
+    if short_name in model2layers:
+        return model2layers[short_name]
+
+    if model_type in _BERTSCORE_NUM_LAYERS_OVERRIDES:
+        return _BERTSCORE_NUM_LAYERS_OVERRIDES[model_type]
+
+    name_lower = model_type.replace("\\", "/").lower()
+    for key in sorted(_BERTSCORE_NUM_LAYERS_OVERRIDES, key=len, reverse=True):
+        key_short = key.split("/")[-1].lower()
+        if key_short in name_lower:
+            return _BERTSCORE_NUM_LAYERS_OVERRIDES[key]
+
+    return None
+
+
+class EvaluationMetrics:
+    
+    @staticmethod
+    def _tokenize_sql(sql: str) -> List[str]:
+        """Basic SQL tokenization for BLEU/ROUGE."""
+        if not sql: return []
+        sql = sql.lower().strip()
+        sql = re.sub(r'[^a-z0-9\s_\*\(\)\,\.\'\"\=\>\<\!\+\-\%]', ' ', sql)
+        return sql.split()
+
+    @staticmethod
+    def compute_bleu(references: List[str], predictions: List[str]) -> float:
+        try:
+            refs = [[EvaluationMetrics._tokenize_sql(r)] for r in references]
+            preds = [EvaluationMetrics._tokenize_sql(p) for p in predictions]
+            smoother = SmoothingFunction().method1
+            score = corpus_bleu(refs, preds, smoothing_function=smoother)
+            return round(score, 4)
+        except Exception as e:
+            logger.error(f"BLEU calculation failed: {e}")
+            return 0.0
+
+    @staticmethod
+    def compute_rouge(references: List[str], predictions: List[str]) -> Dict[str, float]:
+        try:
+            from rouge_score import rouge_scorer
+            scorer = rouge_scorer.RougeScorer(['rouge1', 'rougeL'], use_stemmer=True)
+            rouge1_scores, rougeL_scores = [], []
+            
+            for ref, pred in zip(references, predictions):
+                scores = scorer.score(ref, pred)
+                rouge1_scores.append(scores['rouge1'].fmeasure)
+                rougeL_scores.append(scores['rougeL'].fmeasure)
+                
+            return {
+                "ROUGE-1": round(sum(rouge1_scores) / len(rouge1_scores), 4) if rouge1_scores else 0.0,
+                "ROUGE-L": round(sum(rougeL_scores) / len(rougeL_scores), 4) if rougeL_scores else 0.0
+            }
+        except Exception as e:
+            logger.error(f"ROUGE calculation failed: {e}")
+            return {"ROUGE-1": 0.0, "ROUGE-L": 0.0}
+
+    @staticmethod
+    def compute_bertscore(references: List[str], predictions: List[str], device: str = "cpu", model_type: str = "models/codebert-base") -> float:
+        """CodeBERTScore: BERTScore F1 computed with a code-aware encoder (CodeBERT by default).
+
+        NOTE: `model_type` values that aren't in bert_score's built-in
+        `model2layers` table (which is the case for both "microsoft/codebert-base"
+        and "FacebookAI/roberta-base") need an explicit `num_layers`, or the
+        underlying `bert_score.score()` call raises a KeyError. See
+        `_resolve_bertscore_num_layers` above -- this is what was previously
+        making CodeBERTScore silently return 0.0 for every model category.
+        """
+        try:
+            from bert_score import score
+            valid_pairs = [(r, p) for r, p in zip(references, predictions) if r.strip() and p.strip()]
+            if not valid_pairs:
+                return 0.0
+            refs, preds = zip(*valid_pairs)
+
+            num_layers = _resolve_bertscore_num_layers(model_type)
+            score_kwargs = {"num_layers": num_layers} if num_layers is not None else {}
+
+            with torch.no_grad():
+                P, R, F1 = score(
+                    list(preds), list(refs),
+                    model_type=model_type,
+                    verbose=False,
+                    device=device,
+                    **score_kwargs,
+                )
+            return round(F1.mean().item(), 4)
+        except Exception as e:
+            logger.error(f"BERTScore ({model_type}) calculation failed: {e}")
+            return 0.0
+
+    @staticmethod
+    def _find_top_level_groups(s: str) -> List[tuple]:
+        """Locates every top-level `{...}`/`[...]` span in `s`, skipping over
+        brace/bracket characters that appear inside double-quoted string
+        literals (so e.g. a `$regex: "a{2,3}"` value doesn't confuse the
+        bracket matching)."""
+        groups = []
+        depth = 0
+        start = None
+        in_string = False
+        i, n = 0, len(s)
+        while i < n:
+            ch = s[i]
+            if in_string:
+                if ch == '\\':
+                    i += 2
+                    continue
+                if ch == '"':
+                    in_string = False
+                i += 1
+                continue
+            if ch == '"':
+                in_string = True
+                i += 1
+                continue
+            if ch in '{[':
+                if depth == 0:
+                    start = i
+                depth += 1
+            elif ch in '}]':
+                depth -= 1
+                if depth == 0 and start is not None:
+                    groups.append((start, i + 1))
+                    start = None
+            i += 1
+        return groups
+
+    @staticmethod
+    def _split_top_level(inner: str, sep: str) -> List[str]:
+        """Splits `inner` on `sep`, but only at bracket-depth 0 and never
+        inside a double-quoted string, so nested `{...}`/`[...]`/`(...)`
+        and comma-containing string values aren't split incorrectly."""
+        parts, depth, current = [], 0, []
+        in_string = False
+        i = 0
+        while i < len(inner):
+            ch = inner[i]
+            if in_string:
+                current.append(ch)
+                if ch == '\\' and i + 1 < len(inner):
+                    current.append(inner[i + 1])
+                    i += 1
+                elif ch == '"':
+                    in_string = False
+                i += 1
+                continue
+            if ch == '"':
+                in_string = True
+                current.append(ch)
+                i += 1
+                continue
+            if ch in '{[(':
+                depth += 1
+            elif ch in '}])':
+                depth -= 1
+            if ch == sep and depth == 0:
+                parts.append(''.join(current))
+                current = []
+            else:
+                current.append(ch)
+            i += 1
+        parts.append(''.join(current))
+        return [p for p in parts if p.strip() != '']
+
+    @classmethod
+    def _canon_group(cls, s: str) -> str:
+        s = s.strip()
+        if s.startswith('{') and s.endswith('}'):
+            inner = s[1:-1]
+            items = cls._split_top_level(inner, ',')
+            kvs = []
+            for item in items:
+                if ':' not in item:
+                    return s  # malformed for our purposes -- bail out safely
+                k, v = item.split(':', 1)
+                kvs.append((k.strip(), cls._canon_group(v.strip())))
+            kvs.sort(key=lambda kv: kv[0])
+            return '{' + ','.join(f'{k}:{v}' for k, v in kvs) + '}'
+        elif s.startswith('[') and s.endswith(']'):
+            inner = s[1:-1]
+            items = cls._split_top_level(inner, ',')
+            return '[' + ','.join(cls._canon_group(i.strip()) for i in items) + ']'
+        return s
+
+    @classmethod
+    def _canonicalize_mql(cls, s: str) -> str:
+        try:
+            if s.count('{') != s.count('}') or s.count('[') != s.count(']'):
+                return s
+            groups = cls._find_top_level_groups(s)
+            if not groups:
+                return s
+            out, last = [], 0
+            for gstart, gend in groups:
+                out.append(s[last:gstart])
+                out.append(cls._canon_group(s[gstart:gend]))
+                last = gend
+            out.append(s[last:])
+            return ''.join(out)
+        except Exception:
+            return s
+
+    # @classmethod
+    # def compute_exact_match(cls, gold: str, pred: str) -> float:
+    #     if not gold or not pred:
+    #         return 0.0
+
+    #     def normalize(s):
+    #         s = s.strip().rstrip(';')
+    #         s = s.replace("'", '"')            # Normalize single quotes to double quotes
+    #         s = cls._canonicalize_mql(s)       # Order-independent object-key comparison
+    #         s = re.sub(r'\s+', '', s)          # Remove all remaining whitespace
+    #         return s.lower()
+
+    #     return 1.0 if normalize(gold) == normalize(pred) else 0.0
+
+    @staticmethod
+    def _tokenize_mql(query: str) -> List[str]:
+        """Tokenizer optimized for MongoDB Query Language (MQL) brackets and punctuation."""
+        if not query: return []
+        # Add spaces around brackets, commas, colons, and dots
+        query = re.sub(r'([(){}[\],:.])', r' \1 ', query)
+        return [tok.lower() for tok in query.split() if tok.strip()]
+
+    @staticmethod
+    def compute_response_accuracy(gold: str, pred: str) -> float:
+        """Semantic response accuracy (recall-based token overlap)."""
+        if not gold or not pred: return 0.0
+        g_tokens = set(EvaluationMetrics._tokenize_mql(gold))
+        p_tokens = set(EvaluationMetrics._tokenize_mql(pred))
+        if not g_tokens: return 0.0
+        overlap = len(g_tokens.intersection(p_tokens))
+        return round(overlap / len(g_tokens), 4)

src/evaluation/rag_evaluation.py

@@ -0,0 +1,132 @@
+import logging
+from typing import List, Dict, Any
+import matplotlib.pyplot as plt
+from pathlib import Path
+from evaluation.metrics import EvaluationMetrics
+
+logger = logging.getLogger(__name__)
+
+class RAGEvaluator:
+    """Evaluates RAG pipeline improvements and optimizes Top-K retrieval."""
+    
+    def __init__(self, comparator):
+        self.comparator = comparator
+        self.metrics = EvaluationMetrics()
+        logger.info(f"RAGEvaluator using BERTScore/CodeBERTScore backbone: {getattr(comparator, 'bertscore_model', 'models/codebert-base')}")
+
+    def measure_rag_improvement(self, rag_pipeline, references, queries, lora_preds, top_k=3):
+        logger.info(f"Measuring RAG improvement over baseline LoRA (top_k={top_k})...")
+        rag_preds = []
+        for q in queries:
+            out, _ = rag_pipeline.generate_with_rag(q, top_k=top_k)
+            rag_preds.append(out)
+            
+        bertscore_model = getattr(self.comparator, "bertscore_model", "models/codebert-base")
+
+        rag_metrics = {
+            "BLEU": self.metrics.compute_bleu(references, rag_preds),
+            "BERTScore": self.metrics.compute_bertscore(references, rag_preds, device=self.comparator.device, model_type=bertscore_model)
+            # "Exact_Match_Accuracy": round(sum([self.metrics.compute_exact_match(r, p) for r, p in zip(references, rag_preds)]) / len(references), 4)
+        }
+        
+        baseline_metrics = {
+            "BLEU": self.metrics.compute_bleu(references, lora_preds),
+            "BERTScore": self.metrics.compute_bertscore(references, lora_preds, device=self.comparator.device, model_type=bertscore_model)
+            # "Exact_Match_Accuracy": round(sum([self.metrics.compute_exact_match(r, p) for r, p in zip(references, lora_preds)]) / len(references), 4)
+        }
+        
+        gain = {k: round(rag_metrics[k] - baseline_metrics[k], 4) for k in rag_metrics}
+        
+        return {
+            "metrics": {"RAG_Pipeline": rag_metrics, "Baseline_LoRA": baseline_metrics},
+            "gain_over_baseline": gain,
+            "baseline_used_for_gain": "Baseline_LoRA"
+        }
+
+    def sweep_top_k(self, rag_pipeline, references, queries, baseline_preds, k_values=[1, 2, 3, 5, 8], verbose=True):
+        sweep_results = {}
+        
+        # Pre-compute baseline metrics once to save time
+        baseline_bleu = self.metrics.compute_bleu(references, baseline_preds)
+        baseline_rouge = self.metrics.compute_rouge(references, baseline_preds).get("ROUGE-1", 0.0)
+        
+        for k in k_values:
+            logger.info(f"Sweeping top_k={k}...")
+            preds = []
+            for q in queries:
+                out, _ = rag_pipeline.generate_with_rag(q, top_k=k)
+                preds.append(out)
+                
+            # Use BLEU and ROUGE-1 for sweeping 
+            bleu = self.metrics.compute_bleu(references, preds)
+            rouge1 = self.metrics.compute_rouge(references, preds).get("ROUGE-1", 0.0)
+            
+            metrics = {"BLEU": bleu, "ROUGE-1": rouge1}
+            gain = {
+                "BLEU": round(bleu - baseline_bleu, 4),
+                "ROUGE-1": round(rouge1 - baseline_rouge, 4)
+            }
+            
+            sweep_results[k] = {"metrics": metrics, "gain_over_baseline": gain, "predictions": preds}
+
+        if verbose:
+            print("\n" + self.format_sweep_table(sweep_results))
+
+        return sweep_results
+
+    @staticmethod
+    def format_sweep_table(sweep_results: Dict[int, Dict[str, Any]]) -> str:
+        ks = list(sweep_results.keys())
+        if not ks:
+            return "(no sweep results to display)"
+
+        columns = [
+            ("Top-K", lambda k, r: str(k)),
+            ("BLEU", lambda k, r: f"{r['metrics'].get('BLEU', 0.0):.4f}"),
+            ("ROUGE-1", lambda k, r: f"{r['metrics'].get('ROUGE-1', 0.0):.4f}"),
+            ("BLEU Gain", lambda k, r: f"{r['gain_over_baseline'].get('BLEU', 0.0):+.4f}"),
+            ("ROUGE-1 Gain", lambda k, r: f"{r['gain_over_baseline'].get('ROUGE-1', 0.0):+.4f}"),
+        ]
+
+        rows = [[render(k, sweep_results[k]) for _, render in columns] for k in ks]
+        headers = [name for name, _ in columns]
+
+        widths = [
+            max(len(headers[i]), *(len(row[i]) for row in rows))
+            for i in range(len(columns))
+        ]
+
+        def fmt_row(cells):
+            return "| " + " | ".join(c.center(w) for c, w in zip(cells, widths)) + " |"
+
+        separator = "+-" + "-+-".join("-" * w for w in widths) + "-+"
+
+        lines = ["Top-K Retrieval Sweep Results", separator, fmt_row(headers), separator]
+        lines += [fmt_row(row) for row in rows]
+        lines.append(separator)
+        return "\n".join(lines)
+
+    def best_k(self, sweep_results, metric="BLEU"):
+        best_k_val = -999
+        best_k = None
+        for k, res in sweep_results.items():
+            val = res["gain_over_baseline"].get(metric, -999)
+            if val > best_k_val:
+                best_k_val = val
+                best_k = k
+        return best_k if best_k is not None else list(sweep_results.keys())[0]
+
+    def plot_gain_vs_k(self, sweep_results, save_path="output/plots/topk_gain_sweep.png"):
+        Path(save_path).parent.mkdir(parents=True, exist_ok=True)
+        ks = list(sweep_results.keys())
+        bleu_gains = [sweep_results[k]["gain_over_baseline"].get("BLEU", 0) for k in ks]
+        
+        plt.figure(figsize=(8, 5))
+        plt.plot(ks, bleu_gains, marker='o', linestyle='-', color='b')
+        plt.title("RAG Gain over Baseline by Top-K")
+        plt.xlabel("Top-K Retrieved Contexts")
+        plt.ylabel("BLEU Score Gain")
+        plt.grid(True)
+        plt.savefig(save_path)
+        plt.close()
+        logger.info(f"✅ Saved Top-K sweep plot to {save_path}")

src/evaluation/visualization.py

@@ -0,0 +1,76 @@
+import logging
+import matplotlib.pyplot as plt
+import seaborn as sns
+import pandas as pd
+from pathlib import Path
+from typing import Dict
+
+logger = logging.getLogger(__name__)
+
+class ResultVisualizer:
+    """Visualizes evaluation metrics across Base, LoRA, RAG, and LLM categories."""
+    
+    def __init__(self, output_dir: str = "output/plots"):
+        self.output_dir = Path(output_dir)
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+        sns.set_theme(style="whitegrid", palette="muted")
+
+    def plot_comparison(self, eval_results: Dict[str, Dict[str, float]], save_name: str = "model_comparison.png", bertscore_model_label: str = None):
+        # Flatten dictionary for seaborn
+        data = []
+        for model_name, metrics in eval_results.items():
+            for metric_name, value in metrics.items():
+                data.append({"Model": model_name, "Metric": metric_name, "Score": value})
+                
+        df = pd.DataFrame(data)
+        
+        if df.empty:
+            logger.warning("No data to plot.")
+            return
+
+        plt.figure(figsize=(14, 8))
+        sns.barplot(x="Metric", y="Score", hue="Model", data=df)
+        
+        plt.title("Comprehensive Model Comparison: Base vs LoRA vs RAG vs LLM", fontsize=16, fontweight='bold')
+        if bertscore_model_label:
+            plt.suptitle(f"BERTScore backbone: {bertscore_model_label}", fontsize=9, y=0.93, color="gray")
+        plt.ylabel("Score", fontsize=12)
+        plt.xlabel("Evaluation Metric", fontsize=12)
+        plt.xticks(rotation=45, ha='right')
+        plt.legend(title="Model Category", bbox_to_anchor=(1.05, 1), loc='upper left')
+        plt.tight_layout()
+        
+        save_path = self.output_dir / save_name
+        plt.savefig(save_path, dpi=300)
+        plt.close()
+        logger.info(f"✅ Saved comparison plot to {save_path}")
+
+    def plot_radar_chart(self, eval_results: Dict[str, Dict[str, float]], save_name: str = "radar_comparison.png"):
+        """Generates a radar chart for a holistic view of model capabilities."""
+        import numpy as np
+        
+        categories = list(next(iter(eval_results.values())).keys())
+        N = len(categories)
+        if N == 0: return
+
+        # Compute the angle for each axis
+        angles = [n / float(N) * 2 * 3.14159 for n in range(N)]
+        angles += angles[:1]
+
+        fig, ax = plt.subplots(figsize=(8, 8), subplot_kw=dict(polar=True))
+        
+        for model_name, metrics in eval_results.items():
+            values = list(metrics.values())
+            values += values[:1]
+            ax.plot(angles, values, linewidth=2, label=model_name)
+            ax.fill(angles, values, alpha=0.15)
+
+        ax.set_xticks(angles[:-1])
+        ax.set_xticklabels(categories, fontsize=10)
+        plt.title("Model Capability Radar", size=15, y=1.1, fontweight='bold')
+        plt.legend(loc='upper right', bbox_to_anchor=(1.3, 1.1))
+        
+        save_path = self.output_dir / save_name
+        plt.savefig(save_path, dpi=300, bbox_inches='tight')
+        plt.close()
+        logger.info(f"✅ Saved radar chart to {save_path}")

src/generators/__init__.py

@@ -0,0 +1,11 @@
+from generators.program_generator import GenerationPipeline
+from generators.doc_generator import DocGenerator
+from generators.commit_generator import CommitMessageGenerator
+from generators.text_to_mongo_generator import TextToMongoGenerator
+
+__all__ = [
+    "GenerationPipeline",
+    "DocGenerator",
+    "CommitMessageGenerator",
+    "TextToMongoGenerator"
+]

src/generators/commit_generator.py

@@ -0,0 +1,26 @@
+import logging
+
+logger = logging.getLogger(__name__)
+
+class CommitMessageGenerator:
+    """
+    Generates conventional commit messages based on git diffs.
+    """
+    def __init__(self, pipeline):
+        self.pipeline = pipeline
+
+    def _build_prompt(self, diff: str) -> str:
+        prompt = (
+            "You are an expert software developer. Generate a concise, conventional commit message "
+            "for the following git diff. The message should start with a standard type "
+            "(feat, fix, docs, style, refactor, test, chore) followed by a colon and a short, "
+            "imperative summary of the changes.\n\n"
+            f"Diff:\n{diff}\n\n"
+            "Commit Message:"
+        )
+        return prompt
+
+    def generate_commit_msg(self, diff: str) -> str:
+        """Generates a commit message for the provided git diff."""
+        prompt = self._build_prompt(diff)
+        return self.pipeline.generate(prompt)

src/generators/doc_generator.py

@@ -0,0 +1,47 @@
+import logging
+
+logger = logging.getLogger(__name__)
+
+class DocGenerator:
+    """
+    Generates documentation/docstrings for code snippets using 5-shot prompting.
+    """
+    def __init__(self, pipeline):
+        self.pipeline = pipeline
+
+    def _build_prompt(self, code_snippet: str) -> str:
+        # 5 high-quality examples for few-shot learning
+        examples = [
+            {
+                "code": "def add(a, b):\n    return a + b",
+                "doc": "Adds two numbers together and returns their sum."
+            },
+            {
+                "code": "import os\ndef list_files(directory):\n    return os.listdir(directory)",
+                "doc": "Returns a list containing the names of all files and directories in the specified directory path."
+            },
+            {
+                "code": "class Dog:\n    def __init__(self, name):\n        self.name = name\n    def bark(self):\n        return 'Woof!'",
+                "doc": "A simple Dog class that initializes with a name and provides a method to make the dog bark."
+            },
+            {
+                "code": "def factorial(n):\n    if n == 0:\n        return 1\n    return n * factorial(n-1)",
+                "doc": "Calculates the factorial of a given non-negative integer n using a recursive approach."
+            },
+            {
+                "code": "import json\ndef parse_json(file_path):\n    with open(file_path, 'r') as f:\n        return json.load(f)",
+                "doc": "Reads a JSON file from the specified file path and returns the parsed content as a Python dictionary."
+            }
+        ]
+        
+        prompt = "You are an expert programmer. Generate a concise, accurate, and professional docstring/description for the given code.\n\n"
+        for ex in examples:
+            prompt += f"Code:\n{ex['code']}\nDescription: {ex['doc']}\n\n"
+            
+        prompt += f"Code:\n{code_snippet}\nDescription:"
+        return prompt
+
+    def generate_docstring(self, code_snippet: str) -> str:
+        """Generates a docstring for the provided code snippet."""
+        prompt = self._build_prompt(code_snippet)
+        return self.pipeline.generate(prompt)

src/generators/llm_baseline.py

@@ -0,0 +1,85 @@
+import os
+import logging
+import time
+from typing import Optional
+from google import genai
+
+logger = logging.getLogger(__name__)
+
+class LargeLLMGenerator:
+    """
+    Wrapper for Google Gemini API to serve as an upper-bound baseline for fast, reliable generation.
+    """
+    def __init__(self, config=None, backend: str = "api", api_model: str = "gemini-3.6-flash", local_model_name: str = None):
+        self.backend = backend
+        self.api_model = api_model
+        self.local_model_name = local_model_name
+        self.client = None
+        self.config = config
+        self.api_key_env_name = "GEMINI_API_KEY"
+
+        # Override with config if provided
+        if config and hasattr(config, "llm_baseline"):
+            llm_cfg = config.llm_baseline
+            self.backend = getattr(llm_cfg, "backend", self.backend)
+            self.api_model = getattr(llm_cfg, "api_model", self.api_model)
+            self.local_model_name = getattr(llm_cfg, "local_model", self.local_model_name)
+            self.api_key_env_name = getattr(llm_cfg, "api_key_env", None) or self.api_key_env_name
+
+        self._initialize_client()
+
+    def _initialize_client(self):
+        if self.backend == "api":
+            self._init_gemini()
+        else:
+            logger.info(f"Using local/off-the-shelf LLM backend: {self.local_model_name}")
+
+    def _init_gemini(self):
+        try:
+            api_key = os.getenv(self.api_key_env_name)
+
+            if not api_key:
+                logger.error(
+                    f"API key not found. Set the '{self.api_key_env_name}' environment "
+                    "variable before initializing LargeLLMGenerator."
+                )
+                return
+
+            # Initialize the modern google-genai client
+            self.client = genai.Client(api_key=api_key)
+            logger.info(f"✅ Initialized Gemini client for model: {self.api_model}")
+        except ImportError:
+            logger.error("google-genai package not installed. Run: pip install google-genai")
+        except Exception as e:
+            logger.error(f"Failed to initialize Gemini client: {e}")
+
+    def generate(self, prompt: str, max_tokens: int = 512, temperature: float = 0.2) -> str:
+        if self.backend == "api":
+            return self._generate_api(prompt, max_tokens, temperature)
+        else:
+            logger.warning("Local LLM generation not fully implemented in this baseline wrapper. Returning empty string.")
+            return ""
+
+    def _generate_api(self, prompt: str, max_tokens: int, temperature: float) -> str:
+        if not self.client:
+            return "[API Client Not Initialized]"
+        try:
+            start_time = time.time()
+
+            # Generate content using the google-genai client
+            response = self.client.models.generate_content(
+                model=self.api_model,
+                contents=prompt,
+                config={
+                    "max_output_tokens": max_tokens,
+                    "temperature": temperature
+                }
+            )
+            text = response.text.strip()
+
+            latency = round((time.time() - start_time) * 1000, 2)
+            logger.info(f"Gemini API ({self.api_model}) generated {len(text)} chars in {latency}ms")
+            return text
+        except Exception as e:
+            logger.error(f"API Generation failed: {e}")
+            return f"[API Error: {str(e)}]"

src/generators/program_generator.py

@@ -0,0 +1,70 @@
+import logging
+import torch
+import os
+os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1")
+
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+class GenerationPipeline:
+    def __init__(self, model: Any, tokenizer: Any, config: Any = None):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.config = config
+
+    def _build_generation_kwargs(self) -> dict:
+        """Builds generation kwargs based on config or defaults."""
+        # Safely extract token IDs as standard Python ints to prevent PyTorch warnings
+        eos_id = self.tokenizer.eos_token_id
+        if isinstance(eos_id, torch.Tensor):
+            eos_id = eos_id.item()
+
+        kwargs = {
+            "max_new_tokens": 128,  # Hard cap: MongoDB queries are short
+            "temperature": 0.2,
+            "top_p": 0.95,
+            "top_k": 50,
+            "do_sample": True,
+            "pad_token_id": eos_id,
+            "eos_token_id": eos_id,
+            "repetition_penalty": 1.2,
+            "no_repeat_ngram_size": 3,  # CRITICAL: Prevents repeating blocks of numbers/text
+        }
+        if self.config:
+            kwargs["max_new_tokens"] = min(getattr(self.config, "max_length", 256), 128)
+            kwargs["temperature"] = getattr(self.config, "temperature", 0.2)
+            kwargs["top_p"] = getattr(self.config, "top_p", 0.95)
+            kwargs["top_k"] = getattr(self.config, "top_k", 50)
+        return kwargs
+
+    def generate(self, prompt: str) -> str:
+        """Generates text from a given prompt."""
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
+        with torch.no_grad():
+            outputs = self.model.generate(**inputs, **self._build_generation_kwargs())
+
+        generated_text = self.tokenizer.decode(
+            outputs[0][inputs["input_ids"].shape[1]:],
+            skip_special_tokens=True
+        )
+        generated_text = generated_text.strip()
+
+        generated_text = self._clean_output(generated_text)
+        return generated_text
+
+    @staticmethod
+    def _clean_output(text: str) -> str:
+        """Remove hallucinated refusals or trailing junk after the actual answer."""
+        # Common patterns that indicate the model went off-track
+        stop_patterns = [
+            "I'm sorry", "I can't", "I cannot", "I am sorry",
+            "I apologize", "As an AI", "As a language model",
+            "```", "\n\n", "###", "Human:", "User:", "Assistant:"
+        ]
+        for pattern in stop_patterns:
+            idx = text.find(pattern)
+            if idx > 0:  # Only cut if it's not at the very start
+                text = text[:idx]
+        return text.strip()

src/generators/text_to_mongo_generator.py

@@ -0,0 +1,125 @@
+import json
+import time
+import logging
+from pathlib import Path
+from typing import List, Dict, Any
+from tqdm import tqdm
+
+logger = logging.getLogger(__name__)
+
+
+class TextToMongoGenerator:
+    def __init__(self, pipeline, config: Any = None):
+        self.pipeline = pipeline
+        self.config = config
+
+    def _build_prompt(self, question: str, schema: str = "") -> str:
+        prompt = "### Task: Generate MongoDB Query (MQL)\n"
+        if schema:
+            prompt += f"### Collection Schema / Structure:\n{schema}\n\n"
+        prompt += f"### Question:\n{question}\n\n### MongoDB Query:\n db."
+        return prompt
+
+    def generate_single(self, question: str, schema: str = "") -> Dict[str, Any]:
+        """Generates a single MongoDB query and measures latency."""
+        prompt = self._build_prompt(question, schema)
+        start_time = time.perf_counter()
+        raw_output = self.pipeline.generate(prompt)
+        end_time = time.perf_counter()
+        latency_ms = (end_time - start_time) * 1000
+
+        # Clean up hallucinations and repetition loops
+        generated_query = self._clean_mongo_output(raw_output)
+
+        return {
+            "generated_query": generated_query,
+            "latency_ms": round(latency_ms, 2)
+        }
+
+    @staticmethod
+    def clean_mongo_output(raw_output: str) -> str:
+        stripped = raw_output.strip()
+        if stripped.lower().startswith("db."):
+            # Model re-emitted the prefix itself; don't duplicate it.
+            restored = stripped
+        else:
+            restored = "db." + stripped
+
+        cleaned = (
+            restored.replace("```javascript", "")
+            .replace("```json", "")
+            .replace("```js", "")
+            .replace("```", "")
+            .strip()
+        )
+
+        # Stop at common hallucination markers
+        stop_markers = ["--->", "\n\n", "\nUser:", "\nQuestion:", "```", "###", "Human:"]
+        for marker in stop_markers:
+            idx = cleaned.find(marker)
+            if idx != -1:
+                cleaned = cleaned[:idx]
+
+        # Extract the first valid MongoDB query line
+        lines = cleaned.split('\n')
+        for line in lines:
+            line = line.strip()
+            if line.startswith("db."):
+                return line.rstrip(';').strip()
+
+        # Fallback: find "db." anywhere in the text
+        if "db." in cleaned:
+            idx = cleaned.find("db.")
+            end_idx = len(cleaned)
+            for marker in stop_markers + ["\n"]:
+                pos = cleaned.find(marker, idx)
+                if pos != -1:
+                    end_idx = min(end_idx, pos)
+            return cleaned[idx:end_idx].strip().rstrip(';')
+
+        return cleaned.strip()
+
+    _clean_mongo_output = clean_mongo_output
+
+    def generate_batch(self, dataset: List[Dict[str, Any]], output_path: str = None) -> List[Dict[str, Any]]:
+        results = []
+        logger.info(f"🚀 Starting batch MongoDB query generation for {len(dataset)} examples...")
+
+        for item in tqdm(dataset, desc="Generating MongoDB Queries", unit="query"):
+            query_id = item.get("query_id", f"sample_{len(results)}")
+            database_id = item.get("database_id", "unknown")
+            question = item.get("question", "")
+            schema = item.get("schema", "")
+            gold_query = item.get("generated_mongodb_query", item.get("gold_query", ""))
+            if isinstance(gold_query, list) and len(gold_query) > 0:
+                gold_query = gold_query[0]
+
+            if not question:
+                continue
+
+            gen_result = self.generate_single(question, schema)
+
+            result_entry = {
+                "query_id": query_id,
+                "database_id": database_id,
+                "question": question,
+                "gold_query": gold_query,
+                "generated_query": gen_result["generated_query"],
+                "latency_ms": gen_result["latency_ms"],
+            }
+            results.append(result_entry)
+
+        if output_path:
+            self.save_results(results, output_path)
+
+        return results
+
+    def save_results(self, results: List[Dict[str, Any]], output_path: str):
+        path = Path(output_path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+
+        output_data = {"results": results}
+        with open(path, "w", encoding="utf-8") as f:
+            json.dump(output_data, f, indent=4, ensure_ascii=False)
+
+        logger.info(f"✅ Successfully saved {len(results)} generated queries to {path}")

src/models/load_model.py

@@ -0,0 +1,263 @@
+import os
+import logging
+import torch
+from pathlib import Path
+from tqdm import tqdm
+from peft import LoraConfig, get_peft_model, PeftModel
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from torch.utils.data import DataLoader, Dataset
+
+os.environ["PYTORCH_MPS_HIGH_WATERMARK_RATIO"] = "0.0"
+
+logger = logging.getLogger(__name__)
+
+class SpiderDataset(Dataset):
+    """Formats Spider JSON data into causal LM training examples (Question -> MongoDB Query)."""
+
+    def __init__(self, data, tokenizer, max_length=512):
+        self.data = data
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        item = self.data[idx]
+        question = item.get("question", "")
+        mongodb_query = item.get("generated_mongodb_query", "")
+
+        # Handle if stored as a list
+        if isinstance(mongodb_query, list) and len(mongodb_query) > 0:
+            mongodb_query = mongodb_query[0]
+
+        prompt = f"### Task: Generate MongoDB Query (MQL)\n### Question:\n{question}\n\n### MongoDB Query:\n"
+
+        prompt_enc = self.tokenizer(prompt, truncation=True, max_length=self.max_length, return_tensors="pt")
+        completion_enc = self.tokenizer(mongodb_query, truncation=True, max_length=self.max_length, return_tensors="pt")
+
+        input_ids = torch.cat([prompt_enc.input_ids[0], completion_enc.input_ids[0]], dim=0)
+        attention_mask = torch.cat([prompt_enc.attention_mask[0], completion_enc.attention_mask[0]], dim=0)
+
+        prompt_len = prompt_enc.input_ids.shape[1]
+        labels = torch.cat([torch.full((prompt_len,), -100, dtype=torch.long), completion_enc.input_ids[0]], dim=0)
+
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "labels": labels
+        }
+
+class ModelLoader:
+    def __init__(self, config, device_override=None):
+        self.config = config
+        self.device = device_override or self._detect_device()
+        if self.device == "cuda":
+            self.dtype = torch.float16
+        elif self.device == "mps":
+            self.dtype = torch.float16
+        else:
+            self.dtype = torch.float32
+        self.tokenizer = None
+        self.base_model = None
+        self.lora_model = None
+        
+        logger.info(f"🖥️ Initialized ModelLoader | Device: {self.device} | Dtype: {self.dtype}")
+
+    def _detect_device(self):
+        if torch.cuda.is_available():
+            return "cuda"
+        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+            return "mps"
+        return "cpu"
+
+    def load_tokenizer(self):
+        model_name = getattr(self.config, "model_name", "Qwen/Qwen2.5-Coder-0.5B-Instruct")
+        logger.info(f"🔄 Loading tokenizer: {model_name}")
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+        return self.tokenizer
+
+    def load_base_model(self):
+        model_name = getattr(self.config, "model_name", "Qwen/Qwen2.5-Coder-0.5B-Instruct")
+        logger.info(f"🔄 Loading base model: {model_name}")
+
+        if self.device == "cuda":
+            self.base_model = AutoModelForCausalLM.from_pretrained(
+                model_name,
+                torch_dtype=self.dtype,
+                trust_remote_code=True,
+                low_cpu_mem_usage=True,
+                device_map="auto"
+            )
+        else:
+            self.base_model = AutoModelForCausalLM.from_pretrained(
+                model_name,
+                torch_dtype=self.dtype,
+                trust_remote_code=True,
+                low_cpu_mem_usage=True
+            )
+            self.base_model.to(self.device)
+
+        self.base_model.eval()
+        return self.base_model
+
+    def setup_lora(self, trainable_fraction=1.0, r=16, alpha=32, dropout=0.05):
+        if self.base_model is None:
+            self.load_base_model()
+            
+        target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
+        
+        peft_config = LoraConfig(
+            r=r,
+            lora_alpha=alpha,
+            target_modules=target_modules,
+            lora_dropout=dropout,
+            bias="none",
+            task_type="CAUSAL_LM"
+        )
+        
+        self.lora_model = get_peft_model(self.base_model, peft_config)
+        
+        if trainable_fraction < 1.0:
+            self._apply_trainable_fraction(trainable_fraction)
+        else:
+            logger.info("✅ LoRA configured with 100% trainable parameters.")
+            
+        self.lora_model.print_trainable_parameters()
+        return self.lora_model
+
+    def _apply_trainable_fraction(self, fraction):
+        lora_params = [(n, p) for n, p in self.lora_model.named_parameters() if "lora_" in n and p.requires_grad]
+        total_lora = len(lora_params)
+        num_to_freeze = int(total_lora * (1.0 - fraction))
+        
+        if num_to_freeze > 0:
+            for i, (name, param) in enumerate(lora_params):
+                if i < num_to_freeze:
+                    param.requires_grad = False
+            
+            trainable_count = sum(1 for p in self.lora_model.parameters() if p.requires_grad)
+            logger.info(f"🔒 Froze {num_to_freeze}/{total_lora} LoRA parameters. Remaining trainable: {trainable_count} ({fraction*100}%).")
+
+    def autodetect_saved_model(self, output_dir):
+        output_dir = Path(output_dir)
+        if output_dir.exists() and (output_dir / "adapter_config.json").exists():
+            logger.info(f"✅ Found saved LoRA adapter at {output_dir}. Loading...")
+            if self.base_model is None:
+                self.load_base_model()
+            # NOTE (ZeroGPU compatibility): peft's from_pretrained() infers the load
+            # device via infer_device(), NOT from the base model's actual .device.
+            # On HF Spaces ZeroGPU, `spaces` monkey-patches torch.cuda.is_available()
+            # to report True even before a real GPU is attached to the process, so
+            # infer_device() always says "cuda" here -- which makes peft try to load
+            # the adapter's .safetensors weights with a raw, direct CUDA file load.
+            # That bypasses spaces' safe .to("cuda") staging and crashes with
+            # "No CUDA GPUs are available". Forcing torch_device="cpu" here loads the
+            # adapter safely on CPU; the single .to(self.device) call right after is
+            # a normal nn.Module.to() move, which spaces *does* stage/support.
+            self.lora_model = PeftModel.from_pretrained(self.base_model, str(output_dir), torch_device="cpu")
+            self.lora_model.to(self.device)
+            self.lora_model.eval()
+            return self.lora_model
+        logger.info(f"ℹ️ No saved adapter found at {output_dir}.")
+        return None
+
+    def train_lora(self, train_data, output_dir, epochs=3, batch_size=1, grad_accum=4, lr=2e-4, max_length=512):
+        lr = float(lr) # Ensure lr is a float to prevent PyTorch optimizer TypeError
+        if self.lora_model is None:
+            self.setup_lora()
+            
+        dataset = SpiderDataset(train_data, self.tokenizer, max_length)
+        
+        def collate_fn(batch):
+            max_len = max(len(b["input_ids"]) for b in batch)
+            input_ids, attention_mask, labels = [], [], []
+            pad_id = self.tokenizer.pad_token_id
+            
+            for b in batch:
+                pad_len = max_len - len(b["input_ids"])
+                input_ids.append(b["input_ids"].tolist() + [pad_id] * pad_len)
+                attention_mask.append(b["attention_mask"].tolist() + [0] * pad_len)
+                labels.append(b["labels"].tolist() + [-100] * pad_len)
+                
+            return {
+                "input_ids": torch.tensor(input_ids, dtype=torch.long).to(self.device),
+                "attention_mask": torch.tensor(attention_mask, dtype=torch.long).to(self.device),
+                "labels": torch.tensor(labels, dtype=torch.long).to(self.device)
+            }
+            
+        loader = DataLoader(dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)
+        
+        optimizer = torch.optim.AdamW(
+            filter(lambda p: p.requires_grad, self.lora_model.parameters()), 
+            lr=lr
+        )
+        
+        self.lora_model.train()
+        
+        logger.info(f"🚀 Starting LoRA training | Device: {self.device} | Epochs: {epochs} | Batch: {batch_size} | Grad Accum: {grad_accum} | LR: {lr}")
+        
+        for epoch in range(epochs):
+            epoch_loss = 0.0
+            progress_bar = tqdm(loader, desc=f"Epoch {epoch+1}/{epochs}", leave=False)
+            
+            optimizer.zero_grad()
+            
+            for step, batch in enumerate(progress_bar):
+                outputs = self.lora_model(**batch)
+                loss = outputs.loss / grad_accum
+                loss.backward()
+                epoch_loss += loss.item() * grad_accum
+                
+                if (step + 1) % grad_accum == 0:
+                    optimizer.step()
+                    optimizer.zero_grad()
+                    
+                progress_bar.set_postfix({"loss": f"{loss.item() * grad_accum:.4f}"})
+                
+                # Clear MPS cache periodically to prevent memory leaks on Mac
+                if self.device == "mps" and (step + 1) % 10 == 0:
+                    torch.mps.empty_cache()
+                    
+            optimizer.step()
+            optimizer.zero_grad()
+            
+            avg_loss = epoch_loss / len(loader)
+            logger.info(f"✅ Epoch {epoch+1} completed | Avg Loss: {avg_loss:.4f}")
+            
+        output_dir = Path(output_dir)
+        output_dir.mkdir(parents=True, exist_ok=True)
+        self.lora_model.save_pretrained(str(output_dir))
+        self.tokenizer.save_pretrained(str(output_dir))
+        logger.info(f"💾 LoRA adapter saved to {output_dir}")
+        
+        self.lora_model.eval()
+        return self.lora_model
+
+    def load_models(self, lora_path=None):
+        """Returns a dictionary of models and the tokenizer."""
+        if self.tokenizer is None:
+            self.load_tokenizer()
+        if self.base_model is None:
+            self.load_base_model()
+
+        models = {"base": self.base_model, "lora": None}
+
+        if self.lora_model is not None:
+            logger.info("Reusing already-loaded LoRA model in memory.")
+            models["lora"] = self.lora_model
+        elif lora_path:
+            lora_path = Path(lora_path)
+            if lora_path.exists() and (lora_path / "adapter_config.json").exists():
+                logger.info(f"Loading LoRA model from {lora_path}")
+                self.lora_model = PeftModel.from_pretrained(self.base_model, str(lora_path))
+                self.lora_model.to(self.device)
+                self.lora_model.eval()
+                models["lora"] = self.lora_model
+            else:
+                logger.warning(f"LoRA path {lora_path} does not exist or is invalid. Returning base model as 'lora'.")
+                models["lora"] = self.base_model
+
+        return models, self.tokenizer

src/rag/rag_pipeline.py

@@ -0,0 +1,166 @@
+import logging
+import torch
+from typing import Any, List, Tuple, Dict, Optional
+
+logger = logging.getLogger(__name__)
+
+METADATA_TEXT_FIELDS = ["generated_mongodb_query", "mongodb_query", "code", "sql", "SQL", "question", "text", "schema"]
+
+class RAGPipeline:
+    def __init__(
+        self, 
+        retriever: Any, 
+        generator: Optional[Any] = None, 
+        model: Optional[Any] = None, 
+        tokenizer: Optional[Any] = None, 
+        config: Any = None
+    ):
+        """
+        Args:
+            retriever: An instance of IndexManager or HybridRetriever that has a `.search(query, k)` method.
+            generator: An instance of GenerationPipeline (has `.generate(prompt)`). 
+                       If provided, `model` and `tokenizer` are ignored.
+            model: HuggingFace model (used if generator is None).
+            tokenizer: HuggingFace tokenizer (used if generator is None).
+            config: Generation configuration object.
+        """
+        self.retriever = retriever
+        self.generator = generator
+        self.model = model
+        self.tokenizer = tokenizer
+        self.config = config
+        
+        if self.generator is None and (self.model is None or self.tokenizer is None):
+            raise ValueError("Either `generator` or both `model` and `tokenizer` must be provided.")
+
+    def retrieve_context(self, query: str, top_k: int = 3) -> List[Dict[str, Any]]:
+        """
+        Retrieves relevant context using the hybrid retriever.
+        Returns a list of metadata dictionaries.
+        """
+        try:
+            # The new HybridRetriever/IndexManager returns List[Dict] or List[Tuple[Dict, float]]
+            results = self.retriever.search(query, k=top_k)
+            
+            # Handle different return formats gracefully
+            if results and isinstance(results[0], tuple):
+                return [metadata for metadata, score in results]
+            elif results and isinstance(results[0], dict):
+                return results
+            else:
+                return []
+        except Exception as e:
+            logger.error(f"Error during context retrieval: {e}")
+            return []
+
+    @staticmethod
+    def _metadata_to_text(metadata: Any) -> str:
+        """Extracts the most relevant text field from a metadata dictionary."""
+        if not isinstance(metadata, dict):
+            return str(metadata)
+        
+        for field in METADATA_TEXT_FIELDS:
+            if metadata.get(field):
+                return str(metadata[field])
+        
+        # Fallback: concatenate all string values
+        return " | ".join([f"{k}: {v}" for k, v in metadata.items() if isinstance(v, str)])
+
+    def build_prompt(self, query: str, context_items: List[Dict[str, Any]]) -> str:
+        context_str = ""
+        for i, item in enumerate(context_items):
+            mongodb_query = (
+                    item.get('generated_mongodb_query', '')
+                    or item.get('mongodb_query', '')
+            )
+            question = item.get('question', '')
+
+            if isinstance(mongodb_query, list) and len(mongodb_query) > 0:
+                mongodb_query = mongodb_query[0]
+
+            if mongodb_query:
+                context_str += f"--- Example {i + 1} ---\n"
+                if question:
+                    context_str += f"Question: {question}\n"
+                context_str += f"MongoDB Query: {mongodb_query}\n\n"
+
+        prompt = (
+            "You are an expert MongoDB query generator. Given a question, output ONLY the MongoDB query.\n"
+            "Do not include explanations, apologies, or extra text.\n\n"
+            f"Context Examples:\n{context_str}\n"
+            f"### Task: Generate MongoDB Query (MQL)\n"
+            f"### Question:\n{query}\n\n"
+            f"### MongoDB Query:\ndb."
+        )
+        return prompt
+
+    def _build_generation_kwargs(self) -> dict:
+        eos_id = self.tokenizer.eos_token_id if self.tokenizer else None
+        if isinstance(eos_id, torch.Tensor):
+            eos_id = eos_id.item()
+
+        kwargs = {
+            "max_new_tokens": 128,
+            "temperature": 0.2,
+            "top_p": 0.95,
+            "top_k": 50,
+            "do_sample": True,
+            "pad_token_id": eos_id,
+            "eos_token_id": eos_id,
+            "repetition_penalty": 1.2,
+            "no_repeat_ngram_size": 3,
+        }
+        if self.config:
+            kwargs["max_new_tokens"] = min(getattr(self.config, "max_length", 256), 128)
+            kwargs["temperature"] = getattr(self.config, "temperature", 0.2)
+            kwargs["top_p"] = getattr(self.config, "top_p", 0.95)
+            kwargs["top_k"] = getattr(self.config, "top_k", 50)
+        return kwargs
+
+    @staticmethod
+    def _clean_mongo_output(raw_output: str) -> str:
+        stripped = raw_output.strip()
+        if stripped.lower().startswith("db."):
+            restored = stripped
+        else:
+            restored = "db." + stripped
+
+        cleaned = restored.replace("```javascript", "").replace("```json", "").replace("```", "").strip()
+        stop_markers = ["--->", "\n\n", "\nUser:", "\nQuestion:", "```", "###", "Human:"]
+        for marker in stop_markers:
+            idx = cleaned.find(marker)
+            if idx != -1:
+                cleaned = cleaned[:idx]
+        lines = cleaned.split('\n')
+        for line in lines:
+            if line.strip().startswith("db."):
+                return line.strip().rstrip(';')
+        if "db." in cleaned:
+            idx = cleaned.find("db.")
+            return cleaned[idx:].split('\n')[0].strip().rstrip(';')
+        return cleaned.strip()
+
+    # Update generate_with_rag to clean the output before returning:
+    def generate_with_rag(self, query: str, top_k: int = 3) -> Tuple[str, List[Dict[str, Any]]]:
+        context_metadata = self.retrieve_context(query, top_k=top_k)
+        prompt = self.build_prompt(query, context_metadata)
+
+        if self.generator:
+            generated_text = self.generator.generate(prompt)
+        else:
+            inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
+            gen_kwargs = self._build_generation_kwargs()
+            with torch.no_grad():
+                outputs = self.model.generate(**inputs, **gen_kwargs)
+            generated_text = self.tokenizer.decode(
+                outputs[0][inputs["input_ids"].shape[1]:],
+                skip_special_tokens=True
+            ).strip()
+
+        generated_text = self._clean_mongo_output(generated_text)
+        return generated_text, context_metadata
+
+    def generate_program(self, prompt: str, top_k: int = 3) -> str:
+        """Wrapper for simple program generation using RAG."""
+        generated_text, _ = self.generate_with_rag(query=prompt, top_k=top_k)
+        return generated_text

src/vectorDB/__init__.py

@@ -0,0 +1,13 @@
+from vectorDB.embedder import HFEmbedder
+from vectorDB.semantic_index import SemanticIndex
+from vectorDB.ast_index import ASTIndex
+from vectorDB.hybrid_retriever import HybridRetriever
+from vectorDB.index_manager import IndexManager
+
+__all__ = [
+    "HFEmbedder",
+    "SemanticIndex",
+    "ASTIndex",
+    "HybridRetriever",
+    "IndexManager"
+]