Pvnn/grasp-rag

Efficient context compression to optimize token usage and accuracy in Retrieval-Augmented Generation (RAG) pipelines

0

stars

62

commits

Python

primary language

Jul 9, 2026

updated

context-compression
nlp
rag
token-optimization

README

GRASP

Graph-based Relevance and Attention-driven Span Pruning
for Enhanced Retrieval-Augmented Generation

Python PyTorch FastAPI License Paper

Paper under review Amritha Krishna S · Gopika M · Pavan Raj · Nimisha Abraham · Neethu Subash Department of Computer Science and Engineering, Mar Athanasius College of Engineering (Autonomous), Kothamangalam, India


Overview

Standard Retrieval-Augmented Generation (RAG) pipelines forward entire retrieved passages to the reader LLM with little or no filtering. This inflates token budgets, increases inference latency, and introduces distracting noise — especially in multi-hop scenarios where critical evidence is buried under irrelevant text. Existing compression techniques either fragment multi-sentence evidence (extractive methods) or introduce hallucinations (abstractive methods).

GRASP addresses this with a multi-stage, query-aware context compression framework. It intelligently prunes retrieved content before it reaches the LLM, preserving only semantically cohesive, query-critical evidence spans — not just individual sentences.


Pipeline

GRASP End-to-End Pipeline

The pipeline consists of four sequential stages:


Key Contributions

ContributionDescription
Hybrid RetrieverFuses dense (Contriever-MSMARCO) and sparse (BM25-Okapi) retrieval via Reciprocal Rank Fusion (RRF, κ=10) to maximise initial recall
QUITO-X Sentence-Level FilterAdapts QUITO-X from token-level to sentence-level by introducing boundary tracking between tokens and spans, with max-pooled cross-attention scoring from a frozen Flan-T5-Small model
EP-EXIT Evidence Span GroupingA GraphSeg-inspired graph-based clustering mechanism that groups semantically related neighbouring sentences into coherent evidence spans before relevance classification
Span-Level CompressionExtends EXIT from sentence-level to span-level classification, allowing the model to evaluate groups of related sentences instead of treating them independently
Pareto Frontier ShiftSimultaneously improves accuracyand reduces latency versus the EXIT baseline on single-hop tasks — pushing the accuracy-efficiency trade-off frontier

Architecture

QUITO-X: Coarse Filter

QUITO-X Coarse Filtering Architecture

QUITO-X adapts the Information Bottleneck (IB) framework to sentence-level filtering using the cross-attention mechanism of a frozen Flan-T5-Small (80M parameters):

  1. Tokenise all sentences into a flat token array, tracking sentence boundaries [start_j, end_j)
  2. Chunk into windows of ≤ 510 − |query tokens| tokens; prepend chunk, append query + EOS
  3. Forward through Flan-T5 with output_attentions=True; extract last-layer decoder cross-attention
  4. Apply Gaussian smoothing (σ=2.0) across the token attention array for stability
  5. Compute per-sentence score as max(smoothed_attentions[start_j : end_j])
  6. Retain top-k sentences by score using a dynamic threshold: k = max(k_min, ⌈r × |S_i|⌉), where r=0.8

EP-EXIT: Fine-Grained Span Filter

EP-EXIT Evidence Unit Extraction

EP-EXIT extends EXIT with graph-based evidence unit decomposition:

  1. Embed sentences via all-MiniLM-L6-v2 (SentenceTransformer)

  2. Build similarity graph: Add edge (i, j) iff |i−j| ≤ w (locality window=2) and cosine_sim(i,j) ≥ δ (δ=0.45)

  3. Extract Evidence Units: Each connected component in the graph becomes one evidence span; isolated nodes are singletons

  4. Classify spans: Pass each span through the EXIT-fine-tuned Gemma-2B-it classifier (doubleyyh/exit-gemma-2b, 4-bit NF4 quantised). Relevance score:

    r(u) = P("Yes" | q, d_i, u) / [P("Yes" | q, d_i, u) + P("No" | q, d_i, u)]
    
  5. Retain spans where r(u) ≥ τ (τ=0.5); concatenate retained spans in original document order

Key difference from GraphSeg: GRASP uses connected components (not maximal cliques) for broader, more robust grouping; a locality window for efficiency; and is optimised for evidence span selection in RAG (not generic text segmentation).


Experimental Results

