AI Retrieval Benchmark for Document Analysis

Reproducible benchmark that compares retrieval strategies for LLM analysis of legal text.

It runs the same model prompt against multiple context-construction methods and measures:

• response quality versus a ground-truth pass
• latency
• token usage
• retrieval coverage and hallucination behavior

The benchmark target document is nda_section.md (NDA Section 8).

TL;DR

• Goal: compare retrieval strategies (flat vs semantic vs hierarchical) on the same legal extraction task.
• Entry point: run.py.
• Output: results.json + terminal summary table.
• Local run: BACKEND=ollama python run.py.
• Longleaf run: BACKEND=vllm_server python run.py (recommended — keeps model loaded between runs).

Quick start

# Local
pip install requests transformers torch langchain-experimental langchain-huggingface sentence-transformers
BACKEND=ollama python run.py

# HPC — start model server in one terminal, run benchmark in another
# Terminal 1 (keep open):
vllm serve <model_path> --dtype float16 --port 8000
# Terminal 2:
BACKEND=vllm_server python run.py

# HPC (batch)
sbatch submit.sl

Why this exists

Most "RAG quality" discussions mix retrieval and generation errors together. This project separates them:

• Retrieval quality: did the retriever surface enough source text?
• Extraction quality: given that text, did the LLM answer correctly?

run.py evaluates both with field-level scoring against a ground-truth run.

Repository layout

• run.py — benchmark driver, chunking/retrieval logic, model invocation, scoring, and summary output
• nda_section.md — fixed legal source text under test
• prompt.md — fixed extraction prompt template ({context} placeholder)
• results.json — machine-readable benchmark output
• submit.sl — Slurm job script for Longleaf GPU execution

Benchmark runs

The script currently executes:

1. GROUND TRUTH — Full Section 8

   • passes full source text into context
   • acts as reference output for comparison

2. RUN 2 — BASELINE (flat chunking, BGE-M3 top 10)

   • fixed-size word chunks (approx 500 tokens/chunk)
   • BGE-M3 cosine similarity retrieval, top 10

3. RUN 2.1 — BASELINE + CROSS-ENCODER (flat, BGE-M3→CE top 10)

   • same flat chunks as Run 2
   • BGE-M3 retrieves top 20 candidates; cross-encoder (ms-marco-MiniLM-L6-v2) reranks to top 10

4. RUN 3 — SEMANTIC CHUNKING (topic-shift, BGE-M3 top 10)

   • LangChain SemanticChunker with MiniLM embeddings for splitting
   • BGE-M3 cosine similarity retrieval, top 10

5. RUN 3.1 — SEMANTIC + CROSS-ENCODER (topic-shift, BGE-M3→CE top 10)

   • same semantic chunks as Run 3
   • BGE-M3 retrieves top 20 candidates; cross-encoder reranks to top 10

6. RUN 4 — HICHUNK + AUTO-MERGE (hierarchical, top 10)

   • structure-aware regex chunker identifies section/subsection/clause boundaries
   • BGE-M3 retrieval over leaf nodes, auto-merge promotes siblings to parent context
   • if full HiChunk repo is unavailable, falls back to the structure-aware chunker for local compatibility

7. RUN 4.1 — LLAMAINDEX (Structure-Aware HiChunk + AutoMerge, top 10)

   • same structure_aware_chunk() splits as Run 4, built as LlamaIndex TextNode objects
   • explicit PARENT/CHILD NodeRelationship links; parent text = full leaf concatenation
   • Qdrant in-memory vector store, BGE-M3 embeddings, AutoMergingRetriever

8. RUN 4.2 — HICHUNK + CROSS-ENCODER + AUTO-MERGE (BGE-M3→CE top 10)

   • same HiChunk tree as Run 4
   • BGE-M3 retrieves top 20 leaf candidates; cross-encoder reranks to top 10 before auto-merge

