# zhishanq/quco-rag architecture

> Architecture diagram of the GitHub repository zhishanq/quco-rag, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-06-23). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/zhishanq/quco-rag
- Repository: https://github.com/zhishanq/quco-rag

## Overview

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: [`src/main_quco.py`](https://github.com/zhishanq/quco-rag/blob/main/src/main_quco.py) (QuCo runner), [`src/main_baseline.py`](https://github.com/zhishanq/quco-rag/blob/main/src/main_baseline.py) (Baseline runner), [`src/generate_quco.py`](https://github.com/zhishanq/quco-rag/blob/main/src/generate_quco.py) (QuCo generation), [`src/retriever.py`](https://github.com/zhishanq/quco-rag/blob/main/src/retriever.py) (Retriever)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_g1["Experiment control"]
  node_config_tree["Config tree<br/>config dispatcher"]
  node_main_quco["QuCo runner<br/>entrypoint<br/>[main_quco.py]"]
  node_main_baseline["Baseline runner<br/>entrypoint<br/>[main_baseline.py]"]
  node_evaluate["Evaluator<br/>metrics<br/>[evaluate.py]"]
end

subgraph group_g2["Core pipeline"]
  node_data["Dataset loader<br/>input boundary<br/>[data.py]"]
  node_prompts["Prompt templates<br/>prompt layer"]
  node_gen_quco["QuCo generation<br/>adaptive RAG<br/>[generate_quco.py]"]
  node_gen_baselines["Baseline generation<br/>fixed-policy RAG"]
end

subgraph group_g3["Retrieval services"]
  node_retriever["Retriever<br/>BM25 adapter<br/>[retriever.py]"]
  node_qwen3_retriever["Qwen3 retriever<br/>embedding retriever<br/>[qwen3_retriever.py]"]
  node_gpt_client["GPT API client<br/>LLM boundary<br/>[gpt_api_client.py]"]
  node_es_index[("Elasticsearch wiki<br/>external index")]
  node_infini_gram[("Infini-gram<br/>external stats service")]
end

subgraph group_g4["Infra and artifacts"]
  node_cache[("Cache")]
  node_result_dir[("Run artifacts<br/>experiment outputs")]
  node_prep_index["Index prep<br/>setup script"]
  node_encode_qwen3["Qwen3 encoding<br/>prep pipeline"]
  node_start_es["ES startup<br/>ops script<br/>[start_es.sh]"]
end

node_config_tree -->|"select QuCo run"| node_main_quco
node_config_tree -->|"select baseline run"| node_main_baseline
node_main_quco -->|"load data"| node_data
node_main_baseline -->|"load data"| node_data
node_main_quco -->|"format inputs"| node_prompts
node_main_baseline -->|"format inputs"| node_prompts
node_main_quco -->|"run method"| node_gen_quco
node_main_baseline -->|"run method"| node_gen_baselines
node_gen_quco -->|"retrieve docs"| node_retriever
node_gen_baselines -->|"retrieve docs"| node_retriever
node_gen_quco -->|"score uncertainty"| node_infini_gram
node_gen_quco -->|"cache lookups"| node_cache
node_retriever -->|"query index"| node_es_index
node_retriever -->|"read/write cache"| node_cache
node_qwen3_retriever -->|"read/write cache"| node_cache
node_qwen3_retriever -->|"query index"| node_es_index
node_main_quco -->|"API generation"| node_gpt_client
node_main_baseline -->|"API generation"| node_gpt_client
node_main_quco -->|"write artifacts"| node_result_dir
node_main_baseline -->|"write artifacts"| node_result_dir
node_result_dir -->|"consume artifacts"| node_evaluate
node_start_es -->|"start service"| node_es_index
node_prep_index -->|"build index"| node_es_index
node_encode_qwen3 -->|"prepare embeddings"| node_qwen3_retriever

click node_config_tree "https://github.com/zhishanq/quco-rag/tree/main/config"
click node_main_quco "https://github.com/zhishanq/quco-rag/blob/main/src/main_quco.py"
click node_main_baseline "https://github.com/zhishanq/quco-rag/blob/main/src/main_baseline.py"
click node_evaluate "https://github.com/zhishanq/quco-rag/blob/main/src/evaluate.py"
click node_data "https://github.com/zhishanq/quco-rag/blob/main/src/data.py"
click node_prompts "https://github.com/zhishanq/quco-rag/blob/main/src/prompt_templates.json"
click node_gen_quco "https://github.com/zhishanq/quco-rag/blob/main/src/generate_quco.py"
click node_gen_baselines "https://github.com/zhishanq/quco-rag/blob/main/src/generate_baselines.py"
click node_retriever "https://github.com/zhishanq/quco-rag/blob/main/src/retriever.py"
click node_qwen3_retriever "https://github.com/zhishanq/quco-rag/blob/main/src/qwen3_retriever.py"
click node_gpt_client "https://github.com/zhishanq/quco-rag/blob/main/src/gpt_api_client.py"
click node_prep_index "https://github.com/zhishanq/quco-rag/blob/main/tools/prep_elastic_index_with_progress.py"
click node_encode_qwen3 "https://github.com/zhishanq/quco-rag/blob/main/tools/encode_qwen3_vllm.py"
click node_start_es "https://github.com/zhishanq/quco-rag/blob/main/start_es.sh"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_config_tree,node_main_quco,node_main_baseline,node_evaluate toneBlue
class node_data,node_prompts,node_gen_quco,node_gen_baselines toneAmber
class node_retriever,node_qwen3_retriever,node_gpt_client,node_es_index,node_infini_gram toneMint
class node_cache,node_result_dir,node_prep_index,node_encode_qwen3,node_start_es toneRose
```

## Components

### Experiment control

Config-driven entrypoints and evaluation loop

- [Config tree](https://github.com/zhishanq/quco-rag/tree/main/config): config dispatcher · Model- and dataset-specific experiment settings · `config`
- [QuCo runner](https://github.com/zhishanq/quco-rag/blob/main/src/main_quco.py): entrypoint · Runs QuCo-RAG experiments from config · `src/main_quco.py`
- [Baseline runner](https://github.com/zhishanq/quco-rag/blob/main/src/main_baseline.py): entrypoint · Runs static baseline RAG variants · `src/main_baseline.py`
- [Evaluator](https://github.com/zhishanq/quco-rag/blob/main/src/evaluate.py): metrics · Computes EM/F1 and per-sample analysis · `src/evaluate.py`

### Core pipeline

Dataset loading, prompting, and generation logic

- [Dataset loader](https://github.com/zhishanq/quco-rag/blob/main/src/data.py): input boundary · Loads 2WikiMultihopQA and HotpotQA samples · `src/data.py`
- [Prompt templates](https://github.com/zhishanq/quco-rag/blob/main/src/prompt_templates.json): prompt layer · Centralized method-specific prompt construction · `src/prompt_templates.json`
- [QuCo generation](https://github.com/zhishanq/quco-rag/blob/main/src/generate_quco.py): adaptive RAG · Two-stage retrieval triggering and verification logic · `src/generate_quco.py`
- [Baseline generation](https://github.com/zhishanq/quco-rag/blob/main/src/generate_baselines.py): fixed-policy RAG · 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](https://github.com/zhishanq/quco-rag/blob/main/src/retriever.py): BM25 adapter · Main retrieval abstraction over Elasticsearch wiki index · `src/retriever.py`
- [Qwen3 retriever](https://github.com/zhishanq/quco-rag/blob/main/src/qwen3_retriever.py): embedding retriever · Alternate embedding-based retrieval path · `src/qwen3_retriever.py`
- [GPT API client](https://github.com/zhishanq/quco-rag/blob/main/src/gpt_api_client.py): LLM boundary · API-based model access for GPT-family runs · `src/gpt_api_client.py`
- **Elasticsearch wiki**: external index · Wikipedia BM25 index named wiki
- **Infini-gram**: external stats service · Corpus frequency and co-occurrence source for uncertainty checks

### Infra and artifacts

Index preparation, cache, and run outputs

- **Cache**: read-through cache · Stores entity extraction, Infini-gram, and retrieval results
- **Run artifacts**: experiment outputs · Config snapshots, raw outputs, retrieved docs, and metrics
- [Index prep](https://github.com/zhishanq/quco-rag/blob/main/tools/prep_elastic_index_with_progress.py): setup script · Builds or loads Elasticsearch wiki index · `tools/prep_elastic_index_with_progress.py`
- [Qwen3 encoding](https://github.com/zhishanq/quco-rag/blob/main/tools/encode_qwen3_vllm.py): prep pipeline · Creates embeddings with Qwen3-Embedding via vLLM · `tools/encode_qwen3_vllm.py`
- [ES startup](https://github.com/zhishanq/quco-rag/blob/main/start_es.sh): ops script · Starts Elasticsearch from local or HF assets · `start_es.sh`

## Connections

- 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