Evaluated on three QA benchmarks (50 queries each) using Qwen-2.5:3B as the reader on an NVIDIA RTX-4000 Ada Edition GPU (20GB VRAM).

Comprehensive Results (Table I)

MethodHotpotQA EMHotpotQA F1HotpotQA R-LHotpotQA CompHotpotQA Lat(s)NQ EMNQ F1NQ R-LNQ CompNQ Lat(s)TriviaQA EMTriviaQA F1TriviaQA R-LTriviaQA CompTriviaQA Lat(s)
NoOp (Uncompressed)48.054.6153.940.0%0.4130.027.5428.160.0%0.4382.041.0240.420.0%0.70
RECOMP-Extr34.034.8134.5385.0%0.2732.030.5130.9179.9%0.2760.060.5159.8579.8%0.31
LLMLingua-214.017.9617.3685.9%0.8824.022.5122.4884.8%0.9950.051.6853.6885.8%1.27
EXIT48.053.1452.4087.3%1.7324.024.2924.2269.3%1.7266.065.8466.8081.6%2.91
GRASP (Proposed)42.046.9546.7579.5%1.1934.035.8535.1144.8%1.2370.068.2167.8364.3%2.11

↑ higher is better for EM, F1, R-L · Comp = compression ratio (% tokens removed vs NoOp) · Lat = end-to-end latency

Key Findings