9. RUN 5 — LIGHTRAG (full document, hybrid)

   • full document inserted into LightRAG in a single call
   • LightRAG extracts entities and relationships using the active LLM backend
   • retrieval via hybrid graph traversal (local entity + global relationship)
   • knowledge graph built in a temp dir and discarded after the run

All retrieval runs use the same QUERY (aligned with the prompt question) and BAAI/bge-m3 for embeddings. Only chunking and retrieval strategy vary — this isolates those two variables across runs.

Run matrix (at a glance)

Run	Chunking	Retrieval	Expected behavior
Ground Truth	Full document	—	Reference output
Run 2	Flat (500-token)	BGE-M3 top 10	Flat boundaries cut across sections
Run 2.1	Flat (500-token)	BGE-M3→cross-encoder top 10	Tests whether reranking compensates for flat boundaries
Run 3	Semantic (MiniLM)	BGE-M3 top 10	Better split points, same retrieval
Run 3.1	Semantic (MiniLM)	BGE-M3→cross-encoder top 10	Tests whether reranking improves on semantic splits
Run 4	Hierarchical (regex)	BGE-M3 + auto-merge top 10	Structure-aware, sibling promotion
Run 4.1	Hierarchical (regex, LlamaIndex)	BGE-M3 + auto-merge top 10	Same chunking as Run 4, LlamaIndex/Qdrant stack
Run 4.2	Hierarchical (regex)	BGE-M3→cross-encoder + auto-merge top 10	Tests whether better leaf selection improves merged context
Run 5	Full document	LightRAG hybrid graph traversal	Entity/relation-aware retrieval

Backends

run.py supports four LLM backends, selected via the BACKEND environment variable:

BACKEND	Where	Model	Notes
ollama	Local Mac	qwen2.5:7b	Requires Ollama running locally
vllm_server	Longleaf GPU	Qwen/Qwen2.5-14B-Instruct	Recommended on GPU — connects to a persistent vllm serve process; no reload between runs
vllm	Longleaf GPU	Qwen/Qwen2.5-14B-Instruct	Loads model in-process; reloads on every run
unset (huggingface)	Longleaf GPU	Qwen/Qwen2.5-14B-Instruct	Falls back to transformers.pipeline

Environment variables

