How zhishanq/quco-rag works
1. Repository purpose
- This is an experimental RAG framework centered on QuCo-RAG, a dynamic retrieval-augmented generation method for multi-hop QA.
- The repo also includes several baseline RAG variants and an evaluation pipeline, so the main architecture is: config-driven experiment runner + retrieval infrastructure + LLM generation + scoring.
2. Top-level experiment flow
- A run starts from a config file under config/.
- The entrypoint in src/main_quco.py or src/main_baseline.py loads the dataset, initializes the chosen model, connects to retrieval backends, generates answers, and writes run artifacts to result/.
- src/evaluate.py then consumes those artifacts to compute EM/F1 and per-sample outputs.
3. Configuration as the control plane
- The config/ tree is the main orchestration layer.
- Configs are organized by model family (OLMo, Qwen, Meta-Llama, GPT API models) and by dataset (2WikiMultihopQA, HotpotQA).
- They also encode method choice, retrieval type, caching, prompt settings, output paths, and thresholds, which makes the system highly parameterized and easy to diagram as a “configuration-driven dispatcher.”
4. Dataset ingestion
- src/data.py is the dataset boundary.
- It loads benchmark data from local files under data/, specifically the two supported QA benchmarks.
- The README shows the system often samples a subset for faster experiments, but the underlying unit of processing is a question instance with supporting metadata.
5. QuCo-RAG generation pipeline
- src/generate_quco.py implements the core method.
- The pipeline is two-stage: it decides when to retrieve before generation using corpus-grounded uncertainty, then can verify during generation whether retrieval is needed again.
- The method uses entity extraction plus Infini-gram frequency/co-occurrence checks to detect low-confidence or hallucination-prone cases before issuing retrieval calls.
6. Baseline generation pipeline
- src/generate_baselines.py supports comparison methods such as SR-RAG, FL-RAG, FLARE, DRAGIN, and no-retrieval.
- src/main_baseline.py routes into these modes with the same overall experiment skeleton as QuCo-RAG, but with simpler or fixed retrieval behavior.
- This creates a clean architectural boundary between the adaptive QuCo logic and static baseline policies.
7. Retrieval subsystem
- src/retriever.py is the main retrieval abstraction, with src/qwen3_retriever.py as an alternate embedding-based retriever path.
- The default production retrieval backend is Elasticsearch BM25 over a Wikipedia index named wiki.
- Retrieval is a distinct service boundary: the LLM does not search directly; it asks the retriever, which queries Elasticsearch and returns documents.
8. External search/index infrastructure
- The repo depends on a prebuilt or locally built Elasticsearch 7.17.9 instance.
- Scripts at the root and in tools/ manage index setup: start_es.sh, Start_Elasticsearch_from_hf.sh, Start_Elasticsearch_from_hf_HPC.sh, and tools/prep_elastic_index_with_progress.py.
- The wiki index is a major persistent asset in the architecture; experiments assume it is already available and running before inference starts.
9. Corpus-grounded uncertainty source
- A key architectural dependency is Infini-gram, used to query entity frequency and co-occurrence statistics from a very large pretraining corpus.
- This is not a local index in the repo; it behaves like an external knowledge/statistics service that QuCo-RAG consults during triggering decisions.
- The config fields such as infini_gram_index_name and ngram_threshold_question show that corpus statistics are a first-class control input to the method.
10. Model execution layer
- The system supports both local Hugging Face models and OpenAI API models.
- src/gpt_api_client.py is the boundary for API-based LLM access, while local model execution uses standard model loading from Hugging Face.
- Supported model families in the configs include OLMo, Qwen, Llama, and GPT-4.1/4o/5-chat-latest, so the architecture is explicitly model-agnostic.
11. Prompting and task formatting
- src/prompt_templates.json centralizes prompt variants and method-specific prompt logic.
- Prompt construction sits between dataset loading and model generation, shaping how retrieved evidence and question text are turned into generation inputs.
- This is important for the diagram because it is a separate transformation layer, not part of retrieval or model inference itself.
12. Caching and acceleration path
- The repo has an optional cache layer under data/cache/, and configs can enable it with enable_cache.
- The cache stores results for entity extraction, Infini-gram queries, and wiki retrieval, reducing repeated external calls and Elasticsearch traffic.
- Functionally, the cache sits between QuCo’s decision logic and its external dependencies, acting as a read-through accelerator.
13. Optional alternative retriever encoding path
- tools/encode_qwen3_vllm.py and its shell wrapper provide an alternate path for building embeddings with Qwen3-Embedding-0.6B via vLLM.
- This is not the default path, but it represents an alternate retrieval backend preparation pipeline for users who want embedding-based retrieval instead of BM25.
14. Outputs and evaluation
- Each experiment writes a run directory in result/ containing the config snapshot, raw outputs, retrieved documents, and metrics files.
- src/evaluate.py consumes those outputs to produce EM/F1 and detailed per-example analysis.
- This gives the repo a closed loop: config → run → artifacts → evaluation, which is central for reproducing paper results.
15. Main architecture boundaries
- The clearest boundaries for a diagram are:
- Experiment orchestration: main_quco.py / main_baseline.py
- Data loading: data.py
- Prompting: prompt_templates.json
- LLM access: local HF models or gpt_api_client.py
- Retrieval: retriever.py / qwen3_retriever.py
- External services: Elasticsearch, Infini-gram, OpenAI API
- Persistence: data/cache/ and result/
- The repo’s distinctive feature is that QuCo-RAG makes retrieval decisions from corpus statistics, not just model uncertainty, so the architecture should show that extra decision branch explicitly.
Main entry points: main_quco.py QuCo runner · main_baseline.py Baseline runner · generate_quco.py QuCo generation · retriever.py Retriever
Components 18
Experiment control
Config-driven entrypoints and evaluation loop
- Config tree · Model- and dataset-specific experiment settings
config - QuCo runner · Runs QuCo-RAG experiments from config
src/main_quco.py - Baseline runner · Runs static baseline RAG variants
src/main_baseline.py - Evaluator · Computes EM/F1 and per-sample analysis
src/evaluate.py
Core pipeline
Dataset loading, prompting, and generation logic
- Dataset loader · Loads 2WikiMultihopQA and HotpotQA samples
src/data.py - Prompt templates · Centralized method-specific prompt construction
src/prompt_templates.json - QuCo generation · Two-stage retrieval triggering and verification logic
src/generate_quco.py - Baseline generation · Implements SR-RAG, FL-RAG, FLARE, DRAGIN, and no-retrieval modes
src/generate_baselines.py
Retrieval services
Document retrieval backends and external search dependencies
- Retriever · Main retrieval abstraction over Elasticsearch wiki index
src/retriever.py - Qwen3 retriever · Alternate embedding-based retrieval path
src/qwen3_retriever.py - GPT API client · API-based model access for GPT-family runs
src/gpt_api_client.py - Elasticsearch wiki · Wikipedia BM25 index named wiki
- Infini-gram · Corpus frequency and co-occurrence source for uncertainty checks
Infra and artifacts
Index preparation, cache, and run outputs
- Cache · Stores entity extraction, Infini-gram, and retrieval results
- Run artifacts · Config snapshots, raw outputs, retrieved docs, and metrics
- Index prep · Builds or loads Elasticsearch wiki index
tools/prep_elastic_index_with_progress.py - Qwen3 encoding · Creates embeddings with Qwen3-Embedding via vLLM
tools/encode_qwen3_vllm.py - ES startup · Starts Elasticsearch from local or HF assets
start_es.sh
Connections 24
- Config tree → QuCo runner: select QuCo run · Config selects method, model, dataset, retrieval, and thresholds
- Config tree → Baseline runner: select baseline run · Config routes into baseline variants and model family
- QuCo runner → Dataset loader: load data · Loads benchmark questions and metadata
- Baseline runner → Dataset loader: load data · Loads benchmark questions and metadata
- QuCo runner → Prompt templates: format inputs · Builds prompts from question, evidence, and method settings
- Baseline runner → Prompt templates: format inputs · Builds prompts for baseline generation modes
- QuCo runner → QuCo generation: run method · Executes adaptive retrieval and generation
- Baseline runner → Baseline generation: run method · Executes fixed-policy baseline generation
- QuCo generation → Retriever: retrieve docs · Triggers retrieval when uncertainty or verification requires it
- Baseline generation → Retriever: retrieve docs · Uses retrieval according to baseline policy
- QuCo generation → Infini-gram: score uncertainty · Uses corpus statistics to decide whether to retrieve
- QuCo generation → Cache: cache lookups · Reuses entity, corpus-stat, and retrieval results when enabled
- Retriever → Elasticsearch wiki: query index · Queries Elasticsearch BM25 wiki corpus
- Retriever → Cache: read/write cache · Caches retrieval outputs when enabled
- Qwen3 retriever → Cache: read/write cache · Can reuse cached embedding retrieval outputs
- Qwen3 retriever → Elasticsearch wiki: query index · Alternate retrieval path over indexed corpus
- QuCo runner → GPT API client: API generation · Uses API-based LLMs when configured
- Baseline runner → GPT API client: API generation · Uses API-based LLMs when configured
- QuCo runner → Run artifacts: write artifacts · Persists run outputs and metadata
- Baseline runner → Run artifacts: write artifacts · Persists run outputs and metadata
- Run artifacts → Evaluator: consume artifacts · Evaluation reads saved run outputs
- ES startup → Elasticsearch wiki: start service · Brings up Elasticsearch for retrieval
- Index prep → Elasticsearch wiki: build index · Prepares the wiki index used for BM25 retrieval
- Qwen3 encoding → Qwen3 retriever: prepare embeddings · Produces embeddings for the alternate retrieval backend