Single-hop tasks (NQ & TriviaQA):

  • GRASP achieves +10 EM on Natural Questions (34.0% vs EXIT's 24.0%) — even beating the NoOp upper bound (30.0%)
  • GRASP achieves +4 EM on TriviaQA (70.0% vs EXIT's 66.0%)
  • 27% latency speedup on TriviaQA (2.11s vs EXIT's 2.91s) and 28% speedup on NQ (1.23s vs 1.72s) — despite adding an extra compression stage
  • The speedup is paradoxical: QUITO-X's lightweight 60M-param filter removes ~20% of irrelevant sentences before the expensive 2B-param Gemma classifier, shrinking its payload significantly

Multi-hop task (HotpotQA):

  • GRASP scores 42.0% EM vs EXIT's 48.0% — a known trade-off: graph-based span grouping can remove "bridge" sentences needed to chain multi-hop reasoning steps. This is an identified area for future work.

Case Study

Case Study: GRASP vs RECOMP on HotpotQA

The case study illustrates how GRASP preserves the critical multi-document evidence span needed to correctly answer a HotpotQA multi-hop question, while RECOMP fails to retrieve the necessary cross-document evidence.


Evaluation Metrics

MetricDefinition
Exact Match (EM)1 if the normalised predicted answer contains the normalised gold answer, 0 otherwise
Token F1 (F1)Harmonic mean of token-level precision and recall between prediction and gold
ROUGE-L (R-L)F-measure of the longest common subsequence (LCS) between prediction and gold
Compression (Comp)(1 − compressed_tokens / original_tokens) × 100% relative to NoOp
Latency (Lat)Total wall-clock time for compression + answer generation per query (seconds)

Repository Structure

grasp/
├── app.py                          # FastAPI REST backend (lifespan, auto-indexing)
├── requirements.txt                # Python dependencies
├── .env.example                    # Environment variable template
├── run_evals.sh / run_evals.ps1    # Sequential multi-dataset evaluation runner
├── run_preprocess.sh / .ps1        # Dataset preprocessing runner
│
├── assets/                         # README figures (workflow, QUITO-X, EP-EXIT, case study)
│
├── src/
│   ├── rag_pipeline.py             # End-to-end QueryAwareRAG orchestrator
│   │
│   ├── retrieval/
│   │   ├── retriever.py            # Dense-only Contriever retriever
│   │   └── hybrid_retriever.py     # Hybrid RRF retriever (Contriever + BM25)
│   │
│   ├── compression/
│   │   ├── hybrid_compressor.py    # Two-stage HybridCompressor (QUITO-X → EP-EXIT)
│   │   ├── quitox_filter.py        # Stage 2: QUITO-X coarse filter (Flan-T5-Small)
│   │   ├── ep_exit.py              # Stage 3: EP-EXIT fine filter (graph + Gemma-2B)
│   │   ├── exit_baseline.py        # EXIT baseline compressor (Gemma-2B, 4-bit NF4)
│   │   └── baselines/              # Re-implemented baseline compressors
│   │       ├── exit.py             #   EXIT
│   │       ├── compact.py          #   CompAct
│   │       ├── llmlingua2.py       #   LLMLingua-2
│   │       ├── recomp_abst.py      #   RECOMP Abstractive
│   │       ├── refiner.py          #   Refiner
│   │       └── recomp_extr/        #   RECOMP Extractive
│   │
│   ├── generation/
│   │   ├── reader.py               # Ollama-based RAG reader (Llama-3.1:8B)
│   │   └── gemma_reader.py         # Gemma reader variant
│   │
│   ├── eval/
│   │   ├── eval_pipeline.py        # GenerativeEvaluator (compress → read → score)
│   │   ├── metrics.py              # EM, Token F1, ROUGE-L
│   │   ├── adapters.py             # Unified compressor adapter interface
│   │   ├── interfaces.py           # SearchResult dataclass
│   │   ├── eval_hotpotqa.py        # HotpotQA benchmark script
│   │   ├── eval_nq.py              # Natural Questions benchmark script
│   │   ├── eval_tqa.py             # TriviaQA benchmark script
│   │   ├── eval_asqa.py            # ASQA benchmark script
│   │   └── eval_2wiki.py           # 2WikiMultiHopQA benchmark script
│   │
│   ├── budget_predictor/           # (Experimental) Query-difficulty budget predictor
│   └── data/
│       └── demo_loader.py          # Demo dataset loader for the UI
│
├── scripts/
│   ├── analyze_exit.py             # Step-by-step EXIT pipeline observability tool
│   ├── preprocess_hotpotqa.py      # HotpotQA preprocessing (offline retrieval)
│   ├── preprocess_nq.py            # NQ preprocessing
│   ├── preprocess_asqa.py          # ASQA preprocessing
│   ├── preprocess_2wiki.py         # 2WikiMultiHopQA preprocessing
│   └── preprocess_tqa.py           # TriviaQA preprocessing
│
├── data/                           # Pre-processed datasets (gitignored)
│
└── rag-frontend/                   # React + Vite interactive demo UI

Installation

Prerequisites

RequirementVersionNotes
Python≥ 3.10
CUDA≥ 11.8Required for Gemma-2B (4-bit NF4) and Flan-T5
GPU VRAM≥ 8 GBNVIDIA RTX 4000 Ada or equivalent
RAM≥ 16 GB
Disk≥ 20 GBFor datasets + model checkpoints
OllamaLatestRequired for the reader LLM

Step 1 — Clone & create environment

git clone https://github.com/Pvnn/grasp-rag.git
cd grasp-rag

python -m venv env
# Windows
env\Scripts\activate
# Linux/macOS
source env/bin/activate

Step 2 — Install PyTorch (CUDA 11.8)

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

Step 3 — Install project dependencies

pip install -r requirements.txt
python -m spacy download en_core_web_sm

Step 4 — Configure environment

cp .env.example .env

Edit .env:

HF_TOKEN=hf_xxxxxxxxxxxxxxxxxxxx     # Required: Hugging Face token (Gemma-2B access)

Gemma-2B access: Accept the Gemma model licence on Hugging Face before the first run.

Step 5 — Pull the reader model via Ollama

# Install Ollama from https://ollama.ai/
ollama pull llama3.1:8b

Quick Start

Running the API Server

python app.py

The FastAPI server starts at http://127.0.0.1:8000. On startup it initialises the full GRASP pipeline and loads demo datasets. API docs at http://127.0.0.1:8000/docs.

Running the Interactive UI

cd rag-frontend
npm install
npm run dev

Served at http://localhost:5173. The UI provides:

  • Dataset and query selection with auto-indexing
  • Per-sentence QUITO-X attention score visualisation (retained / pruned)
  • EP-EXIT evidence unit grouping (kept / removed spans)
  • Side-by-side compressed vs. uncompressed answer comparison
  • Live pipeline metrics (compression ratio, token savings, latency)

Programmatic Usage

from dotenv import load_dotenv
import os
from src.rag_pipeline import QueryAwareRAG

load_dotenv()
pipeline = QueryAwareRAG(token=os.getenv("HF_TOKEN"))

# Index your documents (list of dicts with 'title' and 'text' keys)
pipeline.retriever.index_documents(my_documents)

result = pipeline.run(
    query="What instruments on the James Webb Space Telescope are used for observation?",
    top_k=5,
    compare_original=True,   # Runs uncompressed baseline in parallel for comparison
    use_coarse=True,          # Enable QUITO-X stage
    use_fine=True             # Enable EP-EXIT stage
)

print(result['answer'])
print(f"Compression ratio: {result['metrics']['compression']['ratio_chars']:.1f}%")
print(f"Time saved:        {result['metrics']['times']['net_time_saved']:.2f}s")

REST API Reference

EndpointMethodDescription
/datasetsGETList available demo datasets and queries
/dataset/loadPOSTManually load a dataset into the retrieval index
/queryPOSTExecute the full GRASP pipeline on a query

POST /query payload:

{
  "query": "Were Scott Derrickson and Ed Wood of the same nationality?",
  "top_k": 4,
  "compare_original": true,
  "use_coarse": true,
  "use_fine": true
}

Dataset Preparation & Evaluation

1. Preprocess Datasets

# Windows
.\run_preprocess.ps1

# Linux/macOS
bash run_preprocess.sh

Preprocesses all benchmarks with offline top-30 hybrid retrieval into data/<dataset>/. To run a single dataset:

python scripts/preprocess_hotpotqa.py

2. Run Evaluations

# All datasets, 50 samples each
# Windows
.\run_evals.ps1 -N 50

# Linux/macOS
bash run_evals.sh 50

# Single dataset
python -m src.eval.eval_hotpotqa -n 50
python -m src.eval.eval_nq -n 50
python -m src.eval.eval_tqa -n 50

Results are saved to src/eval_results/<dataset>/:

  • final_benchmark_results.csv — aggregate metrics for all compressors
  • details_<compressor>.csv — per-query traces with context, predictions, and scores

Configuration

ParameterDefaultDescription
top_k5Documents to retrieve
coarse_ratio (r)0.8Fraction of sentences retained after QUITO-X
quitox_min_keep2Minimum sentences kept per document
fine_threshold (τ)0.5EP-EXIT inclusion threshold
similarity_threshold (δ)0.45Cosine similarity edge threshold for evidence graph
locality_window (w)2Maximum sentence index distance for graph edge
rrf_k (κ)10RRF smoothing constant
gaussian_sigma (σ)2.0Smoothing for QUITO-X attention scores

Ablation Setup

Each stage can be toggled independently:

# QUITO-X only
result = pipeline.run(query, use_coarse=True, use_fine=False)

# EP-EXIT only
result = pipeline.run(query, use_coarse=False, use_fine=True)

# Full GRASP pipeline
result = pipeline.run(query, use_coarse=True, use_fine=True)

# No compression (NoOp baseline)
result = pipeline.run(query, use_coarse=False, use_fine=False)

Models

ComponentModelParametersQuantisation
Dense Retrieverfacebook/contriever-msmarco110MFP32/FP16
QUITO-X Filtergoogle/flan-t5-small80MFP16
Sentence Embedderall-MiniLM-L6-v222MFP16 (GPU)
EP-EXIT Classifierdoubleyyh/exit-gemma-2b2B4-bit NF4
Reader LLMllama3.1:8b (Ollama)8BQ4_K_M

Limitations & Future Work

  • Multi-hop reasoning: The graph-based span grouping can remove "bridge" sentences needed to connect multi-hop reasoning chains. Future work targets hop-aware weighting in the evidence graph.
  • Fixed compression ratio: QUITO-X uses a static coarse_ratio. A query-difficulty-aware budget predictor (see src/budget_predictor/) is under development.
  • English only: Current models are English-specific. Multilingual dense retrievers and sentence encoders are a direct extension path.
  • End-to-end training: Compression and reader are currently independent; joint training is planned.
  • Scope: Text-based English-language QA only. Abstractive summarisation, multi-modal retrieval, and cross-lingual RAG are out of scope.

Project Context

Department of Computer Science and Engineering Mar Athanasius College of Engineering (Autonomous), Kothamangalam APJ Abdul Kalam Technological University · May 2026

Project Guide: Prof. Neethu Subash · Coordinator: Prof. Nimisha Abraham


Built with ❤️ at Mar Athanasius College of Engineering · Department of Computer Science and Engineering

Contributors

Pvnn

53 commits

Pvnn/grasp-rag

Efficient context compression to optimize token usage and accuracy in Retrieval-Augmented Generation (RAG) pipelines

0

stars

62

commits

Python

primary language

Jul 9, 2026

updated

context-compression
nlp
rag
token-optimization

README

GRASP

Graph-based Relevance and Attention-driven Span Pruning
for Enhanced Retrieval-Augmented Generation

Python PyTorch FastAPI License Paper

Paper under review Amritha Krishna S · Gopika M · Pavan Raj · Nimisha Abraham · Neethu Subash Department of Computer Science and Engineering, Mar Athanasius College of Engineering (Autonomous), Kothamangalam, India


Overview

Standard Retrieval-Augmented Generation (RAG) pipelines forward entire retrieved passages to the reader LLM with little or no filtering. This inflates token budgets, increases inference latency, and introduces distracting noise — especially in multi-hop scenarios where critical evidence is buried under irrelevant text. Existing compression techniques either fragment multi-sentence evidence (extractive methods) or introduce hallucinations (abstractive methods).

GRASP addresses this with a multi-stage, query-aware context compression framework. It intelligently prunes retrieved content before it reaches the LLM, preserving only semantically cohesive, query-critical evidence spans — not just individual sentences.


Pipeline

GRASP End-to-End Pipeline

The pipeline consists of four sequential stages:


Key Contributions

ContributionDescription
Hybrid RetrieverFuses dense (Contriever-MSMARCO) and sparse (BM25-Okapi) retrieval via Reciprocal Rank Fusion (RRF, κ=10) to maximise initial recall
QUITO-X Sentence-Level FilterAdapts QUITO-X from token-level to sentence-level by introducing boundary tracking between tokens and spans, with max-pooled cross-attention scoring from a frozen Flan-T5-Small model
EP-EXIT Evidence Span GroupingA GraphSeg-inspired graph-based clustering mechanism that groups semantically related neighbouring sentences into coherent evidence spans before relevance classification
Span-Level CompressionExtends EXIT from sentence-level to span-level classification, allowing the model to evaluate groups of related sentences instead of treating them independently
Pareto Frontier ShiftSimultaneously improves accuracyand reduces latency versus the EXIT baseline on single-hop tasks — pushing the accuracy-efficiency trade-off frontier

Architecture

QUITO-X: Coarse Filter

QUITO-X Coarse Filtering Architecture

QUITO-X adapts the Information Bottleneck (IB) framework to sentence-level filtering using the cross-attention mechanism of a frozen Flan-T5-Small (80M parameters):

  1. Tokenise all sentences into a flat token array, tracking sentence boundaries [start_j, end_j)
  2. Chunk into windows of ≤ 510 − |query tokens| tokens; prepend chunk, append query + EOS
  3. Forward through Flan-T5 with output_attentions=True; extract last-layer decoder cross-attention
  4. Apply Gaussian smoothing (σ=2.0) across the token attention array for stability
  5. Compute per-sentence score as max(smoothed_attentions[start_j : end_j])
  6. Retain top-k sentences by score using a dynamic threshold: k = max(k_min, ⌈r × |S_i|⌉), where r=0.8

EP-EXIT: Fine-Grained Span Filter

EP-EXIT Evidence Unit Extraction

EP-EXIT extends EXIT with graph-based evidence unit decomposition:

  1. Embed sentences via all-MiniLM-L6-v2 (SentenceTransformer)

  2. Build similarity graph: Add edge (i, j) iff |i−j| ≤ w (locality window=2) and cosine_sim(i,j) ≥ δ (δ=0.45)

  3. Extract Evidence Units: Each connected component in the graph becomes one evidence span; isolated nodes are singletons

  4. Classify spans: Pass each span through the EXIT-fine-tuned Gemma-2B-it classifier (doubleyyh/exit-gemma-2b, 4-bit NF4 quantised). Relevance score:

    r(u) = P("Yes" | q, d_i, u) / [P("Yes" | q, d_i, u) + P("No" | q, d_i, u)]
    
  5. Retain spans where r(u) ≥ τ (τ=0.5); concatenate retained spans in original document order

Key difference from GraphSeg: GRASP uses connected components (not maximal cliques) for broader, more robust grouping; a locality window for efficiency; and is optimised for evidence span selection in RAG (not generic text segmentation).


Experimental Results

Evaluated on three QA benchmarks (50 queries each) using Qwen-2.5:3B as the reader on an NVIDIA RTX-4000 Ada Edition GPU (20GB VRAM).

Comprehensive Results (Table I)

MethodHotpotQA EMHotpotQA F1HotpotQA R-LHotpotQA CompHotpotQA Lat(s)NQ EMNQ F1NQ R-LNQ CompNQ Lat(s)TriviaQA EMTriviaQA F1TriviaQA R-LTriviaQA CompTriviaQA Lat(s)
NoOp (Uncompressed)48.054.6153.940.0%0.4130.027.5428.160.0%0.4382.041.0240.420.0%0.70
RECOMP-Extr34.034.8134.5385.0%0.2732.030.5130.9179.9%0.2760.060.5159.8579.8%0.31
LLMLingua-214.017.9617.3685.9%0.8824.022.5122.4884.8%0.9950.051.6853.6885.8%1.27
EXIT48.053.1452.4087.3%1.7324.024.2924.2269.3%1.7266.065.8466.8081.6%2.91
GRASP (Proposed)42.046.9546.7579.5%1.1934.035.8535.1144.8%1.2370.068.2167.8364.3%2.11

↑ higher is better for EM, F1, R-L · Comp = compression ratio (% tokens removed vs NoOp) · Lat = end-to-end latency

Key Findings

Single-hop tasks (NQ & TriviaQA):

  • GRASP achieves +10 EM on Natural Questions (34.0% vs EXIT's 24.0%) — even beating the NoOp upper bound (30.0%)
  • GRASP achieves +4 EM on TriviaQA (70.0% vs EXIT's 66.0%)
  • 27% latency speedup on TriviaQA (2.11s vs EXIT's 2.91s) and 28% speedup on NQ (1.23s vs 1.72s) — despite adding an extra compression stage
  • The speedup is paradoxical: QUITO-X's lightweight 60M-param filter removes ~20% of irrelevant sentences before the expensive 2B-param Gemma classifier, shrinking its payload significantly

Multi-hop task (HotpotQA):

  • GRASP scores 42.0% EM vs EXIT's 48.0% — a known trade-off: graph-based span grouping can remove "bridge" sentences needed to chain multi-hop reasoning steps. This is an identified area for future work.

Case Study

Case Study: GRASP vs RECOMP on HotpotQA

The case study illustrates how GRASP preserves the critical multi-document evidence span needed to correctly answer a HotpotQA multi-hop question, while RECOMP fails to retrieve the necessary cross-document evidence.


Evaluation Metrics

MetricDefinition
Exact Match (EM)1 if the normalised predicted answer contains the normalised gold answer, 0 otherwise
Token F1 (F1)Harmonic mean of token-level precision and recall between prediction and gold
ROUGE-L (R-L)F-measure of the longest common subsequence (LCS) between prediction and gold
Compression (Comp)(1 − compressed_tokens / original_tokens) × 100% relative to NoOp
Latency (Lat)Total wall-clock time for compression + answer generation per query (seconds)

Repository Structure

grasp/
├── app.py                          # FastAPI REST backend (lifespan, auto-indexing)
├── requirements.txt                # Python dependencies
├── .env.example                    # Environment variable template
├── run_evals.sh / run_evals.ps1    # Sequential multi-dataset evaluation runner
├── run_preprocess.sh / .ps1        # Dataset preprocessing runner
│
├── assets/                         # README figures (workflow, QUITO-X, EP-EXIT, case study)
│
├── src/
│   ├── rag_pipeline.py             # End-to-end QueryAwareRAG orchestrator
│   │
│   ├── retrieval/
│   │   ├── retriever.py            # Dense-only Contriever retriever
│   │   └── hybrid_retriever.py     # Hybrid RRF retriever (Contriever + BM25)
│   │
│   ├── compression/
│   │   ├── hybrid_compressor.py    # Two-stage HybridCompressor (QUITO-X → EP-EXIT)
│   │   ├── quitox_filter.py        # Stage 2: QUITO-X coarse filter (Flan-T5-Small)
│   │   ├── ep_exit.py              # Stage 3: EP-EXIT fine filter (graph + Gemma-2B)
│   │   ├── exit_baseline.py        # EXIT baseline compressor (Gemma-2B, 4-bit NF4)
│   │   └── baselines/              # Re-implemented baseline compressors
│   │       ├── exit.py             #   EXIT
│   │       ├── compact.py          #   CompAct
│   │       ├── llmlingua2.py       #   LLMLingua-2
│   │       ├── recomp_abst.py      #   RECOMP Abstractive
│   │       ├── refiner.py          #   Refiner
│   │       └── recomp_extr/        #   RECOMP Extractive
│   │
│   ├── generation/
│   │   ├── reader.py               # Ollama-based RAG reader (Llama-3.1:8B)
│   │   └── gemma_reader.py         # Gemma reader variant
│   │
│   ├── eval/
│   │   ├── eval_pipeline.py        # GenerativeEvaluator (compress → read → score)
│   │   ├── metrics.py              # EM, Token F1, ROUGE-L
│   │   ├── adapters.py             # Unified compressor adapter interface
│   │   ├── interfaces.py           # SearchResult dataclass
│   │   ├── eval_hotpotqa.py        # HotpotQA benchmark script
│   │   ├── eval_nq.py              # Natural Questions benchmark script
│   │   ├── eval_tqa.py             # TriviaQA benchmark script
│   │   ├── eval_asqa.py            # ASQA benchmark script
│   │   └── eval_2wiki.py           # 2WikiMultiHopQA benchmark script
│   │
│   ├── budget_predictor/           # (Experimental) Query-difficulty budget predictor
│   └── data/
│       └── demo_loader.py          # Demo dataset loader for the UI
│
├── scripts/
│   ├── analyze_exit.py             # Step-by-step EXIT pipeline observability tool
│   ├── preprocess_hotpotqa.py      # HotpotQA preprocessing (offline retrieval)
│   ├── preprocess_nq.py            # NQ preprocessing
│   ├── preprocess_asqa.py          # ASQA preprocessing
│   ├── preprocess_2wiki.py         # 2WikiMultiHopQA preprocessing
│   └── preprocess_tqa.py           # TriviaQA preprocessing
│
├── data/                           # Pre-processed datasets (gitignored)
│
└── rag-frontend/                   # React + Vite interactive demo UI

Installation

Prerequisites

RequirementVersionNotes
Python≥ 3.10
CUDA≥ 11.8Required for Gemma-2B (4-bit NF4) and Flan-T5
GPU VRAM≥ 8 GBNVIDIA RTX 4000 Ada or equivalent
RAM≥ 16 GB
Disk≥ 20 GBFor datasets + model checkpoints
OllamaLatestRequired for the reader LLM

Step 1 — Clone & create environment

git clone https://github.com/Pvnn/grasp-rag.git
cd grasp-rag

python -m venv env
# Windows
env\Scripts\activate
# Linux/macOS
source env/bin/activate

Step 2 — Install PyTorch (CUDA 11.8)

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

Step 3 — Install project dependencies

pip install -r requirements.txt
python -m spacy download en_core_web_sm

Step 4 — Configure environment

cp .env.example .env

Edit .env:

HF_TOKEN=hf_xxxxxxxxxxxxxxxxxxxx     # Required: Hugging Face token (Gemma-2B access)

Gemma-2B access: Accept the Gemma model licence on Hugging Face before the first run.

Step 5 — Pull the reader model via Ollama

# Install Ollama from https://ollama.ai/
ollama pull llama3.1:8b

Quick Start

Running the API Server

python app.py

The FastAPI server starts at http://127.0.0.1:8000. On startup it initialises the full GRASP pipeline and loads demo datasets. API docs at http://127.0.0.1:8000/docs.

Running the Interactive UI

cd rag-frontend
npm install
npm run dev

Served at http://localhost:5173. The UI provides:

  • Dataset and query selection with auto-indexing
  • Per-sentence QUITO-X attention score visualisation (retained / pruned)
  • EP-EXIT evidence unit grouping (kept / removed spans)
  • Side-by-side compressed vs. uncompressed answer comparison
  • Live pipeline metrics (compression ratio, token savings, latency)

Programmatic Usage

from dotenv import load_dotenv
import os
from src.rag_pipeline import QueryAwareRAG

load_dotenv()
pipeline = QueryAwareRAG(token=os.getenv("HF_TOKEN"))

# Index your documents (list of dicts with 'title' and 'text' keys)
pipeline.retriever.index_documents(my_documents)

result = pipeline.run(
    query="What instruments on the James Webb Space Telescope are used for observation?",
    top_k=5,
    compare_original=True,   # Runs uncompressed baseline in parallel for comparison
    use_coarse=True,          # Enable QUITO-X stage
    use_fine=True             # Enable EP-EXIT stage
)

print(result['answer'])
print(f"Compression ratio: {result['metrics']['compression']['ratio_chars']:.1f}%")
print(f"Time saved:        {result['metrics']['times']['net_time_saved']:.2f}s")

REST API Reference

EndpointMethodDescription
/datasetsGETList available demo datasets and queries
/dataset/loadPOSTManually load a dataset into the retrieval index
/queryPOSTExecute the full GRASP pipeline on a query

POST /query payload:

{
  "query": "Were Scott Derrickson and Ed Wood of the same nationality?",
  "top_k": 4,
  "compare_original": true,
  "use_coarse": true,
  "use_fine": true
}

Dataset Preparation & Evaluation

1. Preprocess Datasets

# Windows
.\run_preprocess.ps1

# Linux/macOS
bash run_preprocess.sh

Preprocesses all benchmarks with offline top-30 hybrid retrieval into data/<dataset>/. To run a single dataset:

python scripts/preprocess_hotpotqa.py

2. Run Evaluations

# All datasets, 50 samples each
# Windows
.\run_evals.ps1 -N 50

# Linux/macOS
bash run_evals.sh 50

# Single dataset
python -m src.eval.eval_hotpotqa -n 50
python -m src.eval.eval_nq -n 50
python -m src.eval.eval_tqa -n 50

Results are saved to src/eval_results/<dataset>/:

  • final_benchmark_results.csv — aggregate metrics for all compressors
  • details_<compressor>.csv — per-query traces with context, predictions, and scores

Configuration

ParameterDefaultDescription
top_k5Documents to retrieve
coarse_ratio (r)0.8Fraction of sentences retained after QUITO-X
quitox_min_keep2Minimum sentences kept per document
fine_threshold (τ)0.5EP-EXIT inclusion threshold
similarity_threshold (δ)0.45Cosine similarity edge threshold for evidence graph
locality_window (w)2Maximum sentence index distance for graph edge
rrf_k (κ)10RRF smoothing constant
gaussian_sigma (σ)2.0Smoothing for QUITO-X attention scores

Ablation Setup

Each stage can be toggled independently:

# QUITO-X only
result = pipeline.run(query, use_coarse=True, use_fine=False)

# EP-EXIT only
result = pipeline.run(query, use_coarse=False, use_fine=True)

# Full GRASP pipeline
result = pipeline.run(query, use_coarse=True, use_fine=True)

# No compression (NoOp baseline)
result = pipeline.run(query, use_coarse=False, use_fine=False)

Models

ComponentModelParametersQuantisation
Dense Retrieverfacebook/contriever-msmarco110MFP32/FP16
QUITO-X Filtergoogle/flan-t5-small80MFP16
Sentence Embedderall-MiniLM-L6-v222MFP16 (GPU)
EP-EXIT Classifierdoubleyyh/exit-gemma-2b2B4-bit NF4
Reader LLMllama3.1:8b (Ollama)8BQ4_K_M

Limitations & Future Work

  • Multi-hop reasoning: The graph-based span grouping can remove "bridge" sentences needed to connect multi-hop reasoning chains. Future work targets hop-aware weighting in the evidence graph.
  • Fixed compression ratio: QUITO-X uses a static coarse_ratio. A query-difficulty-aware budget predictor (see src/budget_predictor/) is under development.
  • English only: Current models are English-specific. Multilingual dense retrievers and sentence encoders are a direct extension path.
  • End-to-end training: Compression and reader are currently independent; joint training is planned.
  • Scope: Text-based English-language QA only. Abstractive summarisation, multi-modal retrieval, and cross-lingual RAG are out of scope.

Project Context

Department of Computer Science and Engineering Mar Athanasius College of Engineering (Autonomous), Kothamangalam APJ Abdul Kalam Technological University · May 2026

Project Guide: Prof. Neethu Subash · Coordinator: Prof. Nimisha Abraham


Built with ❤️ at Mar Athanasius College of Engineering · Department of Computer Science and Engineering

Contributors

Pvnn

53 commits

Languages

Python

78.7%

JavaScript

16.5%

PowerShell

2.4%

Shell

1.6%