Variable	Purpose
HF_HOME	Where HuggingFace looks for cached models. Set to work storage (/work/users/...) to avoid filling your home quota with the 29 GB model.
HF_HUB_OFFLINE=1	Prevents huggingface_hub from making any network requests. Required on compute nodes, which have no internet access.
TRANSFORMERS_OFFLINE=1	Same as above but for the transformers library. Both must be set since they are separate libraries.
VLLM_SERVER_URL	URL of the vLLM server (default: http://localhost:8000). Override if serving on a different port.

Setup

Option A: Local Mac workflow (typical development path)

1. Install Python dependencies:

pip install requests transformers torch langchain-experimental langchain-huggingface sentence-transformers \
  llama-index-core llama-index-vector-stores-qdrant llama-index-embeddings-huggingface qdrant-client lightrag-hku

2. Install and run Ollama:

ollama pull qwen2.5:7b

3. Run benchmark:

BACKEND=ollama python run.py

Notes:

• Runs 1–4 work in this mode.
• Run 4.1 requires llama-index-* and qdrant-client — skipped gracefully if not installed.
• Run 5 requires lightrag-hku — skipped gracefully if not installed.

Option B: Longleaf GPU workflow

One-time model download

Compute nodes have no internet access. Download the model from the data transfer node before your first run:

ssh longleaf-xfer.its.unc.edu
module load anaconda
conda activate hichunk
export HF_HOME=/work/users/v/k/vkereka/hf_cache
huggingface-cli download Qwen/Qwen2.5-14B-Instruct

Run the download inside tmux — it pulls ~29 GB and takes 15–20 minutes. Do not run large downloads on the login node; it will be killed by the memory limit.

One-time dependency install

If any packages are missing from the hichunk env, install them on a compute node:

pip install langchain-experimental langchain-huggingface sentence-transformers \
  llama-index-core llama-index-vector-stores-qdrant llama-index-embeddings-huggingface \
  qdrant-client lightrag-hku

Running interactively (recommended: persistent server)

This approach loads the model once and keeps it in GPU memory between runs.

Terminal 1 — start the model server (keep this open):

module load anaconda && module load cuda && conda activate hichunk
export HF_HOME=/work/users/v/k/vkereka/hf_cache
export HF_HUB_OFFLINE=1
export TRANSFORMERS_OFFLINE=1
vllm serve /work/users/v/k/vkereka/hf_cache/hub/models--Qwen--Qwen2.5-14B-Instruct/snapshots/<hash> --dtype float16 --port 8000

Wait for Application startup complete.

Terminal 2 — run the benchmark (repeat as needed):

module load anaconda && conda activate hichunk
cd ~/Projects/vv
BACKEND=vllm_server python run.py

Running interactively (single terminal)

Loads the model fresh on every run — slower but simpler.

srun -n 1 --cpus-per-task=4 --mem=32g -t 1:00:00 -p l40-gpu --qos=gpu_access --gres=gpu:1 --pty bash
module purge && module load anaconda && module load cuda && conda activate hichunk
export HF_HOME=/work/users/v/k/vkereka/hf_cache
export TRANSFORMERS_OFFLINE=1
export HF_HUB_OFFLINE=1
cd ~/Projects/vv
BACKEND=vllm python run.py

Running via batch

submit.sl is pre-configured for the l40-gpu partition with the correct HF_HOME:

sbatch submit.sl

Output and metrics

Results are written to results.json with two top-level keys:

{
  "summary": [ ... ],
  "runs":    [ ... ]
}

summary is a compact table — one row per run with just the headline numbers. runs is the full detail.

Terminal summary table

After every run the script prints:

Run                        Latency  Instruction Following  Isolated Extraction Score  Document Coverage  Halluc
GROUND TRUTH               17.3s          (baseline)                      —                  —               —
RUN 2 — BASELINE           16.2s             85.5%                     84.8%              95.5%              0
RUN 3 — SEMANTIC           98.7s             94.1%                     83.8%              36.4%              0
RUN 4 — HICHUNK            15.4s            100.0%                    100.0%              54.5%              0

Column	What it means
Latency	Wall-clock time for the LLM call
Instruction Following	Average per-field score from the LLM judge across all fields. Fields where the retriever failed and the model correctly said [NOT FOUND] still score 1.0 — so this number can look high even when coverage is low.
Isolated Extraction Score	Same scoring, but fields where the retriever never surfaced the content are excluded entirely. This isolates extraction quality: given what the retriever found, how well did the model use it? This is the primary quality signal.
Document Coverage	Fraction of the document's fields the retriever actually surfaced. A run can score 100% Isolated Extraction but only 50% Document Coverage — meaning it was perfect on what it found, but missed half the document.
Halluc	Count of fields where the model asserted a specific fact not present in the retrieved context.

How to read Instruction Following vs Isolated Extraction Score together: A high Instruction Following score with low Document Coverage means the model behaved correctly but the retriever failed it. A high Isolated Extraction Score with low Document Coverage means the retrieval strategy is precise but narrow — increase top_k to improve coverage.

Per-field scoring categories

Each field in details.judge_results has a category:

Category	Meaning
correct	Answer matches ground truth in legal meaning
correct_gap	Context didn't contain the answer; model correctly said [NOT FOUND]
over_extracted	Correct answer but includes extra text beyond what was asked
under_extracted	Correct but missing key detail or conditions
partial	Got some of the answer but missed parts (e.g. 1 of 6 survival sections)
missing	Context had the answer but model returned something wrong
hallucinated	Model stated a specific fact not present in the retrieved context
parse_error	LLM judge returned malformed JSON; score defaults to 0.0

Full run object fields

Each entry in runs includes:

• run — label
• latency_seconds
• tokens_prompt / tokens_response
• context_used — exact text sent to the model for that run
• output — raw model response
• auto_merge_fired — (Runs 4, 4.1, 4.2) whether auto-merge promoted any retrieved chunk to its parent section
• comparison — scoring results vs ground truth, including per-field judge_results

How to read results.json

Start with summary. This is the same as the terminal table — one row per run. Look at retrieval_coverage first. If coverage is low, the retriever is the problem regardless of what the other scores say. Then look at retrieval_adjusted_score to see how well the model used what it did retrieve.

To investigate a specific run, find it in runs by matching the run label.

Check context_used to see exactly what text the model had access to. If a field scored [NOT FOUND] and you think the answer should have been there, look here first — the answer may simply not have been retrieved.

Check output to see the raw model response before any parsing. If scores look wrong, check whether the model followed the output format correctly.

To find where a run went wrong, go to comparison.details.judge_results. Each entry is one field with its score and a one-line reason. Look for category: "partial", "missing", or "hallucinated" — those are the failures. correct_gap means the retriever missed it, not the model.

retrieval_gaps vs missing_in_run: both count fields that got [NOT FOUND], but retrieval_gaps is the subset where the context genuinely didn't contain the answer. If missing_in_run is high but retrieval_gaps is low, the model had the answer available and still said [NOT FOUND] — that's an extraction failure, not a retrieval failure.

parse_error entries scored 0.0 but may not actually be wrong — the judge produced invalid JSON for those fields. Check the reasoning string; if it starts with Failed to parse judge response:, treat that field's score as unreliable.

How scoring works

compare_to_ground_truth() parses structured field outputs and compares each run to run 1.

The judge prompt in run.py enforces a two-step evaluation:

1. Determine whether the needed information existed in the provided context.
2. Score extraction quality only after checking retrieval sufficiency.

This avoids penalizing a model for missing data it never received, and flags hallucinations when it invents unsupported facts.

Reproducibility notes

• prompt.md and nda_section.md are benchmark fixtures; keep them stable when comparing retrieval strategies.
• Summary trends are more informative than one-off run variance.
• For fair comparisons, keep backend/model, prompt, and source document constant.

What success looks like

• RUN 4 should generally improve coverage vs RUN 3 on deeply nested clauses.
• retrieval_coverage should increase when hierarchical retrieval is working.
• retrieval_adjusted_score helps separate extraction mistakes from retrieval misses.
• Hallucinations should remain low even as context breadth increases.

Customization points

Safe modifications for experiments:

• retrieval parameters (top_k, chunk settings)
• retrieval algorithms/chunkers
• backend model choice

Avoid modifying while benchmarking strategy quality:

• run() interface
• scoring logic in compare_to_ground_truth()
• fixture semantics in prompt.md and nda_section.md

Troubleshooting

• Skipping RUN 4 — sentence-transformers not installed

  • run pip install langchain-experimental langchain-huggingface sentence-transformers

• Cannot find the requested files in the disk cache

  • model hasn't been downloaded yet, or HF_HOME is pointing to the wrong path
  • verify with ls /work/users/v/k/vkereka/hf_cache/hub/

• Killed during model download

  • login nodes have memory limits — use longleaf-xfer.its.unc.edu for large downloads

• Could not find nvcc / FlashInfer build error

  • CUDA toolkit is not loaded — run module load cuda before starting vLLM

• Connection refused with BACKEND=vllm_server

  • the vLLM server in Terminal 1 is not running or hasn't finished starting up yet

• slow or failed model load

  • use BACKEND=ollama locally, or BACKEND=vllm_server on GPU

• unexpected retrieval quality regressions

  • inspect context_used in results.json
  • verify top_k, merge behavior, and chunk boundaries

• no summary differences across runs

  • ensure chunking/retrieval settings actually differ per run