tarzanagh commited on
Commit
8e874f5
·
verified ·
1 Parent(s): 761dfc5

Add code: src, benchmarks, data, requirements, run.sh, README

Browse files
This view is limited to 50 files because it contains too many changes.   See raw diff
Files changed (50) hide show
  1. .gitattributes +6 -0
  2. .gitignore +47 -0
  3. README.md +345 -0
  4. benchmarks/multihop/benchmark_multihop.py +531 -0
  5. benchmarks/run.py +371 -0
  6. benchmarks/summarization/benchmark_summarization.py +673 -0
  7. benchmarks/text2sql/benchmark_text2sql.py +385 -0
  8. benchmarks/text2sql/evaluate_schema.py +213 -0
  9. benchmarks/ultradomain/benchmark_ultradomain.py +548 -0
  10. data/multihop/2wikimultihopqa.json +0 -0
  11. data/multihop/2wikimultihopqa_corpus.json +0 -0
  12. data/multihop/README.md +23 -0
  13. data/multihop/hotpotqa.json +0 -0
  14. data/multihop/hotpotqa_corpus.json +0 -0
  15. data/multihop/musique.json +3 -0
  16. data/multihop/musique_corpus.json +0 -0
  17. data/summarization/README.md +26 -0
  18. data/text2sql/README.md +70 -0
  19. data/text2sql/bird/bird.jsonl +129 -0
  20. data/text2sql/bird/databases/superhero/superhero.sqlite +3 -0
  21. data/text2sql/bird/databases/superhero/superhero_db_summary.json +0 -0
  22. data/text2sql/spider2-lite/bigquery/austin/austin_bigquery_summary.json +39 -0
  23. data/text2sql/spider2-lite/golden_lite_spider_total.json +1443 -0
  24. data/text2sql/spider2-lite/snowflake/AUSTIN/AUSTIN_db_summary.json +39 -0
  25. data/text2sql/spider2-lite/spider2-lite.jsonl +0 -0
  26. data/text2sql/spider2-lite/sqlite/Pagila/Pagila.sqlite +3 -0
  27. data/text2sql/spider2-lite/sqlite/Pagila/Pagila_db_summary.json +0 -0
  28. data/ultradomain/README.md +25 -0
  29. docs/figs/GraphRAG.png +3 -0
  30. docs/figs/LightRAG.png +3 -0
  31. docs/figs/QAFD-RAG.png +3 -0
  32. docs/index.html +452 -0
  33. experiments/edge_weight_ablation.py +401 -0
  34. experiments/query_aware_ablation.py +478 -0
  35. experiments/run_all_ablations.sh +49 -0
  36. experiments/run_alpha_qa_sweep.sh +29 -0
  37. experiments/run_batch_push_ablation.sh +29 -0
  38. experiments/run_batch_push_nvidia.sh +55 -0
  39. experiments/run_decoupled_ablation.sh +32 -0
  40. experiments/run_nvidia_qa_ablation.sh +43 -0
  41. experiments/run_overnight_qa_awareness.sh +136 -0
  42. experiments/run_qa_awareness_v2.sh +34 -0
  43. experiments/run_sim_mode_ablation.sh +29 -0
  44. requirements.txt +46 -0
  45. run.sh +113 -0
  46. src/QAFD_RAG.py +747 -0
  47. src/__init__.py +4 -0
  48. src/answering/__init__.py +42 -0
  49. src/answering/clusters.py +739 -0
  50. src/answering/context.py +322 -0
.gitattributes CHANGED
@@ -33,3 +33,9 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
 
 
 
 
 
 
 
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
36
+ data/multihop/musique.json filter=lfs diff=lfs merge=lfs -text
37
+ data/text2sql/bird/databases/superhero/superhero.sqlite filter=lfs diff=lfs merge=lfs -text
38
+ data/text2sql/spider2-lite/sqlite/Pagila/Pagila.sqlite filter=lfs diff=lfs merge=lfs -text
39
+ docs/figs/GraphRAG.png filter=lfs diff=lfs merge=lfs -text
40
+ docs/figs/LightRAG.png filter=lfs diff=lfs merge=lfs -text
41
+ docs/figs/QAFD-RAG.png filter=lfs diff=lfs merge=lfs -text
.gitignore ADDED
@@ -0,0 +1,47 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Knowledge graphs (large, download from https://huggingface.co/datasets/qafd/kg)
2
+ kg/**
3
+ !kg/
4
+ !kg/ultradomain/
5
+ !kg/ultradomain/.gitkeep
6
+ !kg/multihop/
7
+ !kg/multihop/.gitkeep
8
+ !kg/summarization/
9
+ !kg/summarization/.gitkeep
10
+ !kg/text2sql/
11
+ !kg/text2sql/.gitkeep
12
+
13
+ # Python
14
+ __pycache__/
15
+ *.pyc
16
+ *.pyo
17
+ *.pyd
18
+ *.egg-info/
19
+ dist/
20
+ build/
21
+
22
+ # Environment
23
+ .env
24
+ *.env
25
+ .venv/
26
+ env/
27
+
28
+ # Jupyter
29
+ .ipynb_checkpoints/
30
+
31
+ # IDE
32
+ .vscode/
33
+ .idea/
34
+
35
+ # OS
36
+ .DS_Store
37
+ Thumbs.db
38
+
39
+ # Large datasets (downloaded at runtime)
40
+ data/multihop/narrativeqa*
41
+
42
+ # Run logs and results (regenerable)
43
+ run_all_embeddings.log
44
+ results/
45
+ outputs/
46
+ outputs_ablation/
47
+ experiments/results/
README.md ADDED
@@ -0,0 +1,345 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # QAFD-RAG
2
+
3
+ Official code for ICLR 2026 paper: **[Query-Aware Flow Diffusion for Graph-Based RAG with Retrieval Guarantees](https://openreview.net/forum?id=n28wnc2QTc)**
4
+
5
+ QAFD-RAG uses **query-aware flow diffusion** to retrieve contextually relevant subgraphs from a knowledge graph. Unlike community-based (GraphRAG) or entity-centric (LightRAG) approaches, QAFD-RAG dynamically re-weights edges based on query relevance and propagates flow through the graph to discover multi-hop context with retrieval guarantees.
6
+
7
+ <p align="center"><em>Query: "Introduce Steve Jobs's products in Apple."</em></p>
8
+
9
+ <table align="center"><tr>
10
+ <td align="center"><img src="docs/figs/GraphRAG.png" width="260" alt="GraphRAG"/><br/><em>GraphRAG</em></td>
11
+ <td align="center"><img src="docs/figs/LightRAG.png" width="260" alt="LightRAG"/><br/><em>LightRAG</em></td>
12
+ <td align="center"><img src="docs/figs/QAFD-RAG.png" width="260" alt="QAFD-RAG"/><br/><em>QAFD-RAG</em></td>
13
+ </tr></table>
14
+ <p align="center"><em>QAFD-RAG reweights edges by query relevance, suppressing irrelevant neighborhoods (e.g., Amazon, Apple fruit).</em></p>
15
+
16
+ ## Quick Start
17
+
18
+ ```bash
19
+ # 1. Create conda environment and install dependencies
20
+ conda create -n qafd-rag python=3.10 -y
21
+ conda activate qafd-rag
22
+ pip install -r requirements.txt
23
+
24
+ # 2. Set your OpenAI API key (required for LLM and answer generation)
25
+ export OPENAI_API_KEY="sk-..."
26
+
27
+ # 3. Download pre-built KGs (recommended — saves hours of build time + API costs)
28
+ huggingface-cli download qafd/kg --repo-type dataset --include "multihop/*" --local-dir ./kg
29
+ huggingface-cli download qafd/kg --repo-type dataset --include "ultradomain/*" --local-dir ./kg
30
+
31
+ # 4. Run a benchmark
32
+ # Note: First run for multihop will auto-download the nvidia/NV-Embed-v2
33
+ # embedding model (~8GB, one-time, requires GPU with 16GB+ VRAM).
34
+ # No GPU? See "Rebuild with openai-small" below.
35
+ python benchmarks/run.py --task multihop --dataset musique --questions 10
36
+ python benchmarks/run.py --task ultradomain --dataset mix --questions 10
37
+ ```
38
+
39
+ > **Embeddings:** For best multihop results, download the pre-built KGs which use `nvidia-nv-embed-v2` (requires GPU with 16GB+ VRAM, auto-downloaded on first run). If you don't have a GPU, you can rebuild the KGs from scratch with `openai-small` instead: `python benchmarks/run.py --task multihop --dataset musique --force_build --embedding openai-small`.
40
+ >
41
+ > **Graph Types:** Multihop defaults to **passage-entity** graph (entities + passages + facts as nodes). UltraDomain/text2sql/summarization default to **entity** graph (classic KG). Override with `--graph_type`.
42
+
43
+ ## Two Graph Types
44
+
45
+ QAFD-RAG supports two knowledge graph representations, both using the same Query-Aware Flow Diffusion algorithm:
46
+
47
+ | | **Entity Graph** | **Passage-Entity Graph** |
48
+ |---|---|---|
49
+ | **Nodes** | Entities + relationships | Entities + passages + facts ([Gutiérrez et al., 2024](https://arxiv.org/abs/2405.14831)) |
50
+ | **Extraction** | LLM entity/relationship extraction | OpenIE (NER + triple extraction) |
51
+ | **Edges** | Entity-entity relationships | Fact edges + passage edges + synonymy edges |
52
+ | **QAFD traversal** | Flow reaches entities, passages looked up after | Flow reaches passages directly as graph nodes |
53
+ | **Best for** | General QA, text2sql, summarization | Multi-hop reasoning |
54
+ | **Default tasks** | ultradomain, text2sql, summarization | multihop |
55
+
56
+ Both graph types use query-aware flow diffusion with the same parameters (alpha=2.0, epsilon=0.01, step_size=0.2).
57
+
58
+ ## Project Structure
59
+
60
+ ```
61
+ QAFD-RAG/
62
+ ├── benchmarks/
63
+ │ ├── run.py # Unified benchmark runner (--graph_type, --task)
64
+ │ ├── multihop/ # Multi-hop reasoning
65
+ │ ├── ultradomain/ # General domain QA
66
+ │ ├── text2sql/ # Natural language to SQL
67
+ │ └── summarization/ # Document summarization
68
+ ├── src/
69
+ │ ├── QAFD_RAG.py # Entity graph pipeline (main class)
70
+ │ ├── base.py # Storage interfaces, QueryParam
71
+ │ ├── llm.py # LLM and embedding functions
72
+ │ ├── storage.py # KV, vector, and graph storage
73
+ │ ├── operate.py # Query operations
74
+ │ ├── evaluation.py # Answer evaluation
75
+ │ ├── answering/ # Query processing (entity graph)
76
+ │ │ ├── handler.py # kg_query entry point
77
+ │ │ ├── context.py # Context building (local/global/hybrid)
78
+ │ │ ├── clusters.py # Flow diffusion clustering
79
+ │ │ └── text_units.py # Text chunk retrieval
80
+ │ ├── retrievers/
81
+ │ │ ├── flow_diffusion.py # QAFD on NetworkX (entity graph)
82
+ │ │ └── base.py # Retriever interface
83
+ │ ├── hipporag_pipeline/ # Passage-entity graph pipeline
84
+ │ │ ├── kg_builder.py # OpenIE → igraph (entities + passages + facts)
85
+ │ │ ├── graph_adapter.py # QAFD on igraph (passage-entity graph)
86
+ │ │ ├── retriever.py # Fact reranking → seed selection → QAFD
87
+ │ │ ├── config.py # Pipeline configuration
88
+ │ │ ├── openie.py # NER + triple extraction
89
+ │ │ ├── embedding_store.py # Parquet-backed vector store
90
+ │ │ ├── reranker.py # LLM-based fact reranker
91
+ │ │ ├── prompts.py # Prompt templates
92
+ │ │ ├── benchmark_runner.py # Standalone runner
93
+ │ │ └── utils.py # Helpers
94
+ │ ├── indexing/ # KG construction (entity graph)
95
+ │ │ ├── chunkers.py # Token-based chunking
96
+ │ │ ├── extractors.py # Entity/relationship extraction
97
+ │ │ ├── schema_builder.py # Database schema KG builder
98
+ │ │ └── build_kg.py # CLI KG builder
99
+ │ ├── embedding_models/ # Embedding model implementations
100
+ │ ├── prompts/ # LLM prompt templates
101
+ │ ├── text2sql/ # Text-to-SQL support
102
+ │ └── utils/ # Helpers
103
+ ├── data/ # Datasets (auto-downloaded from HuggingFace)
104
+ ├── kg/ # Knowledge graphs (auto-generated or downloaded)
105
+ ├── docs/figs/ # Figures for README
106
+ ├── ICLR2026/ # Paper source (LaTeX)
107
+ └── results/ # Benchmark results (JSON)
108
+ ```
109
+
110
+ ## Benchmarks
111
+
112
+ | Task | Datasets | Metrics |
113
+ |------|----------|---------|
114
+ | **UltraDomain** | agriculture, biology, cs, finance, legal, math, medicine, mix, music, philosophy, physics | Quality scores (comprehensiveness, diversity, relevance, logicality, coherence) |
115
+ | **Multi-hop QA** | MuSiQue, HotpotQA, 2WikiMultiHopQA | F1, Exact Match |
116
+ | **Text-to-SQL** | Spider2-lite, Bird | Schema retrieval precision/recall |
117
+ | **Summarization** | SQuALITY | BLEU, ROUGE, METEOR, quality scores |
118
+
119
+ ## Usage
120
+
121
+ ### Unified Benchmark Runner
122
+
123
+ ```bash
124
+ python benchmarks/run.py --task <task> --dataset <dataset> [options]
125
+ ```
126
+
127
+ The runner automatically selects the appropriate graph type (passage-entity for multihop, entity for others). Override with `--graph_type`:
128
+
129
+ ```bash
130
+ # Multihop with passage-entity graph (default)
131
+ python benchmarks/run.py --task multihop --dataset musique --questions 100
132
+
133
+ # Multihop with entity graph (override)
134
+ python benchmarks/run.py --task multihop --dataset musique --graph_type entity
135
+
136
+ # Ultradomain
137
+ python benchmarks/run.py --task ultradomain --dataset mix --questions 10
138
+
139
+ # Retrieval only (skip QA)
140
+ python benchmarks/run.py --task multihop --dataset musique --skip_qa
141
+
142
+ # Build KG only
143
+ python benchmarks/run.py --task multihop --dataset musique --build_only
144
+ ```
145
+
146
+ ### Legacy CLI (entity graph only)
147
+
148
+ All benchmarks can also be run through `./run.sh` (always uses entity graph):
149
+
150
+ ```bash
151
+ ./run.sh <task> [options]
152
+ ```
153
+
154
+ ### Knowledge Graphs
155
+
156
+ Pre-built KGs are available at [huggingface.co/datasets/qafd/kg](https://huggingface.co/datasets/qafd/kg). Downloading is **recommended** to avoid hours of build time and API costs.
157
+
158
+ | Benchmark | Embedding | GPU needed? | Notes |
159
+ |-----------|-----------|-------------|-------|
160
+ | **multihop** | `nvidia-nv-embed-v2` | Yes (16GB+ VRAM) | Auto-downloaded from HuggingFace on first run |
161
+ | **ultradomain** | `openai-small` | No | Uses OpenAI API (`OPENAI_API_KEY`) |
162
+ | **text2sql** | `openai-small` | No | Uses OpenAI API (`OPENAI_API_KEY`) |
163
+
164
+ ```bash
165
+ # Download pre-built KGs (recommended)
166
+ huggingface-cli download qafd/kg --repo-type dataset --include "multihop/*" --local-dir ./kg
167
+ huggingface-cli download qafd/kg --repo-type dataset --include "ultradomain/*" --local-dir ./kg
168
+
169
+ # Or download everything
170
+ huggingface-cli download qafd/kg --repo-type dataset --local-dir ./kg
171
+ ```
172
+
173
+ To rebuild KGs with a different embedding (e.g., no GPU available):
174
+
175
+ ```bash
176
+ # Rebuild multihop with openai-small (no GPU needed, uses OpenAI API)
177
+ python benchmarks/run.py --task multihop --dataset musique --force_build --embedding openai-small
178
+ ```
179
+
180
+ ### Run Benchmarks
181
+
182
+ ```bash
183
+ # Multi-hop QA (passage-entity graph by default)
184
+ python benchmarks/run.py --task multihop --dataset musique --questions 100
185
+ python benchmarks/run.py --task multihop --dataset hotpotqa --questions 100
186
+ python benchmarks/run.py --task multihop --dataset 2wikimultihopqa --questions 100
187
+
188
+ # UltraDomain (entity graph by default)
189
+ python benchmarks/run.py --task ultradomain --dataset mix --questions 10
190
+
191
+ # UltraDomain with passage-entity graph
192
+ python benchmarks/run.py --task ultradomain --dataset mix --graph_type passage-entity --questions 10
193
+
194
+ # Text-to-SQL (entity graph)
195
+ python benchmarks/run.py --task text2sql --dataset spider2-lite
196
+
197
+ # Summarization (entity graph)
198
+ python benchmarks/run.py --task summarization --dataset squality
199
+ ```
200
+
201
+ ### Common Options
202
+
203
+ | Option | Description |
204
+ |--------|-------------|
205
+ | `--task TASK` | `multihop`, `ultradomain`, `text2sql`, `summarization` |
206
+ | `--dataset NAME` | Dataset name (e.g., `musique`, `mix`, `spider2-lite`) |
207
+ | `--graph_type TYPE` | `passage-entity` or `entity` (auto-selected by task) |
208
+ | `--questions N` | Number of questions to evaluate |
209
+ | `--build_only` | Build KG only, skip benchmark |
210
+ | `--force_build` | Rebuild KG even if it exists |
211
+ | `--skip_qa` | Run retrieval only, skip QA (passage-entity) |
212
+ | `--embedding MODEL` | Embedding model (see table below) |
213
+ | `--llm MODEL` | LLM model (`gpt-4o-mini`, `gpt-4o`, `gpt-5-nano`, `gpt-5-mini`, `gpt-5`, `gpt-oss-120b`) |
214
+ | `--alpha FLOAT` | QAFD alpha parameter (default: 2.0) |
215
+ | `--epsilon FLOAT` | QAFD convergence threshold (default: 0.01) |
216
+ | `--weight_scheme` | Query-aware edge weighting: `original`, `multiply`, `add` |
217
+
218
+ ## Data
219
+
220
+ Datasets are automatically downloaded from HuggingFace on first run:
221
+
222
+ | Task | Source | Reference |
223
+ |------|--------|-----------|
224
+ | **UltraDomain** | [TommyChien/UltraDomain](https://huggingface.co/datasets/TommyChien/UltraDomain) | agriculture, biology, cs, finance, legal, math, medicine, mix, music, philosophy, physics |
225
+ | **Multi-hop QA** | [osunlp](https://huggingface.co/osunlp) | MuSiQue, HotpotQA, 2WikiMultiHopQA |
226
+ | **Summarization** | [pszemraj/SQuALITY-v1.3](https://huggingface.co/datasets/pszemraj/SQuALITY-v1.3) | SQuALITY |
227
+ | **Text-to-SQL** | Included (`data/text2sql/`) | Spider2-lite (Pagila) + Bird (superhero) with auto-generated DB summaries |
228
+
229
+ ## Configuration
230
+
231
+ ### Embedding Models
232
+
233
+ | Key | Model | Dimensions |
234
+ |-----|-------|-----------|
235
+ | `openai-small` | text-embedding-3-small | 1536 |
236
+ | `openai-large` | text-embedding-3-large | 3072 |
237
+ | `jina-v3` | jinaai/jina-embeddings-v3 | 1024 |
238
+ | `gritlm` | GritLM/GritLM-7B | 4096 |
239
+ | `nvidia-nv-embed-v2` | nvidia/NV-Embed-v2 | 4096 |
240
+
241
+ The default embedding is **`openai-small`** which requires `OPENAI_API_KEY`.
242
+
243
+ Use `--embedding <key>` to select a model:
244
+
245
+ ```bash
246
+ ./run.sh ultradomain --questions 10 --embedding jina-v3
247
+ ```
248
+
249
+ ### Local Embedding Models
250
+
251
+ Local embeddings (`jina-v3`, `gritlm`, `nvidia-nv-embed-v2`) run on your GPU and do not require an API key for embeddings. They are downloaded automatically from HuggingFace on first use.
252
+
253
+ **Requirements:**
254
+ - CUDA-capable GPU with sufficient VRAM (8GB+ recommended)
255
+ - Models are cached in `~/.cache/huggingface/`
256
+
257
+ **Usage with local embeddings (no OpenAI API needed for embeddings):**
258
+
259
+ ```bash
260
+ # Use Jina v3 (1024-dim, lightweight, good quality)
261
+ ./run.sh ultradomain --questions 10 --embedding jina-v3
262
+
263
+ # Use GritLM (4096-dim, unified embedding+generation)
264
+ ./run.sh multihop --dataset musique --questions 10 --embedding gritlm
265
+
266
+ # Use NVIDIA NV-Embed-v2 (4096-dim, 32K context, high quality)
267
+ ./run.sh ultradomain --questions 10 --embedding nvidia-nv-embed-v2
268
+ ```
269
+
270
+ > **Note:** Even with local embeddings, an LLM API key is still required for entity extraction (KG building) and answer generation. Set `OPENAI_API_KEY` or use `--llm gpt-oss-120b` for a free open-source LLM.
271
+
272
+ ### LLM Models
273
+
274
+ `gpt-4o-mini` (default), `gpt-4o`, `gpt-5-nano`, `gpt-5-mini`, `gpt-5`, `gpt-oss-120b` (free, local)
275
+
276
+ ## Python API
277
+
278
+ ```python
279
+ from src import QAFD_RAG, QueryParam
280
+
281
+ # Initialize
282
+ rag = QAFD_RAG(
283
+ working_dir="./my_kg",
284
+ llm_model_name="gpt-4o-mini",
285
+ embedding_model_key="jina-v3",
286
+ )
287
+
288
+ # Index documents
289
+ rag.insert(["Document text 1...", "Document text 2..."])
290
+
291
+ # Query
292
+ answer = rag.query(
293
+ "What is X?",
294
+ param=QueryParam(mode="hybrid"),
295
+ )
296
+ print(answer)
297
+ ```
298
+
299
+ ### QueryParam Options
300
+
301
+ | Parameter | Default | Description |
302
+ |-----------|---------|-------------|
303
+ | `mode` | `"hybrid"` | Retrieval mode: `local`, `global`, `hybrid` |
304
+ | `top_k` | 40 | Number of entities to retrieve |
305
+ | `max_source_nodes` | 40 | Max source nodes for flow diffusion |
306
+ | `min_flow_threshold` | 0.1 | Minimum flow value to include a node |
307
+ | `enable_query_aware_flow_diffusion` | `True` | Use query-aware edge weighting |
308
+ | `alpha` | 50.0 | Flow diffusion initialization factor |
309
+ | `response_type` | `"Multiple Paragraphs"` | Output format |
310
+
311
+ ## How It Works
312
+
313
+ QAFD-RAG processes documents through a two-stage pipeline:
314
+
315
+ **Stage 1 -- Knowledge Graph Construction**
316
+
317
+ *Entity Graph:*
318
+ 1. Documents are split into token-based chunks.
319
+ 2. An LLM extracts entities and relationships from each chunk.
320
+ 3. Entities become nodes, relationships become weighted edges.
321
+
322
+ *Passage-Entity Graph:*
323
+ 1. Documents are split into passages (chunks).
324
+ 2. OpenIE extracts named entities and (subject, predicate, object) triples.
325
+ 3. Entities and passages are both graph nodes, connected by fact edges, passage-entity edges, and synonymy edges (entity pairs with cosine similarity > 0.8).
326
+
327
+ **Stage 2 -- Query-Aware Flow Diffusion**
328
+
329
+ 1. **Seed Selection**: Query is matched to entities/facts via embedding similarity. An LLM reranker filters the most relevant facts.
330
+ 2. **Flow Diffusion**: Mass is injected at seed nodes and propagated through the graph via push-relabel. Edge weights are dynamically adjusted based on each node's similarity to the query.
331
+ 3. **Passage Ranking**: Nodes accumulate importance scores proportional to flow. In the passage-entity graph, passages are ranked directly. In the entity graph, associated text chunks are retrieved from top-ranked entities.
332
+ 4. **LLM Answering**: Top passages are assembled into context and passed to an LLM for answer generation.
333
+
334
+ The flow diffusion algorithm is the key differentiator: rather than simple graph traversal or vector search alone, it combines graph structure with query relevance to find contextually important information that may be several hops away from the initial match.
335
+
336
+ ## Citation
337
+
338
+ ```bibtex
339
+ @inproceedings{zhou2026qafd,
340
+ title={Query-Aware Flow Diffusion for Graph-Based RAG with Retrieval Guarantees},
341
+ author={Zhou, Zhuoping and Ataee Tarzanagh, Davoud and Didari, Sima and Hu, Wenjun and Gutow, Baruch and Verkholyak, Oxana and Faraki, Masoud and Hao, Heng and Moon, Hankyu and Min, Seungjai},
342
+ booktitle={International Conference on Learning Representations (ICLR)},
343
+ year={2026}
344
+ }
345
+ ```
benchmarks/multihop/benchmark_multihop.py ADDED
@@ -0,0 +1,531 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """
3
+ Unified Multi-hop QA Benchmark for QAFD-RAG
4
+ Supports: MuSiQue, HotpotQA, 2WikiMultiHopQA
5
+ """
6
+
7
+ import os
8
+ import sys
9
+ import asyncio
10
+ import time
11
+ import json
12
+ import re
13
+ import string
14
+ import logging
15
+ from typing import List, Dict
16
+ from dataclasses import dataclass, asdict
17
+ from datetime import datetime
18
+ from pathlib import Path
19
+ from collections import Counter
20
+ import numpy as np
21
+
22
+ # Suppress verbose logging
23
+ logging.getLogger("httpx").setLevel(logging.ERROR)
24
+ logging.getLogger("QAFD_RAG").setLevel(logging.ERROR)
25
+ logging.getLogger("nano-vectordb").setLevel(logging.ERROR)
26
+ logging.getLogger("OpenAI").setLevel(logging.ERROR)
27
+
28
+ # Add QAFD-RAG to path
29
+ QAFD_RAG_HOME = str(Path(__file__).parent.parent.parent)
30
+ sys.path.insert(0, QAFD_RAG_HOME)
31
+
32
+ import nest_asyncio
33
+ nest_asyncio.apply()
34
+
35
+ # Dataset configurations
36
+ DATASETS = {
37
+ "musique": {
38
+ "name": "MuSiQue",
39
+ "data_file": "musique.json",
40
+ "corpus_file": "musique_corpus.json",
41
+ "kg_dir": "musique",
42
+ },
43
+ "hotpotqa": {
44
+ "name": "HotpotQA",
45
+ "data_file": "hotpotqa.json",
46
+ "corpus_file": "hotpotqa_corpus.json",
47
+ "kg_dir": "hotpotqa",
48
+ },
49
+ "2wikimultihopqa": {
50
+ "name": "2WikiMultiHopQA",
51
+ "data_file": "2wikimultihopqa.json",
52
+ "corpus_file": "2wikimultihopqa_corpus.json",
53
+ "kg_dir": "2wikimultihopqa",
54
+ },
55
+ }
56
+
57
+
58
+ def print_header(title: str, width: int = 70):
59
+ """Print a formatted header"""
60
+ print(f"\n{'=' * width}")
61
+ print(f" {title}")
62
+ print(f"{'=' * width}")
63
+
64
+
65
+ def print_config(items: Dict[str, str], width: int = 70):
66
+ """Print configuration items"""
67
+ print(f"{'─' * width}")
68
+ for key, value in items.items():
69
+ print(f" {key:<20} {value}")
70
+ print(f"{'─' * width}")
71
+
72
+
73
+ def print_progress(current: int, total: int, prefix: str = "", width: int = 40):
74
+ """Print a progress bar"""
75
+ percent = current / total
76
+ filled = int(width * percent)
77
+ bar = '█' * filled + '░' * (width - filled)
78
+ print(f"\r {prefix} [{bar}] {current}/{total} ({percent*100:.1f}%)", end='', flush=True)
79
+ if current == total:
80
+ print()
81
+
82
+
83
+ def print_metric(name: str, value: float, std: float = None, width: int = 25):
84
+ """Print a metric with optional std"""
85
+ if std is not None:
86
+ print(f" {name:<{width}} {value:.4f} ± {std:.4f}")
87
+ else:
88
+ print(f" {name:<{width}} {value:.4f}")
89
+
90
+
91
+ @dataclass
92
+ class BenchmarkResult:
93
+ """Benchmark result for multi-hop QA"""
94
+ model_name: str
95
+ dataset_name: str
96
+ total_questions: int
97
+ success_count: int
98
+ total_time: float
99
+ kg_build_time: float
100
+ query_time: float
101
+ avg_time_per_question: float
102
+ f1_score_mean: float
103
+ f1_score_std: float
104
+ exact_match_mean: float
105
+ exact_match_std: float
106
+ f1_scores: List[float] = None
107
+ exact_match_scores: List[float] = None
108
+ responses: List[str] = None
109
+ questions: List[str] = None
110
+ gold_answers: List[List[str]] = None
111
+ error_message: str = ""
112
+
113
+
114
+ def normalize_answer(answer: str) -> str:
115
+ """Normalize answer for comparison"""
116
+ def remove_articles(text):
117
+ return re.sub(r"\b(a|an|the)\b", " ", text)
118
+ def white_space_fix(text):
119
+ return " ".join(text.split())
120
+ def remove_punc(text):
121
+ exclude = set(string.punctuation)
122
+ return "".join(ch for ch in text if ch not in exclude)
123
+ def lower(text):
124
+ return text.lower()
125
+ return white_space_fix(remove_articles(remove_punc(lower(answer))))
126
+
127
+
128
+ def compute_f1(gold: str, predicted: str) -> float:
129
+ """Compute F1 score between gold and predicted answers"""
130
+ gold_tokens = normalize_answer(gold).split()
131
+ predicted_tokens = normalize_answer(predicted).split()
132
+ common = Counter(predicted_tokens) & Counter(gold_tokens)
133
+ num_same = sum(common.values())
134
+ if num_same == 0:
135
+ return 0.0
136
+ precision = 1.0 * num_same / len(predicted_tokens) if predicted_tokens else 0.0
137
+ recall = 1.0 * num_same / len(gold_tokens) if gold_tokens else 0.0
138
+ if precision + recall == 0:
139
+ return 0.0
140
+ return 2 * (precision * recall) / (precision + recall)
141
+
142
+
143
+ def compute_exact_match(gold: str, predicted: str) -> float:
144
+ """Compute exact match score"""
145
+ return 1.0 if normalize_answer(gold) == normalize_answer(predicted) else 0.0
146
+
147
+
148
+ def get_gold_answers(samples):
149
+ """Extract gold answers from samples"""
150
+ gold_answers = []
151
+ for sample in samples:
152
+ if 'answer' in sample:
153
+ gold_ans = sample['answer']
154
+ elif 'reference' in sample:
155
+ gold_ans = sample['reference']
156
+ else:
157
+ gold_ans = "Unknown"
158
+ if isinstance(gold_ans, str):
159
+ gold_ans = [gold_ans]
160
+ elif not isinstance(gold_ans, list):
161
+ gold_ans = [str(gold_ans)]
162
+ gold_answers.append(gold_ans)
163
+ return gold_answers
164
+
165
+
166
+ class MultiHopBenchmark:
167
+ def __init__(self, dataset: str, api_key: str, embedding_model: str = "openai-small",
168
+ llm_model: str = "gpt-4o-mini"):
169
+ if dataset not in DATASETS:
170
+ raise ValueError(f"Unknown dataset: {dataset}. Choose from: {list(DATASETS.keys())}")
171
+
172
+ self.dataset = dataset
173
+ self.config = DATASETS[dataset]
174
+ self.api_key = api_key
175
+ self.embedding_model = embedding_model
176
+ self.llm_model = llm_model
177
+ os.environ["OPENAI_API_KEY"] = api_key
178
+ os.environ["OPENAI_API_BASE"] = "https://api.openai.com/v1"
179
+
180
+ def _get_working_dir(self) -> str:
181
+ return os.path.join(QAFD_RAG_HOME, "kg", "multihop", f"{self.llm_model}_{self.embedding_model}_{self.config['kg_dir']}")
182
+
183
+ def _kg_exists(self, working_dir: str) -> bool:
184
+ kg_files = [
185
+ os.path.join(working_dir, "vdb_entities.json"),
186
+ os.path.join(working_dir, "vdb_chunks.json"),
187
+ os.path.join(working_dir, "kv_store_full_docs.json"),
188
+ ]
189
+ return all(os.path.exists(f) for f in kg_files)
190
+
191
+ def _get_llm_func(self):
192
+ from src import llm
193
+ llm_funcs = {
194
+ "gpt-4o-mini": llm.gpt_4o_mini_complete,
195
+ "gpt-4o": llm.gpt_4o_complete,
196
+ "gpt-oss-120b": llm.gpt_oss_120b_complete,
197
+ "gpt-5": llm.gpt_5_complete,
198
+ "gpt-5-mini": llm.gpt_5_mini_complete,
199
+ "gpt-5-nano": llm.gpt_5_nano_complete,
200
+ }
201
+ return llm_funcs.get(self.llm_model, llm.gpt_4o_mini_complete)
202
+
203
+ def _ensure_data_file(self, filename: str) -> str:
204
+ """Return path to data file, downloading from HuggingFace if missing."""
205
+ data_dir = os.path.join(QAFD_RAG_HOME, "data", "multihop")
206
+ filepath = os.path.join(data_dir, filename)
207
+ if not os.path.exists(filepath):
208
+ print(f" Downloading {filename} from HuggingFace...", end=" ", flush=True)
209
+ from huggingface_hub import hf_hub_download
210
+ os.makedirs(data_dir, exist_ok=True)
211
+ hf_hub_download(
212
+ repo_id="osunlp/HippoRAG_2",
213
+ filename=filename,
214
+ repo_type="dataset",
215
+ local_dir=data_dir,
216
+ )
217
+ print("done")
218
+ return filepath
219
+
220
+ def _load_dataset(self) -> List[Dict]:
221
+ dataset_path = self._ensure_data_file(self.config["data_file"])
222
+ with open(dataset_path, 'r', encoding='utf-8') as f:
223
+ samples = json.load(f)
224
+ return samples
225
+
226
+ def _load_corpus(self) -> List[str]:
227
+ corpus_path = self._ensure_data_file(self.config["corpus_file"])
228
+ with open(corpus_path, 'r', encoding='utf-8') as f:
229
+ corpus = json.load(f)
230
+ docs = [f"{doc['title']}\n{doc['text']}" for doc in corpus]
231
+ return docs
232
+
233
+ async def build_kg(self, max_documents: int = None) -> bool:
234
+ """Build KG only (no benchmark)"""
235
+ from src.QAFD_RAG import QAFD_RAG
236
+
237
+ print_header(f"QAFD-RAG Knowledge Graph Builder")
238
+ print_config({
239
+ "Graph Type": "entity",
240
+ "Dataset": self.config['name'],
241
+ "Embedding": self.embedding_model,
242
+ "LLM": self.llm_model,
243
+ "Working Dir": self._get_working_dir()
244
+ })
245
+
246
+ working_dir = self._get_working_dir()
247
+ os.makedirs(working_dir, exist_ok=True)
248
+
249
+ llm_func = self._get_llm_func()
250
+ rag = QAFD_RAG(
251
+ working_dir=working_dir,
252
+ llm_model_func=llm_func,
253
+ llm_model_name=self.llm_model,
254
+ embedding_model_key=self.embedding_model,
255
+ enable_llm_cache=True,
256
+ )
257
+
258
+ print("\n Loading corpus...", end=" ", flush=True)
259
+ docs = self._load_corpus()
260
+ docs_to_process = min(max_documents, len(docs)) if max_documents else len(docs)
261
+ print(f"done ({len(docs)} documents available)")
262
+
263
+ print(f"\n Building KG from {docs_to_process} documents...")
264
+ start_time = time.time()
265
+
266
+ for i, doc in enumerate(docs[:docs_to_process]):
267
+ print_progress(i + 1, docs_to_process, "Progress")
268
+ await rag.ainsert(doc)
269
+
270
+ build_time = time.time() - start_time
271
+
272
+ print_header("Build Complete")
273
+ print(f" Documents processed: {docs_to_process}")
274
+ print(f" Time elapsed: {build_time:.2f}s")
275
+ print(f" Avg per document: {build_time/docs_to_process:.2f}s")
276
+ print(f" Output directory: {working_dir}")
277
+ print()
278
+
279
+ return True
280
+
281
+ async def run_benchmark(self, question_count: int = 100, force_build: bool = False,
282
+ max_documents: int = None, mode: str = "hybrid",
283
+ max_source_nodes: int = 20, min_flow_threshold: float = 0.1,
284
+ alpha: float = 2.0) -> BenchmarkResult:
285
+ """Run benchmark"""
286
+ from src.QAFD_RAG import QAFD_RAG, QueryParam
287
+
288
+ print_header(f"QAFD-RAG Multi-hop QA Benchmark")
289
+ print_config({
290
+ "Graph Type": "entity",
291
+ "Dataset": self.config['name'],
292
+ "Questions": str(question_count),
293
+ "Embedding": self.embedding_model,
294
+ "LLM": self.llm_model,
295
+ "Mode": mode,
296
+ "Max Nodes": str(max_source_nodes),
297
+ "Threshold": str(min_flow_threshold),
298
+ "Alpha": str(alpha)
299
+ })
300
+
301
+ working_dir = self._get_working_dir()
302
+ os.makedirs(working_dir, exist_ok=True)
303
+
304
+ llm_func = self._get_llm_func()
305
+ rag = QAFD_RAG(
306
+ working_dir=working_dir,
307
+ llm_model_func=llm_func,
308
+ llm_model_name=self.llm_model,
309
+ embedding_model_key=self.embedding_model,
310
+ enable_llm_cache=True,
311
+ )
312
+
313
+ # Check if KG exists or needs to be built
314
+ kg_build_time = 0.0
315
+ if self._kg_exists(working_dir) and not force_build:
316
+ print(f"\n Using existing KG: {working_dir}")
317
+ else:
318
+ print("\n Loading corpus...", end=" ", flush=True)
319
+ docs = self._load_corpus()
320
+ docs_to_process = min(max_documents, len(docs)) if max_documents else len(docs)
321
+ print(f"done ({docs_to_process} documents)")
322
+
323
+ print(f" Building KG...")
324
+ start_time = time.time()
325
+ for i, doc in enumerate(docs[:docs_to_process]):
326
+ print_progress(i + 1, docs_to_process, "Progress")
327
+ await rag.ainsert(doc)
328
+ kg_build_time = time.time() - start_time
329
+ print(f" KG built in {kg_build_time:.2f}s")
330
+
331
+ # Load dataset
332
+ print("\n Loading dataset...", end=" ", flush=True)
333
+ samples = self._load_dataset()
334
+ samples = samples[:question_count]
335
+ questions = [s['question'] for s in samples]
336
+ gold_answers = get_gold_answers(samples)
337
+ print(f"done ({len(questions)} questions)")
338
+
339
+ # Run queries
340
+ print(f"\n Running queries...")
341
+ start_time = time.time()
342
+ responses = []
343
+ success_count = 0
344
+
345
+ for i, question in enumerate(questions):
346
+ try:
347
+ print_progress(i + 1, len(questions), "Progress")
348
+ query_param = QueryParam(
349
+ mode=mode,
350
+ max_source_nodes=max_source_nodes,
351
+ min_flow_threshold=min_flow_threshold,
352
+ alpha=alpha,
353
+ response_type="Brief, accurate answer (maximum 14 words)."
354
+ )
355
+ response = await rag.aquery(question, query_param)
356
+ if response and len(response.split()) > 14:
357
+ response = " ".join(response.split()[:14])
358
+ responses.append(response)
359
+ success_count += 1
360
+ except Exception as e:
361
+ responses.append("")
362
+
363
+ query_time = time.time() - start_time
364
+
365
+ # Calculate metrics
366
+ f1_scores = []
367
+ em_scores = []
368
+ for gold_list, predicted in zip(gold_answers, responses):
369
+ if not predicted:
370
+ f1_scores.append(0.0)
371
+ em_scores.append(0.0)
372
+ continue
373
+ f1_scores.append(max(compute_f1(g, predicted) for g in gold_list))
374
+ em_scores.append(max(compute_exact_match(g, predicted) for g in gold_list))
375
+
376
+ result = BenchmarkResult(
377
+ model_name="QAFD_RAG",
378
+ dataset_name=self.dataset,
379
+ total_questions=len(questions),
380
+ success_count=success_count,
381
+ total_time=kg_build_time + query_time,
382
+ kg_build_time=kg_build_time,
383
+ query_time=query_time,
384
+ avg_time_per_question=query_time / len(questions) if questions else 0,
385
+ f1_score_mean=float(np.mean(f1_scores)) if f1_scores else 0,
386
+ f1_score_std=float(np.std(f1_scores)) if f1_scores else 0,
387
+ exact_match_mean=float(np.mean(em_scores)) if em_scores else 0,
388
+ exact_match_std=float(np.std(em_scores)) if em_scores else 0,
389
+ f1_scores=f1_scores,
390
+ exact_match_scores=em_scores,
391
+ responses=responses,
392
+ questions=questions,
393
+ gold_answers=gold_answers,
394
+ )
395
+
396
+ self.print_results(result)
397
+ return result
398
+
399
+ def print_results(self, result: BenchmarkResult):
400
+ """Print benchmark results"""
401
+ print_header(f"Results: {self.config['name']}")
402
+
403
+ print("\n PERFORMANCE")
404
+ print(f" {'─' * 40}")
405
+ print(f" {'Questions':<25} {result.total_questions}")
406
+ print(f" {'Successful':<25} {result.success_count}/{result.total_questions}")
407
+ print(f" {'KG Build Time':<25} {result.kg_build_time:.2f}s")
408
+ print(f" {'Query Time':<25} {result.query_time:.2f}s")
409
+ print(f" {'Avg per Question':<25} {result.avg_time_per_question:.2f}s")
410
+
411
+ print("\n ACCURACY METRICS")
412
+ print(f" {'─' * 40}")
413
+ print_metric("F1 Score", result.f1_score_mean, result.f1_score_std)
414
+ print_metric("Exact Match", result.exact_match_mean, result.exact_match_std)
415
+ print()
416
+
417
+ def save_results(self, result: BenchmarkResult):
418
+ """Save results as two separate files: eval metrics and generated responses"""
419
+ timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
420
+ results_dir = os.path.join(QAFD_RAG_HOME, "results", "multihop", self.dataset)
421
+ os.makedirs(results_dir, exist_ok=True)
422
+
423
+ eval_file = os.path.join(results_dir, f"{self.dataset}_{timestamp}_eval.json")
424
+ output_file = os.path.join(results_dir, f"{self.dataset}_{timestamp}_responses.json")
425
+
426
+ # --- Eval file: metrics and timing ---
427
+ eval_data = {
428
+ "timestamp": datetime.now().isoformat(),
429
+ "graph_type": "entity",
430
+ "model": result.model_name,
431
+ "llm": self.llm_model,
432
+ "embedding": self.embedding_model,
433
+ "dataset": result.dataset_name,
434
+ "performance": {
435
+ "total_questions": result.total_questions,
436
+ "success_count": result.success_count,
437
+ "kg_build_time": result.kg_build_time,
438
+ "query_time": result.query_time,
439
+ "total_time": result.total_time,
440
+ "avg_time_per_question": result.avg_time_per_question,
441
+ },
442
+ "metrics": {
443
+ "f1_score_mean": result.f1_score_mean,
444
+ "f1_score_std": result.f1_score_std,
445
+ "exact_match_mean": result.exact_match_mean,
446
+ "exact_match_std": result.exact_match_std,
447
+ },
448
+ "per_question_f1": result.f1_scores,
449
+ "per_question_em": result.exact_match_scores,
450
+ "error": result.error_message,
451
+ }
452
+
453
+ with open(eval_file, 'w', encoding='utf-8') as f:
454
+ json.dump(eval_data, f, indent=2, ensure_ascii=False)
455
+
456
+ # --- Responses file: questions + generated answers + gold answers ---
457
+ output_entries = []
458
+ if result.responses:
459
+ for i, response in enumerate(result.responses):
460
+ entry = {
461
+ "id": i + 1,
462
+ "question": result.questions[i] if result.questions else "",
463
+ "generated_answer": response,
464
+ "gold_answers": result.gold_answers[i] if result.gold_answers else [],
465
+ }
466
+ output_entries.append(entry)
467
+
468
+ output_data = {
469
+ "timestamp": datetime.now().isoformat(),
470
+ "model": result.model_name,
471
+ "llm": self.llm_model,
472
+ "embedding": self.embedding_model,
473
+ "dataset": result.dataset_name,
474
+ "num_responses": len(output_entries),
475
+ "responses": output_entries,
476
+ }
477
+
478
+ with open(output_file, 'w', encoding='utf-8') as f:
479
+ json.dump(output_data, f, indent=2, ensure_ascii=False)
480
+
481
+ print(f" Eval saved: {eval_file}")
482
+ print(f" Responses saved: {output_file}\n")
483
+
484
+
485
+ async def main():
486
+ import argparse
487
+
488
+ parser = argparse.ArgumentParser(description="QAFD_RAG Multi-hop QA Benchmark")
489
+ parser.add_argument("--dataset", type=str, required=True,
490
+ choices=["musique", "hotpotqa", "2wikimultihopqa"])
491
+ parser.add_argument("--questions", type=int, default=100)
492
+ parser.add_argument("--max-documents", type=int, default=None)
493
+ parser.add_argument("--build", action="store_true")
494
+ parser.add_argument("--force-build", action="store_true")
495
+ parser.add_argument("--embedding", type=str, default="openai-small",
496
+ choices=["openai-small", "openai-large", "jina-v3", "gritlm", "nvidia-nv-embed-v2"])
497
+ parser.add_argument("--llm", type=str, default="gpt-4o-mini",
498
+ choices=["gpt-4o-mini", "gpt-4o", "gpt-oss-120b", "gpt-5", "gpt-5-mini", "gpt-5-nano"])
499
+ parser.add_argument("--mode", type=str, default="hybrid",
500
+ choices=["local", "global", "hybrid"])
501
+ parser.add_argument("--max-source-nodes", type=int, default=20)
502
+ parser.add_argument("--min-flow-threshold", type=float, default=0.1)
503
+ parser.add_argument("--alpha", type=float, default=2.0)
504
+
505
+ args = parser.parse_args()
506
+
507
+ api_key = os.environ.get("OPENAI_API_KEY")
508
+ if not api_key:
509
+ print("ERROR: Set OPENAI_API_KEY environment variable")
510
+ return
511
+
512
+ benchmark = MultiHopBenchmark(args.dataset, api_key, args.embedding, args.llm)
513
+
514
+ if args.build:
515
+ await benchmark.build_kg(max_documents=args.max_documents)
516
+ return
517
+
518
+ result = await benchmark.run_benchmark(
519
+ question_count=args.questions,
520
+ force_build=args.force_build,
521
+ max_documents=args.max_documents,
522
+ mode=args.mode,
523
+ max_source_nodes=args.max_source_nodes,
524
+ min_flow_threshold=args.min_flow_threshold,
525
+ alpha=args.alpha
526
+ )
527
+ benchmark.save_results(result)
528
+
529
+
530
+ if __name__ == "__main__":
531
+ asyncio.run(main())
benchmarks/run.py ADDED
@@ -0,0 +1,371 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """
3
+ Unified Benchmark Runner for QAFD-RAG
4
+ ======================================
5
+
6
+ Supports two graph types (always with Query-Aware Flow Diffusion):
7
+
8
+ - **passage-entity**: Entities + passages + facts as nodes, synonymy edges.
9
+ Flow diffusion reaches passages directly. Default for multihop.
10
+
11
+ - **entity**: Classic KG with entity + relationship nodes.
12
+ Passages are looked up after graph traversal. Default for other tasks.
13
+
14
+ Usage::
15
+
16
+ # Multihop (auto-selects passage-entity graph)
17
+ python benchmarks/run.py --task multihop --dataset musique
18
+
19
+ # Override graph type
20
+ python benchmarks/run.py --task multihop --dataset musique --graph_type entity
21
+
22
+ # Ultradomain (auto-selects entity graph)
23
+ python benchmarks/run.py --task ultradomain --dataset mix
24
+
25
+ # Text2SQL
26
+ python benchmarks/run.py --task text2sql --dataset spider2-lite
27
+
28
+ # Build KG only
29
+ python benchmarks/run.py --task multihop --dataset musique --build_only
30
+ """
31
+
32
+ import argparse
33
+ import asyncio
34
+ import os
35
+ import sys
36
+
37
+ QAFD_RAG_HOME = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
38
+ sys.path.insert(0, QAFD_RAG_HOME)
39
+
40
+ # ── Task → dataset mapping ──────────────────────────────────────────────
41
+
42
+ TASK_DATASETS = {
43
+ "multihop": ["musique", "hotpotqa", "2wikimultihopqa"],
44
+ "ultradomain": [
45
+ "agriculture", "biology", "cs", "finance", "legal",
46
+ "math", "medicine", "mix", "music", "philosophy",
47
+ "physics", "psychology",
48
+ ],
49
+ "text2sql": ["spider2-lite", "bird"],
50
+ "summarization": ["squality"],
51
+ }
52
+
53
+ # ── Default graph type per task ─────────────────────────────────────────
54
+
55
+ TASK_DEFAULT_GRAPH_TYPE = {
56
+ "multihop": "passage-entity",
57
+ "ultradomain": "entity",
58
+ "text2sql": "entity",
59
+ "summarization": "entity",
60
+ }
61
+
62
+ # ── Default QAFD parameters per graph type ──────────────────────────────
63
+ # These come from proven successful runs in each pipeline.
64
+
65
+ GRAPH_TYPE_DEFAULTS = {
66
+ "passage-entity": {
67
+ "alpha": 2.0,
68
+ "epsilon": 0.01,
69
+ "max_iterations": 500,
70
+ "step_size": 0.2,
71
+ "weight_scheme": "original",
72
+ "linking_top_k": 5, # fact seeds
73
+ "passage_node_weight": 0.05,
74
+ "retrieval_top_k": 200,
75
+ },
76
+ "entity": {
77
+ "alpha": 2.0,
78
+ "epsilon": 0.01,
79
+ "max_iterations": 500,
80
+ "step_size": 0.2,
81
+ "weight_scheme": "original",
82
+ "max_source_nodes": 20,
83
+ "min_flow_threshold": 0.1,
84
+ },
85
+ }
86
+
87
+
88
+ def run_passage_entity(args):
89
+ """Run benchmark using passage-entity graph (hipporag_pipeline)."""
90
+ # Bypass src/__init__.py (heavy AWS deps)
91
+ import types as _types
92
+ for _pkg_path in ["src", "src.retrievers", "src.hipporag_pipeline"]:
93
+ if _pkg_path not in sys.modules:
94
+ _m = _types.ModuleType(_pkg_path)
95
+ _m.__path__ = [os.path.join(QAFD_RAG_HOME, *_pkg_path.split("."))]
96
+ _m.__package__ = _pkg_path
97
+ sys.modules[_pkg_path] = _m
98
+
99
+ import importlib.util as _ilu
100
+ def _load_mod(fqn, filepath):
101
+ spec = _ilu.spec_from_file_location(fqn, filepath)
102
+ mod = _ilu.module_from_spec(spec)
103
+ sys.modules[fqn] = mod
104
+ spec.loader.exec_module(mod)
105
+ return mod
106
+
107
+ _src = os.path.join(QAFD_RAG_HOME, "src")
108
+ _load_mod("src.retrievers.base", os.path.join(_src, "retrievers", "base.py"))
109
+ _load_mod("src.retrievers.flow_diffusion", os.path.join(_src, "retrievers", "flow_diffusion.py"))
110
+
111
+ # Import after module setup
112
+ from src.hipporag_pipeline.benchmark_runner import main as hippo_main
113
+
114
+ # Build sys.argv for the sub-module
115
+ sub_argv = [
116
+ "benchmark_runner",
117
+ "--task", args.task,
118
+ "--dataset", args.dataset,
119
+ "--data_dir", os.path.join(QAFD_RAG_HOME, "data", "multihop"),
120
+ "--embedding_model", args.embedding,
121
+ "--llm_model", args.llm,
122
+ "--num_queries", str(args.questions),
123
+ "--qafd_alpha", str(args.alpha),
124
+ "--qafd_epsilon", str(args.epsilon),
125
+ "--qafd_max_iterations", str(args.max_iterations),
126
+ "--qafd_step_size", str(args.step_size),
127
+ "--qafd_weight_scheme", str(args.weight_scheme),
128
+ "--linking_top_k", str(args.linking_top_k),
129
+ "--passage_node_weight", str(args.passage_node_weight),
130
+ "--retrieval_top_k", str(args.retrieval_top_k),
131
+ ]
132
+ if args.skip_qa:
133
+ sub_argv.append("--skip_qa")
134
+ if getattr(args, 'batch_push', False):
135
+ sub_argv.append("--batch_push")
136
+ if args.force_build:
137
+ sub_argv.append("--force_index")
138
+ sub_argv.append("--force_openie")
139
+
140
+ old_argv = sys.argv
141
+ sys.argv = sub_argv
142
+ try:
143
+ hippo_main()
144
+ finally:
145
+ sys.argv = old_argv
146
+
147
+
148
+ def run_entity(args):
149
+ """Run benchmark using entity graph (original QAFD-RAG pipeline)."""
150
+ import nest_asyncio
151
+ nest_asyncio.apply()
152
+
153
+ task = args.task
154
+
155
+ if task == "multihop":
156
+ from benchmarks.multihop.benchmark_multihop import MultiHopBenchmark
157
+
158
+ api_key = os.environ.get("OPENAI_API_KEY")
159
+ if not api_key:
160
+ print("ERROR: Set OPENAI_API_KEY environment variable")
161
+ return
162
+
163
+ benchmark = MultiHopBenchmark(args.dataset, api_key, args.embedding, args.llm)
164
+
165
+ if args.build_only:
166
+ asyncio.run(benchmark.build_kg(max_documents=args.max_documents))
167
+ return
168
+
169
+ result = asyncio.run(benchmark.run_benchmark(
170
+ question_count=args.questions,
171
+ force_build=args.force_build,
172
+ max_documents=args.max_documents,
173
+ mode="hybrid",
174
+ max_source_nodes=args.max_source_nodes,
175
+ min_flow_threshold=args.min_flow_threshold,
176
+ alpha=args.alpha,
177
+ ))
178
+ benchmark.save_results(result)
179
+
180
+ elif task == "ultradomain":
181
+ # Delegate to ultradomain's own argparse
182
+ # Ultradomain expects dataset as "mix.jsonl" format
183
+ ud_dataset = args.dataset if args.dataset.endswith(".jsonl") else f"{args.dataset}.jsonl"
184
+ sub_argv = [
185
+ "benchmark_ultradomain",
186
+ "--dataset", ud_dataset,
187
+ "--questions", str(args.questions),
188
+ "--embedding", args.embedding,
189
+ "--llm", args.llm,
190
+ ]
191
+ if args.force_build:
192
+ sub_argv.append("--force-build")
193
+ if args.build_only:
194
+ sub_argv.append("--build")
195
+ if args.max_documents:
196
+ sub_argv.extend(["--max-documents", str(args.max_documents)])
197
+
198
+ old_argv = sys.argv
199
+ sys.argv = sub_argv
200
+ try:
201
+ from benchmarks.ultradomain.benchmark_ultradomain import main as ultra_main
202
+ asyncio.run(ultra_main())
203
+ finally:
204
+ sys.argv = old_argv
205
+
206
+ elif task == "text2sql":
207
+ sub_argv = ["benchmark_text2sql"]
208
+ if args.max_documents:
209
+ sub_argv.extend(["--max-documents", str(args.max_documents)])
210
+
211
+ old_argv = sys.argv
212
+ sys.argv = sub_argv
213
+ try:
214
+ from benchmarks.text2sql.benchmark_text2sql import main as text2sql_main
215
+ text2sql_main()
216
+ finally:
217
+ sys.argv = old_argv
218
+
219
+ elif task == "summarization":
220
+ sub_argv = [
221
+ "benchmark_summarization",
222
+ "--dataset", args.dataset,
223
+ "--questions", str(args.questions),
224
+ "--embedding", args.embedding,
225
+ "--llm", args.llm,
226
+ ]
227
+ if args.force_build:
228
+ sub_argv.append("--force-build")
229
+ if args.build_only:
230
+ sub_argv.append("--build")
231
+ if args.max_documents:
232
+ sub_argv.extend(["--max-documents", str(args.max_documents)])
233
+
234
+ old_argv = sys.argv
235
+ sys.argv = sub_argv
236
+ try:
237
+ from benchmarks.summarization.benchmark_summarization import main as summ_main
238
+ asyncio.run(summ_main())
239
+ finally:
240
+ sys.argv = old_argv
241
+
242
+ else:
243
+ print(f"ERROR: Unknown task '{task}'")
244
+
245
+
246
+ def main():
247
+ parser = argparse.ArgumentParser(
248
+ description="QAFD-RAG Unified Benchmark Runner",
249
+ formatter_class=argparse.RawDescriptionHelpFormatter,
250
+ epilog="""
251
+ Examples:
252
+ # Multihop with passage-entity graph (default)
253
+ python benchmarks/run.py --task multihop --dataset musique
254
+
255
+ # Multihop with entity graph (override)
256
+ python benchmarks/run.py --task multihop --dataset musique --graph_type entity
257
+
258
+ # Ultradomain with entity graph (default)
259
+ python benchmarks/run.py --task ultradomain --dataset mix
260
+
261
+ # Retrieval only (skip QA)
262
+ python benchmarks/run.py --task multihop --dataset musique --skip_qa
263
+
264
+ # Build KG only
265
+ python benchmarks/run.py --task multihop --dataset musique --build_only
266
+ """,
267
+ )
268
+
269
+ # ── Required ────────────────────────────────────────────────────────
270
+ parser.add_argument("--task", type=str, required=True,
271
+ choices=["multihop", "ultradomain", "text2sql", "summarization"])
272
+ parser.add_argument("--dataset", type=str, required=True,
273
+ help="Dataset name (e.g. musique, hotpotqa, mix, spider2-lite)")
274
+
275
+ # ── Graph type ──────────────────────────────────────────────────────
276
+ parser.add_argument("--graph_type", type=str, default=None,
277
+ choices=["passage-entity", "entity"],
278
+ help="Graph type (default: passage-entity for multihop, entity for others)")
279
+
280
+ # ── Model ──────────────────────────���────────────────────────────────
281
+ parser.add_argument("--llm", type=str, default="gpt-4o-mini")
282
+ parser.add_argument("--embedding", type=str, default=None,
283
+ help="Embedding model (auto-selected per task: "
284
+ "nvidia-nv-embed-v2 for multihop, openai-small for others)")
285
+
286
+ # ── Run control ─────────────────────────────────────────────────────
287
+ parser.add_argument("--questions", type=int, default=100)
288
+ parser.add_argument("--max_documents", type=int, default=None)
289
+ parser.add_argument("--build_only", action="store_true",
290
+ help="Build KG only, skip benchmark")
291
+ parser.add_argument("--force_build", action="store_true",
292
+ help="Rebuild KG even if it exists")
293
+ parser.add_argument("--skip_qa", action="store_true",
294
+ help="Run retrieval only, skip QA (passage-entity only)")
295
+ parser.add_argument("--batch_push", action="store_true",
296
+ help="Batch push-relabel (process all excess nodes per iter)")
297
+
298
+ # ── QAFD parameters (shared) ────────────────────────────────────────
299
+ parser.add_argument("--alpha", type=float, default=None,
300
+ help="QAFD alpha (default: 2.0)")
301
+ parser.add_argument("--epsilon", type=float, default=None)
302
+ parser.add_argument("--max_iterations", type=int, default=None)
303
+ parser.add_argument("--step_size", type=float, default=None)
304
+ parser.add_argument("--weight_scheme", type=str, default=None,
305
+ choices=["original", "multiply", "add"])
306
+
307
+ # ── Passage-entity specific ─────────────────────────────────────────
308
+ parser.add_argument("--linking_top_k", type=int, default=None,
309
+ help="Number of fact seeds (passage-entity only)")
310
+ parser.add_argument("--passage_node_weight", type=float, default=None,
311
+ help="Passage node weight in seed computation (passage-entity only)")
312
+ parser.add_argument("--retrieval_top_k", type=int, default=None,
313
+ help="Number of passages to retrieve (passage-entity only)")
314
+
315
+ # ── Entity graph specific ───────────────────────────────────────────
316
+ parser.add_argument("--max_source_nodes", type=int, default=None,
317
+ help="Max source nodes for flow diffusion (entity only)")
318
+ parser.add_argument("--min_flow_threshold", type=float, default=None,
319
+ help="Min flow threshold for clusters (entity only)")
320
+
321
+ args = parser.parse_args()
322
+
323
+ # ── Resolve graph type ──────────────────────────────────────────────
324
+ if args.graph_type is None:
325
+ args.graph_type = TASK_DEFAULT_GRAPH_TYPE[args.task]
326
+
327
+ # ── Resolve embedding (must match pre-built KGs on HuggingFace) ────
328
+ if args.embedding is None:
329
+ _task_embeddings = {
330
+ "multihop": "nvidia-nv-embed-v2",
331
+ "ultradomain": "openai-small",
332
+ "text2sql": "openai-small",
333
+ "summarization": "openai-small",
334
+ }
335
+ args.embedding = _task_embeddings[args.task]
336
+
337
+ # ── Validate dataset ────────────────────────────────────────────────
338
+ valid = TASK_DATASETS.get(args.task, [])
339
+ if args.dataset not in valid and args.dataset != "all":
340
+ print(f"ERROR: Unknown dataset '{args.dataset}' for task '{args.task}'")
341
+ print(f" Valid: {valid}")
342
+ return
343
+
344
+ # ── Apply graph-type defaults for unset params ──────────────────────
345
+ defaults = GRAPH_TYPE_DEFAULTS[args.graph_type]
346
+ for key, default_val in defaults.items():
347
+ if getattr(args, key, None) is None:
348
+ setattr(args, key, default_val)
349
+
350
+ # ── Print config ────────────────────────────────────────────────────
351
+ print(f"\n{'=' * 70}")
352
+ print(f" QAFD-RAG Benchmark")
353
+ print(f"{'=' * 70}")
354
+ print(f" Task: {args.task}")
355
+ print(f" Dataset: {args.dataset}")
356
+ print(f" Graph type: {args.graph_type}")
357
+ print(f" LLM: {args.llm}")
358
+ print(f" Embedding: {args.embedding}")
359
+ print(f" QAFD alpha: {args.alpha}")
360
+ print(f" Questions: {args.questions}")
361
+ print(f"{'=' * 70}\n")
362
+
363
+ # ── Dispatch ──────────────────���─────────────────────────────────────
364
+ if args.graph_type == "passage-entity":
365
+ run_passage_entity(args)
366
+ else:
367
+ run_entity(args)
368
+
369
+
370
+ if __name__ == "__main__":
371
+ main()
benchmarks/summarization/benchmark_summarization.py ADDED
@@ -0,0 +1,673 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """
3
+ SQuALITY Benchmark for QAFD_RAG
4
+ Multi-reference BLEU/ROUGE/METEOR metrics + LLM quality evaluation
5
+ """
6
+
7
+ import os
8
+ import sys
9
+ import asyncio
10
+ import time
11
+ import json
12
+ import logging
13
+ import nest_asyncio
14
+ from typing import Dict, List
15
+ from dataclasses import dataclass
16
+ from datetime import datetime
17
+ from pathlib import Path
18
+
19
+ # Suppress verbose logging
20
+ logging.getLogger("httpx").setLevel(logging.ERROR)
21
+ logging.getLogger("QAFD_RAG").setLevel(logging.ERROR)
22
+ logging.getLogger("nano-vectordb").setLevel(logging.ERROR)
23
+ logging.getLogger("OpenAI").setLevel(logging.ERROR)
24
+ logging.getLogger("absl").setLevel(logging.ERROR)
25
+
26
+ QAFD_RAG_HOME = str(Path(__file__).parent.parent.parent)
27
+ sys.path.insert(0, QAFD_RAG_HOME)
28
+
29
+ nest_asyncio.apply()
30
+
31
+
32
+ def print_header(title: str, width: int = 70):
33
+ """Print a formatted header"""
34
+ print(f"\n{'=' * width}")
35
+ print(f" {title}")
36
+ print(f"{'=' * width}")
37
+
38
+
39
+ def print_config(items: Dict[str, str], width: int = 70):
40
+ """Print configuration items"""
41
+ print(f"{'─' * width}")
42
+ for key, value in items.items():
43
+ print(f" {key:<20} {value}")
44
+ print(f"{'─' * width}")
45
+
46
+
47
+ def print_progress(current: int, total: int, prefix: str = "", width: int = 40):
48
+ """Print a progress bar"""
49
+ percent = current / total
50
+ filled = int(width * percent)
51
+ bar = '█' * filled + '░' * (width - filled)
52
+ print(f"\r {prefix} [{bar}] {current}/{total} ({percent*100:.1f}%)", end='', flush=True)
53
+ if current == total:
54
+ print()
55
+
56
+ @dataclass
57
+ class BenchmarkResult:
58
+ """Benchmark results"""
59
+ model_name: str
60
+ dataset_name: str
61
+ success: bool = False
62
+ total_time: float = 0.0
63
+ kg_build_time: float = 0.0
64
+ query_time: float = 0.0
65
+ num_questions: int = 0
66
+ metrics: Dict[str, float] = None # BLEU, ROUGE, METEOR
67
+ quality_scores: Dict[str, List[float]] = None # Comprehensiveness, Diversity, etc.
68
+ error_message: str = ""
69
+ responses: List[str] = None # Generated summaries
70
+ questions: List[str] = None # Original questions
71
+ reference_answers: List[List[str]] = None # Reference answers per question
72
+
73
+ class SQuALITYBenchmark:
74
+ def __init__(self, api_key: str, embedding_model: str = "openai-small", llm_model: str = "gpt-4o-mini"):
75
+ self.api_key = api_key
76
+ self.embedding_model = embedding_model
77
+ self.llm_model = llm_model
78
+ os.environ["OPENAI_API_KEY"] = api_key
79
+ os.environ["OPENAI_API_BASE"] = "https://api.openai.com/v1"
80
+
81
+ def _get_llm_func(self):
82
+ """Get LLM function based on model name"""
83
+ from src import llm
84
+ llm_funcs = {
85
+ "gpt-4o-mini": llm.gpt_4o_mini_complete,
86
+ "gpt-4o": llm.gpt_4o_complete,
87
+ "gpt-oss-120b": llm.gpt_oss_120b_complete,
88
+ "gpt-5": llm.gpt_5_complete,
89
+ "gpt-5-mini": llm.gpt_5_mini_complete,
90
+ "gpt-5-nano": llm.gpt_5_nano_complete,
91
+ }
92
+ return llm_funcs.get(self.llm_model, llm.gpt_4o_mini_complete)
93
+
94
+ def _get_working_dir(self, dataset_name: str) -> str:
95
+ """Get working directory for dataset"""
96
+ return os.path.join(QAFD_RAG_HOME, "kg", "summarization", f"{self.llm_model}_{self.embedding_model}_{dataset_name}")
97
+
98
+ def _kg_exists(self, working_dir: str) -> bool:
99
+ """Check if KG already exists"""
100
+ # Check for key files that indicate a built KG
101
+ kg_files = [
102
+ os.path.join(working_dir, "vdb_entities.json"),
103
+ os.path.join(working_dir, "vdb_chunks.json"),
104
+ os.path.join(working_dir, "kv_store_full_docs.json"),
105
+ ]
106
+ return all(os.path.exists(f) for f in kg_files)
107
+
108
+ async def _benchmark_qafd(self, dataset: List[Dict], dataset_name: str,
109
+ question_count: int, max_documents: int = None,
110
+ mode: str = "hybrid", max_source_nodes: int = 40,
111
+ min_flow_threshold: float = 0.01, alpha: float = 5.0,
112
+ evaluator: str = "mini", force_build: bool = False) -> BenchmarkResult:
113
+ """Benchmark QAFD_RAG on SQuALITY"""
114
+ try:
115
+ from src.QAFD_RAG import QAFD_RAG, QueryParam
116
+
117
+ working_dir = self._get_working_dir(dataset_name)
118
+
119
+ llm_func = self._get_llm_func()
120
+ rag = QAFD_RAG(
121
+ working_dir=working_dir,
122
+ llm_model_func=llm_func,
123
+ llm_model_name=self.llm_model,
124
+ embedding_model_key=self.embedding_model,
125
+ enable_llm_cache=True,
126
+ )
127
+
128
+ # Get unique documents and selected questions
129
+ unique_docs = {}
130
+ selected_questions = []
131
+
132
+ for item in dataset:
133
+ passage_id = item["passage_id"]
134
+
135
+ if passage_id not in unique_docs:
136
+ if max_documents and len(unique_docs) >= max_documents:
137
+ continue
138
+ unique_docs[passage_id] = item["document"]
139
+
140
+ if passage_id in unique_docs:
141
+ selected_questions.append(item)
142
+ if len(selected_questions) >= question_count:
143
+ break
144
+
145
+ dataset = selected_questions
146
+
147
+ # Check if KG exists or needs to be built
148
+ kg_build_time = 0.0
149
+ if self._kg_exists(working_dir) and not force_build:
150
+ print(f"\n Using existing KG: {working_dir}")
151
+ else:
152
+ print(f"\n Building KG from {len(unique_docs)} documents...")
153
+ start_time = time.time()
154
+
155
+ for i, (passage_id, document) in enumerate(unique_docs.items()):
156
+ print_progress(i + 1, len(unique_docs), "Progress")
157
+ await rag.ainsert(document)
158
+
159
+ kg_build_time = time.time() - start_time
160
+ print(f" KG built in {kg_build_time:.2f}s")
161
+
162
+ # Query phase
163
+ print(f"\n Running queries...")
164
+ start_time = time.time()
165
+ responses = []
166
+
167
+ for i in range(min(question_count, len(dataset))):
168
+ print_progress(i + 1, min(question_count, len(dataset)), "Progress")
169
+ question = dataset[i]["question"]
170
+
171
+ query_param = QueryParam(
172
+ mode=mode,
173
+ max_source_nodes=max_source_nodes,
174
+ min_flow_threshold=min_flow_threshold,
175
+ alpha=alpha
176
+ )
177
+ response = await rag.aquery(question, query_param)
178
+ responses.append(response)
179
+
180
+ query_time = time.time() - start_time
181
+
182
+ # Calculate metrics with MULTIPLE REFERENCES
183
+ print(f"\n Calculating reference metrics...", end=" ", flush=True)
184
+ metrics = self._calculate_metrics(responses, dataset[:len(responses)])
185
+ print("done")
186
+
187
+ # Calculate quality scores
188
+ print(f" Evaluating quality ({len(responses)} responses)...")
189
+ quality_scores = await self._evaluate_quality(dataset, responses, evaluator)
190
+
191
+ # Collect questions and reference answers for output
192
+ questions_list = [dataset[i]["question"] for i in range(len(responses))]
193
+ references_list = [dataset[i].get("all_answers", [dataset[i]["answer"]]) for i in range(len(responses))]
194
+
195
+ return BenchmarkResult(
196
+ model_name="QAFD_RAG",
197
+ dataset_name=dataset_name,
198
+ success=True,
199
+ total_time=kg_build_time + query_time,
200
+ kg_build_time=kg_build_time,
201
+ query_time=query_time,
202
+ num_questions=len(responses),
203
+ metrics=metrics,
204
+ quality_scores=quality_scores,
205
+ responses=responses,
206
+ questions=questions_list,
207
+ reference_answers=references_list,
208
+ )
209
+
210
+ except Exception as e:
211
+ print(f" ERROR: {str(e)}")
212
+ import traceback
213
+ traceback.print_exc()
214
+ return BenchmarkResult(
215
+ model_name="QAFD_RAG",
216
+ dataset_name=dataset_name,
217
+ success=False,
218
+ error_message=str(e)
219
+ )
220
+
221
+ def _calculate_metrics(self, responses: List[str], dataset: List[Dict]) -> Dict[str, float]:
222
+ """Calculate BLEU, ROUGE, METEOR metrics with MULTIPLE REFERENCES"""
223
+ try:
224
+ from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction
225
+ from nltk.translate.meteor_score import meteor_score
226
+ from nltk.tokenize import word_tokenize
227
+ from rouge_score import rouge_scorer
228
+ import nltk
229
+
230
+ nltk.download('wordnet', quiet=True)
231
+ nltk.download('omw-1.4', quiet=True)
232
+ nltk.download('punkt', quiet=True)
233
+
234
+ scorer = rouge_scorer.RougeScorer(['rouge1', 'rouge2'], use_stemmer=True)
235
+ smoothing = SmoothingFunction().method1
236
+
237
+ bleu1_scores = []
238
+ bleu2_scores = []
239
+ rouge1_scores = []
240
+ rouge2_scores = []
241
+ meteor_scores = []
242
+
243
+ for i, response in enumerate(responses):
244
+ if not response:
245
+ continue
246
+
247
+ # Get ALL reference answers for this question
248
+ all_refs = dataset[i].get("all_answers", [dataset[i]["answer"]])
249
+
250
+ # Tokenize all references
251
+ try:
252
+ ref_tokens_list = [word_tokenize(ref.lower()) for ref in all_refs]
253
+ resp_tokens = word_tokenize(response.lower())
254
+ except:
255
+ ref_tokens_list = [ref.lower().split() for ref in all_refs]
256
+ resp_tokens = response.lower().split()
257
+
258
+ # BLEU with multiple references
259
+ bleu1 = sentence_bleu(ref_tokens_list, resp_tokens, weights=(1,0,0,0), smoothing_function=smoothing)
260
+ bleu2 = sentence_bleu(ref_tokens_list, resp_tokens, weights=(0.5,0.5,0,0), smoothing_function=smoothing)
261
+ bleu1_scores.append(bleu1)
262
+ bleu2_scores.append(bleu2)
263
+
264
+ # ROUGE uses first reference (standard)
265
+ primary_ref = all_refs[0]
266
+ rouge_result = scorer.score(primary_ref, response)
267
+ rouge1_scores.append(rouge_result['rouge1'].fmeasure)
268
+ rouge2_scores.append(rouge_result['rouge2'].fmeasure)
269
+
270
+ # METEOR with multiple references
271
+ try:
272
+ meteor = meteor_score(ref_tokens_list, resp_tokens)
273
+ meteor_scores.append(meteor)
274
+ except:
275
+ meteor_scores.append(0.0)
276
+
277
+ return {
278
+ "bleu_1": sum(bleu1_scores) / len(bleu1_scores) if bleu1_scores else 0,
279
+ "bleu_2": sum(bleu2_scores) / len(bleu2_scores) if bleu2_scores else 0,
280
+ "rouge_1_f1": sum(rouge1_scores) / len(rouge1_scores) if rouge1_scores else 0,
281
+ "rouge_2_f1": sum(rouge2_scores) / len(rouge2_scores) if rouge2_scores else 0,
282
+ "meteor": sum(meteor_scores) / len(meteor_scores) if meteor_scores else 0,
283
+ }
284
+
285
+ except Exception as e:
286
+ print(f" Warning: Metrics calculation failed: {e}")
287
+ import traceback
288
+ traceback.print_exc()
289
+ return {}
290
+
291
+ async def _evaluate_quality(self, dataset: List[Dict], responses: List[str],
292
+ evaluator: str = "mini") -> Dict[str, List[float]]:
293
+ """Evaluate response quality using GPT-4o or GPT-4o-mini"""
294
+ if not responses:
295
+ return {}
296
+
297
+ try:
298
+ if evaluator == "4o":
299
+ from src.llm import gpt_4o_complete as eval_complete
300
+ else:
301
+ from src.llm import gpt_4o_mini_complete as eval_complete
302
+
303
+ criteria = ["comprehensiveness", "diversity", "logicality", "relevance", "coherence"]
304
+ response_scores = {criterion: [] for criterion in criteria}
305
+
306
+ for i in range(len(responses)):
307
+ print_progress(i + 1, len(responses), "Evaluating")
308
+
309
+ question = dataset[i]["question"]
310
+ response = responses[i]
311
+
312
+ if not response:
313
+ for criterion in criteria:
314
+ response_scores[criterion].append(0.0)
315
+ continue
316
+
317
+ response_criterion_scores = {criterion: [] for criterion in criteria}
318
+ for _ in range(5):
319
+ prompt = f"""Evaluate the following response based on five criteria. Rate each from 0-100.
320
+
321
+ Question: {question}
322
+ Response: {response}
323
+
324
+ Please evaluate based on these criteria:
325
+ - Comprehensiveness: How much detail does the answer provide to cover all aspects and details of the question?
326
+ - Diversity: How varied and rich is the answer in providing different perspectives and insights on the question?
327
+ - Logicality: How logically does the answer respond to all parts of the question?
328
+ - Relevance: How relevant is the answer to the question, staying focused and addressing the intended topic or issue?
329
+ - Coherence: How well does the answer maintain internal logical connections between its parts, ensuring a smooth and consistent structure?
330
+
331
+ JSON format:
332
+ {{
333
+ "comprehensiveness": [score],
334
+ "diversity": [score],
335
+ "logicality": [score],
336
+ "relevance": [score],
337
+ "coherence": [score]
338
+ }}"""
339
+
340
+ evaluation_result = await eval_complete(prompt, max_tokens=200)
341
+
342
+ import re
343
+ json_match = re.search(r'\{.*\}', evaluation_result, re.DOTALL)
344
+ if json_match:
345
+ try:
346
+ scores = json.loads(json_match.group())
347
+ for criterion in criteria:
348
+ if criterion in scores:
349
+ score_value = float(scores[criterion])
350
+ if 0 <= score_value <= 100:
351
+ response_criterion_scores[criterion].append(score_value)
352
+ except:
353
+ pass
354
+
355
+ for criterion in criteria:
356
+ if response_criterion_scores[criterion]:
357
+ avg = sum(response_criterion_scores[criterion]) / len(response_criterion_scores[criterion])
358
+ response_scores[criterion].append(avg)
359
+ else:
360
+ response_scores[criterion].append(0.0)
361
+
362
+ return response_scores
363
+ except Exception as e:
364
+ print(f" Warning: Quality evaluation failed: {e}")
365
+ return {}
366
+
367
+ def _download_dataset(self) -> List[Dict]:
368
+ """Download SQuALITY dataset - COLLECTS ALL REFERENCE ANSWERS"""
369
+ try:
370
+ from datasets import load_dataset
371
+ import warnings
372
+ warnings.filterwarnings("ignore", category=FutureWarning)
373
+ logging.getLogger("datasets").setLevel(logging.WARNING)
374
+ except ImportError:
375
+ print("ERROR: Install datasets with: pip install datasets")
376
+ sys.exit(1)
377
+
378
+ dataset = load_dataset("pszemraj/SQuALITY-v1.3", split="train")
379
+
380
+ prepared_data = []
381
+ for i, item in enumerate(dataset):
382
+ document = item.get("document", "")
383
+ questions = item.get("questions", [])
384
+ metadata = item.get("metadata", {})
385
+ passage_id = metadata.get("passage_id", f"passage_{i}") if isinstance(metadata, dict) else f"passage_{i}"
386
+
387
+ if not document or not questions:
388
+ continue
389
+
390
+ for q in questions:
391
+ question_text = q.get("question_text", "")
392
+ responses = q.get("responses", [])
393
+
394
+ if responses and len(responses) > 0:
395
+ # COLLECT ALL REFERENCE ANSWERS (KEY CHANGE!)
396
+ answer_texts = [r.get("response_text", "") for r in responses if r.get("response_text", "")]
397
+
398
+ if question_text and answer_texts:
399
+ prepared_data.append({
400
+ "document": document,
401
+ "question": question_text,
402
+ "answer": answer_texts[0], # Primary answer
403
+ "all_answers": answer_texts, # All references for BLEU
404
+ "passage_id": passage_id
405
+ })
406
+
407
+ return prepared_data
408
+
409
+ async def build_kg(self, dataset_name: str = "squality", max_documents: int = None) -> bool:
410
+ """Build KG only (no benchmark)"""
411
+ print_header("QAFD-RAG Knowledge Graph Builder")
412
+ print_config({
413
+ "Dataset": "SQuALITY",
414
+ "Embedding": self.embedding_model,
415
+ "LLM": self.llm_model,
416
+ "Max Documents": str(max_documents) if max_documents else "all"
417
+ })
418
+
419
+ print("\n Downloading dataset...", end=" ", flush=True)
420
+ dataset = self._download_dataset()
421
+ print("done")
422
+
423
+ try:
424
+ from src.QAFD_RAG import QAFD_RAG
425
+
426
+ working_dir = self._get_working_dir(dataset_name)
427
+ os.makedirs(working_dir, exist_ok=True)
428
+
429
+ llm_func = self._get_llm_func()
430
+ rag = QAFD_RAG(
431
+ working_dir=working_dir,
432
+ llm_model_func=llm_func,
433
+ llm_model_name=self.llm_model,
434
+ embedding_model_key=self.embedding_model,
435
+ enable_llm_cache=True,
436
+ )
437
+
438
+ # Get unique documents
439
+ unique_docs = {}
440
+ for item in dataset:
441
+ passage_id = item["passage_id"]
442
+ if passage_id not in unique_docs:
443
+ if max_documents and len(unique_docs) >= max_documents:
444
+ break
445
+ unique_docs[passage_id] = item["document"]
446
+
447
+ print(f"\n Building KG from {len(unique_docs)} documents...")
448
+ start_time = time.time()
449
+
450
+ for i, (passage_id, document) in enumerate(unique_docs.items()):
451
+ print_progress(i + 1, len(unique_docs), "Progress")
452
+ await rag.ainsert(document)
453
+
454
+ build_time = time.time() - start_time
455
+
456
+ print_header("Build Complete")
457
+ print(f" Documents processed: {len(unique_docs)}")
458
+ print(f" Time elapsed: {build_time:.2f}s")
459
+ print(f" Avg per document: {build_time/len(unique_docs):.2f}s")
460
+ print(f" Output directory: {working_dir}")
461
+ print()
462
+ return True
463
+
464
+ except Exception as e:
465
+ print(f"\n ERROR: {e}")
466
+ return False
467
+
468
+ async def run_benchmark(self, dataset_name: str = "squality", question_count: int = 250, max_documents: int = None,
469
+ mode: str = "hybrid", max_source_nodes: int = 40,
470
+ min_flow_threshold: float = 0.01, alpha: float = 10.0,
471
+ evaluator: str = "mini", force_build: bool = False) -> BenchmarkResult:
472
+ """Run benchmark"""
473
+ print_header("QAFD-RAG SQuALITY Benchmark")
474
+ print_config({
475
+ "Questions": str(question_count),
476
+ "Embedding": self.embedding_model,
477
+ "LLM": self.llm_model,
478
+ "Mode": mode,
479
+ "Max Nodes": str(max_source_nodes),
480
+ "Threshold": str(min_flow_threshold),
481
+ "Alpha": str(alpha),
482
+ "Evaluator": "GPT-4o" if evaluator == "4o" else "GPT-4o-mini"
483
+ })
484
+
485
+ print("\n Downloading dataset...", end=" ", flush=True)
486
+ dataset = self._download_dataset()
487
+ print("done")
488
+
489
+ if len(dataset) < question_count:
490
+ question_count = len(dataset)
491
+
492
+ result = await self._benchmark_qafd(dataset, dataset_name, question_count, max_documents,
493
+ mode, max_source_nodes, min_flow_threshold, alpha,
494
+ evaluator, force_build)
495
+ return result
496
+
497
+ def print_results(self, result: BenchmarkResult):
498
+ """Print results"""
499
+ print_header("Results: SQuALITY")
500
+
501
+ if not result.success:
502
+ print(f"\n Status: FAILED")
503
+ print(f" Error: {result.error_message}")
504
+ return
505
+
506
+ print("\n PERFORMANCE")
507
+ print(f" {'─' * 40}")
508
+ print(f" {'Questions':<25} {result.num_questions}")
509
+ print(f" {'KG Build Time':<25} {result.kg_build_time:.2f}s")
510
+ print(f" {'Query Time':<25} {result.query_time:.2f}s")
511
+ print(f" {'Total Time':<25} {result.total_time:.2f}s")
512
+ if result.num_questions > 0:
513
+ print(f" {'Avg per Question':<25} {result.query_time/result.num_questions:.2f}s")
514
+
515
+ if result.metrics:
516
+ print("\n REFERENCE METRICS")
517
+ print(f" {'─' * 40}")
518
+ for k, v in result.metrics.items():
519
+ print(f" {k:<25} {v*100:.2f}%")
520
+
521
+ if result.quality_scores:
522
+ print("\n QUALITY METRICS")
523
+ print(f" {'─' * 40}")
524
+ import statistics
525
+ for k, v in result.quality_scores.items():
526
+ if v:
527
+ avg = sum(v) / len(v)
528
+ std = statistics.stdev(v) if len(v) > 1 else 0
529
+ print(f" {k:<25} {avg:.2f} ± {std:.2f}")
530
+
531
+ all_scores = []
532
+ for scores in result.quality_scores.values():
533
+ all_scores.extend(scores)
534
+ if all_scores:
535
+ overall = sum(all_scores) / len(all_scores)
536
+ print(f"\n {'Overall Average':<25} {overall:.2f}")
537
+ print()
538
+
539
+ def save_results(self, result: BenchmarkResult, dataset_name: str = "squality"):
540
+ """Save results as two separate files: eval metrics and generated responses"""
541
+ timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
542
+ results_dir = os.path.join(QAFD_RAG_HOME, "results", "summarization", dataset_name)
543
+ os.makedirs(results_dir, exist_ok=True)
544
+
545
+ eval_file = os.path.join(results_dir, f"{dataset_name}_{timestamp}_eval.json")
546
+ output_file = os.path.join(results_dir, f"{dataset_name}_{timestamp}_responses.json")
547
+
548
+ # --- Eval file: metrics, quality scores, timing ---
549
+ quality_averages = {}
550
+ quality_stds = {}
551
+ if result.quality_scores:
552
+ import statistics
553
+ for criterion, scores in result.quality_scores.items():
554
+ if scores:
555
+ quality_averages[criterion] = sum(scores) / len(scores)
556
+ quality_stds[criterion] = statistics.stdev(scores) if len(scores) > 1 else 0.0
557
+
558
+ eval_data = {
559
+ "timestamp": datetime.now().isoformat(),
560
+ "model": result.model_name,
561
+ "llm": self.llm_model,
562
+ "embedding": self.embedding_model,
563
+ "success": result.success,
564
+ "performance": {
565
+ "num_questions": result.num_questions,
566
+ "kg_build_time": result.kg_build_time,
567
+ "query_time": result.query_time,
568
+ "total_time": result.total_time,
569
+ "avg_per_question": result.query_time / result.num_questions if result.num_questions > 0 else 0
570
+ },
571
+ "reference_metrics": result.metrics,
572
+ "quality_scores": quality_averages,
573
+ "quality_stds": quality_stds,
574
+ "quality_raw": result.quality_scores,
575
+ "error": result.error_message
576
+ }
577
+
578
+ with open(eval_file, 'w') as f:
579
+ json.dump(eval_data, f, indent=2)
580
+
581
+ # --- Responses file: questions + generated summaries + references ---
582
+ output_entries = []
583
+ if result.responses:
584
+ for i, response in enumerate(result.responses):
585
+ entry = {
586
+ "id": i + 1,
587
+ "question": result.questions[i] if result.questions else "",
588
+ "generated_summary": response,
589
+ "reference_answers": result.reference_answers[i] if result.reference_answers else [],
590
+ }
591
+ output_entries.append(entry)
592
+
593
+ output_data = {
594
+ "timestamp": datetime.now().isoformat(),
595
+ "model": result.model_name,
596
+ "llm": self.llm_model,
597
+ "embedding": self.embedding_model,
598
+ "num_responses": len(output_entries),
599
+ "responses": output_entries,
600
+ }
601
+
602
+ with open(output_file, 'w') as f:
603
+ json.dump(output_data, f, indent=2, ensure_ascii=False)
604
+
605
+ print(f"\nEval saved: {eval_file}")
606
+ print(f"Responses saved: {output_file}")
607
+
608
+ async def main():
609
+ """Main entry point"""
610
+ import argparse
611
+
612
+ parser = argparse.ArgumentParser(description="SQuALITY Benchmark for QAFD_RAG")
613
+ parser.add_argument("--questions", type=int, default=250,
614
+ help="Number of questions (default: 250)")
615
+ parser.add_argument("--max-documents", type=int, default=None,
616
+ help="Max documents (default: unlimited)")
617
+ parser.add_argument("--mode", type=str, default="hybrid",
618
+ choices=["local", "global", "hybrid"],
619
+ help="Query mode (default: hybrid)")
620
+ parser.add_argument("--max-source-nodes", type=int, default=40,
621
+ help="Max source nodes (default: 40)")
622
+ parser.add_argument("--min-flow-threshold", type=float, default=0.01,
623
+ help="Min flow threshold (default: 0.01)")
624
+ parser.add_argument("--alpha", type=float, default=10.0,
625
+ help="Alpha (default: 10.0)")
626
+ parser.add_argument("--evaluator", type=str, default="mini",
627
+ choices=["mini", "4o"],
628
+ help="Evaluator model: mini (gpt-4o-mini) or 4o (gpt-4o) (default: mini)")
629
+ parser.add_argument("--api-key", type=str,
630
+ help="OpenAI API key")
631
+ parser.add_argument("--build", action="store_true",
632
+ help="Build KG only (no benchmark)")
633
+ parser.add_argument("--force-build", action="store_true",
634
+ help="Force rebuild KG even if it exists")
635
+ parser.add_argument("--embedding", type=str, default="openai-small",
636
+ choices=["openai-small", "openai-large", "jina-v3", "gritlm", "nvidia-nv-embed-v2"],
637
+ help="Embedding model (default: openai-small)")
638
+ parser.add_argument("--llm", type=str, default="gpt-4o-mini",
639
+ choices=["gpt-4o-mini", "gpt-4o", "gpt-oss-120b", "gpt-5", "gpt-5-mini", "gpt-5-nano"],
640
+ help="LLM model (default: gpt-4o-mini)")
641
+ parser.add_argument("--dataset", type=str, default="squality",
642
+ help="Dataset name (default: squality)")
643
+
644
+ args = parser.parse_args()
645
+
646
+ api_key = args.api_key or os.environ.get("OPENAI_API_KEY")
647
+ if not api_key:
648
+ print("ERROR: Set OPENAI_API_KEY")
649
+ return
650
+
651
+ benchmark = SQuALITYBenchmark(api_key, args.embedding, args.llm)
652
+
653
+ # Build-only mode
654
+ if args.build:
655
+ await benchmark.build_kg(dataset_name=args.dataset, max_documents=args.max_documents)
656
+ return
657
+
658
+ result = await benchmark.run_benchmark(
659
+ dataset_name=args.dataset,
660
+ question_count=args.questions,
661
+ max_documents=args.max_documents,
662
+ mode=args.mode,
663
+ max_source_nodes=args.max_source_nodes,
664
+ min_flow_threshold=args.min_flow_threshold,
665
+ alpha=args.alpha,
666
+ evaluator=args.evaluator,
667
+ force_build=args.force_build
668
+ )
669
+ benchmark.print_results(result)
670
+ benchmark.save_results(result, dataset_name=args.dataset)
671
+
672
+ if __name__ == "__main__":
673
+ asyncio.run(main())
benchmarks/text2sql/benchmark_text2sql.py ADDED
@@ -0,0 +1,385 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """
3
+ Text2SQL Benchmark using QAFD-RAG for schema context retrieval.
4
+ Auto-builds KGs from database summaries.
5
+ """
6
+
7
+ import os
8
+ import sys
9
+ import asyncio
10
+ import json
11
+ import logging
12
+ from datetime import datetime
13
+ from typing import Dict, List, Any
14
+ from dataclasses import dataclass, asdict
15
+ from pathlib import Path
16
+
17
+ # Suppress verbose logging
18
+ logging.getLogger("httpx").setLevel(logging.ERROR)
19
+ logging.getLogger("QAFD_RAG").setLevel(logging.ERROR)
20
+ logging.getLogger("nano-vectordb").setLevel(logging.ERROR)
21
+ logging.getLogger("OpenAI").setLevel(logging.ERROR)
22
+
23
+ # Add QAFD-RAG to path
24
+ QAFD_RAG_HOME = Path(__file__).parent.parent.parent
25
+ sys.path.insert(0, str(QAFD_RAG_HOME))
26
+
27
+ import nest_asyncio
28
+ nest_asyncio.apply()
29
+
30
+
31
+ def print_header(title: str, width: int = 70):
32
+ """Print a formatted header"""
33
+ print(f"\n{'=' * width}")
34
+ print(f" {title}")
35
+ print(f"{'=' * width}")
36
+
37
+
38
+ def print_config(items: Dict[str, str], width: int = 70):
39
+ """Print configuration items"""
40
+ print(f"{'─' * width}")
41
+ for key, value in items.items():
42
+ print(f" {key:<20} {value}")
43
+ print(f"{'─' * width}")
44
+
45
+
46
+ def print_progress(current: int, total: int, prefix: str = "", width: int = 40):
47
+ """Print a progress bar"""
48
+ percent = current / total
49
+ filled = int(width * percent)
50
+ bar = '█' * filled + '░' * (width - filled)
51
+ print(f"\r {prefix} [{bar}] {current}/{total} ({percent*100:.1f}%)", end='', flush=True)
52
+ if current == total:
53
+ print()
54
+
55
+
56
+ @dataclass
57
+ class Text2SQLResult:
58
+ instance_id: str
59
+ db: str
60
+ question: str
61
+ create_table: str
62
+ success: bool = True
63
+ error_message: str = ""
64
+ retrieved_tables: List[str] = None
65
+ retrieved_columns: List[str] = None
66
+
67
+
68
+ def parse_schema_from_create_table(create_table: str) -> tuple:
69
+ """Extract table names and table.column pairs from CREATE TABLE output."""
70
+ import re
71
+ tables = []
72
+ columns = []
73
+ current_table = None
74
+ for line in create_table.split('\n'):
75
+ m = re.match(r'^CREATE TABLE\s+`?(\w+)`?\s*\(', line)
76
+ if m:
77
+ current_table = m.group(1)
78
+ tables.append(current_table)
79
+ continue
80
+ if current_table and line.strip().startswith('`'):
81
+ cm = re.match(r'\s*`(\w+)`', line)
82
+ if cm:
83
+ columns.append(f"{current_table}.{cm.group(1)}")
84
+ elif current_table and line.strip().startswith('"'):
85
+ cm = re.match(r'\s*"(\w+)"', line)
86
+ if cm:
87
+ columns.append(f"{current_table}.{cm.group(1)}")
88
+ return tables, columns
89
+
90
+
91
+ def compute_schema_metrics(results: List[Text2SQLResult], golden_path: str) -> dict:
92
+ """Compute table and column recall/precision/F1 against golden annotations."""
93
+ if not os.path.exists(golden_path):
94
+ return {}
95
+
96
+ with open(golden_path, 'r') as f:
97
+ golden = json.load(f)
98
+
99
+ table_metrics = []
100
+ column_metrics = []
101
+
102
+ for r in results:
103
+ if not r.success or r.instance_id not in golden:
104
+ continue
105
+ g = golden[r.instance_id]
106
+ if 'schema_extraction' not in g:
107
+ continue
108
+
109
+ golden_tables = set(t.lower() for t in g['schema_extraction'].get('tables', []))
110
+ golden_cols = set(c.lower() for c in g['schema_extraction'].get('columns', []))
111
+
112
+ retrieved_tables = set(t.lower() for t in (r.retrieved_tables or []))
113
+ retrieved_cols = set(c.lower() for c in (r.retrieved_columns or []))
114
+
115
+ # Table metrics
116
+ if golden_tables:
117
+ tp = len(retrieved_tables & golden_tables)
118
+ precision = tp / len(retrieved_tables) if retrieved_tables else 0
119
+ recall = tp / len(golden_tables)
120
+ f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0
121
+ table_metrics.append({"precision": precision, "recall": recall, "f1": f1})
122
+
123
+ # Column metrics
124
+ if golden_cols:
125
+ tp = len(retrieved_cols & golden_cols)
126
+ precision = tp / len(retrieved_cols) if retrieved_cols else 0
127
+ recall = tp / len(golden_cols)
128
+ f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0
129
+ column_metrics.append({"precision": precision, "recall": recall, "f1": f1})
130
+
131
+ if not table_metrics:
132
+ return {}
133
+
134
+ avg = lambda lst, key: sum(m[key] for m in lst) / len(lst)
135
+ return {
136
+ "table_recall": round(avg(table_metrics, "recall"), 4),
137
+ "table_precision": round(avg(table_metrics, "precision"), 4),
138
+ "table_f1": round(avg(table_metrics, "f1"), 4),
139
+ "column_recall": round(avg(column_metrics, "recall"), 4),
140
+ "column_precision": round(avg(column_metrics, "precision"), 4),
141
+ "column_f1": round(avg(column_metrics, "f1"), 4),
142
+ "num_evaluated": len(table_metrics),
143
+ }
144
+
145
+
146
+ class Text2SQLBenchmark:
147
+ def __init__(self, api_key: str, embedding_model: str = "openai-small", llm_model: str = "gpt-4o-mini"):
148
+ self.api_key = api_key
149
+ self.embedding_model = embedding_model
150
+ self.llm_model = llm_model
151
+ os.environ["OPENAI_API_KEY"] = api_key
152
+ os.environ["OPENAI_API_BASE"] = "https://api.openai.com/v1"
153
+
154
+ def load_questions(self, data_path: str, db_filter: List[str] = None) -> List[Dict]:
155
+ """Load questions from spider2-lite.jsonl"""
156
+ questions = []
157
+ with open(data_path, 'r') as f:
158
+ for line in f:
159
+ item = json.loads(line)
160
+ if db_filter is None or item.get("db") in db_filter:
161
+ questions.append(item)
162
+ return questions
163
+
164
+ async def run_benchmark(self, questions: List[Dict], question_count: int = 5) -> List[Text2SQLResult]:
165
+ """Run text2sql benchmark using the generic runner."""
166
+ from src.text2sql.runner import kg_exists, get_schema_path, ensure_kg, query_kg, get_kg_dir
167
+ from src.text2sql.prompt_parser import parse_qafd_clusters
168
+
169
+ results = []
170
+ total = min(question_count, len(questions))
171
+
172
+ print(f"\n Running queries...")
173
+
174
+ for i, item in enumerate(questions[:question_count]):
175
+ instance_id = item.get("instance_id", f"q{i}")
176
+ db = item.get("db", "unknown")
177
+ question = item.get("question", "")
178
+
179
+ print_progress(i + 1, total, "Progress")
180
+
181
+ # Auto-build KG if schema summary exists (generates from .sqlite if needed)
182
+ if not kg_exists(db):
183
+ from src.text2sql.runner import ensure_db_summary
184
+ schema_path = ensure_db_summary(db) or get_schema_path(db)
185
+ if not schema_path:
186
+ results.append(Text2SQLResult(
187
+ instance_id=instance_id, db=db, question=question,
188
+ create_table="",
189
+ success=False, error_message=f"No DB summary or .sqlite file for {db}"
190
+ ))
191
+ continue
192
+
193
+ try:
194
+ await ensure_kg(db, schema_path,
195
+ embedding_model=self.embedding_model,
196
+ llm_model=self.llm_model)
197
+ except Exception as e:
198
+ results.append(Text2SQLResult(
199
+ instance_id=instance_id, db=db, question=question,
200
+ create_table="",
201
+ success=False, error_message=f"KG build failed: {e}"
202
+ ))
203
+ continue
204
+
205
+ working_dir = get_kg_dir(db, llm_model=self.llm_model, embedding_model=self.embedding_model)
206
+
207
+ try:
208
+ # Get raw clusters (same as CoFD-M pipeline)
209
+ clusters = await query_kg(
210
+ question, working_dir,
211
+ embedding_model=self.embedding_model,
212
+ llm_model=self.llm_model,
213
+ return_raw=True,
214
+ )
215
+
216
+ # Load db_summary for types, PK/FK, sample rows
217
+ schema_data = None
218
+ schema_path = get_schema_path(db)
219
+ if schema_path:
220
+ with open(schema_path, 'r') as sf:
221
+ schema_data = json.load(sf)
222
+
223
+ # Format as CREATE TABLE (with types, PK/FK, constraints, sample rows)
224
+ if isinstance(clusters, list):
225
+ create_table_str = parse_qafd_clusters(
226
+ clusters,
227
+ add_sample_rows=True,
228
+ schema_data=schema_data,
229
+ format_type="create_table",
230
+ )
231
+ else:
232
+ create_table_str = str(clusters)
233
+
234
+ tables, columns = parse_schema_from_create_table(create_table_str)
235
+ results.append(Text2SQLResult(
236
+ instance_id=instance_id, db=db, question=question,
237
+ create_table=create_table_str,
238
+ success=True,
239
+ retrieved_tables=tables,
240
+ retrieved_columns=columns,
241
+ ))
242
+
243
+ except Exception as e:
244
+ results.append(Text2SQLResult(
245
+ instance_id=instance_id, db=db, question=question,
246
+ create_table="",
247
+ success=False, error_message=str(e)
248
+ ))
249
+
250
+ return results
251
+
252
+ def save_results(self, results: List[Text2SQLResult], db_name: str = None,
253
+ benchmark: str = "spider2-lite"):
254
+ """Save results as two separate files: eval metrics and generated responses"""
255
+ timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
256
+ if db_name:
257
+ results_dir = QAFD_RAG_HOME / "results" / "text2sql" / db_name
258
+ else:
259
+ results_dir = QAFD_RAG_HOME / "results" / "text2sql"
260
+ results_dir.mkdir(parents=True, exist_ok=True)
261
+
262
+ eval_file = str(results_dir / f"text2sql_{timestamp}_eval.json")
263
+
264
+ success_count = sum(1 for r in results if r.success)
265
+
266
+ # --- Schema accuracy metrics (if golden file exists) ---
267
+ golden_paths = {
268
+ "spider2-lite": QAFD_RAG_HOME / "data" / "text2sql" / "spider2-lite" / "golden_lite_spider_total.json",
269
+ }
270
+ golden_path = str(golden_paths.get(benchmark, ""))
271
+ schema_metrics = compute_schema_metrics(results, golden_path)
272
+
273
+ # --- Eval file: success/fail stats + schema metrics ---
274
+ eval_data = {
275
+ "timestamp": datetime.now().isoformat(),
276
+ "llm": self.llm_model,
277
+ "embedding": self.embedding_model,
278
+ "total_questions": len(results),
279
+ "success_count": success_count,
280
+ "fail_count": len(results) - success_count,
281
+ "failed_instances": [
282
+ {"instance_id": r.instance_id, "db": r.db, "error": r.error_message}
283
+ for r in results if not r.success
284
+ ],
285
+ }
286
+ if schema_metrics:
287
+ eval_data["schema_metrics"] = schema_metrics
288
+
289
+ with open(eval_file, 'w', encoding='utf-8') as f:
290
+ json.dump(eval_data, f, indent=2, ensure_ascii=False)
291
+
292
+ # --- Responses: one prompts txt per instance ---
293
+ for r in results:
294
+ if not r.success or not r.create_table:
295
+ continue
296
+ prompt_file = results_dir / f"{r.instance_id}_prompts.txt"
297
+ with open(prompt_file, 'w', encoding='utf-8') as f:
298
+ f.write(r.create_table)
299
+
300
+ print(f" Eval saved: {eval_file}")
301
+ print(f" Prompts saved: {results_dir}/<instance_id>_prompts.txt\n")
302
+
303
+
304
+ async def main():
305
+ import argparse
306
+
307
+ parser = argparse.ArgumentParser(description="Text2SQL Benchmark using QAFD_RAG")
308
+ parser.add_argument("--questions", type=int, default=5, help="Number of questions")
309
+ parser.add_argument("--embedding", type=str, default="openai-small",
310
+ choices=["openai-small", "openai-large", "jina-v3", "gritlm", "nvidia-nv-embed-v2"])
311
+ parser.add_argument("--llm", type=str, default="gpt-4o-mini",
312
+ choices=["gpt-4o-mini", "gpt-4o", "gpt-oss-120b", "gpt-5", "gpt-5-mini", "gpt-5-nano"])
313
+ parser.add_argument("--db", type=str, default=None,
314
+ help="Filter by database name (e.g. Pagila, superhero)")
315
+ parser.add_argument("--benchmark", type=str, default="spider2-lite",
316
+ choices=["spider2-lite", "bird"],
317
+ help="Benchmark dataset (default: spider2-lite)")
318
+
319
+ args = parser.parse_args()
320
+
321
+ api_key = os.environ.get("OPENAI_API_KEY")
322
+ if not api_key:
323
+ print("ERROR: Set OPENAI_API_KEY environment variable")
324
+ return
325
+
326
+ benchmark = Text2SQLBenchmark(api_key, args.embedding, args.llm)
327
+
328
+ jsonl_files = {
329
+ "spider2-lite": QAFD_RAG_HOME / "data" / "text2sql" / "spider2-lite" / "spider2-lite.jsonl",
330
+ "bird": QAFD_RAG_HOME / "data" / "text2sql" / "bird" / "bird.jsonl",
331
+ }
332
+ data_path = str(jsonl_files[args.benchmark])
333
+ db_filter = [args.db] if args.db else None
334
+ questions = benchmark.load_questions(data_path, db_filter)
335
+
336
+ print_header("QAFD-RAG Text2SQL Benchmark")
337
+ print_config({
338
+ "Benchmark": args.benchmark,
339
+ "Questions": f"{min(args.questions, len(questions))} / {len(questions)}",
340
+ "DB Filter": str(db_filter[0]) if db_filter else "all",
341
+ "Embedding": args.embedding,
342
+ "LLM": args.llm
343
+ })
344
+
345
+ results = await benchmark.run_benchmark(questions, args.questions)
346
+
347
+ success_count = sum(1 for r in results if r.success)
348
+
349
+ print_header("Results: Text2SQL")
350
+ print(f"\n SUMMARY")
351
+ print(f" {'─' * 40}")
352
+ print(f" {'Successful':<25} {success_count}/{len(results)}")
353
+ print(f" {'Failed':<25} {len(results) - success_count}")
354
+
355
+ # Schema accuracy metrics
356
+ golden_paths = {
357
+ "spider2-lite": QAFD_RAG_HOME / "data" / "text2sql" / "spider2-lite" / "golden_lite_spider_total.json",
358
+ }
359
+ golden_path = str(golden_paths.get(args.benchmark, ""))
360
+ schema_metrics = compute_schema_metrics(results, golden_path)
361
+ if schema_metrics:
362
+ print(f"\n SCHEMA RETRIEVAL ACCURACY")
363
+ print(f" {'─' * 40}")
364
+ print(f" {'Table Recall':<25} {schema_metrics['table_recall']*100:.1f}%")
365
+ print(f" {'Table Precision':<25} {schema_metrics['table_precision']*100:.1f}%")
366
+ print(f" {'Table F1':<25} {schema_metrics['table_f1']*100:.1f}%")
367
+ print(f" {'Column Recall':<25} {schema_metrics['column_recall']*100:.1f}%")
368
+ print(f" {'Column Precision':<25} {schema_metrics['column_precision']*100:.1f}%")
369
+ print(f" {'Column F1':<25} {schema_metrics['column_f1']*100:.1f}%")
370
+
371
+ # Show failed instances
372
+ failed = [r for r in results if not r.success]
373
+ if failed:
374
+ print(f"\n FAILED INSTANCES")
375
+ print(f" {'─' * 40}")
376
+ for r in failed[:5]:
377
+ print(f" {r.instance_id:<20} {r.error_message[:40]}")
378
+ if len(failed) > 5:
379
+ print(f" ... and {len(failed) - 5} more")
380
+ print()
381
+
382
+ benchmark.save_results(results, db_name=args.db, benchmark=args.benchmark)
383
+
384
+ if __name__ == "__main__":
385
+ asyncio.run(main())
benchmarks/text2sql/evaluate_schema.py ADDED
@@ -0,0 +1,213 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """
3
+ Schema Retrieval Precision/Recall Evaluator for QAFD-RAG Text2SQL.
4
+
5
+ Compares retrieved tables/columns against gold standard.
6
+
7
+ Usage:
8
+ python benchmarks/text2sql/evaluate_schema.py --results <benchmark_results.json>
9
+ python benchmarks/text2sql/evaluate_schema.py --results <results.json> --gold <gold.json>
10
+ """
11
+
12
+ import os
13
+ import sys
14
+ import json
15
+ import re
16
+ import argparse
17
+ from pathlib import Path
18
+ from typing import Dict, Set, Tuple
19
+
20
+ QAFD_RAG_HOME = Path(__file__).parent.parent.parent
21
+ DEFAULT_GOLD = QAFD_RAG_HOME / "data" / "text2sql" / "spider2-lite" / "golden_lite_spider_total.json"
22
+
23
+
24
+ def load_json_file(file_path: str) -> Dict:
25
+ with open(file_path, 'r', encoding='utf-8') as f:
26
+ return json.load(f)
27
+
28
+
29
+ def extract_columns_from_gold(data: Dict) -> Dict[str, Set[str]]:
30
+ """Extract columns from gold standard (CoFD-M format)."""
31
+ columns_by_instance = {}
32
+ for instance_id, instance_data in data.items():
33
+ if isinstance(instance_data, dict) and 'schema_extraction' in instance_data:
34
+ schema = instance_data['schema_extraction']
35
+ columns_by_instance[instance_id] = set(schema.get('columns', []))
36
+ else:
37
+ columns_by_instance[instance_id] = set()
38
+ return columns_by_instance
39
+
40
+
41
+ def _parse_create_table(create_table_text: str) -> Tuple[Set[str], Set[str]]:
42
+ """Parse CREATE TABLE statements to extract table and column names."""
43
+ tables = set()
44
+ columns = set()
45
+ current_table = None
46
+ for line in create_table_text.splitlines():
47
+ stripped = line.strip()
48
+ table_match = re.match(r'CREATE\s+TABLE\s+`?(\w+)`?\s*\(', stripped, re.IGNORECASE)
49
+ if table_match:
50
+ current_table = table_match.group(1)
51
+ tables.add(current_table)
52
+ continue
53
+ if current_table and stripped.startswith('`'):
54
+ col_match = re.match(r'`(\w+)`\s+\w+', stripped)
55
+ if col_match:
56
+ columns.add(f"{current_table}.{col_match.group(1)}")
57
+ if stripped.startswith(');'):
58
+ current_table = None
59
+ return tables, columns
60
+
61
+
62
+ def extract_columns_from_results(results_data: Dict) -> Dict[str, Set[str]]:
63
+ """Extract columns from QAFD-RAG benchmark results by parsing create_table field."""
64
+ columns_by_instance = {}
65
+ for result in results_data.get('results', []):
66
+ instance_id = result.get('instance_id', '')
67
+ create_table = result.get('create_table', '')
68
+ if create_table:
69
+ _, cols = _parse_create_table(create_table)
70
+ columns_by_instance[instance_id] = cols
71
+ else:
72
+ columns_by_instance[instance_id] = set()
73
+ return columns_by_instance
74
+
75
+
76
+ def extract_tables_from_gold(data: Dict) -> Dict[str, Set[str]]:
77
+ """Extract tables from gold standard."""
78
+ tables_by_instance = {}
79
+ for instance_id, instance_data in data.items():
80
+ if isinstance(instance_data, dict) and 'schema_extraction' in instance_data:
81
+ schema = instance_data['schema_extraction']
82
+ tables_by_instance[instance_id] = set(schema.get('tables', []))
83
+ else:
84
+ tables_by_instance[instance_id] = set()
85
+ return tables_by_instance
86
+
87
+
88
+ def extract_tables_from_results(results_data: Dict) -> Dict[str, Set[str]]:
89
+ """Extract tables from QAFD-RAG benchmark results by parsing create_table field."""
90
+ tables_by_instance = {}
91
+ for result in results_data.get('results', []):
92
+ instance_id = result.get('instance_id', '')
93
+ create_table = result.get('create_table', '')
94
+ if create_table:
95
+ tbls, _ = _parse_create_table(create_table)
96
+ tables_by_instance[instance_id] = tbls
97
+ else:
98
+ tables_by_instance[instance_id] = set()
99
+ return tables_by_instance
100
+
101
+
102
+ def calculate_precision_recall(predicted: Set[str], actual: Set[str]) -> Tuple[float, float]:
103
+ """Calculate precision and recall."""
104
+ if not predicted and not actual:
105
+ return 1.0, 1.0
106
+ if not predicted:
107
+ return 0.0, 0.0
108
+ if not actual:
109
+ return 0.0, 0.0
110
+
111
+ tp = len(predicted & actual)
112
+ precision = tp / len(predicted) if predicted else 0.0
113
+ recall = tp / len(actual) if actual else 0.0
114
+ return precision, recall
115
+
116
+
117
+ def evaluate(gold: Dict[str, Set[str]], predicted: Dict[str, Set[str]]) -> Dict:
118
+ """Evaluate predicted against gold for all common instances."""
119
+ common = set(gold.keys()) & set(predicted.keys())
120
+ if not common:
121
+ return {'precision': 0.0, 'recall': 0.0, 'f1': 0.0, 'instances': 0, 'details': {}}
122
+
123
+ precisions, recalls = [], []
124
+ details = {}
125
+
126
+ for iid in sorted(common):
127
+ p, r = calculate_precision_recall(predicted[iid], gold[iid])
128
+ precisions.append(p)
129
+ recalls.append(r)
130
+ details[iid] = {
131
+ 'precision': round(p, 4),
132
+ 'recall': round(r, 4),
133
+ 'predicted': sorted(predicted[iid]),
134
+ 'gold': sorted(gold[iid]),
135
+ 'missing': sorted(gold[iid] - predicted[iid]),
136
+ 'extra': sorted(predicted[iid] - gold[iid]),
137
+ }
138
+
139
+ avg_p = sum(precisions) / len(precisions)
140
+ avg_r = sum(recalls) / len(recalls)
141
+ f1 = 2 * avg_p * avg_r / (avg_p + avg_r) if (avg_p + avg_r) > 0 else 0.0
142
+
143
+ return {
144
+ 'precision': round(avg_p, 4),
145
+ 'recall': round(avg_r, 4),
146
+ 'f1': round(f1, 4),
147
+ 'instances': len(common),
148
+ 'details': details,
149
+ }
150
+
151
+
152
+ def main():
153
+ parser = argparse.ArgumentParser(description="Evaluate text2sql schema retrieval")
154
+ parser.add_argument("--results", required=True, help="QAFD-RAG benchmark results JSON")
155
+ parser.add_argument("--gold", default=str(DEFAULT_GOLD), help="Gold standard JSON")
156
+ parser.add_argument("--output", default=None, help="Save detailed results to JSON")
157
+ args = parser.parse_args()
158
+
159
+ print("Loading files...")
160
+ results_data = load_json_file(args.results)
161
+ gold_data = load_json_file(args.gold)
162
+
163
+ # Extract columns and tables
164
+ gold_cols = extract_columns_from_gold(gold_data)
165
+ pred_cols = extract_columns_from_results(results_data)
166
+ gold_tbls = extract_tables_from_gold(gold_data)
167
+ pred_tbls = extract_tables_from_results(results_data)
168
+
169
+ # Evaluate
170
+ col_eval = evaluate(gold_cols, pred_cols)
171
+ tbl_eval = evaluate(gold_tbls, pred_tbls)
172
+
173
+ # Print results
174
+ print(f"\n{'=' * 50}")
175
+ print("SCHEMA RETRIEVAL EVALUATION")
176
+ print(f"{'=' * 50}")
177
+ print(f" Instances evaluated: {col_eval['instances']}")
178
+
179
+ print(f"\n COLUMN RETRIEVAL")
180
+ print(f" {'─' * 40}")
181
+ print(f" Precision: {col_eval['precision']:.4f}")
182
+ print(f" Recall: {col_eval['recall']:.4f}")
183
+ print(f" F1: {col_eval['f1']:.4f}")
184
+
185
+ print(f"\n TABLE RETRIEVAL")
186
+ print(f" {'─' * 40}")
187
+ print(f" Precision: {tbl_eval['precision']:.4f}")
188
+ print(f" Recall: {tbl_eval['recall']:.4f}")
189
+ print(f" F1: {tbl_eval['f1']:.4f}")
190
+
191
+ # Per-instance details
192
+ print(f"\n PER-INSTANCE COLUMN DETAILS")
193
+ print(f" {'─' * 40}")
194
+ for iid, d in col_eval['details'].items():
195
+ status = "OK" if d['recall'] == 1.0 else f"missing: {d['missing']}"
196
+ print(f" {iid:<15} P={d['precision']:.2f} R={d['recall']:.2f} {status}")
197
+
198
+ # Save
199
+ output_path = args.output
200
+ if not output_path:
201
+ output_path = args.results.replace('.json', '_eval.json')
202
+
203
+ output = {
204
+ 'column_evaluation': col_eval,
205
+ 'table_evaluation': tbl_eval,
206
+ }
207
+ with open(output_path, 'w') as f:
208
+ json.dump(output, f, indent=2)
209
+ print(f"\nDetailed results saved to: {output_path}")
210
+
211
+
212
+ if __name__ == "__main__":
213
+ main()
benchmarks/ultradomain/benchmark_ultradomain.py ADDED
@@ -0,0 +1,548 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """
3
+ UltraDomain Benchmark for QAFD-RAG
4
+ Quality evaluation using GPT-4o as judge
5
+ """
6
+
7
+ import os
8
+ import sys
9
+ import asyncio
10
+ import time
11
+ import json
12
+ import logging
13
+ import nest_asyncio
14
+ from pathlib import Path
15
+ from typing import Dict, List, Any, Optional
16
+ from dataclasses import dataclass, asdict
17
+ from datetime import datetime
18
+
19
+ # Suppress verbose logging
20
+ logging.getLogger("httpx").setLevel(logging.ERROR)
21
+ logging.getLogger("QAFD_RAG").setLevel(logging.ERROR)
22
+ logging.getLogger("nano-vectordb").setLevel(logging.ERROR)
23
+ logging.getLogger("OpenAI").setLevel(logging.ERROR)
24
+ logging.getLogger("datasets").setLevel(logging.ERROR)
25
+
26
+ # Add QAFD-RAG to path (so 'src' becomes the package)
27
+ QAFD_RAG_HOME = str(Path(__file__).parent.parent.parent)
28
+ sys.path.insert(0, QAFD_RAG_HOME)
29
+
30
+ nest_asyncio.apply()
31
+
32
+ # Define datasets
33
+ DATASETS = ["mix.jsonl"]
34
+
35
+
36
+ def print_header(title: str, width: int = 70):
37
+ """Print a formatted header"""
38
+ print(f"\n{'=' * width}")
39
+ print(f" {title}")
40
+ print(f"{'=' * width}")
41
+
42
+
43
+ def print_config(items: Dict[str, str], width: int = 70):
44
+ """Print configuration items"""
45
+ print(f"{'─' * width}")
46
+ for key, value in items.items():
47
+ print(f" {key:<20} {value}")
48
+ print(f"{'─' * width}")
49
+
50
+
51
+ def print_progress(current: int, total: int, prefix: str = "", width: int = 40):
52
+ """Print a progress bar"""
53
+ percent = current / total
54
+ filled = int(width * percent)
55
+ bar = '█' * filled + '░' * (width - filled)
56
+ print(f"\r {prefix} [{bar}] {current}/{total} ({percent*100:.1f}%)", end='', flush=True)
57
+ if current == total:
58
+ print()
59
+
60
+ @dataclass
61
+ class BenchmarkResult:
62
+ """Combined performance and quality metrics"""
63
+ model_name: str
64
+ dataset_name: str
65
+ success: bool = False
66
+ total_time: float = 0.0
67
+ quality_scores: Dict[str, List[float]] = None
68
+ responses: List[str] = None
69
+ questions: List[str] = None
70
+ error_message: str = ""
71
+
72
+ class RAGBenchmark:
73
+ def __init__(self, api_key: str, embedding_model: str = "openai-small", llm_model: str = "gpt-4o-mini"):
74
+ self.api_key = api_key
75
+ self.embedding_model = embedding_model
76
+ self.llm_model = llm_model
77
+ os.environ["OPENAI_API_KEY"] = api_key
78
+ os.environ["OPENAI_API_BASE"] = "https://api.openai.com/v1"
79
+
80
+ # Initialize LLM for evaluation
81
+ try:
82
+ from src.llm import gpt_4o_mini_complete
83
+ self.llm_func = gpt_4o_mini_complete
84
+ except ImportError:
85
+ print("Warning: Could not import LLM function for evaluation")
86
+ self.llm_func = None
87
+
88
+
89
+ def _get_working_dir(self, dataset_name: str) -> str:
90
+ """Get working directory for a dataset"""
91
+ return os.path.join(QAFD_RAG_HOME, "kg", "ultradomain", f"{self.llm_model}_{self.embedding_model}_{dataset_name}")
92
+
93
+ def _kg_exists(self, working_dir: str) -> bool:
94
+ """Check if KG already exists"""
95
+ kg_files = [
96
+ os.path.join(working_dir, "vdb_entities.json"),
97
+ os.path.join(working_dir, "vdb_chunks.json"),
98
+ os.path.join(working_dir, "kv_store_full_docs.json"),
99
+ ]
100
+ return all(os.path.exists(f) for f in kg_files)
101
+
102
+ def _get_llm_func(self):
103
+ """Get LLM function based on model name"""
104
+ from src import llm
105
+ llm_funcs = {
106
+ "gpt-4o-mini": llm.gpt_4o_mini_complete,
107
+ "gpt-4o": llm.gpt_4o_complete,
108
+ "gpt-oss-120b": llm.gpt_oss_120b_complete,
109
+ "gpt-5": llm.gpt_5_complete,
110
+ "gpt-5-mini": llm.gpt_5_mini_complete,
111
+ "gpt-5-nano": llm.gpt_5_nano_complete,
112
+ }
113
+ return llm_funcs.get(self.llm_model, llm.gpt_4o_mini_complete)
114
+
115
+ async def _benchmark_qafd(self, dataset, dataset_name: str, question_count: int,
116
+ force_build: bool = False) -> BenchmarkResult:
117
+ """Benchmark QAFD_RAG"""
118
+ try:
119
+ from src import QAFD_RAG, QueryParam
120
+
121
+ working_dir = self._get_working_dir(dataset_name)
122
+ os.makedirs(working_dir, exist_ok=True)
123
+
124
+ # Get LLM function
125
+ llm_func = self._get_llm_func()
126
+
127
+ rag = QAFD_RAG(
128
+ working_dir=working_dir,
129
+ llm_model_func=llm_func,
130
+ llm_model_name=self.llm_model,
131
+ embedding_model_key=self.embedding_model,
132
+ enable_llm_cache=True,
133
+ )
134
+
135
+ # Check if KG exists or needs to be built
136
+ insertion_time = 0.0
137
+ if self._kg_exists(working_dir) and not force_build:
138
+ print(f"\n Using existing KG: {working_dir}")
139
+ else:
140
+ print(f"\n Building KG from {len(dataset)} documents...")
141
+ start_time = time.time()
142
+ for i, item in enumerate(dataset):
143
+ print_progress(i + 1, len(dataset), "Progress")
144
+ context = item.get("context", "")
145
+ if context:
146
+ await rag.ainsert(context)
147
+ insertion_time = time.time() - start_time
148
+ print(f" KG built in {insertion_time:.2f}s")
149
+
150
+ # Query phase
151
+ print(f"\n Running queries...")
152
+ start_time = time.time()
153
+ responses = []
154
+ questions_list = []
155
+ total_questions = min(question_count, len(dataset))
156
+ for i in range(total_questions):
157
+ print_progress(i + 1, total_questions, "Progress")
158
+
159
+ question = dataset[i]["input"]
160
+ questions_list.append(question)
161
+ query_param = QueryParam(
162
+ mode="hybrid",
163
+ max_source_nodes=40,
164
+ min_flow_threshold=0.01
165
+ )
166
+ response = await rag.aquery(question, query_param)
167
+ responses.append(response)
168
+
169
+ query_time = time.time() - start_time
170
+
171
+ return BenchmarkResult(
172
+ model_name="QAFD_RAG",
173
+ dataset_name=dataset_name,
174
+ success=True,
175
+ total_time=insertion_time + query_time,
176
+ responses=responses,
177
+ questions=questions_list,
178
+ )
179
+
180
+ except Exception as e:
181
+ print(f" ERROR: {str(e)}")
182
+ import traceback
183
+ traceback.print_exc()
184
+ return BenchmarkResult(
185
+ model_name="QAFD_RAG",
186
+ dataset_name=dataset_name,
187
+ success=False,
188
+ error_message=str(e)
189
+ )
190
+
191
+ async def _evaluate_quality(self, dataset, responses: List[str], model_name: str, dataset_name: str, question_count: int) -> Dict[str, List[float]]:
192
+ """Evaluate response quality using GPT-4o"""
193
+ if not responses:
194
+ return {}
195
+
196
+ try:
197
+ # Import gpt_4o_complete specifically for evaluation
198
+ from src.llm import gpt_4o_complete
199
+
200
+ criteria = ["comprehensiveness", "diversity", "logicality", "relevance", "coherence"]
201
+ response_scores = {criterion: [] for criterion in criteria}
202
+
203
+ eval_count = min(question_count, len(dataset), len(responses))
204
+
205
+ for i in range(eval_count):
206
+ print_progress(i + 1, eval_count, "Evaluating")
207
+
208
+ query = dataset[i]["input"]
209
+ response = responses[i]
210
+
211
+ # Evaluate this response 5 times
212
+ response_criterion_scores = {criterion: [] for criterion in criteria}
213
+ for eval_round in range(5):
214
+ prompt = f"""Evaluate the following response to a question based on five criteria. Rate each criterion from 0-100.
215
+
216
+ Question: {query}
217
+ Response: {response}
218
+
219
+ Please evaluate based on these criteria:
220
+ - Comprehensiveness: How much detail does the answer provide to cover all aspects and details of the question?
221
+ - Diversity: How varied and rich is the answer in providing different perspectives and insights on the question?
222
+ - Logicality: How logically does the answer respond to all parts of the question?
223
+ - Relevance: How relevant is the answer to the question, staying focused and addressing the intended topic or issue?
224
+ - Coherence: How well does the answer maintain internal logical connections between its parts, ensuring a smooth and consistent structure?
225
+
226
+ Provide scores in JSON format:
227
+ {{
228
+ "comprehensiveness": [score],
229
+ "diversity": [score],
230
+ "logicality": [score],
231
+ "relevance": [score],
232
+ "coherence": [score]
233
+ }}"""
234
+
235
+ evaluation_result = await gpt_4o_complete(prompt, max_tokens=200)
236
+
237
+ import re
238
+ json_match = re.search(r'\{.*\}', evaluation_result, re.DOTALL)
239
+ if json_match:
240
+ try:
241
+ scores = json.loads(json_match.group())
242
+ for criterion in criteria:
243
+ if criterion in scores:
244
+ try:
245
+ score_value = float(scores[criterion])
246
+ if 0 <= score_value <= 100: # Validate score range
247
+ response_criterion_scores[criterion].append(score_value)
248
+ except (ValueError, TypeError):
249
+ continue
250
+ except json.JSONDecodeError:
251
+ continue
252
+
253
+ # Calculate average for this response and add to the list
254
+ for criterion in criteria:
255
+ if response_criterion_scores[criterion]:
256
+ avg_score = sum(response_criterion_scores[criterion]) / len(response_criterion_scores[criterion])
257
+ response_scores[criterion].append(avg_score)
258
+ else:
259
+ response_scores[criterion].append(0.0)
260
+
261
+ return response_scores
262
+
263
+ except Exception as e:
264
+ print(f"Quality evaluation failed: {e}")
265
+ import traceback
266
+ traceback.print_exc()
267
+ return {}
268
+
269
+ async def build_kg(self, max_documents: int = None) -> bool:
270
+ """Build KG only (no benchmark)"""
271
+ try:
272
+ from datasets import load_dataset
273
+ from src import QAFD_RAG
274
+ except ImportError as e:
275
+ print(f"ERROR: Missing dependency: {e}")
276
+ return False
277
+
278
+ for dataset_file in DATASETS:
279
+ dataset_name = dataset_file.replace('.jsonl', '')
280
+
281
+ print_header("QAFD-RAG Knowledge Graph Builder")
282
+ print_config({
283
+ "Graph Type": "entity",
284
+ "Dataset": f"UltraDomain ({dataset_name})",
285
+ "Embedding": self.embedding_model,
286
+ "LLM": self.llm_model,
287
+ "Max Documents": str(max_documents) if max_documents else "all"
288
+ })
289
+
290
+ try:
291
+ print("\n Loading dataset...", end=" ", flush=True)
292
+ dataset = load_dataset("TommyChien/UltraDomain", data_files=dataset_file, split="train")
293
+ print(f"done ({len(dataset)} samples)")
294
+ except Exception as e:
295
+ print(f"failed: {e}")
296
+ return False
297
+
298
+ working_dir = self._get_working_dir(dataset_name)
299
+ os.makedirs(working_dir, exist_ok=True)
300
+
301
+ llm_func = self._get_llm_func()
302
+ rag = QAFD_RAG(
303
+ working_dir=working_dir,
304
+ llm_model_func=llm_func,
305
+ llm_model_name=self.llm_model,
306
+ embedding_model_key=self.embedding_model,
307
+ enable_llm_cache=True,
308
+ )
309
+
310
+ # Determine how many documents to process
311
+ docs_to_process = len(dataset)
312
+ if max_documents:
313
+ docs_to_process = min(max_documents, len(dataset))
314
+
315
+ print(f"\n Building KG from {docs_to_process} documents...")
316
+ start_time = time.time()
317
+
318
+ for i, item in enumerate(dataset):
319
+ if i >= docs_to_process:
320
+ break
321
+ print_progress(i + 1, docs_to_process, "Progress")
322
+ context = item.get("context", "")
323
+ if context:
324
+ await rag.ainsert(context)
325
+
326
+ build_time = time.time() - start_time
327
+
328
+ print_header("Build Complete")
329
+ print(f" Documents processed: {docs_to_process}")
330
+ print(f" Time elapsed: {build_time:.2f}s")
331
+ print(f" Avg per document: {build_time/docs_to_process:.2f}s")
332
+ print(f" Output directory: {working_dir}")
333
+ print()
334
+
335
+ return True
336
+
337
+ async def run_benchmark(self, question_count: int = 10, force_build: bool = False) -> Dict[str, List[BenchmarkResult]]:
338
+ """Run benchmark on QAFD_RAG for all datasets"""
339
+ try:
340
+ from datasets import load_dataset
341
+ except ImportError:
342
+ print("ERROR: 'datasets' library not found. Install with: pip install datasets")
343
+ return {}
344
+
345
+ all_results = {}
346
+ models = [("QAFD_RAG", self._benchmark_qafd)]
347
+
348
+ for dataset_file in DATASETS:
349
+ dataset_name = dataset_file.replace('.jsonl', '')
350
+
351
+ print_header("QAFD-RAG UltraDomain Benchmark")
352
+ print_config({
353
+ "Graph Type": "entity",
354
+ "Dataset": dataset_name,
355
+ "Questions": str(question_count),
356
+ "Embedding": self.embedding_model,
357
+ "LLM": self.llm_model
358
+ })
359
+
360
+ try:
361
+ print("\n Loading dataset...", end=" ", flush=True)
362
+ dataset = load_dataset("TommyChien/UltraDomain", data_files=dataset_file, split="train")
363
+ print(f"done ({len(dataset)} samples)")
364
+ except Exception as e:
365
+ print(f"failed: {e}")
366
+ continue
367
+
368
+ results = []
369
+
370
+ for model_name, benchmark_func in models:
371
+ result = await benchmark_func(dataset, dataset_name, question_count, force_build)
372
+
373
+ # Add quality evaluation if successful
374
+ if result.success:
375
+ print(f"\n Evaluating quality ({len(result.responses)} responses)...")
376
+ result.quality_scores = await self._evaluate_quality(
377
+ dataset, result.responses, model_name, dataset_name, question_count
378
+ )
379
+
380
+ results.append(result)
381
+
382
+ all_results[dataset_name] = results
383
+ self.print_results({dataset_name: results})
384
+
385
+ return all_results
386
+
387
+ def print_results(self, all_results: Dict[str, List[BenchmarkResult]]):
388
+ """Print benchmark results"""
389
+ import statistics
390
+
391
+ for dataset_name, results in all_results.items():
392
+ print_header(f"Results: UltraDomain ({dataset_name})")
393
+
394
+ for result in results:
395
+ if not result.success:
396
+ print(f"\n Status: FAILED")
397
+ print(f" Error: {result.error_message}")
398
+ continue
399
+
400
+ print("\n PERFORMANCE")
401
+ print(f" {'─' * 40}")
402
+ print(f" {'Total Time':<25} {result.total_time:.2f}s")
403
+
404
+ if result.quality_scores:
405
+ scores = result.quality_scores
406
+
407
+ def safe_avg_std(score_list):
408
+ if not score_list:
409
+ return 0.0, 0.0
410
+ avg = sum(score_list) / len(score_list)
411
+ try:
412
+ std = statistics.stdev(score_list) if len(score_list) > 1 else 0.0
413
+ except statistics.StatisticsError:
414
+ std = 0.0
415
+ return avg, std
416
+
417
+ print("\n QUALITY METRICS")
418
+ print(f" {'─' * 40}")
419
+
420
+ compreh_avg, compreh_std = safe_avg_std(scores.get('comprehensiveness', []))
421
+ diversity_avg, diversity_std = safe_avg_std(scores.get('diversity', []))
422
+ logical_avg, logical_std = safe_avg_std(scores.get('logicality', []))
423
+ relevance_avg, relevance_std = safe_avg_std(scores.get('relevance', []))
424
+ coherence_avg, coherence_std = safe_avg_std(scores.get('coherence', []))
425
+
426
+ print(f" {'Comprehensiveness':<25} {compreh_avg:.2f} ± {compreh_std:.2f}")
427
+ print(f" {'Diversity':<25} {diversity_avg:.2f} ± {diversity_std:.2f}")
428
+ print(f" {'Logicality':<25} {logical_avg:.2f} ± {logical_std:.2f}")
429
+ print(f" {'Relevance':<25} {relevance_avg:.2f} ± {relevance_std:.2f}")
430
+ print(f" {'Coherence':<25} {coherence_avg:.2f} ± {coherence_std:.2f}")
431
+
432
+ overall_avg = (compreh_avg + diversity_avg + logical_avg + relevance_avg + coherence_avg) / 5
433
+ print(f"\n {'Overall Average':<25} {overall_avg:.2f}")
434
+ print()
435
+
436
+ def save_results(self, all_results: Dict[str, List[BenchmarkResult]]):
437
+ """Save results as two separate files: eval metrics and generated responses"""
438
+ import statistics
439
+
440
+ timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
441
+ results_dir = os.path.join(QAFD_RAG_HOME, "results", "ultradomain")
442
+ os.makedirs(results_dir, exist_ok=True)
443
+
444
+ for dataset_name, results in all_results.items():
445
+ eval_file = os.path.join(results_dir, f"{dataset_name}_{timestamp}_eval.json")
446
+ output_file = os.path.join(results_dir, f"{dataset_name}_{timestamp}_responses.json")
447
+
448
+ for result in results:
449
+ # --- Eval file: metrics and timing ---
450
+ quality_summary = {}
451
+ if result.quality_scores:
452
+ for criterion, scores in result.quality_scores.items():
453
+ if scores:
454
+ avg = sum(scores) / len(scores)
455
+ std = statistics.stdev(scores) if len(scores) > 1 else 0.0
456
+ quality_summary[criterion] = {"mean": avg, "std": std}
457
+
458
+ eval_data = {
459
+ "timestamp": datetime.now().isoformat(),
460
+ "graph_type": "entity",
461
+ "model": result.model_name,
462
+ "llm": self.llm_model,
463
+ "embedding": self.embedding_model,
464
+ "dataset": dataset_name,
465
+ "success": result.success,
466
+ "performance": {
467
+ "total_time": result.total_time,
468
+ },
469
+ "quality_scores": quality_summary,
470
+ "quality_raw": result.quality_scores,
471
+ "error": result.error_message,
472
+ }
473
+
474
+ with open(eval_file, 'w', encoding='utf-8') as f:
475
+ json.dump(eval_data, f, indent=2, ensure_ascii=False)
476
+
477
+ # --- Responses file: questions + generated answers ---
478
+ output_entries = []
479
+ if result.responses:
480
+ for i, response in enumerate(result.responses):
481
+ entry = {
482
+ "id": i + 1,
483
+ "question": result.questions[i] if result.questions else "",
484
+ "generated_answer": response,
485
+ }
486
+ output_entries.append(entry)
487
+
488
+ output_data = {
489
+ "timestamp": datetime.now().isoformat(),
490
+ "model": result.model_name,
491
+ "llm": self.llm_model,
492
+ "embedding": self.embedding_model,
493
+ "dataset": dataset_name,
494
+ "num_responses": len(output_entries),
495
+ "responses": output_entries,
496
+ }
497
+
498
+ with open(output_file, 'w', encoding='utf-8') as f:
499
+ json.dump(output_data, f, indent=2, ensure_ascii=False)
500
+
501
+ print(f" Eval saved: {eval_file}")
502
+ print(f" Responses saved: {output_file}\n")
503
+
504
+ async def main():
505
+ """Main entry point"""
506
+ import argparse
507
+
508
+ parser = argparse.ArgumentParser(description="QAFD_RAG UltraDomain Benchmark")
509
+ parser.add_argument("--dataset", type=str, default="mix.jsonl", help="Dataset file (default: mix.jsonl)")
510
+ parser.add_argument("--questions", type=int, default=10, help="Number of questions to benchmark (default: 10)")
511
+ parser.add_argument("--api-key", type=str, help="OpenAI API key (or set OPENAI_API_KEY env var)")
512
+ parser.add_argument("--embedding", type=str, default="openai-small",
513
+ choices=["openai-small", "openai-large", "jina-v3", "gritlm", "nvidia-nv-embed-v2"],
514
+ help="Embedding model (default: openai-small)")
515
+ parser.add_argument("--llm", type=str, default="gpt-4o-mini",
516
+ choices=["gpt-4o-mini", "gpt-4o", "gpt-oss-120b", "gpt-5", "gpt-5-mini", "gpt-5-nano"],
517
+ help="LLM model for response generation (default: gpt-4o-mini)")
518
+ parser.add_argument("--build", action="store_true",
519
+ help="Build KG only (no benchmark)")
520
+ parser.add_argument("--force-build", action="store_true",
521
+ help="Force rebuild KG even if exists")
522
+ parser.add_argument("--max-documents", type=int, default=None,
523
+ help="Max documents for KG building (default: all)")
524
+
525
+ args = parser.parse_args()
526
+
527
+ # Get API key from argument or environment
528
+ api_key = args.api_key or os.environ.get("OPENAI_API_KEY")
529
+ if not api_key:
530
+ print("ERROR: OpenAI API key not provided. Set OPENAI_API_KEY environment variable or use --api-key")
531
+ return
532
+
533
+ # Update dataset if provided
534
+ if args.dataset:
535
+ DATASETS[0] = args.dataset
536
+
537
+ benchmark = RAGBenchmark(api_key, embedding_model=args.embedding, llm_model=args.llm)
538
+
539
+ # Build-only mode
540
+ if args.build:
541
+ await benchmark.build_kg(max_documents=args.max_documents)
542
+ return
543
+
544
+ all_results = await benchmark.run_benchmark(question_count=args.questions, force_build=args.force_build)
545
+ benchmark.save_results(all_results)
546
+
547
+ if __name__ == "__main__":
548
+ asyncio.run(main())
data/multihop/2wikimultihopqa.json ADDED
The diff for this file is too large to render. See raw diff
 
data/multihop/2wikimultihopqa_corpus.json ADDED
The diff for this file is too large to render. See raw diff
 
data/multihop/README.md ADDED
@@ -0,0 +1,23 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Multi-hop QA Data
2
+
3
+ Datasets: MuSiQue, HotpotQA, 2WikiMultiHopQA (included in this directory).
4
+
5
+ **Graph type:** passage-entity (default) or entity (`--graph_type entity`)
6
+
7
+ ## Pre-built KGs
8
+
9
+ Pre-built KGs use `nvidia-nv-embed-v2` (requires GPU, 16GB+ VRAM, auto-downloaded):
10
+
11
+ ```bash
12
+ huggingface-cli download qafd/kg --repo-type dataset --include "multihop/*" --local-dir ./kg
13
+ ```
14
+
15
+ ## Usage
16
+
17
+ ```bash
18
+ # Uses downloaded KG, or builds from scratch if not found
19
+ python benchmarks/run.py --task multihop --dataset musique --questions 100
20
+
21
+ # Rebuild with openai-small (no GPU needed)
22
+ python benchmarks/run.py --task multihop --dataset musique --force_build --embedding openai-small
23
+ ```
data/multihop/hotpotqa.json ADDED
The diff for this file is too large to render. See raw diff
 
data/multihop/hotpotqa_corpus.json ADDED
The diff for this file is too large to render. See raw diff
 
data/multihop/musique.json ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:98ed4e21d3076532f6388d42320fb809599c63a0d8dffca8ece5e41922be6b46
3
+ size 12543629
data/multihop/musique_corpus.json ADDED
The diff for this file is too large to render. See raw diff
 
data/summarization/README.md ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Summarization Data
2
+
3
+ The summarization benchmark (SQuALITY) data is loaded automatically from HuggingFace at runtime:
4
+
5
+ ```
6
+ Dataset: pszemraj/SQuALITY-v1.3
7
+ Split: train
8
+ ```
9
+
10
+ No manual download required. The benchmark will fetch the data on first run via the `datasets` library.
11
+
12
+ ## Format
13
+
14
+ Each item contains:
15
+ - `document`: The full source document text
16
+ - `questions`: List of question objects, each with:
17
+ - `question_text`: The question
18
+ - `responses`: List of response objects with `response_text`
19
+ - `metadata.passage_id`: Unique passage identifier
20
+
21
+ ## Usage
22
+
23
+ ```bash
24
+ ./run.sh summarization --questions 50
25
+ ./run.sh summarization --build --max-documents 10
26
+ ```
data/text2sql/README.md ADDED
@@ -0,0 +1,70 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Text-to-SQL Data
2
+
3
+ This directory contains the data for the Text-to-SQL benchmark, following the
4
+ [Spider2-lite](https://github.com/xlang-ai/Spider2) directory layout.
5
+
6
+ ## Structure
7
+
8
+ ```
9
+ data/text2sql/
10
+ ├── spider2-lite.jsonl # Benchmark questions (instance_id, db, question)
11
+ ├── README.md
12
+ └── databases/
13
+ ├── sqlite/ # Local SQLite databases
14
+ │ └── <DB_Name>/
15
+ │ ├── <DB_Name>.sqlite # SQLite database file
16
+ │ └── <DB_Name>_db_summary.json # Auto-generated schema summary
17
+
18
+ ├── bigquery/ # BigQuery dataset schemas
19
+ │ └── <dataset_group>/
20
+ │ └── <dataset_group>_bigquery_summary.json # Schema summary
21
+
22
+ └── snowflake/ # Snowflake database schemas
23
+ └── <DATABASE>/
24
+ └── <DATABASE>_db_summary.json # Schema summary
25
+ ```
26
+
27
+ ## Sample Databases
28
+
29
+ One example db_summary is included per backend type:
30
+
31
+ | Backend | Database | Summary file | Description |
32
+ |-----------|----------|-------------------------------------------|------------------------------|
33
+ | sqlite | Pagila | `Pagila/Pagila_db_summary.json` | DVD rental store (16 tables) |
34
+ | bigquery | austin | `austin/austin_bigquery_summary.json` | Austin 311 service requests |
35
+ | snowflake | AUSTIN | `AUSTIN/AUSTIN_db_summary.json` | Austin 311 service requests |
36
+
37
+ ## Adding a New SQLite Database
38
+
39
+ 1. Place the `.sqlite` file:
40
+ ```
41
+ data/text2sql/databases/sqlite/MyDB/MyDB.sqlite
42
+ ```
43
+ 2. The schema summary is auto-generated on first run, or generate manually:
44
+ ```bash
45
+ python -m src.indexing.extract_db_summary \
46
+ --db-path data/text2sql/databases/sqlite/MyDB/MyDB.sqlite
47
+ ```
48
+
49
+ ## Adding a BigQuery Database
50
+
51
+ 1. Generate the summary using the BigQuery extractor:
52
+ ```bash
53
+ python -m src.indexing.extract_db_summary_bigquery \
54
+ --datasets project_id.dataset_id \
55
+ --output-dir data/text2sql/databases/bigquery/my_dataset/
56
+ ```
57
+
58
+ ## Adding a Snowflake Database
59
+
60
+ 1. Generate the summary using the Snowflake extractor:
61
+ ```bash
62
+ python -m src.indexing.extract_db_summary_snowflake \
63
+ --databases MY_DATABASE \
64
+ --output-dir data/text2sql/databases/snowflake/MY_DATABASE/
65
+ ```
66
+
67
+ ## Full Spider2-lite Data
68
+
69
+ To run the full benchmark, clone [Spider2](https://github.com/xlang-ai/Spider2)
70
+ and symlink or copy the databases into `databases/`.
data/text2sql/bird/bird.jsonl ADDED
@@ -0,0 +1,129 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {"instance_id": "bird_717", "db": "superhero", "question": "Please list all the superpowers of 3-D Man.", "external_knowledge": "superhero.md", "evidence": "3-D Man refers to superhero_name = '3-D Man'; superpowers refers to power_name", "SQL": "SELECT T3.power_name FROM superhero AS T1 INNER JOIN hero_power AS T2 ON T1.id = T2.hero_id INNER JOIN superpower AS T3 ON T2.power_id = T3.id WHERE T1.superhero_name = '3-D Man'", "difficulty": "simple"}
2
+ {"instance_id": "bird_718", "db": "superhero", "question": "How many superheroes have the super power of \"Super Strength\"?", "external_knowledge": "superhero.md", "evidence": "super power of \"Super Strength\" refers to power_name = 'Super Strength'", "SQL": "SELECT COUNT(T1.hero_id) FROM hero_power AS T1 INNER JOIN superpower AS T2 ON T1.power_id = T2.id WHERE T2.power_name = 'Super Strength'", "difficulty": "simple"}
3
+ {"instance_id": "bird_719", "db": "superhero", "question": "Among the superheroes with the super power of \"Super Strength\", how many of them have a height of over 200cm?", "external_knowledge": "superhero.md", "evidence": "super power of \"Super Strength\" refers to power_name = 'Super Strength'; a height of over 200cm refers to height_cm > 200", "SQL": "SELECT COUNT(T1.id) FROM superhero AS T1 INNER JOIN hero_power AS T2 ON T1.id = T2.hero_id INNER JOIN superpower AS T3 ON T2.power_id = T3.id WHERE T3.power_name = 'Super Strength' AND T1.height_cm > 200", "difficulty": "moderate"}
4
+ {"instance_id": "bird_720", "db": "superhero", "question": "Please list the full names of all the superheroes with over 15 super powers.", "external_knowledge": "superhero.md", "evidence": "15 super powers refers to COUNT(full_name) > 15", "SQL": "SELECT DISTINCT T1.full_name FROM superhero AS T1 INNER JOIN hero_power AS T2 ON T1.id = T2.hero_id GROUP BY T1.full_name HAVING COUNT(T2.power_id) > 15", "difficulty": "simple"}
5
+ {"instance_id": "bird_721", "db": "superhero", "question": "How many superheroes have blue eyes?", "external_knowledge": "superhero.md", "evidence": "blue eyes refers to colour = 'Blue' and eye_colour_id = colour.id;", "SQL": "SELECT COUNT(T1.id) FROM superhero AS T1 INNER JOIN colour AS T2 ON T1.eye_colour_id = T2.id WHERE T2.colour = 'Blue'", "difficulty": "simple"}
6
+ {"instance_id": "bird_722", "db": "superhero", "question": "What is the colour of Apocalypse's skin?", "external_knowledge": "superhero.md", "evidence": "Apocalypse refers to superhero_name = 'Apocalypse'; colour of skin refers to colour where skin_colour_id = colour.id", "SQL": "SELECT T2.colour FROM superhero AS T1 INNER JOIN colour AS T2 ON T1.skin_colour_id = T2.id WHERE T1.superhero_name = 'Apocalypse'", "difficulty": "simple"}
7
+ {"instance_id": "bird_723", "db": "superhero", "question": "Among the superheroes with blue eyes, how many of them have the super power of \"Agility\"?", "external_knowledge": "superhero.md", "evidence": "blue eyes refers to colour = 'Blue' and eye_colour_id = colour.id; super power of \"Agility\" refers to power_name = 'Agility'", "SQL": "SELECT COUNT(T1.id) FROM superhero AS T1 INNER JOIN hero_power AS T2 ON T1.id = T2.hero_id INNER JOIN superpower AS T3 ON T2.power_id = T3.id INNER JOIN colour AS T4 ON T1.eye_colour_id = T4.id WHERE T3.power_name = 'Agility' AND T4.colour = 'Blue'", "difficulty": "moderate"}
8
+ {"instance_id": "bird_724", "db": "superhero", "question": "Please list the superhero names of all the superheroes that have blue eyes and blond hair.", "external_knowledge": "superhero.md", "evidence": "blue eyes refers to colour = 'Blue' and eye_colour_id = colour.id; blond hair refers to colour = 'Blond' and hair_colour_id = colour.id; super power of \"Agility\" refers to power_name = 'Agility'", "SQL": "SELECT T1.superhero_name FROM superhero AS T1 INNER JOIN colour AS T2 ON T1.eye_colour_id = T2.id INNER JOIN colour AS T3 ON T1.hair_colour_id = T3.id WHERE T2.colour = 'Blue' AND T3.colour = 'Blond'", "difficulty": "challenging"}
9
+ {"instance_id": "bird_725", "db": "superhero", "question": "How many superheroes are published by Marvel Comics?", "external_knowledge": "superhero.md", "evidence": "published by Marvel Comics refers to publisher_name = 'Marvel Comics'", "SQL": "SELECT COUNT(T1.id) FROM superhero AS T1 INNER JOIN publisher AS T2 ON T1.publisher_id = T2.id WHERE T2.publisher_name = 'Marvel Comics'", "difficulty": "simple"}
10
+ {"instance_id": "bird_726", "db": "superhero", "question": "Rank heroes published by Marvel Comics by their height in descending order.", "external_knowledge": "superhero.md", "evidence": "name refers to superhero_name; the tallest hero refers to MAX(height_cm); published by Marvel Comics refers to publisher_name = 'Marvel Comics'", "SQL": "SELECT superhero_name, height_cm, RANK() OVER (ORDER BY height_cm DESC) AS HeightRank FROM superhero INNER JOIN publisher ON superhero.publisher_id = publisher.id WHERE publisher.publisher_name = 'Marvel Comics'", "difficulty": "moderate"}
11
+ {"instance_id": "bird_727", "db": "superhero", "question": "Who is the publisher of Sauron?", "external_knowledge": "superhero.md", "evidence": "the publisher refers to publisher_name; Sauron refers to superhero_name = 'Sauron'", "SQL": "SELECT T2.publisher_name FROM superhero AS T1 INNER JOIN publisher AS T2 ON T1.publisher_id = T2.id WHERE T1.superhero_name = 'Sauron'", "difficulty": "simple"}
12
+ {"instance_id": "bird_728", "db": "superhero", "question": "Rank superheroes from Marvel Comics by their eye color popularity, starting with the most common color.", "external_knowledge": "superhero.md", "evidence": "the superheroes from Marvel Comics refers to publisher_name = 'Marvel Comics'; most common color refers to COUNT(superhero.id) DESC;", "SQL": "SELECT colour.colour AS EyeColor, COUNT(superhero.id) AS Count, RANK() OVER (ORDER BY COUNT(superhero.id) DESC) AS PopularityRank FROM superhero INNER JOIN colour ON superhero.eye_colour_id = colour.id INNER JOIN publisher ON superhero.publisher_id = publisher.id WHERE publisher.publisher_name = 'Marvel Comics' GROUP BY colour.colour", "difficulty": "moderate"}
13
+ {"instance_id": "bird_729", "db": "superhero", "question": "What is the average height of the superheroes from Marvel Comics?", "external_knowledge": "superhero.md", "evidence": "superheroes from Marvel Comics refers to publisher_name = 'Marvel Comics'; average height of the superheroes refers to AVG(height_cm)", "SQL": "SELECT AVG(T1.height_cm) FROM superhero AS T1 INNER JOIN publisher AS T2 ON T1.publisher_id = T2.id WHERE T2.publisher_name = 'Marvel Comics'", "difficulty": "simple"}
14
+ {"instance_id": "bird_730", "db": "superhero", "question": "List the superheroes from Marvel Comics who have the super power of 'Super Strength'.", "external_knowledge": "superhero.md", "evidence": "the superheroes from Marvel Comics refers to publisher_name = 'Marvel Comics'; super power of \"Super Strength\" refers to power_name = 'Super Strength';", "SQL": "SELECT superhero_name FROM superhero AS T1 WHERE EXISTS (SELECT 1 FROM hero_power AS T2 INNER JOIN superpower AS T3 ON T2.power_id = T3.id WHERE T3.power_name = 'Super Strength' AND T1.id = T2.hero_id)AND EXISTS (SELECT 1 FROM publisher AS T4 WHERE T4.publisher_name = 'Marvel Comics' AND T1.publisher_id = T4.id)", "difficulty": "challenging"}
15
+ {"instance_id": "bird_731", "db": "superhero", "question": "How many superheroes did DC Comics publish?", "external_knowledge": "superhero.md", "evidence": "superheroes that DC Comics published refers to publisher_name = 'DC Comics'", "SQL": "SELECT COUNT(T1.id) FROM superhero AS T1 INNER JOIN publisher AS T2 ON T1.publisher_id = T2.id WHERE T2.publisher_name = 'DC Comics'", "difficulty": "simple"}
16
+ {"instance_id": "bird_732", "db": "superhero", "question": "Which publisher published the slowest superhero?", "external_knowledge": "superhero.md", "evidence": "the slowest superhero refers to attribute_name = 'Speed' where MIN(attribute_value); publisher refers to publisher_name", "SQL": "SELECT T2.publisher_name FROM superhero AS T1 INNER JOIN publisher AS T2 ON T1.publisher_id = T2.id INNER JOIN hero_attribute AS T3 ON T1.id = T3.hero_id INNER JOIN attribute AS T4 ON T3.attribute_id = T4.id WHERE T4.attribute_name = 'Speed' ORDER BY T3.attribute_value LIMIT 1", "difficulty": "moderate"}
17
+ {"instance_id": "bird_733", "db": "superhero", "question": "How many gold-eyed superheroes did Marvel Comics publish?", "external_knowledge": "superhero.md", "evidence": "gold-eyed refers to colour = 'Gold' where eye_colour_id = colour.id; superheroes that Marvel Comics published refers to publisher_name = 'Marvel Comics'", "SQL": "SELECT COUNT(T1.id) FROM superhero AS T1 INNER JOIN publisher AS T2 ON T1.publisher_id = T2.id INNER JOIN colour AS T3 ON T1.eye_colour_id = T3.id WHERE T2.publisher_name = 'Marvel Comics' AND T3.colour = 'Gold'", "difficulty": "moderate"}
18
+ {"instance_id": "bird_734", "db": "superhero", "question": "What is the publisher's name of Blue Beetle II?", "external_knowledge": "superhero.md", "evidence": "Blue Beetle II refers to superhero_name = 'Blue Beetle II'", "SQL": "SELECT T2.publisher_name FROM superhero AS T1 INNER JOIN publisher AS T2 ON T1.publisher_id = T2.id WHERE T1.superhero_name = 'Blue Beetle II'", "difficulty": "simple"}
19
+ {"instance_id": "bird_735", "db": "superhero", "question": "How many superheroes with blonde hair are there?", "external_knowledge": "superhero.md", "evidence": "superheroes with blonde hair refers to colour = 'Blond' where hair_colour_id = colour.id", "SQL": "SELECT COUNT(T1.id) FROM superhero AS T1 INNER JOIN colour AS T2 ON T1.hair_colour_id = T2.id WHERE T2.colour = 'Blond'", "difficulty": "simple"}
20
+ {"instance_id": "bird_736", "db": "superhero", "question": "Who is the dumbest superhero?", "external_knowledge": "superhero.md", "evidence": "the dumbest superhero refers to MIN(attribute_value) where attribute_name = 'Intelligence'", "SQL": "SELECT T1.superhero_name FROM superhero AS T1 INNER JOIN hero_attribute AS T2 ON T1.id = T2.hero_id INNER JOIN attribute AS T3 ON T2.attribute_id = T3.id WHERE T3.attribute_name = 'Intelligence' ORDER BY T2.attribute_value LIMIT 1", "difficulty": "moderate"}
21
+ {"instance_id": "bird_737", "db": "superhero", "question": "What is Copycat's race?", "external_knowledge": "superhero.md", "evidence": "Copycat is the superhero_name;", "SQL": "SELECT T2.race FROM superhero AS T1 INNER JOIN race AS T2 ON T1.race_id = T2.id WHERE T1.superhero_name = 'Copycat'", "difficulty": "simple"}
22
+ {"instance_id": "bird_738", "db": "superhero", "question": "Which superheroes have a durability attribute value of less than 50?", "external_knowledge": "superhero.md", "evidence": "durability of less than 50 refers to attribute_name = 'Durability' AND attribute_value < 50", "SQL": "SELECT superhero_name FROM superhero AS T1 WHERE EXISTS (SELECT 1 FROM hero_attribute AS T2 INNER JOIN attribute AS T3 ON T2.attribute_id = T3.id WHERE T3.attribute_name = 'Durability' AND T2.attribute_value < 50 AND T1.id = T2.hero_id)", "difficulty": "simple"}
23
+ {"instance_id": "bird_739", "db": "superhero", "question": "What are the names of the superheroes with the power of death touch?", "external_knowledge": "superhero.md", "evidence": "name of superheroes refers to refers to superhero_name; the power of death touch refers to power_name = 'Death Touch'", "SQL": "SELECT T1.superhero_name FROM superhero AS T1 INNER JOIN hero_power AS T2 ON T1.id = T2.hero_id INNER JOIN superpower AS T3 ON T2.power_id = T3.id WHERE T3.power_name = 'Death Touch'", "difficulty": "moderate"}
24
+ {"instance_id": "bird_740", "db": "superhero", "question": "How many female superheroes have a strength value of 100?", "external_knowledge": "superhero.md", "evidence": "female refers to gender = 'Female'; strength value of 100 refers to attribute_name = 'Strength' AND attribute_value = 100", "SQL": "SELECT COUNT(T1.id) FROM superhero AS T1 INNER JOIN hero_attribute AS T2 ON T1.id = T2.hero_id INNER JOIN attribute AS T3 ON T2.attribute_id = T3.id INNER JOIN gender AS T4 ON T1.gender_id = T4.id WHERE T3.attribute_name = 'Strength' AND T2.attribute_value = 100 AND T4.gender = 'Female'", "difficulty": "moderate"}
25
+ {"instance_id": "bird_741", "db": "superhero", "question": "What is the name of the superhero that has the most powers?", "external_knowledge": "superhero.md", "evidence": "name of the superhero refers to superhero_name; superhero that has the most powers refers to MAX(COUNT(superhero_name))", "SQL": "SELECT T1.superhero_name FROM superhero AS T1 INNER JOIN hero_power AS T2 ON T1.id = T2.hero_id GROUP BY T1.superhero_name ORDER BY COUNT(T2.hero_id) DESC LIMIT 1", "difficulty": "simple"}
26
+ {"instance_id": "bird_742", "db": "superhero", "question": "How many vampire superheroes are there?", "external_knowledge": "superhero.md", "evidence": "vampire superheroes refers to race = 'Vampire'", "SQL": "SELECT COUNT(T1.superhero_name) FROM superhero AS T1 INNER JOIN race AS T2 ON T1.race_id = T2.id WHERE T2.race = 'Vampire'", "difficulty": "simple"}
27
+ {"instance_id": "bird_743", "db": "superhero", "question": "What is the percentage of superheroes who act in their own self-interest or make decisions based on their own moral code? Indicate how many of the said superheroes were published by Marvel Comics.", "external_knowledge": "superhero.md", "evidence": "published by Marvel Comics refers to publisher_name = 'Marvel Comics'; superheroes who act in their own self-interest or make decisions based on their own moral code refers to alignment = 'Bad'; calculation = MULTIPLY(DIVIDE(SUM(alignment = 'Bad); count(id)), 100)", "SQL": "SELECT (CAST(COUNT(*) AS REAL) * 100 / (SELECT COUNT(*) FROM superhero)), CAST(SUM(CASE WHEN T2.publisher_name = 'Marvel Comics' THEN 1 ELSE 0 END) AS REAL) FROM superhero AS T1 INNER JOIN publisher AS T2 ON T1.publisher_id = T2.id INNER JOIN alignment AS T3 ON T3.id = T1.alignment_id WHERE T3.alignment = 'Bad'", "difficulty": "challenging"}
28
+ {"instance_id": "bird_744", "db": "superhero", "question": "Between DC and Marvel Comics, which publisher has published more superheroes? Find the difference in the number of superheroes they have published.", "external_knowledge": "superhero.md", "evidence": "DC refers to publisher_name = 'DC Comics'; Marvel Comics refers to publisher_name = 'Marvel Comics'; calculation = SUBTRACT(SUM(publisher_name = 'Marvel Comics'), SUM(publisher_name = 'DC Comics'))", "SQL": "SELECT SUM(CASE WHEN T2.publisher_name = 'Marvel Comics' THEN 1 ELSE 0 END) - SUM(CASE WHEN T2.publisher_name = 'DC Comics' THEN 1 ELSE 0 END) FROM superhero AS T1 INNER JOIN publisher AS T2 ON T1.publisher_id = T2.id", "difficulty": "challenging"}
29
+ {"instance_id": "bird_745", "db": "superhero", "question": "Give the publisher ID of Star Trek.", "external_knowledge": "superhero.md", "evidence": "Star Trek is the publisher_name;", "SQL": "SELECT id FROM publisher WHERE publisher_name = 'Star Trek'", "difficulty": "simple"}
30
+ {"instance_id": "bird_746", "db": "superhero", "question": "Calculate the average attribute value of all superheroes.", "external_knowledge": "superhero.md", "evidence": "average attribute value of all superheroes refers to AVG(attribute_value)", "SQL": "SELECT AVG(attribute_value) FROM hero_attribute", "difficulty": "simple"}
31
+ {"instance_id": "bird_747", "db": "superhero", "question": "What is the total number of superheroes without full name?", "external_knowledge": "superhero.md", "evidence": "superheroes without full name refers to full_name IS NULL", "SQL": "SELECT COUNT(id) FROM superhero WHERE full_name IS NULL", "difficulty": "simple"}
32
+ {"instance_id": "bird_748", "db": "superhero", "question": "What is the eye colour of superhero with superhero ID 75?", "external_knowledge": "superhero.md", "evidence": "eye colour refers to colour where eye_colour_id = colour.id;", "SQL": "SELECT T2.colour FROM superhero AS T1 INNER JOIN colour AS T2 ON T1.eye_colour_id = T2.id WHERE T1.id = 75", "difficulty": "simple"}
33
+ {"instance_id": "bird_749", "db": "superhero", "question": "Provide the superpowers of the superhero called Deathlok.", "external_knowledge": "superhero.md", "evidence": "superpowers refers to power_name; Deathlok refers to superhero_name = 'Deathlok'", "SQL": "SELECT T3.power_name FROM superhero AS T1 INNER JOIN hero_power AS T2 ON T1.id = T2.hero_id INNER JOIN superpower AS T3 ON T2.power_id = T3.id WHERE T1.superhero_name = 'Deathlok'", "difficulty": "simple"}
34
+ {"instance_id": "bird_750", "db": "superhero", "question": "What is the average weight of all female superheroes?", "external_knowledge": "superhero.md", "evidence": "female refers to gender = 'Female'; average weight refers to AVG(weight_kg)", "SQL": "SELECT AVG(T1.weight_kg) FROM superhero AS T1 INNER JOIN gender AS T2 ON T1.gender_id = T2.id WHERE T2.gender = 'Female'", "difficulty": "simple"}
35
+ {"instance_id": "bird_751", "db": "superhero", "question": "List down at least five superpowers of male superheroes.", "external_knowledge": "superhero.md", "evidence": "male refers to gender = 'Male'; superpowers refers to power_name;", "SQL": "SELECT T3.power_name FROM superhero AS T1 INNER JOIN hero_power AS T2 ON T1.id = T2.hero_id INNER JOIN superpower AS T3 ON T3.id = T2.power_id INNER JOIN gender AS T4 ON T4.id = T1.gender_id WHERE T4.gender = 'Male' LIMIT 5", "difficulty": "moderate"}
36
+ {"instance_id": "bird_752", "db": "superhero", "question": "Give the name of the alien superheroes.", "external_knowledge": "superhero.md", "evidence": "alien superheroes refers to race = 'Alien'; name of superhero refers to superhero_name;", "SQL": "SELECT T1.superhero_name FROM superhero AS T1 INNER JOIN race AS T2 ON T1.race_id = T2.id WHERE T2.race = 'Alien'", "difficulty": "simple"}
37
+ {"instance_id": "bird_753", "db": "superhero", "question": "Among the superheroes with height from 170 to 190, list the names of the superheroes with no eye color.", "external_knowledge": "superhero.md", "evidence": "height from 170 to 190 refers to height_cm BETWEEN 170 AND 190; no eye color refers to colour = 'No Colour'", "SQL": "SELECT DISTINCT T1.superhero_name FROM superhero AS T1 INNER JOIN colour AS T2 ON T1.eye_colour_id = T2.id WHERE T1.height_cm BETWEEN 170 AND 190 AND T2.colour = 'No Colour'", "difficulty": "moderate"}
38
+ {"instance_id": "bird_754", "db": "superhero", "question": "What is the superpower of hero ID 56?", "external_knowledge": "superhero.md", "evidence": "superpower refers to hero_power", "SQL": "SELECT T2.power_name FROM hero_power AS T1 INNER JOIN superpower AS T2 ON T1.power_id = T2.id WHERE T1.hero_id = 56", "difficulty": "simple"}
39
+ {"instance_id": "bird_755", "db": "superhero", "question": "List down at least five full name of Demi-God superheroes.", "external_knowledge": "superhero.md", "evidence": "Demi-God superheroes refers to race = 'Demi-God'", "SQL": "SELECT T1.full_name FROM superhero AS T1 INNER JOIN race AS T2 ON T1.race_id = T2.id WHERE T2.race = 'Demi-God'", "difficulty": "simple"}
40
+ {"instance_id": "bird_756", "db": "superhero", "question": "How many bad superheroes are there?", "external_knowledge": "superhero.md", "evidence": "bad superheroes refers to alignment_id = Bad", "SQL": "SELECT COUNT(T1.id) FROM superhero AS T1 INNER JOIN alignment AS T2 ON T1.alignment_id = T2.id WHERE T2.alignment = 'Bad'", "difficulty": "simple"}
41
+ {"instance_id": "bird_757", "db": "superhero", "question": "Identify the race of the superhero who weighed 169 kg.", "external_knowledge": "superhero.md", "evidence": "weighed 169 kg refers to weight_kg = 169", "SQL": "SELECT T2.race FROM superhero AS T1 INNER JOIN race AS T2 ON T1.race_id = T2.id WHERE T1.weight_kg = 169", "difficulty": "simple"}
42
+ {"instance_id": "bird_758", "db": "superhero", "question": "Provide the hair colour of the human superhero who is 185 cm tall.", "external_knowledge": "superhero.md", "evidence": "185 cm tall refers to height_cm = 185; human superhero refers to race = 'human'; hair colour refers to colour where hair_colour_id = colour.id;", "SQL": "SELECT DISTINCT T3.colour FROM superhero AS T1 INNER JOIN race AS T2 ON T1.race_id = T2.id INNER JOIN colour AS T3 ON T1.hair_colour_id = T3.id WHERE T1.height_cm = 185 AND T2.race = 'Human'", "difficulty": "moderate"}
43
+ {"instance_id": "bird_759", "db": "superhero", "question": "What is the eye clolour of the heaviest superhero?", "external_knowledge": "superhero.md", "evidence": "the heaviest superhero refers to MAX(weight_kg); eye colour refers to colour where eye_colour_id = colour.id;", "SQL": "SELECT T2.colour FROM superhero AS T1 INNER JOIN colour AS T2 ON T1.eye_colour_id = T2.id ORDER BY T1.weight_kg DESC LIMIT 1", "difficulty": "simple"}
44
+ {"instance_id": "bird_760", "db": "superhero", "question": "In superheroes with height between 150 to 180, what is the percentage of heroes published by Marvel Comics?", "external_knowledge": "superhero.md", "evidence": "height between 150 to 180 refers to height_cm BETWEEN 150 AND 180; heroes published by Marvel Comics refers to publisher_name = 'Marvel Comics'; calculation = MULTIPLY(DIVIDE(SUM(publisher.id = 13)), COUNT(publisher.id), 100)", "SQL": "SELECT CAST(COUNT(CASE WHEN T2.publisher_name = 'Marvel Comics' THEN 1 ELSE NULL END) AS REAL) * 100 / COUNT(T1.id) FROM superhero AS T1 INNER JOIN publisher AS T2 ON T1.publisher_id = T2.id WHERE T1.height_cm BETWEEN 150 AND 180", "difficulty": "challenging"}
45
+ {"instance_id": "bird_761", "db": "superhero", "question": "Among the male superheroes, list the super hero names of superheroes with weight greater than the 79% average weight of all superheroes.", "external_knowledge": "superhero.md", "evidence": "super hero names refers to superhero_name;male superheros refers to gender = 'Male';Calculation = weight_kg > MULTIPLY(AVG(weight_kg), 0.79)", "SQL": "SELECT T1.superhero_name FROM superhero AS T1 INNER JOIN gender AS T2 ON T1.gender_id = T2.id WHERE T2.gender = 'Male' AND T1.weight_kg * 100 > ( SELECT AVG(weight_kg) FROM superhero ) * 79", "difficulty": "moderate"}
46
+ {"instance_id": "bird_762", "db": "superhero", "question": "Which power do superheroes have the most of?", "external_knowledge": "superhero.md", "evidence": "power that superheroes have the most refers to MAX(COUNT(power_name))", "SQL": "SELECT T2.power_name FROM hero_power AS T1 INNER JOIN superpower AS T2 ON T1.power_id = T2.id GROUP BY T2.power_name ORDER BY COUNT(T1.hero_id) DESC LIMIT 1", "difficulty": "simple"}
47
+ {"instance_id": "bird_763", "db": "superhero", "question": "Indicate the attribute value of superhero Abomination.", "external_knowledge": "superhero.md", "evidence": "Abomination refers to superhero_name = 'Abomination';", "SQL": "SELECT T2.attribute_value FROM superhero AS T1 INNER JOIN hero_attribute AS T2 ON T1.id = T2.hero_id WHERE T1.superhero_name = 'Abomination'", "difficulty": "simple"}
48
+ {"instance_id": "bird_764", "db": "superhero", "question": "What are the superpowers of heroes with ID 1?", "external_knowledge": "superhero.md", "evidence": "superpowers refers to power_name; heroes with ID 1 refers to hero_id = 1;", "SQL": "SELECT DISTINCT T2.power_name FROM hero_power AS T1 INNER JOIN superpower AS T2 ON T1.power_id = T2.id WHERE T1.hero_id = 1", "difficulty": "simple"}
49
+ {"instance_id": "bird_765", "db": "superhero", "question": "How many heroes have stealth power?", "external_knowledge": "superhero.md", "evidence": "stealth power refers to power_name = 'stealth';", "SQL": "SELECT COUNT(T1.hero_id) FROM hero_power AS T1 INNER JOIN superpower AS T2 ON T1.power_id = T2.id WHERE T2.power_name = 'Stealth'", "difficulty": "simple"}
50
+ {"instance_id": "bird_766", "db": "superhero", "question": "What is the hero's full name with the highest attribute in strength?", "external_knowledge": "superhero.md", "evidence": "highest attribute in strength refers to MAX(attribute_value) WHERE attribute_name = 'strength';", "SQL": "SELECT T1.full_name FROM superhero AS T1 INNER JOIN hero_attribute AS T2 ON T1.id = T2.hero_id INNER JOIN attribute AS T3 ON T2.attribute_id = T3.id WHERE T3.attribute_name = 'Strength' ORDER BY T2.attribute_value DESC LIMIT 1", "difficulty": "moderate"}
51
+ {"instance_id": "bird_767", "db": "superhero", "question": "What is the average of superheroes with no skin colour?", "external_knowledge": "superhero.md", "evidence": "average = DIVIDE(COUNT(superhero.id), SUM(skin_colour_id = 1)); no skin colour refers to skin_colour_id WHERE colour.id = 1;", "SQL": "SELECT CAST(COUNT(*) AS REAL) / SUM(CASE WHEN T2.id = 1 THEN 1 ELSE 0 END) FROM superhero AS T1 INNER JOIN colour AS T2 ON T1.skin_colour_id = T2.id", "difficulty": "simple"}
52
+ {"instance_id": "bird_768", "db": "superhero", "question": "How many superheroes were published by Dark Horse Comics?", "external_knowledge": "superhero.md", "evidence": "published by Dark Horse Comics refers to publisher_name = 'Dark Horse Comics';", "SQL": "SELECT COUNT(T1.id) FROM superhero AS T1 INNER JOIN publisher AS T2 ON T1.publisher_id = T2.id WHERE T2.publisher_name = 'Dark Horse Comics'", "difficulty": "simple"}
53
+ {"instance_id": "bird_769", "db": "superhero", "question": "Which superhero has the most durability published by Dark Horse Comics?", "external_knowledge": "superhero.md", "evidence": "which superhero refers to superhero_name; most durability refers to MAX(attribute_value) WHERE attribute_name = 'durability'; published by Dark Horse Comics refers to publisher_name = 'Dark Horse Comics';", "SQL": "SELECT T1.superhero_name FROM superhero AS T1 INNER JOIN hero_attribute AS T2 ON T1.id = T2.hero_id INNER JOIN attribute AS T3 ON T3.id = T2.attribute_id INNER JOIN publisher AS T4 ON T4.id = T1.publisher_id WHERE T4.publisher_name = 'Dark Horse Comics' AND T3.attribute_name = 'Durability' ORDER BY T2.attribute_value DESC LIMIT 1", "difficulty": "challenging"}
54
+ {"instance_id": "bird_770", "db": "superhero", "question": "What is the eyes colour of Abraham Sapien?", "external_knowledge": "superhero.md", "evidence": "eye colour refers to colour.colour where eye_colour_id = colour.id; Abraham Sapien is the full name of superhero;", "SQL": "SELECT T2.colour FROM superhero AS T1 INNER JOIN colour AS T2 ON T1.eye_colour_id = T2.id WHERE T1.full_name = 'Abraham Sapien'", "difficulty": "simple"}
55
+ {"instance_id": "bird_771", "db": "superhero", "question": "List the name of superheroes with flight power.", "external_knowledge": "superhero.md", "evidence": "name of superheroes refers to superhero_name; flight power refers to power_name = 'Flight';", "SQL": "SELECT T1.superhero_name FROM superhero AS T1 INNER JOIN hero_power AS T2 ON T1.id = T2.hero_id INNER JOIN superpower AS T3 ON T2.power_id = T3.id WHERE T3.power_name = 'Flight'", "difficulty": "simple"}
56
+ {"instance_id": "bird_772", "db": "superhero", "question": "List the eyes, hair and skin colour of all female superheroes published by Dark Horse Comics.", "external_knowledge": "superhero.md", "evidence": "eyes refers to eye_colour_id; hair refers to hair_colour_id; skin colour refers to skin_colour_id; female superheroes refers to gender = 'Female'; published by Dark Horse Comics refers to publisher_name = 'Dark Horse Comics';", "SQL": "SELECT T1.eye_colour_id, T1.hair_colour_id, T1.skin_colour_id FROM superhero AS T1 INNER JOIN publisher AS T2 ON T2.id = T1.publisher_id INNER JOIN gender AS T3 ON T3.id = T1.gender_id WHERE T2.publisher_name = 'Dark Horse Comics' AND T3.gender = 'Female'", "difficulty": "challenging"}
57
+ {"instance_id": "bird_773", "db": "superhero", "question": "Which superhero has the same eyes, hair and skin colour? Indicate the publisher of the superhero.", "external_knowledge": "superhero.md", "evidence": "which superhero refers to superhero_name; the same eyes, hair and skin colour refers to hair_colour_id = skin_colour_id AND hair_colour_id = eye_colour_id; publisher refers to publisher_name;", "SQL": "SELECT T1.superhero_name, T2.publisher_name FROM superhero AS T1 INNER JOIN publisher AS T2 ON T1.publisher_id = T2.id WHERE T1.eye_colour_id = T1.hair_colour_id AND T1.eye_colour_id = T1.skin_colour_id", "difficulty": "challenging"}
58
+ {"instance_id": "bird_774", "db": "superhero", "question": "Which group does superhero A-Bomb belong to?", "external_knowledge": "superhero.md", "evidence": "group refers to race; A-Bomb refers to superhero_name = 'A-Bomb';", "SQL": "SELECT T2.race FROM superhero AS T1 INNER JOIN race AS T2 ON T1.race_id = T2.id WHERE T1.superhero_name = 'A-Bomb'", "difficulty": "simple"}
59
+ {"instance_id": "bird_775", "db": "superhero", "question": "What is the percentage of blue female superheroes among all female superheroes?", "external_knowledge": "superhero.md", "evidence": "percentage = MULTIPLY(DIVIDE(SUM(colour = 'Blue' WHERE gender = 'Female'), COUNT(gender = 'Female')), 100); blue refers to the color = 'Blue' WHERE skin_colour_id = colour.id; female refers to gender = 'Female';", "SQL": "SELECT CAST(COUNT(CASE WHEN T3.colour = 'Blue' THEN T1.id ELSE NULL END) AS REAL) * 100 / COUNT(T1.id) FROM superhero AS T1 INNER JOIN gender AS T2 ON T1.gender_id = T2.id INNER JOIN colour AS T3 ON T1.skin_colour_id = T3.id WHERE T2.gender = 'Female'", "difficulty": "challenging"}
60
+ {"instance_id": "bird_776", "db": "superhero", "question": "Provide the hero name and race of Charles Chandler.", "external_knowledge": "superhero.md", "evidence": "hero name refers to superhero_name; Charles Chandler is the full name of superhero;", "SQL": "SELECT T1.superhero_name, T2.race FROM superhero AS T1 INNER JOIN race AS T2 ON T1.race_id = T2.id WHERE T1.full_name = 'Charles Chandler'", "difficulty": "simple"}
61
+ {"instance_id": "bird_777", "db": "superhero", "question": "What is the gender of Agent 13 hero?", "external_knowledge": "superhero.md", "evidence": "Agent 13 hero refers to superhero_name = 'Agent 13';", "SQL": "SELECT T2.gender FROM superhero AS T1 INNER JOIN gender AS T2 ON T1.gender_id = T2.id WHERE T1.superhero_name = 'Agent 13'", "difficulty": "simple"}
62
+ {"instance_id": "bird_778", "db": "superhero", "question": "Provide superheroes' names who have the adaptation power.", "external_knowledge": "superhero.md", "evidence": "adaptation power refers to power_name = 'Adaptation';", "SQL": "SELECT T1.superhero_name FROM superhero AS T1 INNER JOIN hero_power AS T2 ON T1.id = T2.hero_id INNER JOIN superpower AS T3 ON T2.power_id = T3.id WHERE T3.power_name = 'Adaptation'", "difficulty": "simple"}
63
+ {"instance_id": "bird_779", "db": "superhero", "question": "How many powers does Amazo hero have?", "external_knowledge": "superhero.md", "evidence": "Amazo hero refers to superhero_name = 'Amazo';", "SQL": "SELECT COUNT(T1.power_id) FROM hero_power AS T1 INNER JOIN superhero AS T2 ON T1.hero_id = T2.id WHERE T2.superhero_name = 'Amazo'", "difficulty": "simple"}
64
+ {"instance_id": "bird_780", "db": "superhero", "question": "List the powers of Hunter Zolomon.", "external_knowledge": "superhero.md", "evidence": "Hunter Zolomon is the full name of superhero; list the powers refers to power_name;", "SQL": "SELECT T3.power_name FROM superhero AS T1 INNER JOIN hero_power AS T2 ON T1.id = T2.hero_id INNER JOIN superpower AS T3 ON T2.power_id = T3.id WHERE T1.full_name = 'Hunter Zolomon'", "difficulty": "simple"}
65
+ {"instance_id": "bird_781", "db": "superhero", "question": "Provide the heights of the heroes whose eye colours are amber.", "external_knowledge": "superhero.md", "evidence": "heights of the heroes refers to height_cm; eye colours are amber refers to colour.colour = 'Amber' WHERE eye_colour_id = colour.id;", "SQL": "SELECT T1.height_cm FROM superhero AS T1 INNER JOIN colour AS T2 ON T1.eye_colour_id = T2.id WHERE T2.colour = 'Amber'", "difficulty": "simple"}
66
+ {"instance_id": "bird_782", "db": "superhero", "question": "List the heroes' names whose eyes and hair colours are both black.", "external_knowledge": "superhero.md", "evidence": "heroes' names refers to superhero_name; eyes and hair colours are both black refers to eye_colour_id AND hair_colour_id WHERE colour.colour = 'Black';", "SQL": "SELECT T1.superhero_name FROM superhero AS T1 INNER JOIN colour AS T2 ON T1.eye_colour_id = T2.id AND T1.hair_colour_id = T2.id WHERE T2.colour = 'Black'", "difficulty": "moderate"}
67
+ {"instance_id": "bird_783", "db": "superhero", "question": "Provide the eye colours of the heroes whose skin colours are gold.", "external_knowledge": "superhero.md", "evidence": "skin colours are gold refers to colour.colour = 'Gold' WHERE skin_colour_id = colour.id;", "SQL": "SELECT T2.colour FROM superhero AS T1 INNER JOIN colour AS T2 ON T1.eye_colour_id = T2.id INNER JOIN colour AS T3 ON T1.skin_colour_id = T3.id WHERE T3.colour = 'Gold'", "difficulty": "simple"}
68
+ {"instance_id": "bird_784", "db": "superhero", "question": "Provide the full names of vampire heroes.", "external_knowledge": "superhero.md", "evidence": "vampire heroes refers to race = 'Vampire';", "SQL": "SELECT T1.full_name FROM superhero AS T1 INNER JOIN race AS T2 ON T1.race_id = T2.id WHERE T2.race = 'Vampire'", "difficulty": "simple"}
69
+ {"instance_id": "bird_785", "db": "superhero", "question": "Describe the names of neutral alignment superheroes.", "external_knowledge": "superhero.md", "evidence": "names of superheroes refers to superhero_name; neutral alignment refers to alignment = 'Neutral';", "SQL": "SELECT T1.superhero_name FROM superhero AS T1 INNER JOIN alignment AS T2 ON T1.alignment_id = T2.id WHERE T2.alignment = 'Neutral'", "difficulty": "simple"}
70
+ {"instance_id": "bird_786", "db": "superhero", "question": "How many heroes have the highest attribute value in strength?", "external_knowledge": "superhero.md", "evidence": "highest attribute value in strength refers to MAX(attribute_value) WHERE attribute_name = 'Strength';", "SQL": "SELECT COUNT(T1.hero_id) FROM hero_attribute AS T1 INNER JOIN attribute AS T2 ON T1.attribute_id = T2.id WHERE T2.attribute_name = 'Strength' AND T1.attribute_value = ( SELECT MAX(attribute_value) FROM hero_attribute )", "difficulty": "moderate"}
71
+ {"instance_id": "bird_787", "db": "superhero", "question": "What are the race and alignment of Cameron Hicks?", "external_knowledge": "superhero.md", "evidence": "Cameron Hicks refers to superhero_name = 'Cameron Hicks';", "SQL": "SELECT T2.race, T3.alignment FROM superhero AS T1 INNER JOIN race AS T2 ON T1.race_id = T2.id INNER JOIN alignment AS T3 ON T1.alignment_id = T3.id WHERE T1.superhero_name = 'Cameron Hicks'", "difficulty": "simple"}
72
+ {"instance_id": "bird_788", "db": "superhero", "question": "How many percent of female heroes were published by Marvel Comics?", "external_knowledge": "superhero.md", "evidence": "percent = MULTIPLY(DIVIDE(SUM(gender = 'Female' WHERE publisher_name = 'Marvel Comics'), COUNT(publisher_name = 'Marvel Comics')), 100); female heroes refers to gender = 'Female'; Marvel Comics refers to publisher_name = 'Marvel Comics';", "SQL": "SELECT CAST(COUNT(CASE WHEN T2.publisher_name = 'Marvel Comics' THEN 1 ELSE NULL END) AS REAL) * 100 / COUNT(T1.id) FROM superhero AS T1 INNER JOIN publisher AS T2 ON T1.publisher_id = T2.id INNER JOIN gender AS T3 ON T1.gender_id = T3.id WHERE T3.gender = 'Female'", "difficulty": "challenging"}
73
+ {"instance_id": "bird_789", "db": "superhero", "question": "Find the average weight of the heroes who are aliens.", "external_knowledge": "superhero.md", "evidence": "average = AVG(weight_kg); aliens refers to race = 'Alien';", "SQL": "SELECT CAST(SUM(T1.weight_kg) AS REAL) / COUNT(T1.id) FROM superhero AS T1 INNER JOIN race AS T2 ON T1.race_id = T2.id WHERE T2.race = 'Alien'", "difficulty": "simple"}
74
+ {"instance_id": "bird_790", "db": "superhero", "question": "Calculate the difference between Emil Blonsky's weight and Charles Chandler's weight.", "external_knowledge": "superhero.md", "evidence": "difference = SUBTRACT(SUM(weight_kg WHERE full_name = 'Emil Blonsky'), SUM(weight_kg WHERE full_name = 'Charles Chandler')); Emil Blonsky is the full name of superhero; Charles Chandler is the full name of superhero;", "SQL": "SELECT ( SELECT weight_kg FROM superhero WHERE full_name LIKE 'Emil Blonsky' ) - ( SELECT weight_kg FROM superhero WHERE full_name LIKE 'Charles Chandler' ) AS CALCULATE", "difficulty": "moderate"}
75
+ {"instance_id": "bird_791", "db": "superhero", "question": "Calculate the average height for each superhero.", "external_knowledge": "superhero.md", "evidence": "average = DIVIDE(SUM(height_cm), COUNT(all heros));", "SQL": "SELECT CAST(SUM(height_cm) AS REAL) / COUNT(id) FROM superhero", "difficulty": "simple"}
76
+ {"instance_id": "bird_792", "db": "superhero", "question": "What is Abomination's superpower?", "external_knowledge": "superhero.md", "evidence": "Abomination refers to superhero_name = 'Abomination'; superpower refers to power_name;", "SQL": "SELECT T3.power_name FROM superhero AS T1 INNER JOIN hero_power AS T2 ON T1.id = T2.hero_id INNER JOIN superpower AS T3 ON T2.power_id = T3.id WHERE T1.superhero_name = 'Abomination'", "difficulty": "simple"}
77
+ {"instance_id": "bird_793", "db": "superhero", "question": "Among the superheroes with the race of god/eternal, how many of them are male", "external_knowledge": "superhero.md", "evidence": "race \"god/eternal\" refers to race_id = 21; male refers to gender.id = 1", "SQL": "SELECT COUNT(*) FROM superhero AS T1 INNER JOIN race AS T2 ON T1.race_id = T2.id INNER JOIN gender AS T3 ON T3.id = T1.gender_id WHERE T1.race_id = 21 AND T1.gender_id = 1", "difficulty": "simple"}
78
+ {"instance_id": "bird_794", "db": "superhero", "question": "Which hero was the fastest?", "external_knowledge": "superhero.md", "evidence": "which hero refers to superhero_name; fastest refers to MAX(attribute_value) WHERE attribute_name = 'Speed';", "SQL": "SELECT T1.superhero_name FROM superhero AS T1 INNER JOIN hero_attribute AS T2 ON T1.id = T2.hero_id INNER JOIN attribute AS T3 ON T2.attribute_id = T3.id WHERE T3.attribute_name = 'Speed' ORDER BY T2.attribute_value DESC LIMIT 1", "difficulty": "moderate"}
79
+ {"instance_id": "bird_795", "db": "superhero", "question": "How many superheroes have a neutral alignment?", "external_knowledge": "superhero.md", "evidence": "neutral alignment refers to alignment_id = 3;", "SQL": "SELECT COUNT(T1.id) FROM superhero AS T1 INNER JOIN alignment AS T2 ON T1.alignment_id = T2.id WHERE T2.alignment = 'Neutral'", "difficulty": "simple"}
80
+ {"instance_id": "bird_796", "db": "superhero", "question": "State all of 3-D Man's attributes along with their values.", "external_knowledge": "superhero.md", "evidence": "3-D Man is the superhero_name. attributes refers to attribute_name; values refers to attribute_value;", "SQL": "SELECT T3.attribute_name, T2.attribute_value FROM superhero AS T1 INNER JOIN hero_attribute AS T2 ON T1.id = T2.hero_id INNER JOIN attribute AS T3 ON T2.attribute_id = T3.id WHERE T1.superhero_name = '3-D Man'", "difficulty": "moderate"}
81
+ {"instance_id": "bird_797", "db": "superhero", "question": "Which superheroes have blue eyes with brown hair?", "external_knowledge": "superhero.md", "evidence": "which superheroes refers to superhero_name; blue eyes refers to color = 'Blue' and color.id = eye_colour_id; brown hair refers to color = 'Brown' and color.id = hair_colour_id;", "SQL": "SELECT T1.superhero_name FROM superhero AS T1 INNER JOIN colour AS T2 ON T1.eye_colour_id = T2.id INNER JOIN colour AS T3 ON T1.hair_colour_id = T3.id WHERE T2.colour = 'Blue' AND T3.colour = 'Brown'", "difficulty": "moderate"}
82
+ {"instance_id": "bird_798", "db": "superhero", "question": "What is the publisher for Hawkman, Karate Kid and Speedy?", "external_knowledge": "superhero.md", "evidence": "publisher refers to publisher_name; Hawkman refers to superhero_name = 'Hawkman'; Karate Kid refers to superhero_name = 'Karate Kid'; Speedy refers to superhero_name = 'Speedy';", "SQL": "SELECT T2.publisher_name FROM superhero AS T1 INNER JOIN publisher AS T2 ON T1.publisher_id = T2.id WHERE T1.superhero_name IN ('Hawkman', 'Karate Kid', 'Speedy')", "difficulty": "moderate"}
83
+ {"instance_id": "bird_799", "db": "superhero", "question": "How many superheroes didn't have any publisher?", "external_knowledge": "superhero.md", "evidence": "didn't have any publisher refers to publisher.id = 1;", "SQL": "SELECT COUNT(T1.id) FROM superhero AS T1 INNER JOIN publisher AS T2 ON T1.publisher_id = T2.id WHERE T2.id = 1", "difficulty": "simple"}
84
+ {"instance_id": "bird_800", "db": "superhero", "question": "Calculate the percentage of superheroes with blue eyes.", "external_knowledge": "superhero.md", "evidence": "percentage = MULTIPLY(DIVIDE(SUM(superhero_name WHERE color = 'Blue'), COUNT(superhero_name)), 100.0); blue eyes refers to color = 'Blue' and color.id = eye_colour_id = 7;", "SQL": "SELECT CAST(COUNT(CASE WHEN T2.colour = 'Blue' THEN 1 ELSE NULL END) AS REAL) * 100 / COUNT(T1.id) FROM superhero AS T1 INNER JOIN colour AS T2 ON T1.eye_colour_id = T2.id", "difficulty": "moderate"}
85
+ {"instance_id": "bird_801", "db": "superhero", "question": "Find the ratio between male superheroes and female superheroes.", "external_knowledge": "superhero.md", "evidence": "ratio = DIVIDE(SUM(gender_id = 1) / SUM(gender_id = 2)); male superheroes refers to gender = 'Female'; female superheroes refers to gender = 'Male';", "SQL": "SELECT CAST(COUNT(CASE WHEN T2.gender = 'Male' THEN T1.id ELSE NULL END) AS REAL) / COUNT(CASE WHEN T2.gender = 'Female' THEN T1.id ELSE NULL END) FROM superhero AS T1 INNER JOIN gender AS T2 ON T1.gender_id = T2.id", "difficulty": "moderate"}
86
+ {"instance_id": "bird_802", "db": "superhero", "question": "Who is the tallest superhero?", "external_knowledge": "superhero.md", "evidence": "who refers to superhero_name; tallest superhero refers to MAX(height_cm);", "SQL": "SELECT superhero_name FROM superhero ORDER BY height_cm DESC LIMIT 1", "difficulty": "simple"}
87
+ {"instance_id": "bird_803", "db": "superhero", "question": "What is the power ID of cryokinesis?", "external_knowledge": "superhero.md", "evidence": "power ID refers to superpower.id; cryokinesis refers to power_name = 'cryokinesis';", "SQL": "SELECT id FROM superpower WHERE power_name = 'Cryokinesis'", "difficulty": "simple"}
88
+ {"instance_id": "bird_804", "db": "superhero", "question": "Provide the name of superhero with superhero ID 294.", "external_knowledge": "superhero.md", "evidence": "name of superhero refers to superhero_name; superhero ID 294 refers to superhero.id = 294;", "SQL": "SELECT superhero_name FROM superhero WHERE id = 294", "difficulty": "simple"}
89
+ {"instance_id": "bird_805", "db": "superhero", "question": "List the full names of superheroes with missing weight.", "external_knowledge": "superhero.md", "evidence": "missing weight refers to weight_kg = 0 OR weight_kg = NULL;", "SQL": "SELECT DISTINCT full_name FROM superhero WHERE full_name IS NOT NULL AND (weight_kg IS NULL OR weight_kg = 0)", "difficulty": "simple"}
90
+ {"instance_id": "bird_806", "db": "superhero", "question": "Provide the eye colour of the superhero who has Karen Beecher-Duncan as their full name.", "external_knowledge": "superhero.md", "evidence": "eye colour refers to colour.colour where eye_colour_id = colour.id; Karen Beecher-Duncan is the full name of superhero;", "SQL": "SELECT T2.colour FROM superhero AS T1 INNER JOIN colour AS T2 ON T1.eye_colour_id = T2.id WHERE T1.full_name = 'Karen Beecher-Duncan'", "difficulty": "simple"}
91
+ {"instance_id": "bird_807", "db": "superhero", "question": "What is the superpowers of the superhero has Helen Parr as their full name?", "external_knowledge": "superhero.md", "evidence": "superpowers refers to power_name; Helen Parr is the full name of superhero;", "SQL": "SELECT T3.power_name FROM superhero AS T1 INNER JOIN hero_power AS T2 ON T1.id = T2.hero_id INNER JOIN superpower AS T3 ON T2.power_id = T3.id WHERE T1.full_name = 'Helen Parr'", "difficulty": "simple"}
92
+ {"instance_id": "bird_808", "db": "superhero", "question": "Find the race of the superhero who weighs 108kg and is 188cm tall.", "external_knowledge": "superhero.md", "evidence": "weighs 108kg refers to weight_kg = 108; 188cm tall refers to height_cm = 188;", "SQL": "SELECT DISTINCT T2.race FROM superhero AS T1 INNER JOIN race AS T2 ON T1.race_id = T2.id WHERE T1.weight_kg = 108 AND T1.height_cm = 188", "difficulty": "simple"}
93
+ {"instance_id": "bird_809", "db": "superhero", "question": "What is the publisher name of the superhero ID 38?", "external_knowledge": "superhero.md", "evidence": "superhero ID 38 refers to superhero.id = 38;", "SQL": "SELECT T2.publisher_name FROM superhero AS T1 INNER JOIN publisher AS T2 ON T1.publisher_id = T2.id WHERE T1.id = 38", "difficulty": "simple"}
94
+ {"instance_id": "bird_810", "db": "superhero", "question": "What is the race of the superhero with maximum attribute value?", "external_knowledge": "superhero.md", "evidence": "maximum attribute value refers to MAX(attribute_value);", "SQL": "SELECT T3.race FROM superhero AS T1 INNER JOIN hero_attribute AS T2 ON T1.id = T2.hero_id INNER JOIN race AS T3 ON T1.race_id = T3.id ORDER BY T2.attribute_value DESC LIMIT 1", "difficulty": "simple"}
95
+ {"instance_id": "bird_811", "db": "superhero", "question": "Give the alignment and superpowers of the superhero named Atom IV.", "external_knowledge": "superhero.md", "evidence": "superpowers refers to power_name;", "SQL": "SELECT T4.alignment, T3.power_name FROM superhero AS T1 INNER JOIN hero_power AS T2 ON T1.id = T2.hero_id INNER JOIN superpower AS T3 ON T3.id = T2.power_id INNER JOIN alignment AS T4 ON T1.alignment_id = T4.id WHERE T1.superhero_name = 'Atom IV'", "difficulty": "simple"}
96
+ {"instance_id": "bird_812", "db": "superhero", "question": "List down at least five full names of superheroes with blue eyes.", "external_knowledge": "superhero.md", "evidence": "blue eyes refers to colour.colour = 'Blue' WHERE eye_colour_id = colour.id; Name of superheroes refers to superhero_name;", "SQL": "SELECT T1.superhero_name FROM superhero AS T1 INNER JOIN colour AS T2 ON T1.eye_colour_id = T2.id WHERE T2.colour = 'Blue' LIMIT 5", "difficulty": "simple"}
97
+ {"instance_id": "bird_813", "db": "superhero", "question": "Calculate the average attribute value of all neutral superheroes.", "external_knowledge": "superhero.md", "evidence": "average = AVG(attribute_value); neutral superheroes refers to alignment_id = 3;", "SQL": "SELECT AVG(T1.attribute_value) FROM hero_attribute AS T1 INNER JOIN superhero AS T2 ON T1.hero_id = T2.id INNER JOIN alignment AS T3 ON T2.alignment_id = T3.id WHERE T3.alignment = 'Neutral'", "difficulty": "simple"}
98
+ {"instance_id": "bird_814", "db": "superhero", "question": "List the skin colour of the superheroes with 100 attribute value.", "external_knowledge": "superhero.md", "evidence": "skin colour refers to colour.colour where skin_colour_id = colour.id; 100 attribute value refers to attribute_value = 100;", "SQL": "SELECT DISTINCT T2.colour FROM superhero AS T1 INNER JOIN colour AS T2 ON T1.skin_colour_id = T2.id INNER JOIN hero_attribute AS T3 ON T1.id = T3.hero_id WHERE T3.attribute_value = 100", "difficulty": "moderate"}
99
+ {"instance_id": "bird_815", "db": "superhero", "question": "Count the good female superheroes.", "external_knowledge": "superhero.md", "evidence": "good refers to alignment.id = 1; female refers to gender.id = 2;", "SQL": "SELECT COUNT(T1.id) FROM superhero AS T1 INNER JOIN alignment AS T2 ON T1.alignment_id = T2.id INNER JOIN gender AS T3 ON T1.gender_id = T3.id WHERE T2.alignment = 'Good' AND T3.gender = 'Female'", "difficulty": "simple"}
100
+ {"instance_id": "bird_816", "db": "superhero", "question": "Provide the names of superheroes with attribute value between 75 to 80.", "external_knowledge": "superhero.md", "evidence": "names of superheroes refers to superhero_name; attribute value between 75 to 80 refers to attribute_value BETWEEN 75 AND 80;", "SQL": "SELECT T1.superhero_name FROM superhero AS T1 INNER JOIN hero_attribute AS T2 ON T1.id = T2.hero_id WHERE T2.attribute_value BETWEEN 75 AND 80", "difficulty": "simple"}
101
+ {"instance_id": "bird_817", "db": "superhero", "question": "Give the race of the blue-haired male superhero.", "external_knowledge": "superhero.md", "evidence": "blue-haired refers to colour.colour = 'blue' WHERE hair_colour_id = colour.id; male refers to gender = 'male';", "SQL": "SELECT T3.race FROM superhero AS T1 INNER JOIN colour AS T2 ON T1.hair_colour_id = T2.id INNER JOIN race AS T3 ON T1.race_id = T3.id INNER JOIN gender AS T4 ON T1.gender_id = T4.id WHERE T2.colour = 'Blue' AND T4.gender = 'Male'", "difficulty": "moderate"}
102
+ {"instance_id": "bird_818", "db": "superhero", "question": "Among the bad superheroes, what is the percentage of female superheroes?", "external_knowledge": "superhero.md", "evidence": "bad superheroes refers to alignment.id = 2; percentage = MULTIPLY(DIVIDE(SUM(gender.id = 2 WHERE alignment.id = 2), COUNT(alignment.id = 2)), 100.0); female refers to gender.id = 2;", "SQL": "SELECT CAST(COUNT(CASE WHEN T3.gender = 'Female' THEN T1.id ELSE NULL END) AS REAL) * 100 / COUNT(T1.id) FROM superhero AS T1 INNER JOIN alignment AS T2 ON T1.alignment_id = T2.id INNER JOIN gender AS T3 ON T1.gender_id = T3.id WHERE T2.alignment = 'Bad'", "difficulty": "challenging"}
103
+ {"instance_id": "bird_819", "db": "superhero", "question": "In superheroes with missing weight data, calculate the difference between the number of superheroes with blue eyes and no eye color.", "external_knowledge": "superhero.md", "evidence": "missing weight data refers to weight_kg = 0 OR T1.weight_kg = NULL; difference = SUBTRACT(SUM(colour.id = 7), SUM(colour.id = 1)); blue eyes refers to eye_colour_id WHERE colour.id = 7; no eye color refers to eye_colour_id WHERE colour.id = 1;", "SQL": "SELECT SUM(CASE WHEN T2.id = 7 THEN 1 ELSE 0 END) - SUM(CASE WHEN T2.id = 1 THEN 1 ELSE 0 END) FROM superhero AS T1 INNER JOIN colour AS T2 ON T1.eye_colour_id = T2.id WHERE T1.weight_kg = 0 OR T1.weight_kg is NULL", "difficulty": "challenging"}
104
+ {"instance_id": "bird_820", "db": "superhero", "question": "How strong is the Hulk?", "external_knowledge": "superhero.md", "evidence": "how strong refers to attribute_value WHERE attribute_name = 'Strength'; the Hulk refers to superhero_name = 'Hulk';", "SQL": "SELECT T2.attribute_value FROM superhero AS T1 INNER JOIN hero_attribute AS T2 ON T1.id = T2.hero_id INNER JOIN attribute AS T3 ON T2.attribute_id = T3.id WHERE T1.superhero_name = 'Hulk' AND T3.attribute_name = 'Strength'", "difficulty": "moderate"}
105
+ {"instance_id": "bird_821", "db": "superhero", "question": "List down Ajax's superpowers.", "external_knowledge": "superhero.md", "evidence": "Ajax refers to superhero_name = 'Ajax'; superpowers refers to power_name;", "SQL": "SELECT T3.power_name FROM superhero AS T1 INNER JOIN hero_power AS T2 ON T1.id = T2.hero_id INNER JOIN superpower AS T3 ON T2.power_id = T3.id WHERE T1.superhero_name = 'Ajax'", "difficulty": "simple"}
106
+ {"instance_id": "bird_822", "db": "superhero", "question": "How many green-skinned villains are there in the superhero universe?", "external_knowledge": "superhero.md", "evidence": "green-skinned refers to colour.colour = 'Green' WHERE skin_colour_id = colour.id; villains refers to alignment = 'Bad';", "SQL": "SELECT COUNT(T1.id) FROM superhero AS T1 INNER JOIN alignment AS T2 ON T1.alignment_id = T2.id INNER JOIN colour AS T3 ON T1.skin_colour_id = T3.id WHERE T2.alignment = 'Bad' AND T3.colour = 'Green'", "difficulty": "moderate"}
107
+ {"instance_id": "bird_823", "db": "superhero", "question": "How many female superheroes are in Marvel Comics?", "external_knowledge": "superhero.md", "evidence": "female refers to gender = 'Female'; Marvel Comics refers to publisher_name = 'Marvel Comics';", "SQL": "SELECT COUNT(T1.id) FROM superhero AS T1 INNER JOIN publisher AS T2 ON T1.publisher_id = T2.id INNER JOIN gender AS T3 ON T1.gender_id = T3.id WHERE T2.publisher_name = 'Marvel Comics' AND T3.gender = 'Female'", "difficulty": "moderate"}
108
+ {"instance_id": "bird_824", "db": "superhero", "question": "Identify superheroes who can control wind and list their names in alphabetical order.", "external_knowledge": "superhero.md", "evidence": "superheroes refers to superhero_name; can control wind refers to power_name = 'Wind Control';", "SQL": "SELECT T1.superhero_name FROM superhero AS T1 INNER JOIN hero_power AS T2 ON T1.id = T2.hero_id INNER JOIN superpower AS T3 ON T2.power_id = T3.id WHERE T3.power_name = 'Wind Control' ORDER BY T1.superhero_name", "difficulty": "moderate"}
109
+ {"instance_id": "bird_825", "db": "superhero", "question": "Identify the gender of the superhero who has the ability of Phoenix Force.", "external_knowledge": "superhero.md", "evidence": "ability of Phoenix Force refers to power_name = 'Phoenix Force';", "SQL": "SELECT T4.gender FROM superhero AS T1 INNER JOIN hero_power AS T2 ON T1.id = T2.hero_id INNER JOIN superpower AS T3 ON T2.power_id = T3.id INNER JOIN gender AS T4 ON T1.gender_id = T4.id WHERE T3.power_name = 'Phoenix Force'", "difficulty": "moderate"}
110
+ {"instance_id": "bird_826", "db": "superhero", "question": "Identify the heaviest superhero in DC Comics.", "external_knowledge": "superhero.md", "evidence": "heaviest refers to MAX(weight_kg); DC Comics refers to publisher_name = 'DC Comics'; superhero refers to superhero_name;", "SQL": "SELECT T1.superhero_name FROM superhero AS T1 INNER JOIN publisher AS T2 ON T1.publisher_id = T2.id WHERE T2.publisher_name = 'DC Comics' ORDER BY T1.weight_kg DESC LIMIT 1", "difficulty": "simple"}
111
+ {"instance_id": "bird_827", "db": "superhero", "question": "What is the average height of a non-human superhero in Dark Horse Comics?", "external_knowledge": "superhero.md", "evidence": "average = AVG(height_cm); non-human superhero refers to race <> 'Human'; Dark Horse Comics refers to publisher_name = 'Dark Horse Comics';", "SQL": "SELECT AVG(T1.height_cm) FROM superhero AS T1 INNER JOIN publisher AS T2 ON T1.publisher_id = T2.id INNER JOIN race AS T3 ON T1.race_id = T3.id WHERE T2.publisher_name = 'Dark Horse Comics' AND T3.race != 'Human'", "difficulty": "moderate"}
112
+ {"instance_id": "bird_828", "db": "superhero", "question": "Count the fastest superheroes.", "external_knowledge": "superhero.md", "evidence": "fastest refers to attribute_value = 100 WHERE attribute_name = 'Speed';", "SQL": "SELECT COUNT(T3.superhero_name) FROM hero_attribute AS T1 INNER JOIN attribute AS T2 ON T1.attribute_id = T2.id INNER JOIN superhero AS T3 ON T1.hero_id = T3.id WHERE T2.attribute_name = 'Speed' AND T1.attribute_value = 100", "difficulty": "simple"}
113
+ {"instance_id": "bird_829", "db": "superhero", "question": "Which publisher created more superheroes: DC or Marvel Comics? Find the difference in the number of superheroes.", "external_knowledge": "superhero.md", "evidence": "DC refers to publisher_name = 'DC Comics'; Marvel Comics refers to publisher_name = 'Marvel Comics'; difference = SUBTRACT(SUM(publisher_name = 'DC Comics'), SUM(publisher_name = 'Marvel Comics'));", "SQL": "SELECT SUM(CASE WHEN T2.publisher_name = 'DC Comics' THEN 1 ELSE 0 END) - SUM(CASE WHEN T2.publisher_name = 'Marvel Comics' THEN 1 ELSE 0 END) FROM superhero AS T1 INNER JOIN publisher AS T2 ON T1.publisher_id = T2.id", "difficulty": "challenging"}
114
+ {"instance_id": "bird_830", "db": "superhero", "question": "Identify the weakest attribute of the Black Panther.", "external_knowledge": "superhero.md", "evidence": "weakest attribute refers to attribute_name WHERE MIN(attribute_value); Black Panther refers to superhero_name = 'Black Panther';", "SQL": "SELECT T3.attribute_name FROM superhero AS T1 INNER JOIN hero_attribute AS T2 ON T1.id = T2.hero_id INNER JOIN attribute AS T3 ON T2.attribute_id = T3.id WHERE T1.superhero_name = 'Black Panther' ORDER BY T2.attribute_value ASC LIMIT 1", "difficulty": "moderate"}
115
+ {"instance_id": "bird_831", "db": "superhero", "question": "What is Abomination's eye colour?", "external_knowledge": "superhero.md", "evidence": "Abomination refers to superhero_name = 'Abomination'; eye colour refers to colour.colour where eye_colour_id = colour.id;", "SQL": "SELECT T2.colour FROM superhero AS T1 INNER JOIN colour AS T2 ON T1.eye_colour_id = T2.id WHERE T1.superhero_name = 'Abomination'", "difficulty": "simple"}
116
+ {"instance_id": "bird_832", "db": "superhero", "question": "Name the tallest superhero.", "external_knowledge": "superhero.md", "evidence": "tallest superhero refers to MAX(height_cm);", "SQL": "SELECT superhero_name FROM superhero ORDER BY height_cm DESC LIMIT 1", "difficulty": "simple"}
117
+ {"instance_id": "bird_833", "db": "superhero", "question": "Name the superhero, otherwise known as Charles Chandler.", "external_knowledge": "superhero.md", "evidence": "name the superhero refers to superhero_name; Charles Chandler is the full name of superhero;", "SQL": "SELECT superhero_name FROM superhero WHERE full_name = 'Charles Chandler'", "difficulty": "simple"}
118
+ {"instance_id": "bird_834", "db": "superhero", "question": "Among all superheroes created by George Lucas, identify the percentage of female superheroes.", "external_knowledge": "superhero.md", "evidence": "created by George Lucas refers to publisher_name = 'George Lucas'; percentage = MULTIPLY(DIVIDE(SUM(gender = 'Female' WHERE publisher_name = 'George Lucas'), COUNT(publisher_name = 'George Lucas')), 100.0); female refers to gender = 'Female';", "SQL": "SELECT CAST(COUNT(CASE WHEN T3.gender = 'Female' THEN 1 ELSE NULL END) AS REAL) * 100 / COUNT(T1.id) FROM superhero AS T1 INNER JOIN publisher AS T2 ON T1.publisher_id = T2.id INNER JOIN gender AS T3 ON T1.gender_id = T3.id WHERE T2.publisher_name = 'George Lucas'", "difficulty": "challenging"}
119
+ {"instance_id": "bird_835", "db": "superhero", "question": "Among all superheroes in Marvel Comics, identify the percentage of 'good' superheroes.", "external_knowledge": "superhero.md", "evidence": "Marvel Comics refers to publisher_name = 'Marvel Comics'; percentage = MULTIPLY(DIVIDE(SUM(alignment = 'Good' WHERE publisher_name = 'Marvel Comics'), COUNT(publisher_name = 'Marvel Comics')), 100.0); good superheroes refers to alignment = 'Good';", "SQL": "SELECT CAST(COUNT(CASE WHEN T3.alignment = 'Good' THEN T1.id ELSE NULL END) AS REAL) * 100 / COUNT(T1.id) FROM superhero AS T1 INNER JOIN publisher AS T2 ON T1.publisher_id = T2.id INNER JOIN alignment AS T3 ON T1.alignment_id = T3.id WHERE T2.publisher_name = 'Marvel Comics'", "difficulty": "challenging"}
120
+ {"instance_id": "bird_836", "db": "superhero", "question": "What is the total number of superheroes that have John as their first name?", "external_knowledge": "superhero.md", "evidence": "have John as their first name refers to full_name LIKE 'John%';", "SQL": "SELECT COUNT(id) FROM superhero WHERE full_name LIKE 'John%'", "difficulty": "simple"}
121
+ {"instance_id": "bird_837", "db": "superhero", "question": "Give the hero ID of superhero with the lowest attribute value.", "external_knowledge": "superhero.md", "evidence": "lowest attribute value refers to MIN(attribute_value);", "SQL": "SELECT hero_id FROM hero_attribute WHERE attribute_value = ( SELECT MIN(attribute_value) FROM hero_attribute )", "difficulty": "simple"}
122
+ {"instance_id": "bird_838", "db": "superhero", "question": "Provide the full name of the superhero named Alien.", "external_knowledge": "superhero.md", "evidence": "", "SQL": "SELECT full_name FROM superhero WHERE superhero_name = 'Alien'", "difficulty": "simple"}
123
+ {"instance_id": "bird_839", "db": "superhero", "question": "In superheroes with weight less than 100, list the full name of the superheroes with brown eyes.", "external_knowledge": "superhero.md", "evidence": "weight less than 100 refers to weight_kg < 100", "SQL": "SELECT T1.full_name FROM superhero AS T1 INNER JOIN colour AS T2 ON T1.eye_colour_id = T2.id WHERE T1.weight_kg < 100 AND T2.colour = 'Brown'", "difficulty": "simple"}
124
+ {"instance_id": "bird_840", "db": "superhero", "question": "List the attribute value of the superhero named Aquababy.", "external_knowledge": "superhero.md", "evidence": "", "SQL": "SELECT T2.attribute_value FROM superhero AS T1 INNER JOIN hero_attribute AS T2 ON T1.id = T2.hero_id WHERE T1.superhero_name = 'Aquababy'", "difficulty": "simple"}
125
+ {"instance_id": "bird_841", "db": "superhero", "question": "Provide the weight and race of the superhero with superhero ID 40.", "external_knowledge": "superhero.md", "evidence": "weight refers to weight_kg; superhero ID 40 refers to superhero.id = 40;", "SQL": "SELECT T1.weight_kg, T2.race FROM superhero AS T1 INNER JOIN race AS T2 ON T1.race_id = T2.id WHERE T1.id = 40", "difficulty": "simple"}
126
+ {"instance_id": "bird_842", "db": "superhero", "question": "Calculate the average height of all neutral superheroes.", "external_knowledge": "superhero.md", "evidence": "", "SQL": "SELECT AVG(T1.height_cm) FROM superhero AS T1 INNER JOIN alignment AS T2 ON T1.alignment_id = T2.id WHERE T2.alignment = 'Neutral'", "difficulty": "simple"}
127
+ {"instance_id": "bird_843", "db": "superhero", "question": "List the hero ID of superheroes have intellegence as their power.", "external_knowledge": "superhero.md", "evidence": "hero ID refers to superhero.id; have intelligence as their power refers to power_name = 'Intelligence';", "SQL": "SELECT T1.hero_id FROM hero_power AS T1 INNER JOIN superpower AS T2 ON T1.power_id = T2.id WHERE T2.power_name = 'Intelligence'", "difficulty": "simple"}
128
+ {"instance_id": "bird_844", "db": "superhero", "question": "Give the eye colour of Blackwulf.", "external_knowledge": "superhero.md", "evidence": "eye colour refers to colour.colour where eye_colour_id = colour.id; Blackwulf refers to superhero_name = 'Blackwulf';", "SQL": "SELECT T2.colour FROM superhero AS T1 INNER JOIN colour AS T2 ON T1.eye_colour_id = T2.id WHERE T1.superhero_name = 'Blackwulf'", "difficulty": "simple"}
129
+ {"instance_id": "bird_845", "db": "superhero", "question": "List the power of superheroes with height greater than 80% of the average height of all superheroes.", "external_knowledge": "superhero.md", "evidence": "power of superheroes refers to power_name; height greater than 80% of the average height of all superheroes = height_cm > MULTIPLY(AVG(height_cm), 0.8);", "SQL": "SELECT T3.power_name FROM superhero AS T1 INNER JOIN hero_power AS T2 ON T1.id = T2.hero_id INNER JOIN superpower AS T3 ON T2.power_id = T3.id WHERE T1.height_cm * 100 > ( SELECT AVG(height_cm) FROM superhero ) * 80", "difficulty": "moderate"}
data/text2sql/bird/databases/superhero/superhero.sqlite ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:75e94a2c3236ee3bb2c01fb97a1c4b4c1c269bcefd4eab1d04be323d2d0825b1
3
+ size 237568
data/text2sql/bird/databases/superhero/superhero_db_summary.json ADDED
The diff for this file is too large to render. See raw diff
 
data/text2sql/spider2-lite/bigquery/austin/austin_bigquery_summary.json ADDED
@@ -0,0 +1,39 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "tables": {
3
+ "bigquery-public-data.austin_311.311_service_requests": {
4
+ "name": "bigquery-public-data.austin_311.311_service_requests",
5
+ "columns": [
6
+ {"name": "unique_key", "type": "STRING", "is_primary_key": true, "pk_origin": "heuristic", "is_foreign_key": false, "description": "The service request tracking number."},
7
+ {"name": "complaint_description", "type": "STRING", "is_primary_key": false, "is_foreign_key": false, "description": "Service request type"},
8
+ {"name": "source", "type": "STRING", "is_primary_key": false, "is_foreign_key": false, "description": "Contact method Service request was received from."},
9
+ {"name": "status", "type": "STRING", "is_primary_key": false, "is_foreign_key": false, "description": "Service request status."},
10
+ {"name": "status_change_date", "type": "TIMESTAMP", "is_primary_key": false, "is_foreign_key": false, "description": "Date of last Service request status change."},
11
+ {"name": "created_date", "type": "TIMESTAMP", "is_primary_key": false, "is_foreign_key": false, "description": "Date Service request was created."},
12
+ {"name": "last_update_date", "type": "TIMESTAMP", "is_primary_key": false, "is_foreign_key": false, "description": "Date Service request was updated."},
13
+ {"name": "close_date", "type": "TIMESTAMP", "is_primary_key": false, "is_foreign_key": false, "description": "Date Service request was closed."},
14
+ {"name": "incident_address", "type": "STRING", "is_primary_key": false, "is_foreign_key": false, "description": "Service location of Service request."},
15
+ {"name": "street_number", "type": "STRING", "is_primary_key": false, "is_foreign_key": false, "description": "Parsed location information. Street number."},
16
+ {"name": "street_name", "type": "STRING", "is_primary_key": false, "is_foreign_key": false, "description": "Parsed location information. Street name."},
17
+ {"name": "city", "type": "STRING", "is_primary_key": false, "is_foreign_key": false, "description": "Parsed location information. City."},
18
+ {"name": "incident_zip", "type": "INT64", "is_primary_key": false, "is_foreign_key": false, "description": "Parsed location information. Zip code."},
19
+ {"name": "county", "type": "STRING", "is_primary_key": false, "is_foreign_key": false, "description": "Parsed location information. County."},
20
+ {"name": "state_plane_x_coordinate", "type": "STRING", "is_primary_key": false, "is_foreign_key": false, "description": "State plane X coordinate."},
21
+ {"name": "state_plane_y_coordinate", "type": "FLOAT64", "is_primary_key": false, "is_foreign_key": false, "description": "State plane Y coordinate."},
22
+ {"name": "latitude", "type": "FLOAT64", "is_primary_key": false, "is_foreign_key": false, "description": "Service request location latitude coordinate."},
23
+ {"name": "longitude", "type": "FLOAT64", "is_primary_key": false, "is_foreign_key": false, "description": "Service request location longitude coordinate."},
24
+ {"name": "location", "type": "STRING", "is_primary_key": false, "is_foreign_key": false, "description": "Service request location latitude and longitude coordinates."},
25
+ {"name": "council_district_code", "type": "INT64", "is_primary_key": false, "is_foreign_key": false, "description": "Council district corresponding to Service request location."},
26
+ {"name": "map_page", "type": "STRING", "is_primary_key": false, "is_foreign_key": false, "description": "Service request location corresponding map page."},
27
+ {"name": "map_tile", "type": "STRING", "is_primary_key": false, "is_foreign_key": false, "description": "Service request location corresponding map tile."}
28
+ ],
29
+ "samples": {
30
+ "unique_key": ["18-00004837", "18-00004838"],
31
+ "complaint_description": ["Loose Dog", "Animal Control - Assistance Request"],
32
+ "status": ["Closed", "Closed"],
33
+ "city": ["Austin", "Austin"]
34
+ },
35
+ "foreign_keys": []
36
+ }
37
+ },
38
+ "relationships": []
39
+ }
data/text2sql/spider2-lite/golden_lite_spider_total.json ADDED
@@ -0,0 +1,1443 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "local003": {
3
+ "instance_id": "local003",
4
+ "schema_extraction": {
5
+ "tables": [
6
+ "customers",
7
+ "order_items",
8
+ "orders"
9
+ ],
10
+ "columns": [
11
+ "customers.customer_unique_id",
12
+ "customers.customer_id",
13
+ "order_items.price",
14
+ "order_items.order_id",
15
+ "orders.order_purchase_timestamp",
16
+ "orders.customer_id",
17
+ "orders.order_status",
18
+ "orders.order_id"
19
+ ]
20
+ }
21
+ },
22
+ "local004": {
23
+ "instance_id": "local004",
24
+ "schema_extraction": {
25
+ "tables": [
26
+ "customers",
27
+ "order_payments",
28
+ "orders"
29
+ ],
30
+ "columns": [
31
+ "customers.customer_unique_id",
32
+ "customers.customer_id",
33
+ "order_payments.payment_value",
34
+ "order_payments.order_id",
35
+ "orders.customer_id",
36
+ "orders.order_id",
37
+ "orders.order_purchase_timestamp"
38
+ ]
39
+ }
40
+ },
41
+ "local008": {
42
+ "instance_id": "local008",
43
+ "schema_extraction": {
44
+ "tables": [
45
+ "batting",
46
+ "player"
47
+ ],
48
+ "columns": [
49
+ "batting.player_id",
50
+ "batting.g",
51
+ "batting.r",
52
+ "batting.h",
53
+ "batting.hr",
54
+ "player.name_given",
55
+ "player.player_id"
56
+ ]
57
+ }
58
+ },
59
+ "local009": {
60
+ "instance_id": "local009",
61
+ "schema_extraction": {
62
+ "tables": [
63
+ "airports_data",
64
+ "flights"
65
+ ],
66
+ "columns": [
67
+ "airports_data.city",
68
+ "airports_data.airport_code",
69
+ "airports_data.coordinates",
70
+ "flights.flight_id",
71
+ "flights.departure_airport",
72
+ "flights.arrival_airport"
73
+ ]
74
+ }
75
+ },
76
+ "local010": {
77
+ "instance_id": "local010",
78
+ "schema_extraction": {
79
+ "tables": [
80
+ "airports_data",
81
+ "flights"
82
+ ],
83
+ "columns": [
84
+ "airports_data.city",
85
+ "airports_data.airport_code",
86
+ "airports_data.coordinates",
87
+ "flights.flight_id",
88
+ "flights.departure_airport",
89
+ "flights.arrival_airport"
90
+ ]
91
+ }
92
+ },
93
+ "local015": {
94
+ "instance_id": "local015",
95
+ "schema_extraction": {
96
+ "tables": [
97
+ "collisions",
98
+ "parties"
99
+ ],
100
+ "columns": [
101
+ "collisions.case_id",
102
+ "collisions.motorcyclist_killed_count",
103
+ "collisions.motorcycle_collision",
104
+ "parties.party_safety_equipment_1",
105
+ "parties.party_safety_equipment_2",
106
+ "parties.case_id",
107
+ "parties.party_age"
108
+ ]
109
+ }
110
+ },
111
+ "local017": {
112
+ "instance_id": "local017",
113
+ "schema_extraction": {
114
+ "tables": [
115
+ "collisions"
116
+ ],
117
+ "columns": [
118
+ "collisions.pcf_violation_category",
119
+ "collisions.case_id",
120
+ "collisions.collision_date"
121
+ ]
122
+ }
123
+ },
124
+ "local019": {
125
+ "instance_id": "local019",
126
+ "schema_extraction": {
127
+ "tables": [
128
+ "Promotions",
129
+ "Cards",
130
+ "Locations",
131
+ "Events",
132
+ "Matches",
133
+ "Belts",
134
+ "Wrestlers"
135
+ ],
136
+ "columns": [
137
+ "Promotions.id",
138
+ "Promotions.name",
139
+ "Cards.id",
140
+ "Cards.location_id",
141
+ "Cards.event_id",
142
+ "Cards.promotion_id",
143
+ "Locations.name",
144
+ "Locations.id",
145
+ "Events.name",
146
+ "Events.id",
147
+ "Matches.duration",
148
+ "Matches.win_type",
149
+ "Matches.title_id",
150
+ "Matches.winner_id",
151
+ "Matches.loser_id",
152
+ "Matches.card_id",
153
+ "Belts.name",
154
+ "Belts.id",
155
+ "Wrestlers.name",
156
+ "Wrestlers.id"
157
+ ]
158
+ }
159
+ },
160
+ "local026": {
161
+ "instance_id": "local026",
162
+ "schema_extraction": {
163
+ "tables": [
164
+ "player",
165
+ "ball_by_ball",
166
+ "batsman_scored",
167
+ "extra_runs"
168
+ ],
169
+ "columns": [
170
+ "player.player_name",
171
+ "player.player_id",
172
+ "ball_by_ball.match_id",
173
+ "ball_by_ball.bowler",
174
+ "ball_by_ball.over_id",
175
+ "ball_by_ball.innings_no",
176
+ "batsman_scored.match_id",
177
+ "batsman_scored.over_id",
178
+ "batsman_scored.ball_id",
179
+ "batsman_scored.innings_no",
180
+ "batsman_scored.runs_scored",
181
+ "extra_runs.match_id",
182
+ "extra_runs.over_id",
183
+ "extra_runs.ball_id",
184
+ "extra_runs.innings_no",
185
+ "extra_runs.extra_runs"
186
+ ]
187
+ }
188
+ },
189
+ "local021": {
190
+ "instance_id": "local021",
191
+ "schema_extraction": {
192
+ "tables": [
193
+ "ball_by_ball",
194
+ "batsman_scored"
195
+ ],
196
+ "columns": [
197
+ "ball_by_ball.match_id",
198
+ "ball_by_ball.striker",
199
+ "ball_by_ball.innings_no",
200
+ "ball_by_ball.ball_id",
201
+ "ball_by_ball.over_id",
202
+ "batsman_scored.runs_scored",
203
+ "batsman_scored.innings_no",
204
+ "batsman_scored.ball_id",
205
+ "batsman_scored.match_id",
206
+ "batsman_scored.over_id"
207
+ ]
208
+ }
209
+ },
210
+ "local022": {
211
+ "instance_id": "local022",
212
+ "schema_extraction": {
213
+ "tables": [
214
+ "player",
215
+ "match",
216
+ "player_match",
217
+ "ball_by_ball",
218
+ "batsman_scored"
219
+ ],
220
+ "columns": [
221
+ "player.player_name",
222
+ "player.player_id",
223
+ "match.match_id",
224
+ "match.team_1",
225
+ "match.team_2",
226
+ "match.match_winner",
227
+ "player_match.team_id",
228
+ "player_match.player_id",
229
+ "player_match.match_id",
230
+ "ball_by_ball.striker",
231
+ "ball_by_ball.match_id",
232
+ "ball_by_ball.innings_no",
233
+ "ball_by_ball.ball_id",
234
+ "ball_by_ball.over_id",
235
+ "batsman_scored.runs_scored",
236
+ "batsman_scored.innings_no",
237
+ "batsman_scored.ball_id",
238
+ "batsman_scored.match_id",
239
+ "batsman_scored.over_id"
240
+ ]
241
+ }
242
+ },
243
+ "local023": {
244
+ "instance_id": "local023",
245
+ "schema_extraction": {
246
+ "tables": [
247
+ "player",
248
+ "match",
249
+ "ball_by_ball",
250
+ "batsman_scored"
251
+ ],
252
+ "columns": [
253
+ "player.player_name",
254
+ "player.player_id",
255
+ "match.match_id",
256
+ "match.season_id",
257
+ "ball_by_ball.striker",
258
+ "ball_by_ball.match_id",
259
+ "ball_by_ball.innings_no",
260
+ "ball_by_ball.ball_id",
261
+ "ball_by_ball.over_id",
262
+ "batsman_scored.runs_scored",
263
+ "batsman_scored.innings_no",
264
+ "batsman_scored.ball_id",
265
+ "batsman_scored.match_id",
266
+ "batsman_scored.over_id"
267
+ ]
268
+ }
269
+ },
270
+ "local024": {
271
+ "instance_id": "local024",
272
+ "schema_extraction": {
273
+ "tables": [
274
+ "player",
275
+ "ball_by_ball",
276
+ "batsman_scored"
277
+ ],
278
+ "columns": [
279
+ "player.country_name",
280
+ "player.player_id",
281
+ "ball_by_ball.striker",
282
+ "ball_by_ball.match_id",
283
+ "ball_by_ball.innings_no",
284
+ "ball_by_ball.ball_id",
285
+ "ball_by_ball.over_id",
286
+ "batsman_scored.runs_scored",
287
+ "batsman_scored.innings_no",
288
+ "batsman_scored.ball_id",
289
+ "batsman_scored.match_id",
290
+ "batsman_scored.over_id"
291
+ ]
292
+ }
293
+ },
294
+ "local028": {
295
+ "instance_id": "local028",
296
+ "schema_extraction": {
297
+ "tables": [
298
+ "olist_orders"
299
+ ],
300
+ "columns": [
301
+ "olist_orders.order_delivered_customer_date",
302
+ "olist_orders.order_status"
303
+ ]
304
+ }
305
+ },
306
+ "local031": {
307
+ "instance_id": "local031",
308
+ "schema_extraction": {
309
+ "tables": [
310
+ "olist_orders"
311
+ ],
312
+ "columns": [
313
+ "olist_orders.order_status",
314
+ "olist_orders.order_delivered_customer_date"
315
+ ]
316
+ }
317
+ },
318
+ "local029": {
319
+ "instance_id": "local029",
320
+ "schema_extraction": {
321
+ "tables": [
322
+ "olist_customers",
323
+ "olist_orders",
324
+ "olist_order_payments"
325
+ ],
326
+ "columns": [
327
+ "olist_customers.customer_unique_id",
328
+ "olist_customers.customer_city",
329
+ "olist_customers.customer_state",
330
+ "olist_customers.customer_id",
331
+ "olist_orders.order_id",
332
+ "olist_orders.customer_id",
333
+ "olist_orders.order_status",
334
+ "olist_order_payments.payment_value",
335
+ "olist_order_payments.order_id"
336
+ ]
337
+ }
338
+ },
339
+ "local030": {
340
+ "instance_id": "local030",
341
+ "schema_extraction": {
342
+ "tables": [
343
+ "olist_customers",
344
+ "olist_orders",
345
+ "olist_order_payments"
346
+ ],
347
+ "columns": [
348
+ "olist_customers.customer_city",
349
+ "olist_customers.customer_id",
350
+ "olist_orders.customer_id",
351
+ "olist_orders.order_id",
352
+ "olist_orders.order_status",
353
+ "olist_order_payments.payment_value",
354
+ "olist_order_payments.order_id"
355
+ ]
356
+ }
357
+ },
358
+ "local034": {
359
+ "instance_id": "local034",
360
+ "schema_extraction": {
361
+ "tables": [
362
+ "olist_order_items",
363
+ "olist_products",
364
+ "olist_order_payments"
365
+ ],
366
+ "columns": [
367
+ "olist_order_items.product_id",
368
+ "olist_order_items.order_id",
369
+ "olist_products.product_category_name",
370
+ "olist_products.product_id",
371
+ "olist_order_payments.payment_type",
372
+ "olist_order_payments.order_id"
373
+ ]
374
+ }
375
+ },
376
+ "local038": {
377
+ "instance_id": "local038",
378
+ "schema_extraction": {
379
+ "tables": [
380
+ "actor",
381
+ "language",
382
+ "category",
383
+ "film",
384
+ "film_actor",
385
+ "film_category"
386
+ ],
387
+ "columns": [
388
+ "actor.actor_id",
389
+ "actor.first_name",
390
+ "actor.last_name",
391
+ "language.language_id",
392
+ "language.name",
393
+ "category.category_id",
394
+ "category.name",
395
+ "film.film_id",
396
+ "film.language_id",
397
+ "film.length",
398
+ "film.rating",
399
+ "film.release_year",
400
+ "film_actor.actor_id",
401
+ "film_actor.film_id",
402
+ "film_category.film_id",
403
+ "film_category.category_id"
404
+ ]
405
+ }
406
+ },
407
+ "local039": {
408
+ "instance_id": "local039",
409
+ "schema_extraction": {
410
+ "tables": [
411
+ "city",
412
+ "address",
413
+ "category",
414
+ "customer",
415
+ "film_category",
416
+ "inventory",
417
+ "rental"
418
+ ],
419
+ "columns": [
420
+ "city.city_id",
421
+ "city.city",
422
+ "address.address_id",
423
+ "address.city_id",
424
+ "category.name",
425
+ "category.category_id",
426
+ "customer.customer_id",
427
+ "customer.address_id",
428
+ "film_category.film_id",
429
+ "film_category.category_id",
430
+ "inventory.inventory_id",
431
+ "inventory.film_id",
432
+ "rental.customer_id",
433
+ "rental.inventory_id",
434
+ "rental.return_date",
435
+ "rental.rental_date"
436
+ ]
437
+ }
438
+ },
439
+ "local041": {
440
+ "instance_id": "local041",
441
+ "schema_extraction": {
442
+ "tables": [
443
+ "trees"
444
+ ],
445
+ "columns": [
446
+ "trees.boroname",
447
+ "trees.health"
448
+ ]
449
+ }
450
+ },
451
+ "local049": {
452
+ "instance_id": "local049",
453
+ "schema_extraction": {
454
+ "tables": [
455
+ "companies_dates",
456
+ "companies_industries"
457
+ ],
458
+ "columns": [
459
+ "companies_dates.company_id",
460
+ "companies_dates.date_joined",
461
+ "companies_industries.industry",
462
+ "companies_industries.company_id"
463
+ ]
464
+ }
465
+ },
466
+ "local054": {
467
+ "instance_id": "local054",
468
+ "schema_extraction": {
469
+ "tables": [
470
+ "albums",
471
+ "artists",
472
+ "customers",
473
+ "invoices",
474
+ "invoice_items",
475
+ "tracks"
476
+ ],
477
+ "columns": [
478
+ "albums.ArtistId",
479
+ "albums.AlbumId",
480
+ "artists.Name",
481
+ "artists.ArtistId",
482
+ "customers.CustomerId",
483
+ "customers.FirstName",
484
+ "invoices.InvoiceId",
485
+ "invoices.CustomerId",
486
+ "invoice_items.TrackId",
487
+ "invoice_items.InvoiceId",
488
+ "invoice_items.UnitPrice",
489
+ "invoice_items.Quantity",
490
+ "tracks.AlbumId",
491
+ "tracks.TrackId"
492
+ ]
493
+ }
494
+ },
495
+ "local055": {
496
+ "instance_id": "local055",
497
+ "schema_extraction": {
498
+ "tables": [
499
+ "albums",
500
+ "artists",
501
+ "customers",
502
+ "invoices",
503
+ "invoice_items",
504
+ "tracks"
505
+ ],
506
+ "columns": [
507
+ "albums.AlbumId",
508
+ "albums.ArtistId",
509
+ "artists.ArtistId",
510
+ "artists.Name",
511
+ "customers.CustomerId",
512
+ "invoices.InvoiceId",
513
+ "invoices.CustomerId",
514
+ "invoice_items.TrackId",
515
+ "invoice_items.InvoiceId",
516
+ "invoice_items.Quantity",
517
+ "invoice_items.UnitPrice",
518
+ "tracks.TrackId",
519
+ "tracks.AlbumId"
520
+ ]
521
+ }
522
+ },
523
+ "local198": {
524
+ "instance_id": "local198",
525
+ "schema_extraction": {
526
+ "tables": [
527
+ "customers",
528
+ "invoices"
529
+ ],
530
+ "columns": [
531
+ "customers.Country",
532
+ "invoices.BillingCountry",
533
+ "invoices.Total"
534
+ ]
535
+ }
536
+ },
537
+ "local056": {
538
+ "instance_id": "local056",
539
+ "schema_extraction": {
540
+ "tables": [
541
+ "customer",
542
+ "payment"
543
+ ],
544
+ "columns": [
545
+ "customer.last_name",
546
+ "customer.customer_id",
547
+ "customer.first_name",
548
+ "payment.amount",
549
+ "payment.customer_id",
550
+ "payment.payment_date"
551
+ ]
552
+ }
553
+ },
554
+ "local058": {
555
+ "instance_id": "local058",
556
+ "schema_extraction": {
557
+ "tables": [
558
+ "hardware_dim_product",
559
+ "hardware_fact_sales_monthly"
560
+ ],
561
+ "columns": [
562
+ "hardware_dim_product.segment",
563
+ "hardware_dim_product.product_code",
564
+ "hardware_fact_sales_monthly.fiscal_year",
565
+ "hardware_fact_sales_monthly.product_code"
566
+ ]
567
+ }
568
+ },
569
+ "local059": {
570
+ "instance_id": "local059",
571
+ "schema_extraction": {
572
+ "tables": [
573
+ "hardware_dim_product",
574
+ "hardware_fact_sales_monthly"
575
+ ],
576
+ "columns": [
577
+ "hardware_dim_product.product_code",
578
+ "hardware_dim_product.division",
579
+ "hardware_fact_sales_monthly.sold_quantity",
580
+ "hardware_fact_sales_monthly.product_code",
581
+ "hardware_fact_sales_monthly.date"
582
+ ]
583
+ }
584
+ },
585
+ "local071": {
586
+ "instance_id": "local071",
587
+ "schema_extraction": {
588
+ "tables": [
589
+ "cities"
590
+ ],
591
+ "columns": [
592
+ "cities.country_code_2",
593
+ "cities.insert_date"
594
+ ]
595
+ }
596
+ },
597
+ "local066": {
598
+ "instance_id": "local066",
599
+ "schema_extraction": {
600
+ "tables": [
601
+ "pizza_toppings",
602
+ "pizza_recipes"
603
+ ],
604
+ "columns": [
605
+ "pizza_toppings.topping_name",
606
+ "pizza_toppings.topping_id",
607
+ "pizza_recipes.pizza_id",
608
+ "pizza_recipes.toppings"
609
+ ]
610
+ }
611
+ },
612
+ "local065": {
613
+ "instance_id": "local065",
614
+ "schema_extraction": {
615
+ "tables": [
616
+ "pizza_clean_runner_orders",
617
+ "pizza_clean_customer_orders"
618
+ ],
619
+ "columns": [
620
+ "pizza_clean_runner_orders.order_id",
621
+ "pizza_clean_runner_orders.cancellation",
622
+ "pizza_clean_customer_orders.order_id",
623
+ "pizza_clean_customer_orders.pizza_id",
624
+ "pizza_clean_customer_orders.extras"
625
+ ]
626
+ }
627
+ },
628
+ "local064": {
629
+ "instance_id": "local064",
630
+ "schema_extraction": {
631
+ "tables": [
632
+ "customer_transactions"
633
+ ],
634
+ "columns": [
635
+ "customer_transactions.customer_id",
636
+ "customer_transactions.txn_date",
637
+ "customer_transactions.txn_amount",
638
+ "customer_transactions.txn_type"
639
+ ]
640
+ }
641
+ },
642
+ "local299": {
643
+ "instance_id": "local299",
644
+ "schema_extraction": {
645
+ "tables": [
646
+ "customer_transactions"
647
+ ],
648
+ "columns": [
649
+ "customer_transactions.customer_id",
650
+ "customer_transactions.txn_date",
651
+ "customer_transactions.txn_amount",
652
+ "customer_transactions.txn_type"
653
+ ]
654
+ }
655
+ },
656
+ "local300": {
657
+ "instance_id": "local300",
658
+ "schema_extraction": {
659
+ "tables": [
660
+ "customer_transactions"
661
+ ],
662
+ "columns": [
663
+ "customer_transactions.customer_id",
664
+ "customer_transactions.txn_date",
665
+ "customer_transactions.txn_amount",
666
+ "customer_transactions.txn_type"
667
+ ]
668
+ }
669
+ },
670
+ "local075": {
671
+ "instance_id": "local075",
672
+ "schema_extraction": {
673
+ "tables": [
674
+ "shopping_cart_events",
675
+ "shopping_cart_page_hierarchy"
676
+ ],
677
+ "columns": [
678
+ "shopping_cart_events.event_type",
679
+ "shopping_cart_events.visit_id",
680
+ "shopping_cart_page_hierarchy.page_id",
681
+ "shopping_cart_page_hierarchy.page_name",
682
+ "shopping_cart_page_hierarchy.product_id"
683
+ ]
684
+ }
685
+ },
686
+ "local078": {
687
+ "instance_id": "local078",
688
+ "schema_extraction": {
689
+ "tables": [
690
+ "interest_metrics",
691
+ "interest_map"
692
+ ],
693
+ "columns": [
694
+ "interest_metrics.interest_id",
695
+ "interest_metrics.month_year",
696
+ "interest_metrics.composition",
697
+ "interest_map.interest_name",
698
+ "interest_map.id"
699
+ ]
700
+ }
701
+ },
702
+ "local085": {
703
+ "instance_id": "local085",
704
+ "schema_extraction": {
705
+ "tables": [
706
+ "orders"
707
+ ],
708
+ "columns": [
709
+ "orders.employeeid",
710
+ "orders.orderid",
711
+ "orders.shippeddate",
712
+ "orders.requireddate"
713
+ ]
714
+ }
715
+ },
716
+ "local099": {
717
+ "instance_id": "local099",
718
+ "schema_extraction": {
719
+ "tables": [
720
+ "Person",
721
+ "M_Director",
722
+ "M_Cast"
723
+ ],
724
+ "columns": [
725
+ "Person.PID",
726
+ "Person.Name",
727
+ "M_Director.PID",
728
+ "M_Director.MID",
729
+ "M_Cast.PID",
730
+ "M_Cast.MID"
731
+ ]
732
+ }
733
+ },
734
+ "local131": {
735
+ "instance_id": "local131",
736
+ "schema_extraction": {
737
+ "tables": [
738
+ "Musical_Preferences",
739
+ "Musical_Styles"
740
+ ],
741
+ "columns": [
742
+ "Musical_Preferences.StyleID",
743
+ "Musical_Preferences.PreferenceSeq",
744
+ "Musical_Styles.StyleName",
745
+ "Musical_Styles.StyleID"
746
+ ]
747
+ }
748
+ },
749
+ "local132": {
750
+ "instance_id": "local132",
751
+ "schema_extraction": {
752
+ "tables": [
753
+ "Customers",
754
+ "Entertainer_Styles",
755
+ "Entertainers",
756
+ "Musical_Preferences"
757
+ ],
758
+ "columns": [
759
+ "Customers.CustLastName",
760
+ "Customers.CustomerID",
761
+ "Entertainer_Styles.EntertainerID",
762
+ "Entertainer_Styles.StyleID",
763
+ "Entertainer_Styles.StyleStrength",
764
+ "Entertainers.EntStageName",
765
+ "Entertainers.EntertainerID",
766
+ "Musical_Preferences.CustomerID",
767
+ "Musical_Preferences.StyleID",
768
+ "Musical_Preferences.PreferenceSeq"
769
+ ]
770
+ }
771
+ },
772
+ "local157": {
773
+ "instance_id": "local157",
774
+ "schema_extraction": {
775
+ "tables": [
776
+ "bitcoin_prices"
777
+ ],
778
+ "columns": [
779
+ "bitcoin_prices.ticker",
780
+ "bitcoin_prices.market_date",
781
+ "bitcoin_prices.volume"
782
+ ]
783
+ }
784
+ },
785
+ "local163": {
786
+ "instance_id": "local163",
787
+ "schema_extraction": {
788
+ "tables": [
789
+ "university_faculty"
790
+ ],
791
+ "columns": [
792
+ "university_faculty.FacRank",
793
+ "university_faculty.FacFirstName",
794
+ "university_faculty.FacLastName",
795
+ "university_faculty.FacSalary"
796
+ ]
797
+ }
798
+ },
799
+ "local193": {
800
+ "instance_id": "local193",
801
+ "schema_extraction": {
802
+ "tables": [
803
+ "payment"
804
+ ],
805
+ "columns": [
806
+ "payment.customer_id",
807
+ "payment.payment_date",
808
+ "payment.amount"
809
+ ]
810
+ }
811
+ },
812
+ "local194": {
813
+ "instance_id": "local194",
814
+ "schema_extraction": {
815
+ "tables": [
816
+ "film",
817
+ "film_actor",
818
+ "inventory",
819
+ "payment",
820
+ "rental"
821
+ ],
822
+ "columns": [
823
+ "film.film_id",
824
+ "film.title",
825
+ "film_actor.actor_id",
826
+ "film_actor.film_id",
827
+ "inventory.film_id",
828
+ "inventory.inventory_id",
829
+ "payment.amount",
830
+ "payment.rental_id",
831
+ "rental.rental_id",
832
+ "rental.inventory_id"
833
+ ]
834
+ }
835
+ },
836
+ "local195": {
837
+ "instance_id": "local195",
838
+ "schema_extraction": {
839
+ "tables": [
840
+ "actor",
841
+ "customer",
842
+ "film",
843
+ "film_actor",
844
+ "inventory",
845
+ "payment",
846
+ "rental"
847
+ ],
848
+ "columns": [
849
+ "actor.actor_id",
850
+ "actor.first_name",
851
+ "actor.last_name",
852
+ "customer.customer_id",
853
+ "film.film_id",
854
+ "film.title",
855
+ "film_actor.actor_id",
856
+ "film_actor.film_id",
857
+ "inventory.film_id",
858
+ "inventory.inventory_id",
859
+ "payment.amount",
860
+ "payment.rental_id",
861
+ "payment.customer_id",
862
+ "rental.inventory_id",
863
+ "rental.rental_id"
864
+ ]
865
+ }
866
+ },
867
+ "local196": {
868
+ "instance_id": "local196",
869
+ "schema_extraction": {
870
+ "tables": [
871
+ "inventory",
872
+ "payment",
873
+ "rental"
874
+ ],
875
+ "columns": [
876
+ "inventory.inventory_id",
877
+ "inventory.film_id",
878
+ "payment.customer_id",
879
+ "payment.amount",
880
+ "rental.customer_id",
881
+ "rental.inventory_id",
882
+ "rental.rental_date"
883
+ ]
884
+ }
885
+ },
886
+ "local197": {
887
+ "instance_id": "local197",
888
+ "schema_extraction": {
889
+ "tables": [
890
+ "payment"
891
+ ],
892
+ "columns": [
893
+ "payment.customer_id",
894
+ "payment.amount",
895
+ "payment.payment_date"
896
+ ]
897
+ }
898
+ },
899
+ "local199": {
900
+ "instance_id": "local199",
901
+ "schema_extraction": {
902
+ "tables": [
903
+ "staff",
904
+ "rental"
905
+ ],
906
+ "columns": [
907
+ "staff.store_id",
908
+ "staff.staff_id",
909
+ "rental.rental_id",
910
+ "rental.staff_id",
911
+ "rental.rental_date"
912
+ ]
913
+ }
914
+ },
915
+ "local202": {
916
+ "instance_id": "local202",
917
+ "schema_extraction": {
918
+ "tables": [
919
+ "alien_data"
920
+ ],
921
+ "columns": [
922
+ "alien_data.state",
923
+ "alien_data.age",
924
+ "alien_data.aggressive"
925
+ ]
926
+ }
927
+ },
928
+ "local209": {
929
+ "instance_id": "local209",
930
+ "schema_extraction": {
931
+ "tables": [
932
+ "deliveries",
933
+ "stores",
934
+ "orders"
935
+ ],
936
+ "columns": [
937
+ "deliveries.delivery_order_id",
938
+ "deliveries.delivery_status",
939
+ "stores.store_name",
940
+ "stores.store_id",
941
+ "orders.order_id",
942
+ "orders.store_id",
943
+ "orders.delivery_order_id"
944
+ ]
945
+ }
946
+ },
947
+ "local210": {
948
+ "instance_id": "local210",
949
+ "schema_extraction": {
950
+ "tables": [
951
+ "hubs",
952
+ "stores",
953
+ "orders"
954
+ ],
955
+ "columns": [
956
+ "hubs.hub_name",
957
+ "hubs.hub_id",
958
+ "stores.store_id",
959
+ "stores.hub_id",
960
+ "orders.store_id",
961
+ "orders.order_created_month",
962
+ "orders.order_status"
963
+ ]
964
+ }
965
+ },
966
+ "local218": {
967
+ "instance_id": "local218",
968
+ "schema_extraction": {
969
+ "tables": [
970
+ "Match",
971
+ "Team"
972
+ ],
973
+ "columns": [
974
+ "Match.season",
975
+ "Match.away_team_goal",
976
+ "Match.home_team_api_id",
977
+ "Match.away_team_api_id",
978
+ "Match.home_team_goal",
979
+ "Team.id",
980
+ "Team.team_long_name",
981
+ "Team.team_api_id"
982
+ ]
983
+ }
984
+ },
985
+ "local219": {
986
+ "instance_id": "local219",
987
+ "schema_extraction": {
988
+ "tables": [
989
+ "Player",
990
+ "Match",
991
+ "League",
992
+ "Team"
993
+ ],
994
+ "columns": [
995
+ "Player.player_name",
996
+ "Player.player_api_id",
997
+ "Match.id",
998
+ "Match.season",
999
+ "Match.match_api_id",
1000
+ "Match.home_team_goal",
1001
+ "Match.away_team_goal",
1002
+ "Match.goal",
1003
+ "Match.card",
1004
+ "Match.league_id",
1005
+ "Match.home_team_api_id",
1006
+ "Match.away_team_api_id",
1007
+ "Match.home_player_1",
1008
+ "Match.home_player_2",
1009
+ "Match.home_player_3",
1010
+ "Match.home_player_4",
1011
+ "Match.home_player_5",
1012
+ "Match.home_player_6",
1013
+ "Match.home_player_7",
1014
+ "Match.home_player_8",
1015
+ "Match.home_player_9",
1016
+ "Match.home_player_10",
1017
+ "Match.home_player_11",
1018
+ "Match.away_player_1",
1019
+ "Match.away_player_2",
1020
+ "Match.away_player_3",
1021
+ "Match.away_player_4",
1022
+ "Match.away_player_5",
1023
+ "Match.away_player_6",
1024
+ "Match.away_player_7",
1025
+ "Match.away_player_8",
1026
+ "Match.away_player_9",
1027
+ "Match.away_player_10",
1028
+ "Match.away_player_11",
1029
+ "League.name",
1030
+ "League.id",
1031
+ "Team.team_long_name",
1032
+ "Team.team_api_id"
1033
+ ]
1034
+ }
1035
+ },
1036
+ "local229": {
1037
+ "instance_id": "local229",
1038
+ "schema_extraction": {
1039
+ "tables": [
1040
+ "ball_by_ball",
1041
+ "batsman_scored"
1042
+ ],
1043
+ "columns": [
1044
+ "ball_by_ball.match_id",
1045
+ "ball_by_ball.non_striker",
1046
+ "ball_by_ball.striker",
1047
+ "ball_by_ball.innings_no",
1048
+ "ball_by_ball.ball_id",
1049
+ "ball_by_ball.over_id",
1050
+ "batsman_scored.runs_scored",
1051
+ "batsman_scored.innings_no",
1052
+ "batsman_scored.ball_id",
1053
+ "batsman_scored.match_id",
1054
+ "batsman_scored.over_id"
1055
+ ]
1056
+ }
1057
+ },
1058
+ "local244": {
1059
+ "instance_id": "local244",
1060
+ "schema_extraction": {
1061
+ "tables": [
1062
+ "InvoiceLine",
1063
+ "Track"
1064
+ ],
1065
+ "columns": [
1066
+ "InvoiceLine.TrackId",
1067
+ "InvoiceLine.UnitPrice",
1068
+ "InvoiceLine.Quantity",
1069
+ "Track.TrackId",
1070
+ "Track.Milliseconds"
1071
+ ]
1072
+ }
1073
+ },
1074
+ "local263": {
1075
+ "instance_id": "local263",
1076
+ "schema_extraction": {
1077
+ "tables": [
1078
+ "model_score",
1079
+ "model"
1080
+ ],
1081
+ "columns": [
1082
+ "model_score.name",
1083
+ "model_score.version",
1084
+ "model_score.step",
1085
+ "model_score.test_score",
1086
+ "model_score.model",
1087
+ "model.L1_model",
1088
+ "model.name",
1089
+ "model.version"
1090
+ ]
1091
+ }
1092
+ },
1093
+ "local269": {
1094
+ "instance_id": "local269",
1095
+ "schema_extraction": {
1096
+ "tables": [
1097
+ "packaging",
1098
+ "packaging_relations"
1099
+ ],
1100
+ "columns": [
1101
+ "packaging.id",
1102
+ "packaging_relations.contains_id",
1103
+ "packaging_relations.packaging_id",
1104
+ "packaging_relations.qty"
1105
+ ]
1106
+ }
1107
+ },
1108
+ "local274": {
1109
+ "instance_id": "local274",
1110
+ "schema_extraction": {
1111
+ "tables": [
1112
+ "products",
1113
+ "picking_line"
1114
+ ],
1115
+ "columns": [
1116
+ "products.name",
1117
+ "products.id",
1118
+ "picking_line.product_id",
1119
+ "picking_line.location_id",
1120
+ "picking_line.qty",
1121
+ "picking_line.order_id"
1122
+ ]
1123
+ }
1124
+ },
1125
+ "local283": {
1126
+ "instance_id": "local283",
1127
+ "schema_extraction": {
1128
+ "tables": [
1129
+ "Match",
1130
+ "League",
1131
+ "Country",
1132
+ "Team"
1133
+ ],
1134
+ "columns": [
1135
+ "Match.id",
1136
+ "Match.season",
1137
+ "Match.stage",
1138
+ "Match.date",
1139
+ "Match.home_team_goal",
1140
+ "Match.away_team_goal",
1141
+ "Match.country_id",
1142
+ "Match.league_id",
1143
+ "Match.home_team_api_id",
1144
+ "Match.away_team_api_id",
1145
+ "League.name",
1146
+ "League.id",
1147
+ "Country.name",
1148
+ "Country.id",
1149
+ "Team.team_long_name",
1150
+ "Team.team_api_id"
1151
+ ]
1152
+ }
1153
+ },
1154
+ "local284": {
1155
+ "instance_id": "local284",
1156
+ "schema_extraction": {
1157
+ "tables": [
1158
+ "veg_loss_rate_df"
1159
+ ],
1160
+ "columns": [
1161
+ "veg_loss_rate_df.loss_rate_%"
1162
+ ]
1163
+ }
1164
+ },
1165
+ "local285": {
1166
+ "instance_id": "local285",
1167
+ "schema_extraction": {
1168
+ "tables": [
1169
+ "veg_txn_df",
1170
+ "veg_loss_rate_df",
1171
+ "veg_cat",
1172
+ "veg_whsle_df"
1173
+ ],
1174
+ "columns": [
1175
+ "veg_txn_df.item_code",
1176
+ "veg_txn_df.txn_date",
1177
+ "veg_txn_df.qty_sold(kg)",
1178
+ "veg_txn_df.unit_selling_px_rmb/kg",
1179
+ "veg_loss_rate_df.item_code",
1180
+ "veg_loss_rate_df.loss_rate_%",
1181
+ "veg_cat.category_code",
1182
+ "veg_cat.category_name",
1183
+ "veg_cat.item_code",
1184
+ "veg_whsle_df.whsle_px_rmb-kg",
1185
+ "veg_whsle_df.whsle_date",
1186
+ "veg_whsle_df.item_code"
1187
+ ]
1188
+ }
1189
+ },
1190
+ "local301": {
1191
+ "instance_id": "local301",
1192
+ "schema_extraction": {
1193
+ "tables": [
1194
+ "cleaned_weekly_sales"
1195
+ ],
1196
+ "columns": [
1197
+ "cleaned_weekly_sales.week_date",
1198
+ "cleaned_weekly_sales.sales"
1199
+ ]
1200
+ }
1201
+ },
1202
+ "local329": {
1203
+ "instance_id": "local329",
1204
+ "schema_extraction": {
1205
+ "tables": [
1206
+ "form_log"
1207
+ ],
1208
+ "columns": [
1209
+ "form_log.session",
1210
+ "form_log.path",
1211
+ "form_log.stamp"
1212
+ ]
1213
+ }
1214
+ },
1215
+ "local358": {
1216
+ "instance_id": "local358",
1217
+ "schema_extraction": {
1218
+ "tables": [
1219
+ "mst_users"
1220
+ ],
1221
+ "columns": [
1222
+ "mst_users.user_id",
1223
+ "mst_users.birth_date"
1224
+ ]
1225
+ }
1226
+ },
1227
+ "local360": {
1228
+ "instance_id": "local360",
1229
+ "schema_extraction": {
1230
+ "tables": [
1231
+ "activity_log"
1232
+ ],
1233
+ "columns": [
1234
+ "activity_log.session",
1235
+ "activity_log.stamp",
1236
+ "activity_log.path",
1237
+ "activity_log.search_type"
1238
+ ]
1239
+ }
1240
+ },
1241
+ "local336": {
1242
+ "instance_id": "local336",
1243
+ "schema_extraction": {
1244
+ "tables": [
1245
+ "pit_stops",
1246
+ "results",
1247
+ "lap_times_ext",
1248
+ "races_ext",
1249
+ "retirements",
1250
+ "lap_positions"
1251
+ ],
1252
+ "columns": [
1253
+ "pit_stops.driver_id",
1254
+ "pit_stops.lap",
1255
+ "pit_stops.milliseconds",
1256
+ "pit_stops.race_id",
1257
+ "results.race_id",
1258
+ "results.driver_id",
1259
+ "results.grid",
1260
+ "lap_times_ext.lap",
1261
+ "lap_times_ext.race_id",
1262
+ "lap_times_ext.driver_id",
1263
+ "lap_times_ext.running_milliseconds",
1264
+ "races_ext.race_id",
1265
+ "races_ext.is_pit_data_available",
1266
+ "retirements.driver_id",
1267
+ "retirements.race_id",
1268
+ "retirements.lap",
1269
+ "lap_positions.race_id",
1270
+ "lap_positions.driver_id",
1271
+ "lap_positions.lap",
1272
+ "lap_positions.position"
1273
+ ]
1274
+ }
1275
+ },
1276
+ "local309": {
1277
+ "instance_id": "local309",
1278
+ "schema_extraction": {
1279
+ "tables": [
1280
+ "constructors",
1281
+ "drivers",
1282
+ "races",
1283
+ "results"
1284
+ ],
1285
+ "columns": [
1286
+ "constructors.name",
1287
+ "constructors.constructor_id",
1288
+ "drivers.surname",
1289
+ "drivers.driver_id",
1290
+ "drivers.forename",
1291
+ "races.year",
1292
+ "races.race_id",
1293
+ "results.points",
1294
+ "results.race_id",
1295
+ "results.driver_id",
1296
+ "results.constructor_id"
1297
+ ]
1298
+ }
1299
+ },
1300
+ "local311": {
1301
+ "instance_id": "local311",
1302
+ "schema_extraction": {
1303
+ "tables": [
1304
+ "constructors",
1305
+ "drivers",
1306
+ "races",
1307
+ "results"
1308
+ ],
1309
+ "columns": [
1310
+ "constructors.name",
1311
+ "constructors.constructor_id",
1312
+ "drivers.forename",
1313
+ "drivers.surname",
1314
+ "drivers.driver_id",
1315
+ "races.year",
1316
+ "races.race_id",
1317
+ "results.points",
1318
+ "results.race_id",
1319
+ "results.driver_id",
1320
+ "results.constructor_id"
1321
+ ]
1322
+ }
1323
+ },
1324
+ "local354": {
1325
+ "instance_id": "local354",
1326
+ "schema_extraction": {
1327
+ "tables": [
1328
+ "races",
1329
+ "results"
1330
+ ],
1331
+ "columns": [
1332
+ "races.year",
1333
+ "races.round",
1334
+ "races.race_id",
1335
+ "results.driver_id",
1336
+ "results.constructor_id",
1337
+ "results.race_id"
1338
+ ]
1339
+ }
1340
+ },
1341
+ "local355": {
1342
+ "instance_id": "local355",
1343
+ "schema_extraction": {
1344
+ "tables": [
1345
+ "results",
1346
+ "driver_standings_ext",
1347
+ "races_ext"
1348
+ ],
1349
+ "columns": [
1350
+ "results.race_id",
1351
+ "results.driver_id",
1352
+ "results.constructor_id",
1353
+ "driver_standings_ext.race_id",
1354
+ "driver_standings_ext.driver_id",
1355
+ "races_ext.race_id",
1356
+ "races_ext.year",
1357
+ "races_ext.round"
1358
+ ]
1359
+ }
1360
+ },
1361
+ "local074": {
1362
+ "instance_id": "local074",
1363
+ "schema_extraction": {
1364
+ "tables": [
1365
+ "customer_transactions"
1366
+ ],
1367
+ "columns": [
1368
+ "customer_transactions.customer_id",
1369
+ "customer_transactions.txn_date",
1370
+ "customer_transactions.txn_amount"
1371
+ ]
1372
+ }
1373
+ },
1374
+ "local077": {
1375
+ "instance_id": "local077",
1376
+ "schema_extraction": {
1377
+ "tables": [
1378
+ "interest_metrics",
1379
+ "interest_map"
1380
+ ],
1381
+ "columns": [
1382
+ "interest_metrics.month_year",
1383
+ "interest_metrics.interest_id",
1384
+ "interest_metrics.composition",
1385
+ "interest_metrics.index_value",
1386
+ "interest_metrics._year",
1387
+ "interest_metrics._month",
1388
+ "interest_map.id",
1389
+ "interest_map.interest_name"
1390
+ ]
1391
+ }
1392
+ },
1393
+ "local114": {
1394
+ "instance_id": "local114",
1395
+ "schema_extraction": {
1396
+ "tables": [
1397
+ "web_orders",
1398
+ "web_sales_reps",
1399
+ "web_region"
1400
+ ],
1401
+ "columns": [
1402
+ "web_orders.id",
1403
+ "web_orders.account_id",
1404
+ "web_orders.total_amt_usd",
1405
+ "web_sales_reps.id",
1406
+ "web_sales_reps.name",
1407
+ "web_sales_reps.region_id",
1408
+ "web_region.id",
1409
+ "web_region.name"
1410
+ ]
1411
+ }
1412
+ },
1413
+ "local141": {
1414
+ "instance_id": "local141",
1415
+ "schema_extraction": {
1416
+ "tables": [
1417
+ "salesorderheader",
1418
+ "SalesPersonQuotaHistory"
1419
+ ],
1420
+ "columns": [
1421
+ "salesorderheader.salespersonid",
1422
+ "salesorderheader.orderdate",
1423
+ "salesorderheader.totaldue",
1424
+ "SalesPersonQuotaHistory.BusinessEntityID",
1425
+ "SalesPersonQuotaHistory.QuotaDate",
1426
+ "SalesPersonQuotaHistory.SalesQuota"
1427
+ ]
1428
+ }
1429
+ },
1430
+ "local277": {
1431
+ "instance_id": "local277",
1432
+ "schema_extraction": {
1433
+ "tables": [
1434
+ "monthly_sales"
1435
+ ],
1436
+ "columns": [
1437
+ "monthly_sales.product_id",
1438
+ "monthly_sales.mth",
1439
+ "monthly_sales.qty"
1440
+ ]
1441
+ }
1442
+ }
1443
+ }
data/text2sql/spider2-lite/snowflake/AUSTIN/AUSTIN_db_summary.json ADDED
@@ -0,0 +1,39 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "tables": {
3
+ "AUSTIN.AUSTIN_311._311_SERVICE_REQUESTS": {
4
+ "name": "AUSTIN.AUSTIN_311._311_SERVICE_REQUESTS",
5
+ "columns": [
6
+ {"name": "unique_key", "type": "VARCHAR(16777216)", "is_primary_key": true, "pk_origin": "heuristic", "is_foreign_key": false},
7
+ {"name": "complaint_description", "type": "VARCHAR(16777216)", "is_primary_key": false, "is_foreign_key": false},
8
+ {"name": "source", "type": "VARCHAR(16777216)", "is_primary_key": false, "is_foreign_key": false},
9
+ {"name": "status", "type": "VARCHAR(16777216)", "is_primary_key": false, "is_foreign_key": false},
10
+ {"name": "status_change_date", "type": "NUMBER(38,0)", "is_primary_key": false, "is_foreign_key": false},
11
+ {"name": "created_date", "type": "NUMBER(38,0)", "is_primary_key": false, "is_foreign_key": false},
12
+ {"name": "last_update_date", "type": "NUMBER(38,0)", "is_primary_key": false, "is_foreign_key": false},
13
+ {"name": "close_date", "type": "NUMBER(38,0)", "is_primary_key": false, "is_foreign_key": false},
14
+ {"name": "incident_address", "type": "VARCHAR(16777216)", "is_primary_key": false, "is_foreign_key": false},
15
+ {"name": "street_number", "type": "VARCHAR(16777216)", "is_primary_key": false, "is_foreign_key": false},
16
+ {"name": "street_name", "type": "VARCHAR(16777216)", "is_primary_key": false, "is_foreign_key": false},
17
+ {"name": "city", "type": "VARCHAR(16777216)", "is_primary_key": false, "is_foreign_key": false},
18
+ {"name": "incident_zip", "type": "VARCHAR(16777216)", "is_primary_key": false, "is_foreign_key": false},
19
+ {"name": "county", "type": "VARCHAR(16777216)", "is_primary_key": false, "is_foreign_key": false},
20
+ {"name": "state_plane_x_coordinate", "type": "VARCHAR(16777216)", "is_primary_key": false, "is_foreign_key": false},
21
+ {"name": "state_plane_y_coordinate", "type": "FLOAT", "is_primary_key": false, "is_foreign_key": false},
22
+ {"name": "latitude", "type": "FLOAT", "is_primary_key": false, "is_foreign_key": false},
23
+ {"name": "longitude", "type": "FLOAT", "is_primary_key": false, "is_foreign_key": false},
24
+ {"name": "location", "type": "VARCHAR(16777216)", "is_primary_key": false, "is_foreign_key": false},
25
+ {"name": "council_district_code", "type": "NUMBER(38,0)", "is_primary_key": false, "is_foreign_key": false},
26
+ {"name": "map_page", "type": "VARCHAR(16777216)", "is_primary_key": false, "is_foreign_key": false},
27
+ {"name": "map_tile", "type": "VARCHAR(16777216)", "is_primary_key": false, "is_foreign_key": false}
28
+ ],
29
+ "samples": {
30
+ "unique_key": ["18-00004837", "18-00004838"],
31
+ "complaint_description": ["Loose Dog", "Animal Control - Assistance Request"],
32
+ "status": ["Closed", "Closed"],
33
+ "city": ["Austin", "Austin"]
34
+ },
35
+ "foreign_keys": []
36
+ }
37
+ },
38
+ "relationships": []
39
+ }
data/text2sql/spider2-lite/spider2-lite.jsonl ADDED
The diff for this file is too large to render. See raw diff
 
data/text2sql/spider2-lite/sqlite/Pagila/Pagila.sqlite ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:dbca3837b9965d75fbc2b4af1b2daec57d841f95942e656cea44db997e05da88
3
+ size 5828608
data/text2sql/spider2-lite/sqlite/Pagila/Pagila_db_summary.json ADDED
The diff for this file is too large to render. See raw diff
 
data/ultradomain/README.md ADDED
@@ -0,0 +1,25 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # UltraDomain Data
2
+
3
+ Datasets: agriculture, biology, cooking, cs, finance, legal, literature, mathematics, mix, music, philosophy, physics
4
+
5
+ Loaded from [TommyChien/UltraDomain](https://huggingface.co/datasets/TommyChien/UltraDomain) at runtime.
6
+
7
+ **Graph type:** entity (default) or passage-entity (`--graph_type passage-entity`)
8
+
9
+ ## Pre-built KGs
10
+
11
+ Pre-built KGs use `openai-small` embedding (requires `OPENAI_API_KEY`):
12
+
13
+ ```bash
14
+ huggingface-cli download qafd/kg --repo-type dataset --include "ultradomain/*" --local-dir ./kg
15
+ ```
16
+
17
+ ## Usage
18
+
19
+ ```bash
20
+ # Uses downloaded KG, or builds from scratch if not found
21
+ python benchmarks/run.py --task ultradomain --dataset mix --questions 10
22
+
23
+ # Use passage-entity graph instead
24
+ python benchmarks/run.py --task ultradomain --dataset mix --graph_type passage-entity --questions 10
25
+ ```
docs/figs/GraphRAG.png ADDED

Git LFS Details

  • SHA256: 90583dd2d903777a429bb0f29a0ba34a042f397f4554aab18b7a1a86530987ae
  • Pointer size: 131 Bytes
  • Size of remote file: 554 kB
docs/figs/LightRAG.png ADDED

Git LFS Details

  • SHA256: 7ec3175d79dc6ed7f47fe6360aecef98d8adec29474fbbc88085eafc5b84b5db
  • Pointer size: 131 Bytes
  • Size of remote file: 472 kB
docs/figs/QAFD-RAG.png ADDED

Git LFS Details

  • SHA256: 5964565ec81d739886e43af7099247a21fcc6e1ccd20c33395360d204edd0dc9
  • Pointer size: 131 Bytes
  • Size of remote file: 395 kB
docs/index.html ADDED
@@ -0,0 +1,452 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ <!DOCTYPE html>
2
+ <html lang="en">
3
+ <head>
4
+ <meta charset="UTF-8" />
5
+ <meta name="viewport" content="width=device-width, initial-scale=1.0" />
6
+ <title>QAFD-RAG: Query-Aware Flow Diffusion for Graph-Based RAG</title>
7
+ <link rel="preconnect" href="https://fonts.googleapis.com" />
8
+ <link href="https://fonts.googleapis.com/css2?family=DM+Sans:ital,opsz,wght@0,9..40,400;0,9..40,500;0,9..40,600;0,9..40,700&family=JetBrains+Mono:wght@400;500&display=swap" rel="stylesheet" />
9
+ <style>
10
+ :root {
11
+ --navy-900: #0a1628; --navy-800: #0f2140; --navy-700: #1e3a5f;
12
+ --blue-600: #2563eb; --blue-500: #3b82f6; --blue-400: #60a5fa;
13
+ --blue-300: #93c5fd; --blue-200: #bfdbfe; --blue-100: #dbeafe; --blue-50: #eff6ff;
14
+ --slate-900: #0f172a; --slate-800: #1e293b; --slate-700: #334155;
15
+ --slate-500: #64748b; --slate-400: #94a3b8; --slate-300: #cbd5e1;
16
+ --slate-200: #e2e8f0; --slate-100: #f1f5f9; --slate-50: #f8fafc;
17
+ --white: #ffffff; --radius: 10px; --radius-lg: 16px;
18
+ }
19
+ * { margin: 0; padding: 0; box-sizing: border-box; }
20
+ body { font-family: 'DM Sans', -apple-system, sans-serif; color: var(--slate-800); line-height: 1.6; background: var(--white); -webkit-font-smoothing: antialiased; }
21
+
22
+ .hero {
23
+ background: linear-gradient(155deg, #d6e6f7 0%, #dbeafe 40%, #e8f1fc 70%, #eff6ff 100%);
24
+ color: var(--slate-800); padding: 36px 20px 40px; text-align: center;
25
+ position: relative; overflow: hidden;
26
+ }
27
+ .hero::before {
28
+ content: ''; position: absolute; top: -50%; left: -50%; width: 200%; height: 200%;
29
+ background: radial-gradient(ellipse at 30% 50%, rgba(37,99,235,0.04) 0%, transparent 60%),
30
+ radial-gradient(ellipse at 70% 30%, rgba(59,130,246,0.03) 0%, transparent 50%);
31
+ pointer-events: none;
32
+ }
33
+ .hero-content { position: relative; z-index: 1; }
34
+
35
+ .graph-wrap {
36
+ margin: 0 auto 16px; width: 480px; height: 280px;
37
+ animation: logoFadeIn 1s ease-out;
38
+ }
39
+ @keyframes logoFadeIn { from { opacity:0; transform: scale(0.9) translateY(10px); } to { opacity:1; transform: scale(1) translateY(0); } }
40
+ @keyframes waterFlow { 0% { stroke-dashoffset: var(--len); opacity:0.85; } 100% { stroke-dashoffset:0; opacity:0.85; } }
41
+ @keyframes pipeFadeIn { 0% { opacity:0; } 100% { opacity: var(--pipe-op, 0.7); } }
42
+ @keyframes sheenFadeIn { 0% { opacity:0; } 100% { opacity: var(--sheen-op, 0.4); } }
43
+ @keyframes nodeAppear { 0% { transform:scale(0); opacity:0; } 60% { transform:scale(1.2); opacity:1; } 100% { transform:scale(1); opacity:1; } }
44
+ @keyframes fadeUp { from { opacity:0; transform:translateY(16px); } to { opacity:1; transform:translateY(0); } }
45
+ @keyframes seedPulse { 0%,100% { r: 14; } 50% { r: 17; } }
46
+
47
+
48
+
49
+
50
+
51
+ .btn-group { display:flex; gap:10px; justify-content:center; flex-wrap:wrap; animation:fadeUp 0.6s ease-out 0.6s both; }
52
+ .btn { display:inline-flex; align-items:center; gap:7px; padding:10px 22px; border-radius:50px; font-size:0.9rem; font-weight:600; text-decoration:none; transition:all 0.2s; font-family:inherit; }
53
+ .btn:hover { transform:translateY(-2px); box-shadow:0 6px 20px rgba(0,0,0,0.2); }
54
+ .btn-primary { background:var(--blue-600); color:#fff; }
55
+ .btn-outline { background:rgba(30,58,138,0.06); color:var(--slate-800); border:1.5px solid var(--slate-300); }
56
+ .btn-outline:hover { border-color:var(--blue-500); background:rgba(37,99,235,0.08); }
57
+ .btn svg { width:17px; height:17px; flex-shrink:0; }
58
+
59
+ .container { max-width:940px; margin:0 auto; padding:0 24px; }
60
+ section { padding:64px 0; }
61
+ section.alt { background:var(--slate-50); }
62
+ section h2 { font-size:1.7rem; font-weight:700; margin-bottom:10px; text-align:center; letter-spacing:-0.3px; }
63
+ section .section-sub { text-align:center; color:var(--slate-500); font-size:1rem; margin-bottom:32px; max-width:600px; margin-left:auto; margin-right:auto; }
64
+ .about-text { font-size:1.02rem; line-height:1.85; color:var(--slate-500); text-align:center; max-width:680px; margin:0 auto 36px; }
65
+ .duo-grid { display:grid; grid-template-columns:1fr 1fr; gap:20px; }
66
+ .card { background:var(--white); border:1px solid var(--slate-200); border-radius:var(--radius-lg); padding:28px; transition:box-shadow 0.2s; }
67
+ .card:hover { box-shadow:0 4px 20px rgba(0,0,0,0.06); }
68
+ .card h3 { color:var(--blue-600); font-size:1.05rem; font-weight:700; margin-bottom:10px; }
69
+ .card p, .card li { font-size:0.93rem; color:var(--slate-500); line-height:1.7; }
70
+ .card ol { padding-left:18px; } .card li { margin-bottom:6px; } .card li strong { color:var(--slate-800); }
71
+ .figure-row { display:flex; gap:18px; justify-content:center; flex-wrap:wrap; margin-top:8px; }
72
+ .figure-card { text-align:center; background:var(--white); border:1px solid var(--slate-200); border-radius:var(--radius-lg); padding:16px; flex:1; min-width:220px; max-width:280px; transition:box-shadow 0.2s, transform 0.2s; }
73
+ .figure-card:hover { box-shadow:0 6px 24px rgba(0,0,0,0.07); transform:translateY(-2px); }
74
+ .figure-card img { width:100%; border-radius:6px; }
75
+ .figure-card .caption { margin-top:10px; font-size:0.88rem; font-weight:700; color:var(--slate-700); }
76
+ .figure-card .caption small { font-weight:500; color:var(--slate-400); }
77
+ .bench-grid { display:grid; grid-template-columns:repeat(auto-fit,minmax(190px,1fr)); gap:14px; }
78
+ .bench-item { background:var(--white); border:1px solid var(--slate-200); border-radius:var(--radius-lg); padding:22px 18px; text-align:center; transition:box-shadow 0.2s; }
79
+ .bench-item:hover { box-shadow:0 4px 16px rgba(0,0,0,0.06); }
80
+ .bench-item .task { font-weight:700; font-size:1rem; color:var(--blue-600); margin-bottom:4px; }
81
+ .bench-item .dataset { font-size:0.85rem; color:var(--slate-500); }
82
+ .bench-item .metrics { font-size:0.78rem; color:var(--slate-400); margin-top:6px; }
83
+ .models-text { text-align:center; color:var(--slate-500); font-size:1rem; line-height:2; }
84
+ .models-text strong { color:var(--slate-700); }
85
+ pre { background:var(--slate-900); color:var(--slate-300); border-radius:var(--radius-lg); padding:24px; overflow-x:auto; font-size:0.85rem; line-height:1.65; margin-top:16px; }
86
+ code { font-family:'JetBrains Mono','SF Mono',monospace; }
87
+ pre .comment { color:var(--slate-500); } pre .string { color:#7dd3fc; } pre .keyword { color:var(--blue-400); }
88
+ section h3 { font-size:1.1rem; font-weight:700; margin:28px 0 8px; color:var(--slate-700); }
89
+ footer { background:var(--slate-900); color:var(--slate-400); padding:36px 20px; text-align:center; font-size:0.85rem; }
90
+ footer a { color:var(--blue-300); text-decoration:none; } footer a:hover { text-decoration:underline; }
91
+ footer .sep { margin:0 8px; opacity:0.4; }
92
+ @media (max-width:768px) {
93
+ .hero h1 { font-size:2.2rem; }
94
+ .graph-wrap { width:100%; max-width:480px; height:auto; aspect-ratio:420/280; }
95
+ .duo-grid { grid-template-columns:1fr; }
96
+ .figure-card { max-width:100%; }
97
+ }
98
+ </style>
99
+ </head>
100
+ <body>
101
+ <header class="hero">
102
+ <div class="hero-content">
103
+ <div class="graph-wrap">
104
+ <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 420 280" width="100%" height="100%">
105
+ <defs>
106
+ <radialGradient id="src" cx="40%" cy="32%"><stop offset="0%" stop-color="#1E40AF"/><stop offset="45%" stop-color="#1E3A8A"/><stop offset="100%" stop-color="#172554"/></radialGradient>
107
+ <radialGradient id="srcSp" cx="34%" cy="26%"><stop offset="0%" stop-color="#FFF" stop-opacity="0.9"/><stop offset="18%" stop-color="#93C5FD" stop-opacity="0.3"/><stop offset="100%" stop-color="#FFF" stop-opacity="0"/></radialGradient>
108
+ <radialGradient id="srcGl" cx="50%" cy="50%" r="50%"><stop offset="0%" stop-color="#3B82F6" stop-opacity="0.5"/><stop offset="50%" stop-color="#1E3A8A" stop-opacity="0.15"/><stop offset="100%" stop-color="#172554" stop-opacity="0"/></radialGradient>
109
+ <radialGradient id="s1" cx="38%" cy="28%"><stop offset="0%" stop-color="#60A5FA"/><stop offset="45%" stop-color="#2563EB"/><stop offset="100%" stop-color="#1E3A8A"/></radialGradient>
110
+ <radialGradient id="s1sp" cx="32%" cy="24%"><stop offset="0%" stop-color="#FFF" stop-opacity="0.85"/><stop offset="22%" stop-color="#BFDBFE" stop-opacity="0.2"/><stop offset="100%" stop-color="#FFF" stop-opacity="0"/></radialGradient>
111
+ <radialGradient id="m1" cx="38%" cy="28%"><stop offset="0%" stop-color="#93C5FD"/><stop offset="50%" stop-color="#3B82F6"/><stop offset="100%" stop-color="#2563EB"/></radialGradient>
112
+ <radialGradient id="m1sp" cx="32%" cy="24%"><stop offset="0%" stop-color="#FFF" stop-opacity="0.7"/><stop offset="20%" stop-color="#FFF" stop-opacity="0.12"/><stop offset="100%" stop-color="#FFF" stop-opacity="0"/></radialGradient>
113
+ <radialGradient id="w1" cx="38%" cy="28%"><stop offset="0%" stop-color="#BFDBFE"/><stop offset="50%" stop-color="#93C5FD"/><stop offset="100%" stop-color="#60A5FA"/></radialGradient>
114
+ <radialGradient id="w1sp" cx="32%" cy="24%"><stop offset="0%" stop-color="#FFF" stop-opacity="0.75"/><stop offset="18%" stop-color="#FFF" stop-opacity="0.12"/><stop offset="100%" stop-color="#FFF" stop-opacity="0"/></radialGradient>
115
+ <radialGradient id="st1" cx="38%" cy="28%"><stop offset="0%" stop-color="#DBEAFE"/><stop offset="50%" stop-color="#BFDBFE"/><stop offset="100%" stop-color="#93C5FD"/></radialGradient>
116
+ <radialGradient id="lf" cx="38%" cy="28%"><stop offset="0%" stop-color="#EFF6FF"/><stop offset="50%" stop-color="#DBEAFE"/><stop offset="100%" stop-color="#BFDBFE"/></radialGradient>
117
+ <linearGradient id="pSh" x1="0%" y1="0%" x2="0%" y2="100%"><stop offset="0%" stop-color="#FFF" stop-opacity="0.25"/><stop offset="40%" stop-color="#93C5FD" stop-opacity="0.08"/><stop offset="100%" stop-color="#FFF" stop-opacity="0"/></linearGradient>
118
+ <filter id="sB" x="-80%" y="-80%" width="260%" height="260%"><feGaussianBlur in="SourceGraphic" stdDeviation="7" result="b1"/><feGaussianBlur in="SourceGraphic" stdDeviation="3" result="b2"/><feMerge><feMergeNode in="b1"/><feMergeNode in="b2"/></feMerge></filter>
119
+ <filter id="nG" x="-40%" y="-40%" width="180%" height="180%"><feGaussianBlur stdDeviation="2.5" result="b"/><feMerge><feMergeNode in="b"/><feMergeNode in="SourceGraphic"/></feMerge></filter>
120
+ <filter id="sG" x="-40%" y="-40%" width="180%" height="180%"><feGaussianBlur stdDeviation="4" result="b"/><feMerge><feMergeNode in="b"/><feMergeNode in="SourceGraphic"/></feMerge></filter>
121
+ </defs>
122
+
123
+ <!-- ========= BACKGROUND KNOWLEDGE GRAPH (always visible, gray) ========= -->
124
+ <g opacity="0.15" stroke="#94A3B8" stroke-width="0.7" fill="none">
125
+ <line x1="15" y1="60" x2="40" y2="30"/><line x1="40" y1="30" x2="75" y2="22"/>
126
+ <line x1="15" y1="60" x2="10" y2="110"/><line x1="10" y1="110" x2="25" y2="160"/>
127
+ <line x1="25" y1="160" x2="12" y2="210"/><line x1="12" y1="210" x2="40" y2="250"/>
128
+ <line x1="40" y1="30" x2="50" y2="75"/><line x1="50" y1="75" x2="30" y2="120"/>
129
+ <line x1="30" y1="120" x2="55" y2="155"/><line x1="55" y1="155" x2="25" y2="160"/>
130
+ <line x1="75" y1="22" x2="95" y2="55"/><line x1="95" y1="55" x2="75" y2="95"/>
131
+ <line x1="75" y1="95" x2="50" y2="75"/><line x1="95" y1="55" x2="130" y2="40"/>
132
+ <line x1="75" y1="95" x2="110" y2="120"/><line x1="55" y1="155" x2="85" y2="180"/>
133
+ <line x1="85" y1="180" x2="65" y2="220"/><line x1="65" y1="220" x2="40" y2="250"/>
134
+ <line x1="85" y1="180" x2="120" y2="200"/><line x1="120" y1="200" x2="100" y2="240"/>
135
+ <line x1="100" y1="240" x2="130" y2="265"/><line x1="350" y1="22" x2="385" y2="15"/>
136
+ <line x1="385" y1="15" x2="410" y2="40"/><line x1="410" y1="40" x2="405" y2="85"/>
137
+ <line x1="405" y1="85" x2="410" y2="135"/><line x1="410" y1="135" x2="400" y2="180"/>
138
+ <line x1="400" y1="180" x2="408" y2="230"/><line x1="408" y1="230" x2="390" y2="260"/>
139
+ <line x1="350" y1="22" x2="370" y2="60"/><line x1="370" y1="60" x2="405" y2="85"/>
140
+ <line x1="370" y1="60" x2="355" y2="100"/><line x1="355" y1="100" x2="375" y2="140"/>
141
+ <line x1="375" y1="140" x2="400" y2="180"/><line x1="375" y1="140" x2="350" y2="175"/>
142
+ <line x1="350" y1="175" x2="370" y2="215"/><line x1="370" y1="215" x2="390" y2="260"/>
143
+ <line x1="350" y1="175" x2="320" y2="200"/><line x1="320" y1="200" x2="340" y2="240"/>
144
+ <line x1="340" y1="240" x2="308" y2="260"/><line x1="130" y1="40" x2="170" y2="18"/>
145
+ <line x1="170" y1="18" x2="220" y2="12"/><line x1="220" y1="12" x2="275" y2="18"/>
146
+ <line x1="275" y1="18" x2="320" y2="12"/><line x1="320" y1="12" x2="350" y2="22"/>
147
+ <line x1="130" y1="265" x2="175" y2="270"/><line x1="175" y1="270" x2="220" y2="265"/>
148
+ <line x1="220" y1="265" x2="260" y2="255"/><line x1="260" y1="255" x2="308" y2="260"/>
149
+
150
+
151
+
152
+ <line x1="320" y1="200" x2="290" y2="220"/>
153
+ <line x1="290" y1="220" x2="260" y2="255"/>
154
+ <line x1="130" y1="40" x2="158" y2="92"/><line x1="110" y1="120" x2="158" y2="148"/>
155
+ <line x1="120" y1="200" x2="158" y2="175"/><line x1="355" y1="100" x2="325" y2="118"/>
156
+ <line x1="375" y1="140" x2="335" y2="165"/><line x1="275" y1="18" x2="270" y2="45"/>
157
+ <line x1="170" y1="18" x2="210" y2="60"/>
158
+ <line x1="262" y1="92" x2="262" y2="148"/><line x1="210" y1="60" x2="262" y2="92"/>
159
+ <line x1="158" y1="92" x2="210" y2="60"/><line x1="262" y1="148" x2="210" y2="175"/>
160
+ <line x1="210" y1="175" x2="158" y2="148"/><line x1="158" y1="92" x2="158" y2="148"/>
161
+ <line x1="262" y1="92" x2="305" y2="68"/><line x1="305" y1="68" x2="325" y2="118"/>
162
+ <line x1="325" y1="118" x2="262" y2="148"/><line x1="158" y1="92" x2="112" y2="62"/>
163
+ <line x1="112" y1="62" x2="92" y2="35"/><line x1="210" y1="60" x2="210" y2="25"/>
164
+ <line x1="210" y1="175" x2="210" y2="218"/><line x1="210" y1="218" x2="190" y2="250"/>
165
+ <line x1="325" y1="118" x2="335" y2="165"/><line x1="305" y1="68" x2="322" y2="35"/>
166
+ <line x1="158" y1="148" x2="118" y2="178"/>
167
+ </g>
168
+ <g fill="#B0BEC5" opacity="0.22">
169
+ <circle cx="15" cy="60" r="2.5"/><circle cx="40" cy="30" r="2.2"/><circle cx="75" cy="22" r="2.2"/>
170
+ <circle cx="10" cy="110" r="2"/><circle cx="25" cy="160" r="2.2"/><circle cx="12" cy="210" r="2"/>
171
+ <circle cx="40" cy="250" r="2"/><circle cx="50" cy="75" r="2.2"/><circle cx="30" cy="120" r="2"/>
172
+ <circle cx="55" cy="155" r="2.2"/><circle cx="95" cy="55" r="2.5"/><circle cx="75" cy="95" r="2.2"/>
173
+ <circle cx="110" cy="120" r="2"/><circle cx="130" cy="40" r="2.2"/><circle cx="85" cy="180" r="2"/>
174
+ <circle cx="65" cy="220" r="2"/><circle cx="120" cy="200" r="2.2"/><circle cx="100" cy="240" r="2"/>
175
+ <circle cx="130" cy="265" r="2"/><circle cx="170" cy="18" r="2.2"/><circle cx="220" cy="12" r="2.2"/>
176
+ <circle cx="275" cy="18" r="2"/><circle cx="320" cy="12" r="2"/><circle cx="350" cy="22" r="2.2"/>
177
+ <circle cx="385" cy="15" r="2"/><circle cx="410" cy="40" r="2"/><circle cx="405" cy="85" r="2.2"/>
178
+ <circle cx="410" cy="135" r="2"/><circle cx="400" cy="180" r="2.2"/><circle cx="408" cy="230" r="2"/>
179
+ <circle cx="390" cy="260" r="2"/><circle cx="370" cy="60" r="2.5"/><circle cx="355" cy="100" r="2.5"/>
180
+ <circle cx="375" cy="140" r="2.2"/><circle cx="350" cy="175" r="2.2"/><circle cx="370" cy="215" r="2"/>
181
+ <circle cx="320" cy="200" r="2.2"/><circle cx="340" cy="240" r="2"/><circle cx="308" cy="260" r="2"/>
182
+ <circle cx="290" cy="220" r="2"/><circle cx="260" cy="255" r="2"/><circle cx="175" cy="270" r="2"/>
183
+ <circle cx="220" cy="265" r="2"/><circle cx="270" cy="45" r="2"/><circle cx="210" cy="60" r="2"/>
184
+ <circle cx="325" cy="118" r="3"/><circle cx="305" cy="68" r="2.8"/>
185
+ <circle cx="262" cy="148" r="2.8"/><circle cx="262" cy="92" r="2.8"/>
186
+ <circle cx="210" cy="175" r="2.5"/><circle cx="158" cy="92" r="2.5"/>
187
+ <circle cx="158" cy="148" r="2.2"/><circle cx="112" cy="62" r="2.2"/>
188
+ <circle cx="92" cy="35" r="2"/><circle cx="210" cy="25" r="2"/>
189
+ <circle cx="210" cy="218" r="2.2"/><circle cx="190" cy="250" r="2"/>
190
+ <circle cx="335" cy="165" r="2"/><circle cx="118" cy="178" r="2"/>
191
+ <circle cx="322" cy="35" r="2"/>
192
+ </g>
193
+
194
+ <!--
195
+ PHASE 1: BUILD — water shoots from source, pipe+node appear where it lands
196
+ PHASE 2: LOOP — after build, continuous flow from source outward
197
+
198
+ Each element uses <animate> with begin= for precise timing.
199
+ Pipe: starts invisible, fades in when water arrives
200
+ Node: starts invisible, pops when pipe reaches it
201
+ Water: draws from source outward using stroke-dashoffset
202
+ -->
203
+
204
+ <!-- ======= SOURCE (0s) ======= -->
205
+ <circle cx="325" cy="118" r="20" fill="url(#srcGl)" filter="url(#sB)" opacity="0"><animate attributeName="opacity" to="1" dur="0.4s" begin="0.2s" fill="freeze"/></circle>
206
+ <g filter="url(#sG)" opacity="0"><animate attributeName="opacity" to="1" dur="0.4s" begin="0.2s" fill="freeze"/>
207
+ <circle cx="325" cy="118" r="14" fill="url(#src)"/><circle cx="325" cy="118" r="14" fill="url(#srcSp)"/>
208
+ <ellipse cx="320" cy="111" rx="4.5" ry="3" fill="#FFF" opacity="0.5"/>
209
+ </g>
210
+
211
+ <!-- ======= HOP 1: Source→G(305,68) and Source→D(262,148) ======= -->
212
+ <!-- Water Source→G -->
213
+ <path d="M325,118 C318,100 312,85 305,68" fill="none" stroke="#7DD3FC" stroke-width="5.5" stroke-linecap="round" stroke-dasharray="55" stroke-dashoffset="55" opacity="0"><animate attributeName="stroke-dashoffset" to="0" dur="0.6s" begin="0.5s" fill="freeze"/><animate attributeName="opacity" to="0.85" dur="0.1s" begin="0.5s" fill="freeze"/></path>
214
+ <!-- Pipe Source→G fades in -->
215
+ <path d="M305,68 C312,85 318,100 325,118" fill="none" stroke="#172554" stroke-width="10" stroke-linecap="round" opacity="0"><animate attributeName="opacity" to="0.9" dur="0.3s" begin="0.7s" fill="freeze"/></path>
216
+ <path d="M305,68 C312,85 318,100 325,118" fill="none" stroke="url(#pSh)" stroke-width="4.5" stroke-linecap="round" opacity="0"><animate attributeName="opacity" to="0.55" dur="0.3s" begin="0.75s" fill="freeze"/></path>
217
+ <!-- G node pops -->
218
+ <g filter="url(#nG)" opacity="0"><animate attributeName="opacity" to="1" dur="0.3s" begin="0.85s" fill="freeze"/>
219
+ <circle cx="305" cy="68" r="9" fill="url(#s1)"/><circle cx="305" cy="68" r="9" fill="url(#s1sp)"/>
220
+ </g>
221
+
222
+ <!-- Water Source→D -->
223
+ <path d="M325,118 C305,128 282,138 262,148" fill="none" stroke="#7DD3FC" stroke-width="5" stroke-linecap="round" stroke-dasharray="75" stroke-dashoffset="75" opacity="0"><animate attributeName="stroke-dashoffset" to="0" dur="0.7s" begin="0.55s" fill="freeze"/><animate attributeName="opacity" to="0.85" dur="0.1s" begin="0.55s" fill="freeze"/></path>
224
+ <path d="M262,148 C282,138 305,128 325,118" fill="none" stroke="#172554" stroke-width="9.5" stroke-linecap="round" opacity="0"><animate attributeName="opacity" to="0.88" dur="0.3s" begin="0.8s" fill="freeze"/></path>
225
+ <path d="M262,148 C282,138 305,128 325,118" fill="none" stroke="url(#pSh)" stroke-width="4.2" stroke-linecap="round" opacity="0"><animate attributeName="opacity" to="0.52" dur="0.3s" begin="0.85s" fill="freeze"/></path>
226
+ <g filter="url(#nG)" opacity="0"><animate attributeName="opacity" to="1" dur="0.3s" begin="0.95s" fill="freeze"/>
227
+ <circle cx="262" cy="148" r="10" fill="url(#s1)"/><circle cx="262" cy="148" r="10" fill="url(#s1sp)"/>
228
+ </g>
229
+
230
+ <!-- ======= HOP 2: G→C and D→C (converging) ======= -->
231
+ <path d="M305,68 C290,72 275,82 262,92" fill="none" stroke="#7DD3FC" stroke-width="5" stroke-linecap="round" stroke-dasharray="55" stroke-dashoffset="55" opacity="0"><animate attributeName="stroke-dashoffset" to="0" dur="0.6s" begin="1.0s" fill="freeze"/><animate attributeName="opacity" to="0.85" dur="0.1s" begin="1.0s" fill="freeze"/></path>
232
+ <path d="M262,92 C275,82 290,72 305,68" fill="none" stroke="#172554" stroke-width="9.5" stroke-linecap="round" opacity="0"><animate attributeName="opacity" to="0.86" dur="0.3s" begin="1.2s" fill="freeze"/></path>
233
+ <path d="M262,92 C275,82 290,72 305,68" fill="none" stroke="url(#pSh)" stroke-width="4.2" stroke-linecap="round" opacity="0"><animate attributeName="opacity" to="0.5" dur="0.3s" begin="1.25s" fill="freeze"/></path>
234
+
235
+ <path d="M262,148 C262,130 262,110 262,92" fill="none" stroke="#7DD3FC" stroke-width="5" stroke-linecap="round" stroke-dasharray="56" stroke-dashoffset="56" opacity="0"><animate attributeName="stroke-dashoffset" to="0" dur="0.55s" begin="1.05s" fill="freeze"/><animate attributeName="opacity" to="0.85" dur="0.1s" begin="1.05s" fill="freeze"/></path>
236
+ <path d="M262,92 C262,110 262,130 262,148" fill="none" stroke="#172554" stroke-width="9.5" stroke-linecap="round" opacity="0"><animate attributeName="opacity" to="0.84" dur="0.3s" begin="1.25s" fill="freeze"/></path>
237
+ <path d="M262,92 C262,110 262,130 262,148" fill="none" stroke="url(#pSh)" stroke-width="4.2" stroke-linecap="round" opacity="0"><animate attributeName="opacity" to="0.5" dur="0.3s" begin="1.3s" fill="freeze"/></path>
238
+
239
+ <g filter="url(#nG)" opacity="0"><animate attributeName="opacity" to="1" dur="0.3s" begin="1.35s" fill="freeze"/>
240
+ <circle cx="262" cy="92" r="10" fill="url(#s1)"/><circle cx="262" cy="92" r="10" fill="url(#s1sp)"/>
241
+ </g>
242
+
243
+ <!-- ======= HOP 3: C→B→A (6-ring top) and D→E→F (6-ring bottom) ======= -->
244
+ <!-- C→B -->
245
+ <path d="M262,92 C245,78 228,68 210,60" fill="none" stroke="#93C5FD" stroke-width="4.2" stroke-linecap="round" stroke-dasharray="62" stroke-dashoffset="62" opacity="0"><animate attributeName="stroke-dashoffset" to="0" dur="0.6s" begin="1.5s" fill="freeze"/><animate attributeName="opacity" to="0.85" dur="0.1s" begin="1.5s" fill="freeze"/></path>
246
+ <path d="M210,60 C228,68 245,78 262,92" fill="none" stroke="#1E3A8A" stroke-width="8.5" stroke-linecap="round" opacity="0"><animate attributeName="opacity" to="0.72" dur="0.3s" begin="1.7s" fill="freeze"/></path>
247
+ <path d="M210,60 C228,68 245,78 262,92" fill="none" stroke="url(#pSh)" stroke-width="3.8" stroke-linecap="round" opacity="0"><animate attributeName="opacity" to="0.45" dur="0.3s" begin="1.75s" fill="freeze"/></path>
248
+ <g filter="url(#nG)" opacity="0"><animate attributeName="opacity" to="1" dur="0.3s" begin="1.85s" fill="freeze"/>
249
+ <circle cx="210" cy="60" r="9" fill="url(#s1)"/><circle cx="210" cy="60" r="9" fill="url(#s1sp)"/>
250
+ </g>
251
+
252
+ <!-- B→A -->
253
+ <path d="M210,60 C192,68 175,80 158,92" fill="none" stroke="#93C5FD" stroke-width="3.8" stroke-linecap="round" stroke-dasharray="65" stroke-dashoffset="65" opacity="0"><animate attributeName="stroke-dashoffset" to="0" dur="0.6s" begin="1.9s" fill="freeze"/><animate attributeName="opacity" to="0.85" dur="0.1s" begin="1.9s" fill="freeze"/></path>
254
+ <path d="M158,92 C175,80 192,68 210,60" fill="none" stroke="#1E3A8A" stroke-width="8" stroke-linecap="round" opacity="0"><animate attributeName="opacity" to="0.7" dur="0.3s" begin="2.1s" fill="freeze"/></path>
255
+ <path d="M158,92 C175,80 192,68 210,60" fill="none" stroke="url(#pSh)" stroke-width="3.5" stroke-linecap="round" opacity="0"><animate attributeName="opacity" to="0.45" dur="0.3s" begin="2.15s" fill="freeze"/></path>
256
+ <g filter="url(#nG)" opacity="0"><animate attributeName="opacity" to="1" dur="0.3s" begin="2.25s" fill="freeze"/>
257
+ <circle cx="158" cy="92" r="9" fill="url(#s1)"/><circle cx="158" cy="92" r="9" fill="url(#s1sp)"/>
258
+ </g>
259
+
260
+ <!-- D→E -->
261
+ <path d="M262,148 C245,158 228,168 210,175" fill="none" stroke="#93C5FD" stroke-width="3.8" stroke-linecap="round" stroke-dasharray="62" stroke-dashoffset="62" opacity="0"><animate attributeName="stroke-dashoffset" to="0" dur="0.6s" begin="1.5s" fill="freeze"/><animate attributeName="opacity" to="0.85" dur="0.1s" begin="1.5s" fill="freeze"/></path>
262
+ <path d="M210,175 C228,168 245,158 262,148" fill="none" stroke="#1E3A8A" stroke-width="7.5" stroke-linecap="round" opacity="0"><animate attributeName="opacity" to="0.68" dur="0.3s" begin="1.7s" fill="freeze"/></path>
263
+ <path d="M210,175 C228,168 245,158 262,148" fill="none" stroke="url(#pSh)" stroke-width="3.2" stroke-linecap="round" opacity="0"><animate attributeName="opacity" to="0.42" dur="0.3s" begin="1.75s" fill="freeze"/></path>
264
+ <g filter="url(#nG)" opacity="0"><animate attributeName="opacity" to="1" dur="0.3s" begin="1.85s" fill="freeze"/>
265
+ <circle cx="210" cy="175" r="8" fill="url(#m1)"/><circle cx="210" cy="175" r="8" fill="url(#m1sp)"/>
266
+ </g>
267
+
268
+ <!-- E→F -->
269
+ <path d="M210,175 C192,168 175,158 158,148" fill="none" stroke="#BFDBFE" stroke-width="2.8" stroke-linecap="round" stroke-dasharray="65" stroke-dashoffset="65" opacity="0"><animate attributeName="stroke-dashoffset" to="0" dur="0.6s" begin="1.9s" fill="freeze"/><animate attributeName="opacity" to="0.85" dur="0.1s" begin="1.9s" fill="freeze"/></path>
270
+ <path d="M158,148 C175,158 192,168 210,175" fill="none" stroke="#1D4ED8" stroke-width="5.5" stroke-linecap="round" opacity="0"><animate attributeName="opacity" to="0.5" dur="0.3s" begin="2.1s" fill="freeze"/></path>
271
+ <g filter="url(#nG)" opacity="0"><animate attributeName="opacity" to="1" dur="0.3s" begin="2.25s" fill="freeze"/>
272
+ <circle cx="158" cy="148" r="7" fill="url(#m1)"/><circle cx="158" cy="148" r="7" fill="url(#m1sp)"/>
273
+ </g>
274
+
275
+ <!-- F→A ring close -->
276
+ <path d="M158,148 C158,130 158,112 158,92" fill="none" stroke="#BFDBFE" stroke-width="2.2" stroke-linecap="round" stroke-dasharray="56" stroke-dashoffset="56" opacity="0"><animate attributeName="stroke-dashoffset" to="0" dur="0.5s" begin="2.3s" fill="freeze"/><animate attributeName="opacity" to="0.85" dur="0.1s" begin="2.3s" fill="freeze"/></path>
277
+ <path d="M158,92 C158,112 158,130 158,148" fill="none" stroke="#1D4ED8" stroke-width="5" stroke-linecap="round" opacity="0"><animate attributeName="opacity" to="0.48" dur="0.3s" begin="2.5s" fill="freeze"/></path>
278
+
279
+ <!-- ======= HOP 4: BRANCHES ======= -->
280
+ <!-- A→M1 -->
281
+ <path d="M158,92 C142,80 128,70 112,62" fill="none" stroke="#93C5FD" stroke-width="3.2" stroke-linecap="round" stroke-dasharray="52" stroke-dashoffset="52" opacity="0"><animate attributeName="stroke-dashoffset" to="0" dur="0.5s" begin="2.4s" fill="freeze"/><animate attributeName="opacity" to="0.85" dur="0.1s" begin="2.4s" fill="freeze"/></path>
282
+ <path d="M112,62 C128,70 142,80 158,92" fill="none" stroke="#1E3A8A" stroke-width="6.5" stroke-linecap="round" opacity="0"><animate attributeName="opacity" to="0.58" dur="0.3s" begin="2.55s" fill="freeze"/></path>
283
+ <path d="M112,62 C128,70 142,80 158,92" fill="none" stroke="url(#pSh)" stroke-width="2.8" stroke-linecap="round" opacity="0"><animate attributeName="opacity" to="0.38" dur="0.3s" begin="2.6s" fill="freeze"/></path>
284
+ <g filter="url(#nG)" opacity="0"><animate attributeName="opacity" to="1" dur="0.3s" begin="2.7s" fill="freeze"/>
285
+ <circle cx="112" cy="62" r="7.5" fill="url(#m1)"/><circle cx="112" cy="62" r="7.5" fill="url(#m1sp)"/>
286
+ </g>
287
+ <!-- M1→leaf -->
288
+ <path d="M112,62 C104,52 98,42 92,35" fill="none" stroke="#DBEAFE" stroke-width="1.5" stroke-linecap="round" stroke-dasharray="35" stroke-dashoffset="35" opacity="0"><animate attributeName="stroke-dashoffset" to="0" dur="0.4s" begin="2.8s" fill="freeze"/><animate attributeName="opacity" to="0.85" dur="0.1s" begin="2.8s" fill="freeze"/></path>
289
+ <path d="M92,35 C98,42 104,52 112,62" fill="none" stroke="#3B82F6" stroke-width="2.8" stroke-linecap="round" opacity="0"><animate attributeName="opacity" to="0.25" dur="0.3s" begin="2.95s" fill="freeze"/></path>
290
+ <circle cx="92" cy="35" r="3" fill="url(#lf)" opacity="0"><animate attributeName="opacity" to="1" dur="0.3s" begin="3.05s" fill="freeze"/></circle>
291
+
292
+ <!-- B→M2 up -->
293
+ <path d="M210,60 C210,48 210,35 210,25" fill="none" stroke="#BFDBFE" stroke-width="2" stroke-linecap="round" stroke-dasharray="35" stroke-dashoffset="35" opacity="0"><animate attributeName="stroke-dashoffset" to="0" dur="0.4s" begin="2.1s" fill="freeze"/><animate attributeName="opacity" to="0.85" dur="0.1s" begin="2.1s" fill="freeze"/></path>
294
+ <path d="M210,25 C210,35 210,48 210,60" fill="none" stroke="#2563EB" stroke-width="4" stroke-linecap="round" opacity="0"><animate attributeName="opacity" to="0.35" dur="0.3s" begin="2.25s" fill="freeze"/></path>
295
+ <circle cx="210" cy="25" r="5" fill="url(#w1)" opacity="0"><animate attributeName="opacity" to="1" dur="0.3s" begin="2.35s" fill="freeze"/></circle>
296
+
297
+ <!-- E→M3 down -->
298
+ <path d="M210,175 C210,190 210,205 210,218" fill="none" stroke="#BFDBFE" stroke-width="2" stroke-linecap="round" stroke-dasharray="43" stroke-dashoffset="43" opacity="0"><animate attributeName="stroke-dashoffset" to="0" dur="0.4s" begin="2.1s" fill="freeze"/><animate attributeName="opacity" to="0.85" dur="0.1s" begin="2.1s" fill="freeze"/></path>
299
+ <path d="M210,218 C210,205 210,190 210,175" fill="none" stroke="#2563EB" stroke-width="4" stroke-linecap="round" opacity="0"><animate attributeName="opacity" to="0.38" dur="0.3s" begin="2.25s" fill="freeze"/></path>
300
+ <circle cx="210" cy="218" r="5.5" fill="url(#w1)" opacity="0"><animate attributeName="opacity" to="1" dur="0.3s" begin="2.35s" fill="freeze"/></circle>
301
+ <!-- M3→leaf -->
302
+ <path d="M210,218 C202,232 196,242 190,250" fill="none" stroke="#DBEAFE" stroke-width="1.3" stroke-linecap="round" stroke-dasharray="38" stroke-dashoffset="38" opacity="0"><animate attributeName="stroke-dashoffset" to="0" dur="0.4s" begin="2.5s" fill="freeze"/><animate attributeName="opacity" to="0.85" dur="0.1s" begin="2.5s" fill="freeze"/></path>
303
+ <path d="M190,250 C196,242 202,232 210,218" fill="none" stroke="#93C5FD" stroke-width="1.8" stroke-linecap="round" opacity="0"><animate attributeName="opacity" to="0.16" dur="0.3s" begin="2.65s" fill="freeze"/></path>
304
+ <circle cx="190" cy="250" r="2.5" fill="url(#lf)" opacity="0"><animate attributeName="opacity" to="1" dur="0.3s" begin="2.75s" fill="freeze"/></circle>
305
+
306
+ <!-- Source→ext -->
307
+ <path d="M325,118 C328,132 332,150 335,165" fill="none" stroke="#BFDBFE" stroke-width="1.8" stroke-linecap="round" stroke-dasharray="50" stroke-dashoffset="50" opacity="0"><animate attributeName="stroke-dashoffset" to="0" dur="0.5s" begin="2.0s" fill="freeze"/><animate attributeName="opacity" to="0.85" dur="0.1s" begin="2.0s" fill="freeze"/></path>
308
+ <path d="M335,165 C332,150 328,132 325,118" fill="none" stroke="#2563EB" stroke-width="3.2" stroke-linecap="round" opacity="0"><animate attributeName="opacity" to="0.28" dur="0.3s" begin="2.15s" fill="freeze"/></path>
309
+ <circle cx="335" cy="165" r="4" fill="url(#w1)" opacity="0"><animate attributeName="opacity" to="1" dur="0.3s" begin="2.3s" fill="freeze"/></circle>
310
+
311
+ <!-- STARVED -->
312
+ <path d="M158,148 C142,158 130,168 118,178" fill="none" stroke="#DBEAFE" stroke-width="1" stroke-linecap="round" stroke-dasharray="45" stroke-dashoffset="45" opacity="0"><animate attributeName="stroke-dashoffset" to="0" dur="0.7s" begin="3.0s" fill="freeze"/><animate attributeName="opacity" to="0.85" dur="0.1s" begin="3.0s" fill="freeze"/></path>
313
+ <path d="M118,178 C130,168 142,158 158,148" fill="none" stroke="#60A5FA" stroke-width="2" stroke-linecap="round" opacity="0"><animate attributeName="opacity" to="0.13" dur="0.3s" begin="3.2s" fill="freeze"/></path>
314
+ <circle cx="118" cy="178" r="3.5" fill="url(#st1)" opacity="0"><animate attributeName="opacity" to="1" dur="0.3s" begin="3.3s" fill="freeze"/></circle>
315
+
316
+ <path d="M305,68 C310,55 316,42 322,35" fill="none" stroke="#DBEAFE" stroke-width="1" stroke-linecap="round" stroke-dasharray="40" stroke-dashoffset="40" opacity="0"><animate attributeName="stroke-dashoffset" to="0" dur="0.7s" begin="3.1s" fill="freeze"/><animate attributeName="opacity" to="0.85" dur="0.1s" begin="3.1s" fill="freeze"/></path>
317
+ <path d="M322,35 C316,42 310,55 305,68" fill="none" stroke="#60A5FA" stroke-width="2.5" stroke-linecap="round" opacity="0"><animate attributeName="opacity" to="0.18" dur="0.3s" begin="3.3s" fill="freeze"/></path>
318
+ <circle cx="322" cy="35" r="3.5" fill="url(#st1)" opacity="0"><animate attributeName="opacity" to="1" dur="0.3s" begin="3.4s" fill="freeze"/></circle>
319
+
320
+ <!-- ======= PHASE 2: CONTINUOUS WATER LOOP from source outward (4s+) ======= -->
321
+ <g fill="none" stroke-linecap="round" opacity="0">
322
+ <animate attributeName="opacity" to="1" dur="0.5s" begin="4s" fill="freeze"/>
323
+ <!-- 5-ring -->
324
+ <path d="M325,118 C318,100 312,85 305,68" stroke="#7DD3FC" stroke-width="5.5" stroke-dasharray="12 8"><animate attributeName="stroke-dashoffset" values="0;-40" dur="1s" begin="4s" repeatCount="indefinite"/></path>
325
+ <path d="M325,118 C305,128 282,138 262,148" stroke="#7DD3FC" stroke-width="5" stroke-dasharray="12 8"><animate attributeName="stroke-dashoffset" values="0;-40" dur="1.2s" begin="4s" repeatCount="indefinite"/></path>
326
+ <path d="M305,68 C290,72 275,82 262,92" stroke="#7DD3FC" stroke-width="5" stroke-dasharray="12 8"><animate attributeName="stroke-dashoffset" values="0;-40" dur="1.1s" begin="4s" repeatCount="indefinite"/></path>
327
+ <path d="M262,148 C262,130 262,110 262,92" stroke="#7DD3FC" stroke-width="5" stroke-dasharray="12 8"><animate attributeName="stroke-dashoffset" values="0;-40" dur="1s" begin="4s" repeatCount="indefinite"/></path>
328
+ <!-- 6-ring top (outward from C) -->
329
+ <path d="M262,92 C245,78 228,68 210,60" stroke="#93C5FD" stroke-width="4" stroke-dasharray="10 8"><animate attributeName="stroke-dashoffset" values="0;-36" dur="1.4s" begin="4s" repeatCount="indefinite"/></path>
330
+ <path d="M210,60 C192,68 175,80 158,92" stroke="#93C5FD" stroke-width="3.8" stroke-dasharray="10 8"><animate attributeName="stroke-dashoffset" values="0;-36" dur="1.5s" begin="4s" repeatCount="indefinite"/></path>
331
+ <!-- 6-ring bottom (outward from D) -->
332
+ <path d="M262,148 C245,158 228,168 210,175" stroke="#93C5FD" stroke-width="3.5" stroke-dasharray="10 8"><animate attributeName="stroke-dashoffset" values="0;-36" dur="1.3s" begin="4s" repeatCount="indefinite"/></path>
333
+ <path d="M210,175 C192,168 175,158 158,148" stroke="#BFDBFE" stroke-width="2.5" stroke-dasharray="8 8"><animate attributeName="stroke-dashoffset" values="0;-32" dur="1.6s" begin="4s" repeatCount="indefinite"/></path>
334
+ <path d="M158,148 C158,130 158,112 158,92" stroke="#BFDBFE" stroke-width="2" stroke-dasharray="8 8"><animate attributeName="stroke-dashoffset" values="0;-32" dur="1.8s" begin="4s" repeatCount="indefinite"/></path>
335
+ <!-- Branches (outward) -->
336
+ <path d="M158,92 C142,80 128,70 112,62" stroke="#93C5FD" stroke-width="3" stroke-dasharray="8 7"><animate attributeName="stroke-dashoffset" values="0;-30" dur="1.4s" begin="4s" repeatCount="indefinite"/></path>
337
+ <path d="M210,60 C210,48 210,35 210,25" stroke="#BFDBFE" stroke-width="2" stroke-dasharray="6 7"><animate attributeName="stroke-dashoffset" values="0;-26" dur="1.6s" begin="4s" repeatCount="indefinite"/></path>
338
+ <path d="M210,175 C210,190 210,205 210,218" stroke="#BFDBFE" stroke-width="2" stroke-dasharray="6 7"><animate attributeName="stroke-dashoffset" values="0;-26" dur="1.5s" begin="4s" repeatCount="indefinite"/></path>
339
+ <path d="M325,118 C328,132 332,150 335,165" stroke="#BFDBFE" stroke-width="1.5" stroke-dasharray="5 6"><animate attributeName="stroke-dashoffset" values="0;-22" dur="1.7s" begin="4s" repeatCount="indefinite"/></path>
340
+ </g>
341
+ </svg>
342
+ </div>
343
+ <h1 style="font-size:2.8rem;font-weight:700;letter-spacing:-1px;margin-bottom:24px;color:var(--navy-900);animation:fadeUp 0.6s ease-out 0.3s both;">QAFD-RAG</h1>
344
+ <div class="btn-group">
345
+ <a class="btn btn-primary" href="https://openreview.net/pdf?id=n28wnc2QTc"><svg viewBox="0 0 24 24" fill="currentColor"><path d="M14 2H6a2 2 0 0 0-2 2v16a2 2 0 0 0 2 2h12a2 2 0 0 0 2-2V8l-6-6zm-1 2l5 5h-5V4zM6 20V4h5v7h7v9H6z"/></svg> Paper</a>
346
+ <span class="btn btn-outline" style="cursor:default;opacity:0.7;"><svg viewBox="0 0 24 24" fill="currentColor"><path d="M12 0C5.37 0 0 5.37 0 12c0 5.3 3.438 9.8 8.205 11.385.6.113.82-.258.82-.577 0-.285-.01-1.04-.015-2.04-3.338.724-4.042-1.61-4.042-1.61-.546-1.385-1.335-1.755-1.335-1.755-1.087-.744.084-.729.084-.729 1.205.084 1.838 1.236 1.838 1.236 1.07 1.835 2.809 1.305 3.495.998.108-.776.417-1.305.76-1.605-2.665-.3-5.466-1.332-5.466-5.93 0-1.31.465-2.38 1.235-3.22-.135-.303-.54-1.523.105-3.176 0 0 1.005-.322 3.3 1.23.96-.267 1.98-.399 3-.405 1.02.006 2.04.138 3 .405 2.28-1.552 3.285-1.23 3.285-1.23.645 1.653.24 2.873.12 3.176.765.84 1.23 1.91 1.23 3.22 0 4.61-2.805 5.625-5.475 5.92.42.36.81 1.096.81 2.22 0 1.605-.015 2.896-.015 3.286 0 .315.21.69.825.57C20.565 21.795 24 17.295 24 12 24 5.37 18.63 0 12 0z"/></svg> Code (Coming Soon)</span>
347
+ <a class="btn btn-outline" href="https://huggingface.co/datasets/qafd/kg"><svg xmlns="http://www.w3.org/2000/svg" width="17" height="17" viewBox="0 0 95 88" fill="none"><path fill="#FFD21E" d="M47.21 76.5a34.75 34.75 0 1 0 0-69.5 34.75 34.75 0 0 0 0 69.5Z"/><path fill="#FF9D0B" d="M81.96 41.75a34.75 34.75 0 1 0-69.5 0 34.75 34.75 0 0 0 69.5 0Zm-73.5 0a38.75 38.75 0 1 1 77.5 0 38.75 38.75 0 0 1-77.5 0Z"/><path fill="#3A3B45" d="M58.5 32.3c1.28.44 1.78 3.06 3.07 2.38a5 5 0 1 0-6.76-2.07c.61 1.15 2.55-.72 3.7-.32ZM34.95 32.3c-1.28.44-1.79 3.06-3.07 2.38a5 5 0 1 1 6.76-2.07c-.61 1.15-2.56-.72-3.7-.32Z"/><path fill="#FF323D" d="M46.96 56.29c9.83 0 13-8.76 13-13.26 0-2.34-1.57-1.6-4.09-.36-2.33 1.15-5.46 2.74-8.9 2.74-7.19 0-13-6.88-13-2.38s3.16 13.26 13 13.26Z"/><path fill="#3A3B45" fill-rule="evenodd" d="M39.43 54a8.7 8.7 0 0 1 5.3-4.49c.4-.12.81.57 1.24 1.28.4.68.82 1.37 1.24 1.37.45 0 .9-.68 1.33-1.35.45-.7.89-1.38 1.32-1.25a8.61 8.61 0 0 1 5 4.17c3.73-2.94 5.1-7.74 5.1-10.7 0-2.34-1.57-1.6-4.09-.36l-.14.07c-2.31 1.15-5.39 2.67-8.77 2.67s-6.45-1.52-8.77-2.67c-2.6-1.29-4.23-2.1-4.23.29 0 3.05 1.46 8.06 5.47 10.97Z" clip-rule="evenodd"/><path fill="#FF9D0B" d="M70.71 37a3.25 3.25 0 1 0 0-6.5 3.25 3.25 0 0 0 0 6.5ZM24.21 37a3.25 3.25 0 1 0 0-6.5 3.25 3.25 0 0 0 0 6.5ZM17.52 48c-1.62 0-3.06.66-4.07 1.87a5.97 5.97 0 0 0-1.33 3.76 7.1 7.1 0 0 0-1.94-.3c-1.55 0-2.95.59-3.94 1.66a5.8 5.8 0 0 0-.8 7 5.3 5.3 0 0 0-1.79 2.82c-.24.9-.48 2.8.8 4.74a5.22 5.22 0 0 0-.37 5.02c1.02 2.32 3.57 4.14 8.52 6.1 3.07 1.22 5.89 2 5.91 2.01a44.33 44.33 0 0 0 10.93 1.6c5.86 0 10.05-1.8 12.46-5.34 3.88-5.69 3.33-10.9-1.7-15.92-2.77-2.78-4.62-6.87-5-7.77-.78-2.66-2.84-5.62-6.25-5.62a5.7 5.7 0 0 0-4.6 2.46c-1-1.26-1.98-2.25-2.86-2.82A7.4 7.4 0 0 0 17.52 48Zm0 4c.51 0 1.14.22 1.82.65 2.14 1.36 6.25 8.43 7.76 11.18.5.92 1.37 1.31 2.14 1.31 1.55 0 2.75-1.53.15-3.48-3.92-2.93-2.55-7.72-.68-8.01.08-.02.17-.02.24-.02 1.7 0 2.45 2.93 2.45 2.93s2.2 5.52 5.98 9.3c3.77 3.77 3.97 6.8 1.22 10.83-1.88 2.75-5.47 3.58-9.16 3.58-3.81 0-7.73-.9-9.92-1.46-.11-.03-13.45-3.8-11.76-7 .28-.54.75-.76 1.34-.76 2.38 0 6.7 3.54 8.57 3.54.41 0 .7-.17.83-.6.79-2.85-12.06-4.05-10.98-8.17.2-.73.71-1.02 1.44-1.02 3.14 0 10.2 5.53 11.68 5.53.11 0 .2-.03.24-.1.74-1.2.33-2.04-4.9-5.2-5.21-3.16-8.88-5.06-6.8-7.33.24-.26.58-.38 1-.38 3.17 0 10.66 6.82 10.66 6.82s2.02 2.1 3.25 2.1c.28 0 .52-.1.68-.38.86-1.46-8.06-8.22-8.56-11.01-.34-1.9.24-2.85 1.31-2.85Z"/><path fill="#FFD21E" d="M38.6 76.69c2.75-4.04 2.55-7.07-1.22-10.84-3.78-3.77-5.98-9.3-5.98-9.3s-.82-3.2-2.69-2.9c-1.87.3-3.24 5.08.68 8.01 3.91 2.93-.78 4.92-2.29 2.17-1.5-2.75-5.62-9.82-7.76-11.18-2.13-1.35-3.63-.6-3.13 2.2.5 2.79 9.43 9.55 8.56 11-.87 1.47-3.93-1.71-3.93-1.71s-9.57-8.71-11.66-6.44c-2.08 2.27 1.59 4.17 6.8 7.33 5.23 3.16 5.64 4 4.9 5.2-.75 1.2-12.28-8.53-13.36-4.4-1.08 4.11 11.77 5.3 10.98 8.15-.8 2.85-9.06-5.38-10.74-2.18-1.7 3.21 11.65 6.98 11.76 7.01 4.3 1.12 15.25 3.49 19.08-2.12Z"/><path fill="#FF9D0B" d="M77.4 48c1.62 0 3.07.66 4.07 1.87a5.97 5.97 0 0 1 1.33 3.76 7.1 7.1 0 0 1 1.95-.3c1.55 0 2.95.59 3.94 1.66a5.8 5.8 0 0 1 .8 7 5.3 5.3 0 0 1 1.78 2.82c.24.9.48 2.8-.8 4.74a5.22 5.22 0 0 1 .37 5.02c-1.02 2.32-3.57 4.14-8.51 6.1-3.08 1.22-5.9 2-5.92 2.01a44.33 44.33 0 0 1-10.93 1.6c-5.86 0-10.05-1.8-12.46-5.34-3.88-5.69-3.33-10.9 1.7-15.92 2.78-2.78 4.63-6.87 5.01-7.77.78-2.66 2.83-5.62 6.24-5.62a5.7 5.7 0 0 1 4.6 2.46c1-1.26 1.98-2.25 2.87-2.82A7.4 7.4 0 0 1 77.4 48Zm0 4c-.51 0-1.13.22-1.82.65-2.13 1.36-6.25 8.43-7.76 11.18a2.43 2.43 0 0 1-2.14 1.31c-1.54 0-2.75-1.53-.14-3.48 3.91-2.93 2.54-7.72.67-8.01a1.54 1.54 0 0 0-.24-.02c-1.7 0-2.45 2.93-2.45 2.93s-2.2 5.52-5.97 9.3c-3.78 3.77-3.98 6.8-1.22 10.83 1.87 2.75 5.47 3.58 9.15 3.58 3.82 0 7.73-.9 9.93-1.46.1-.03 13.45-3.8 11.76-7-.29-.54-.75-.76-1.34-.76-2.38 0-6.71 3.54-8.57 3.54-.42 0-.71-.17-.83-.6-.8-2.85 12.05-4.05 10.97-8.17-.19-.73-.7-1.02-1.44-1.02-3.14 0-10.2 5.53-11.68 5.53-.1 0-.19-.03-.23-.1-.74-1.2-.34-2.04 4.88-5.2 5.23-3.16 8.9-5.06 6.8-7.33-.23-.26-.57-.38-.98-.38-3.18 0-10.67 6.82-10.67 6.82s-2.02 2.1-3.24 2.1a.74.74 0 0 1-.68-.38c-.87-1.46 8.05-8.22 8.55-11.01.34-1.9-.24-2.85-1.31-2.85Z"/><path fill="#FFD21E" d="M56.33 76.69c-2.75-4.04-2.56-7.07 1.22-10.84 3.77-3.77 5.97-9.3 5.97-9.3s.82-3.2 2.7-2.9c1.86.3 3.23 5.08-.68 8.01-3.92 2.93.78 4.92 2.28 2.17 1.51-2.75 5.63-9.82 7.76-11.18 2.13-1.35 3.64-.6 3.13 2.2-.5 2.79-9.42 9.55-8.55 11 .86 1.47 3.92-1.71 3.92-1.71s9.58-8.71 11.66-6.44c2.08 2.27-1.58 4.17-6.8 7.33-5.23 3.16-5.63 4-4.9 5.2.75 1.2 12.28-8.53 13.36-4.4 1.08 4.11-11.76 5.3-10.97 8.15.8 2.85 9.05-5.38 10.74-2.18 1.69 3.21-11.65 6.98-11.76 7.01-4.31 1.12-15.26 3.49-19.08-2.12Z"/></svg> Data</a>
348
+ </div>
349
+ </div>
350
+ </header>
351
+
352
+ <section>
353
+ <div class="container">
354
+ <h2 style="font-size:1.55rem;font-weight:700;color:var(--slate-800);margin-bottom:16px;">Query-Aware Flow Diffusion for Graph-Based<br/>Retrieval-Augmented Generation</h2>
355
+ <p class="about-text"><strong style="color:var(--blue-600);">QAFD-RAG</strong> is a graph-based retrieval-augmented generation framework that uses <strong>query-aware flow diffusion</strong> to retrieve contextually relevant subgraphs from a knowledge graph with <strong>statistical retrieval guarantees</strong>.</p>
356
+ <div class="duo-grid">
357
+ <div class="card"><h3>The Problem</h3><p>Existing graph-based RAG methods rely on <strong>static retrieval</strong>&mdash;community detection (GraphRAG) or one-hop entity lookup (LightRAG)&mdash;that ignore query context and miss relevant multi-hop connections.</p></div>
358
+ <div class="card"><h3>Our Solution</h3><p>QAFD-RAG <strong>dynamically re-weights</strong> graph edges based on query relevance and <strong>propagates flow</strong> through the knowledge graph to discover multi-hop context, assembling a query-specific subgraph for the LLM.</p></div>
359
+ </div>
360
+ </div>
361
+ </section>
362
+
363
+ <section class="alt">
364
+ <div class="container">
365
+ <h2>Approach Comparison</h2>
366
+ <p class="section-sub">How different retrieval strategies explore a knowledge graph</p>
367
+ <div class="figure-row">
368
+ <div class="figure-card"><img src="figs/GraphRAG.png" alt="GraphRAG"/><div class="caption">GraphRAG<br/><small>Community Detection</small></div></div>
369
+ <div class="figure-card"><img src="figs/LightRAG.png" alt="LightRAG"/><div class="caption">LightRAG<br/><small>One-Hop Entity-Centric</small></div></div>
370
+ <div class="figure-card"><img src="figs/QAFD-RAG.png" alt="QAFD-RAG"/><div class="caption">QAFD-RAG (Ours)<br/><small>Query-Aware Flow Diffusion</small></div></div>
371
+ </div>
372
+ <p style="margin-top:24px;font-size:0.92rem;color:var(--slate-500);line-height:1.75;max-width:860px;margin-left:auto;margin-right:auto;">Comparison on Wikipedia pages (Apple fruit, Apple Inc., Amazon River, Amazon.com). Query: <em>"Introduce Steve Jobs's products in Apple."</em> <strong style="color:var(--slate-700);">GraphRAG</strong> retrieves entire communities, mixing relevant nodes with irrelevant ones. <strong style="color:var(--slate-700);">LightRAG</strong> focuses on 1-hop neighborhoods. <strong style="color:var(--slate-700);">QAFD-RAG</strong> reweights edges by query meaning, suppressing irrelevant neighborhoods.</p>
373
+ </div>
374
+ </section>
375
+
376
+ <section>
377
+ <div class="container">
378
+ <h2>How It Works</h2>
379
+ <p class="section-sub">A two-stage pipeline from raw documents to grounded LLM answers</p>
380
+ <div class="duo-grid">
381
+ <div class="card"><h3>Stage 1 &mdash; Knowledge Graph Construction</h3><ol><li><strong>Chunking</strong> &mdash; Documents split into overlapping token-based chunks.</li><li><strong>Entity Extraction</strong> &mdash; LLM identifies entities and relationships.</li><li><strong>Graph Assembly</strong> &mdash; Entities become nodes, relationships become weighted edges.</li></ol></div>
382
+ <div class="card"><h3>Stage 2 &mdash; Query-Aware Retrieval</h3><ol><li><strong>Entity Matching</strong> &mdash; Query keywords matched to KG entities via vector similarity.</li><li><strong>Flow Diffusion</strong> &mdash; Mass propagates with query-aware edge re-weighting.</li><li><strong>Context Assembly</strong> &mdash; Top-ranked nodes form context for the LLM.</li></ol></div>
383
+ </div>
384
+ </div>
385
+ </section>
386
+
387
+ <section class="alt">
388
+ <div class="container">
389
+ <h2>Benchmarks</h2>
390
+ <p class="section-sub">Evaluated across diverse retrieval and generation tasks</p>
391
+ <div class="bench-grid">
392
+ <div class="bench-item"><div class="task">UltraDomain</div><div class="dataset">Diverse Domain QA</div><div class="metrics">Comprehensiveness · Diversity · Relevance · Logicality · Coherence</div></div>
393
+ <div class="bench-item"><div class="task">Multi-hop QA</div><div class="dataset">MuSiQue · HotpotQA · 2WikiMultiHopQA</div><div class="metrics">F1 · Exact Match</div></div>
394
+ <div class="bench-item"><div class="task">Text-to-SQL</div><div class="dataset">Spider2-lite (Pagila, etc.)</div><div class="metrics">Schema Retrieval Accuracy</div></div>
395
+ <div class="bench-item"><div class="task">Summarization</div><div class="dataset">SQuALITY</div><div class="metrics">BLEU · ROUGE · METEOR</div></div>
396
+ </div>
397
+ </div>
398
+ </section>
399
+
400
+ <section>
401
+ <div class="container">
402
+ <h2>Supported Models</h2>
403
+ <p class="models-text"><strong>Embeddings:</strong> openai-small · openai-large · jina-v3 · gritlm · nvidia-nv-embed-v2<br/><strong>LLMs:</strong> gpt-4o-mini · gpt-4o · gpt-5-nano · gpt-5-mini · gpt-5 · gpt-oss-120b</p>
404
+ </div>
405
+ </section>
406
+
407
+ <section class="alt">
408
+ <div class="container">
409
+ <h2>Quick Start</h2>
410
+ <pre><code><span class="comment"># 1. Install dependencies</span>
411
+ pip install -r requirements.txt
412
+ <span class="comment"># 2. Set your OpenAI API key</span>
413
+ export OPENAI_API_KEY=<span class="string">"sk-..."</span>
414
+ <span class="comment"># 3. Build a knowledge graph</span>
415
+ ./run.sh ultradomain --build --max-documents 100
416
+ <span class="comment"># 4. Run a benchmark</span>
417
+ ./run.sh ultradomain --questions 10</code></pre>
418
+ <h3>Python API</h3>
419
+ <pre><code><span class="keyword">from</span> src <span class="keyword">import</span> QAFD_RAG, QueryParam
420
+ rag = QAFD_RAG(
421
+ working_dir=<span class="string">"./my_kg"</span>,
422
+ llm_model_name=<span class="string">"gpt-4o-mini"</span>,
423
+ embedding_model_key=<span class="string">"jina-v3"</span>,
424
+ )
425
+ <span class="comment"># Index documents</span>
426
+ rag.insert([<span class="string">"Document text 1..."</span>, <span class="string">"Document text 2..."</span>])
427
+ <span class="comment"># Query</span>
428
+ answer = rag.query(<span class="string">"What is X?"</span>, param=QueryParam(mode=<span class="string">"hybrid"</span>))
429
+ print(answer)</code></pre>
430
+ </div>
431
+ </section>
432
+
433
+ <section>
434
+ <div class="container">
435
+ <h2>Citation</h2>
436
+ <pre><code>@inproceedings{zhou2026qafd,
437
+ title={Query-Aware Flow Diffusion for Graph-Based RAG with Retrieval Guarantees},
438
+ author={Zhuoping Zhou and Davoud Ataee Tarzanagh and Sima Didari and Wenjun Hu
439
+ and Baruch Gutow and Oxana Verkholyak and Masoud Faraki and Heng Hao
440
+ and Hankyu Moon and Seungjai Min},
441
+ booktitle={International Conference on Learning Representations (ICLR)},
442
+ year={2026}
443
+ }</code></pre>
444
+ </div>
445
+ </section>
446
+
447
+ <footer>
448
+ <p>QAFD-RAG <span class="sep">·</span><a href="https://openreview.net/pdf?id=n28wnc2QTc">Paper</a> <span class="sep">·</span><a href="https://github.com/Tarzanagh/QAFD-RAG">GitHub</a> <span class="sep">·</span><a href="https://huggingface.co/datasets/qafd/kg">HuggingFace</a></p>
449
+ <p style="margin-top:8px;opacity:0.6;">Contact: tarzanagh [at] gmail [dot] com</p>
450
+ </footer>
451
+ </body>
452
+ </html>
experiments/edge_weight_ablation.py ADDED
@@ -0,0 +1,401 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Edge Weight Ablation Study for QAFD-RAG
3
+ ========================================
4
+
5
+ Runs all experiments from a single script to avoid reloading the KG each time.
6
+
7
+ Experiments:
8
+ 1. Query-aware vs query-agnostic
9
+ 2. Weight scheme comparison (original/multiply/add)
10
+ 3. (a, b) sensitivity sweep for Hybrid
11
+ 4. Qualitative per-query diagnostics
12
+ """
13
+
14
+ import json
15
+ import logging
16
+ import os
17
+ import sys
18
+ import time
19
+ from typing import Dict, List, Tuple
20
+
21
+ import numpy as np
22
+
23
+ # ---------------------------------------------------------------------------
24
+ # Setup paths
25
+ # ---------------------------------------------------------------------------
26
+ _project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
27
+ sys.path.insert(0, _project_root)
28
+
29
+ import types as _types
30
+ for _pkg_path in ["src", "src.retrievers", "src.hipporag_pipeline"]:
31
+ if _pkg_path not in sys.modules:
32
+ _m = _types.ModuleType(_pkg_path)
33
+ _m.__path__ = [os.path.join(_project_root, *_pkg_path.split("."))]
34
+ _m.__package__ = _pkg_path
35
+ sys.modules[_pkg_path] = _m
36
+
37
+ import importlib.util as _ilu
38
+ def _load_mod(fqn, filepath):
39
+ spec = _ilu.spec_from_file_location(fqn, filepath)
40
+ mod = _ilu.module_from_spec(spec)
41
+ sys.modules[fqn] = mod
42
+ spec.loader.exec_module(mod)
43
+ return mod
44
+
45
+ _src = os.path.join(_project_root, "src")
46
+ _load_mod("src.retrievers.base", os.path.join(_src, "retrievers", "base.py"))
47
+ _load_mod("src.retrievers.flow_diffusion", os.path.join(_src, "retrievers", "flow_diffusion.py"))
48
+
49
+ from src.hipporag_pipeline.config import HippoRAGConfig
50
+ from src.hipporag_pipeline.embedding_store import EmbeddingModelWrapper
51
+ from src.hipporag_pipeline.kg_builder import KGBuilder
52
+ from src.hipporag_pipeline.openie import OpenIE
53
+ from src.hipporag_pipeline.reranker import FactReranker
54
+ from src.hipporag_pipeline.retriever import HippoRAGRetriever
55
+ from src.hipporag_pipeline.graph_adapter import run_igraph_qafd
56
+ from src.hipporag_pipeline.benchmark_runner import (
57
+ get_gold_docs, get_gold_answers, recall_at_k,
58
+ exact_match, f1_score, run_qa, _openai_embed, _openai_complete,
59
+ )
60
+
61
+ logging.basicConfig(level=logging.WARNING)
62
+ logger = logging.getLogger(__name__)
63
+
64
+ # ---------------------------------------------------------------------------
65
+ # Load KG once
66
+ # ---------------------------------------------------------------------------
67
+
68
+ def load_everything(dataset="musique", num_queries=10):
69
+ """Load KG, embeddings, data — once for all experiments."""
70
+ api_key = os.environ.get("OPENAI_API_KEY", "")
71
+
72
+ config = HippoRAGConfig(
73
+ llm_model="gpt-4o-mini",
74
+ embedding_model_key="openai-small",
75
+ dataset=dataset,
76
+ save_dir="outputs",
77
+ )
78
+
79
+ import asyncio
80
+
81
+ async def llm_func(prompt, system_prompt=None, history_messages=[], **kwargs):
82
+ return await _openai_complete(
83
+ model="gpt-4o-mini", prompt=prompt,
84
+ system_prompt=system_prompt, history_messages=history_messages,
85
+ api_key=api_key, **kwargs
86
+ )
87
+
88
+ async def embed_func(texts):
89
+ return await _openai_embed(texts, model="text-embedding-3-small", api_key=api_key)
90
+
91
+ embedding_model = EmbeddingModelWrapper(embed_func, batch_size=16)
92
+ openie = OpenIE(llm_func)
93
+ builder = KGBuilder(config, embedding_model, openie)
94
+
95
+ # Load data
96
+ data_dir = os.path.join(_project_root, "data", "multihop")
97
+ with open(os.path.join(data_dir, f"{dataset}_corpus.json")) as f:
98
+ corpus = json.load(f)
99
+ docs = [f"{d['title']}\n{d['text']}" for d in corpus]
100
+
101
+ with open(os.path.join(data_dir, f"{dataset}.json")) as f:
102
+ samples = json.load(f)[:num_queries]
103
+
104
+ queries = [s["question"] for s in samples]
105
+ gold_answers = get_gold_answers(samples)
106
+ gold_docs = get_gold_docs(samples, dataset)
107
+
108
+ # Build / load KG
109
+ builder.index(docs)
110
+
111
+ # Build retriever
112
+ reranker = FactReranker(llm_func)
113
+ retriever = HippoRAGRetriever(
114
+ config=config,
115
+ embedding_model=embedding_model,
116
+ reranker=reranker,
117
+ graph=builder.graph,
118
+ chunk_embedding_store=builder.chunk_embedding_store,
119
+ entity_embedding_store=builder.entity_embedding_store,
120
+ fact_embedding_store=builder.fact_embedding_store,
121
+ openie_results_path=builder.openie_results_path,
122
+ )
123
+ retriever.prepare()
124
+
125
+ return config, retriever, queries, gold_answers, gold_docs, llm_func
126
+
127
+
128
+ def run_retrieval_with_params(
129
+ retriever: HippoRAGRetriever,
130
+ queries: List[str],
131
+ gold_docs: List[List[str]],
132
+ gold_answers,
133
+ weight_scheme: str = "original",
134
+ hybrid_a: float = 1.0,
135
+ hybrid_b: float = 0.5,
136
+ query_aware: bool = True,
137
+ ) -> Dict:
138
+ """Run retrieval with specific edge weight params. Returns metrics dict."""
139
+ from src.hipporag_pipeline.utils import compute_mdhash_id, min_max_normalize
140
+
141
+ retriever._encode_queries(queries)
142
+ k_list = [1, 2, 5, 10, 20, 50, 100, 200]
143
+ all_retrieved = []
144
+ qafd_times = []
145
+ convergence_iters = []
146
+
147
+ for q in queries:
148
+ fact_scores = retriever._get_fact_scores(q)
149
+ top_indices, top_facts, _ = retriever._rerank_facts(q, fact_scores)
150
+
151
+ if len(top_facts) == 0:
152
+ sorted_ids, sorted_scores = retriever._dense_passage_retrieval(q)
153
+ all_retrieved.append([
154
+ retriever.chunk_store.get_row(retriever.passage_node_keys[idx])["content"]
155
+ for idx in sorted_ids[:200]
156
+ ])
157
+ qafd_times.append(0.0)
158
+ convergence_iters.append(-1) # DPR fallback
159
+ continue
160
+
161
+ # Compute seed weights (same as retriever._graph_search)
162
+ n_nodes = retriever.graph.vcount()
163
+ phrase_weights = np.zeros(n_nodes)
164
+ passage_weights = np.zeros(n_nodes)
165
+ number_of_occurs = np.zeros(n_nodes)
166
+
167
+ for rank, f in enumerate(top_facts):
168
+ subj, obj = f[0].lower(), f[2].lower()
169
+ fs = fact_scores[top_indices[rank]] if fact_scores.ndim > 0 else float(fact_scores)
170
+ for phrase in [subj, obj]:
171
+ pk = compute_mdhash_id(phrase, prefix="entity-")
172
+ pid = retriever.node_name_to_vertex_idx.get(pk)
173
+ if pid is not None:
174
+ wfs = fs
175
+ num_chunks = len(retriever.ent_node_to_chunk_ids.get(pk, set()))
176
+ if num_chunks > 0:
177
+ wfs /= num_chunks
178
+ phrase_weights[pid] += wfs
179
+ number_of_occurs[pid] += 1
180
+
181
+ nonzero = number_of_occurs > 0
182
+ phrase_weights[nonzero] /= number_of_occurs[nonzero]
183
+
184
+ dpr_ids, dpr_scores = retriever._dense_passage_retrieval(q)
185
+ norm_dpr = min_max_normalize(dpr_scores)
186
+ pw = retriever.config.passage_node_weight
187
+ for i, did in enumerate(dpr_ids.tolist()):
188
+ pk = retriever.passage_node_keys[did]
189
+ pid = retriever.node_name_to_vertex_idx.get(pk)
190
+ if pid is not None:
191
+ passage_weights[pid] = norm_dpr[i] * pw
192
+
193
+ node_weights = phrase_weights + passage_weights
194
+ if np.sum(node_weights) == 0:
195
+ sorted_ids, sorted_scores = dpr_ids, dpr_scores
196
+ all_retrieved.append([
197
+ retriever.chunk_store.get_row(retriever.passage_node_keys[idx])["content"]
198
+ for idx in sorted_ids[:200]
199
+ ])
200
+ qafd_times.append(0.0)
201
+ convergence_iters.append(-1)
202
+ continue
203
+
204
+ # Build node embeddings
205
+ if not hasattr(retriever, '_node_emb_dict') or retriever._node_emb_dict is None:
206
+ retriever._node_emb_dict = {}
207
+ for i, nk in enumerate(retriever.entity_node_keys):
208
+ if i < len(retriever.entity_embeddings):
209
+ retriever._node_emb_dict[nk] = retriever.entity_embeddings[i]
210
+ for i, nk in enumerate(retriever.passage_node_keys):
211
+ if i < len(retriever.passage_embeddings):
212
+ retriever._node_emb_dict[nk] = retriever.passage_embeddings[i]
213
+
214
+ query_emb = retriever._query_emb_fact.get(q) if query_aware else None
215
+
216
+ t0 = time.time()
217
+ sorted_ids, sorted_scores = run_igraph_qafd(
218
+ graph=retriever.graph,
219
+ node_name_to_idx=retriever.node_name_to_vertex_idx,
220
+ passage_node_idxs=retriever.passage_node_idxs,
221
+ source_weights=node_weights,
222
+ node_embeddings=retriever._node_emb_dict if query_aware else {},
223
+ query_embedding=query_emb,
224
+ alpha=retriever.config.qafd_alpha,
225
+ epsilon=retriever.config.qafd_epsilon,
226
+ max_iterations=retriever.config.qafd_max_iterations,
227
+ step_size=retriever.config.qafd_step_size,
228
+ weight_scheme=weight_scheme,
229
+ hybrid_a=hybrid_a,
230
+ hybrid_b=hybrid_b,
231
+ use_node_degree=retriever.config.qafd_use_node_degree,
232
+ random_seed=retriever.config.qafd_random_seed,
233
+ )
234
+ elapsed = time.time() - t0
235
+ qafd_times.append(elapsed)
236
+
237
+ top_docs = [
238
+ retriever.chunk_store.get_row(retriever.passage_node_keys[idx])["content"]
239
+ for idx in sorted_ids[:200]
240
+ ]
241
+ all_retrieved.append(top_docs)
242
+
243
+ # Compute metrics
244
+ recall_metrics = recall_at_k(gold_docs, all_retrieved, k_list) if gold_docs else {}
245
+
246
+ return {
247
+ "recall": recall_metrics,
248
+ "avg_qafd_time": np.mean([t for t in qafd_times if t > 0]) if any(t > 0 for t in qafd_times) else 0,
249
+ "dpr_fallback_count": sum(1 for t in convergence_iters if t == -1),
250
+ "qafd_queries": sum(1 for t in qafd_times if t > 0),
251
+ }
252
+
253
+
254
+ # ===========================================================================
255
+ # Main experiments
256
+ # ===========================================================================
257
+
258
+ def main():
259
+ import argparse
260
+ parser = argparse.ArgumentParser()
261
+ parser.add_argument("--dataset", default="musique")
262
+ parser.add_argument("--num_queries", type=int, default=10)
263
+ args = parser.parse_args()
264
+
265
+ print("=" * 70)
266
+ print(" QAFD-RAG Edge Weight Ablation Study")
267
+ print(f" Dataset: {args.dataset}, Queries: {args.num_queries}")
268
+ print("=" * 70)
269
+
270
+ config, retriever, queries, gold_answers, gold_docs, llm_func = load_everything(
271
+ args.dataset, args.num_queries
272
+ )
273
+
274
+ results = {}
275
+
276
+ # ── Experiment 1: Query-aware vs Query-agnostic ──────────────────
277
+ print("\n[1/4] Query-aware vs Query-agnostic ablation")
278
+
279
+ print(" Running: Hybrid (a=1, b=0.5) — query-aware ...")
280
+ r = run_retrieval_with_params(
281
+ retriever, queries, gold_docs, gold_answers,
282
+ weight_scheme="original", hybrid_a=1.0, hybrid_b=0.5, query_aware=True,
283
+ )
284
+ results["hybrid_query_aware"] = r
285
+ print(f" R@10={r['recall'].get('Recall@10', 0):.4f} R@100={r['recall'].get('Recall@100', 0):.4f} QAFD={r['avg_qafd_time']:.3f}s DPR_fallback={r['dpr_fallback_count']}")
286
+
287
+ print(" Running: Query-agnostic (b=0) ...")
288
+ r = run_retrieval_with_params(
289
+ retriever, queries, gold_docs, gold_answers,
290
+ weight_scheme="original", hybrid_a=1.0, hybrid_b=0.0, query_aware=False,
291
+ )
292
+ results["query_agnostic"] = r
293
+ print(f" R@10={r['recall'].get('Recall@10', 0):.4f} R@100={r['recall'].get('Recall@100', 0):.4f} QAFD={r['avg_qafd_time']:.3f}s DPR_fallback={r['dpr_fallback_count']}")
294
+
295
+ # ── Experiment 2: Weight scheme comparison ───────────────────────
296
+ print("\n[2/4] Weight scheme comparison")
297
+
298
+ for scheme in ["original", "multiply", "add"]:
299
+ print(f" Running: {scheme} ...")
300
+ r = run_retrieval_with_params(
301
+ retriever, queries, gold_docs, gold_answers,
302
+ weight_scheme=scheme, query_aware=True,
303
+ )
304
+ results[f"scheme_{scheme}"] = r
305
+ print(f" R@10={r['recall'].get('Recall@10', 0):.4f} R@100={r['recall'].get('Recall@100', 0):.4f} QAFD={r['avg_qafd_time']:.3f}s")
306
+
307
+ # ── Experiment 3: (a, b) sensitivity sweep ───────────────────────
308
+ print("\n[3/4] (a, b) sensitivity sweep for Hybrid")
309
+
310
+ a_values = [0.5, 1.0, 2.0]
311
+ b_values = [0.0, 0.1, 0.25, 0.5, 1.0]
312
+
313
+ sweep_results = {}
314
+ for a in a_values:
315
+ for b in b_values:
316
+ label = f"a={a},b={b}"
317
+ r = run_retrieval_with_params(
318
+ retriever, queries, gold_docs, gold_answers,
319
+ weight_scheme="original", hybrid_a=a, hybrid_b=b, query_aware=(b > 0),
320
+ )
321
+ sweep_results[label] = r
322
+ r10 = r['recall'].get('Recall@10', 0)
323
+ r100 = r['recall'].get('Recall@100', 0)
324
+ print(f" {label:>15} R@10={r10:.4f} R@100={r100:.4f} QAFD={r['avg_qafd_time']:.3f}s")
325
+
326
+ results["sweep"] = sweep_results
327
+
328
+ # ── Experiment 4: Per-query diagnostics ──────────────────────────
329
+ print("\n[4/4] Per-query diagnostics (first 5 queries)")
330
+
331
+ # Compare query-aware vs agnostic per query
332
+ diag_queries = queries[:5]
333
+ diag_results = []
334
+
335
+ for qi, q in enumerate(diag_queries):
336
+ # Query-aware
337
+ r_aware = run_retrieval_with_params(
338
+ retriever, [q], [gold_docs[qi]], [gold_answers[qi]],
339
+ weight_scheme="original", hybrid_a=1.0, hybrid_b=0.5, query_aware=True,
340
+ )
341
+ # Query-agnostic
342
+ r_agnostic = run_retrieval_with_params(
343
+ retriever, [q], [gold_docs[qi]], [gold_answers[qi]],
344
+ weight_scheme="original", hybrid_a=1.0, hybrid_b=0.0, query_aware=False,
345
+ )
346
+
347
+ r10_aware = r_aware['recall'].get('Recall@10', 0)
348
+ r10_agnostic = r_agnostic['recall'].get('Recall@10', 0)
349
+ delta = r10_aware - r10_agnostic
350
+
351
+ status = "HELPS" if delta > 0 else ("HURTS" if delta < 0 else "SAME")
352
+ print(f" Q{qi}: R@10 aware={r10_aware:.3f} agnostic={r10_agnostic:.3f} delta={delta:+.3f} [{status}]")
353
+ print(f" Q: {q[:80]}...")
354
+
355
+ diag_results.append({
356
+ "query": q,
357
+ "r10_aware": r10_aware,
358
+ "r10_agnostic": r10_agnostic,
359
+ "delta": delta,
360
+ "status": status,
361
+ })
362
+
363
+ results["diagnostics"] = diag_results
364
+
365
+ # ── Summary ──────────────────────────────────────────────────────
366
+ print("\n" + "=" * 70)
367
+ print(" SUMMARY")
368
+ print("=" * 70)
369
+
370
+ print("\n Ablation: Query-aware vs Query-agnostic")
371
+ print(f" Query-aware (Hybrid a=1,b=0.5): R@10={results['hybrid_query_aware']['recall'].get('Recall@10',0):.4f} R@100={results['hybrid_query_aware']['recall'].get('Recall@100',0):.4f}")
372
+ print(f" Query-agnostic (b=0): R@10={results['query_agnostic']['recall'].get('Recall@10',0):.4f} R@100={results['query_agnostic']['recall'].get('Recall@100',0):.4f}")
373
+
374
+ print("\n Weight scheme comparison")
375
+ for scheme in ["original", "multiply", "add"]:
376
+ r = results[f"scheme_{scheme}"]
377
+ print(f" {scheme:>10}: R@10={r['recall'].get('Recall@10',0):.4f} R@100={r['recall'].get('Recall@100',0):.4f}")
378
+
379
+ print("\n Best (a,b) from sweep:")
380
+ best_label = max(sweep_results, key=lambda k: sweep_results[k]['recall'].get('Recall@10', 0))
381
+ best = sweep_results[best_label]
382
+ print(f" {best_label}: R@10={best['recall'].get('Recall@10',0):.4f} R@100={best['recall'].get('Recall@100',0):.4f}")
383
+
384
+ print("\n Per-query diagnostics:")
385
+ helps = sum(1 for d in diag_results if d["status"] == "HELPS")
386
+ hurts = sum(1 for d in diag_results if d["status"] == "HURTS")
387
+ same = sum(1 for d in diag_results if d["status"] == "SAME")
388
+ print(f" Query-awareness HELPS: {helps}/{len(diag_results)}, HURTS: {hurts}/{len(diag_results)}, SAME: {same}/{len(diag_results)}")
389
+
390
+ # Save
391
+ out_dir = os.path.join(_project_root, "experiments", "results")
392
+ os.makedirs(out_dir, exist_ok=True)
393
+ out_path = os.path.join(out_dir, f"edge_weight_ablation_{args.dataset}.json")
394
+ with open(out_path, "w") as f:
395
+ json.dump(results, f, indent=2, default=str)
396
+ print(f"\n Results saved to {out_path}")
397
+ print("=" * 70)
398
+
399
+
400
+ if __name__ == "__main__":
401
+ main()
experiments/query_aware_ablation.py ADDED
@@ -0,0 +1,478 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Ablation: Query-Aware vs Query-Agnostic Flow Diffusion
3
+ =======================================================
4
+
5
+ Condition A (Query-Agnostic): b=0 → w̄ = H_sim(h(u), h(v)) · a
6
+ Condition B (Query-Aware): b=0.25 → w̄ = H_sim(h(u), h(v)) · (a + b·(sim(u,q)+sim(v,q)))
7
+
8
+ Measures:
9
+ - Downstream task quality (Recall@K, F1, EM)
10
+ - Subgraph size (nodes with nonzero flow)
11
+ - Leakage ratio (mass at irrelevant vs relevant nodes)
12
+ - Convergence iterations
13
+ - Per-query qualitative diagnostics
14
+ """
15
+
16
+ import json
17
+ import logging
18
+ import os
19
+ import sys
20
+ import time
21
+ from collections import defaultdict
22
+ from typing import Dict, List, Optional, Tuple
23
+
24
+ import numpy as np
25
+
26
+ # ---------------------------------------------------------------------------
27
+ _project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
28
+ sys.path.insert(0, _project_root)
29
+
30
+ import types as _types
31
+ for _pkg_path in ["src", "src.retrievers", "src.hipporag_pipeline"]:
32
+ if _pkg_path not in sys.modules:
33
+ _m = _types.ModuleType(_pkg_path)
34
+ _m.__path__ = [os.path.join(_project_root, *_pkg_path.split("."))]
35
+ _m.__package__ = _pkg_path
36
+ sys.modules[_pkg_path] = _m
37
+
38
+ import importlib.util as _ilu
39
+ def _load_mod(fqn, filepath):
40
+ spec = _ilu.spec_from_file_location(fqn, filepath)
41
+ mod = _ilu.module_from_spec(spec)
42
+ sys.modules[fqn] = mod
43
+ spec.loader.exec_module(mod)
44
+ return mod
45
+
46
+ _src = os.path.join(_project_root, "src")
47
+ _load_mod("src.retrievers.base", os.path.join(_src, "retrievers", "base.py"))
48
+ _load_mod("src.retrievers.flow_diffusion", os.path.join(_src, "retrievers", "flow_diffusion.py"))
49
+
50
+ from src.hipporag_pipeline.config import HippoRAGConfig
51
+ from src.hipporag_pipeline.embedding_store import EmbeddingModelWrapper
52
+ from src.hipporag_pipeline.kg_builder import KGBuilder
53
+ from src.hipporag_pipeline.openie import OpenIE
54
+ from src.hipporag_pipeline.reranker import FactReranker
55
+ from src.hipporag_pipeline.retriever import HippoRAGRetriever
56
+ from src.hipporag_pipeline.graph_adapter import IGraphQAFD
57
+ from src.hipporag_pipeline.benchmark_runner import (
58
+ get_gold_docs, get_gold_answers, recall_at_k,
59
+ exact_match, f1_score, run_qa, _openai_embed, _openai_complete,
60
+ )
61
+ from src.hipporag_pipeline.utils import compute_mdhash_id, min_max_normalize
62
+
63
+ logging.basicConfig(level=logging.WARNING)
64
+
65
+
66
+ # ===========================================================================
67
+ # Data + KG loading (once per dataset)
68
+ # ===========================================================================
69
+
70
+ def load_dataset_and_kg(dataset="musique", num_queries=10):
71
+ api_key = os.environ.get("OPENAI_API_KEY", "")
72
+ config = HippoRAGConfig(
73
+ llm_model="gpt-4o-mini",
74
+ embedding_model_key="openai-small",
75
+ dataset=dataset,
76
+ save_dir="outputs",
77
+ )
78
+ import asyncio
79
+
80
+ async def llm_func(prompt, system_prompt=None, history_messages=[], **kwargs):
81
+ return await _openai_complete(
82
+ model="gpt-4o-mini", prompt=prompt,
83
+ system_prompt=system_prompt, history_messages=history_messages,
84
+ api_key=api_key, **kwargs
85
+ )
86
+
87
+ async def embed_func(texts):
88
+ return await _openai_embed(texts, model="text-embedding-3-small", api_key=api_key)
89
+
90
+ embedding_model = EmbeddingModelWrapper(embed_func, batch_size=16)
91
+ openie = OpenIE(llm_func)
92
+ builder = KGBuilder(config, embedding_model, openie)
93
+
94
+ data_dir = os.path.join(_project_root, "data", "multihop")
95
+ with open(os.path.join(data_dir, f"{dataset}_corpus.json")) as f:
96
+ corpus = json.load(f)
97
+ docs = [f"{d['title']}\n{d['text']}" for d in corpus]
98
+
99
+ with open(os.path.join(data_dir, f"{dataset}.json")) as f:
100
+ samples = json.load(f)[:num_queries]
101
+
102
+ queries = [s["question"] for s in samples]
103
+ gold_answers = get_gold_answers(samples)
104
+ gold_docs = get_gold_docs(samples, dataset)
105
+
106
+ builder.index(docs)
107
+
108
+ reranker = FactReranker(llm_func)
109
+ retriever = HippoRAGRetriever(
110
+ config=config, embedding_model=embedding_model, reranker=reranker,
111
+ graph=builder.graph,
112
+ chunk_embedding_store=builder.chunk_embedding_store,
113
+ entity_embedding_store=builder.entity_embedding_store,
114
+ fact_embedding_store=builder.fact_embedding_store,
115
+ openie_results_path=builder.openie_results_path,
116
+ )
117
+ retriever.prepare()
118
+
119
+ # Cache node embeddings
120
+ retriever._node_emb_dict = {}
121
+ for i, nk in enumerate(retriever.entity_node_keys):
122
+ if i < len(retriever.entity_embeddings):
123
+ retriever._node_emb_dict[nk] = retriever.entity_embeddings[i]
124
+ for i, nk in enumerate(retriever.passage_node_keys):
125
+ if i < len(retriever.passage_embeddings):
126
+ retriever._node_emb_dict[nk] = retriever.passage_embeddings[i]
127
+
128
+ return config, retriever, queries, gold_answers, gold_docs, llm_func
129
+
130
+
131
+ # ===========================================================================
132
+ # Core: run QAFD with detailed diagnostics
133
+ # ===========================================================================
134
+
135
+ def run_qafd_detailed(
136
+ retriever: HippoRAGRetriever,
137
+ query: str,
138
+ gold_doc_set: set,
139
+ alpha: float,
140
+ weight_scheme: str,
141
+ hybrid_a: float,
142
+ hybrid_b: float,
143
+ query_aware: bool,
144
+ ) -> Dict:
145
+ """Run QAFD on a single query, return detailed metrics."""
146
+
147
+ retriever._encode_queries([query])
148
+ fact_scores = retriever._get_fact_scores(query)
149
+ top_indices, top_facts, _ = retriever._rerank_facts(query, fact_scores)
150
+
151
+ # DPR fallback
152
+ if len(top_facts) == 0:
153
+ sorted_ids, sorted_scores = retriever._dense_passage_retrieval(query)
154
+ top_docs = [
155
+ retriever.chunk_store.get_row(retriever.passage_node_keys[idx])["content"]
156
+ for idx in sorted_ids[:200]
157
+ ]
158
+ return {
159
+ "method": "DPR_fallback",
160
+ "top_docs": top_docs,
161
+ "subgraph_size": 0,
162
+ "leakage_ratio": 1.0,
163
+ "convergence_iters": -1,
164
+ "qafd_time": 0.0,
165
+ "flow_at_relevant": 0.0,
166
+ "flow_at_irrelevant": 0.0,
167
+ }
168
+
169
+ # Compute seed weights
170
+ n_nodes = retriever.graph.vcount()
171
+ phrase_weights = np.zeros(n_nodes)
172
+ passage_weights = np.zeros(n_nodes)
173
+ number_of_occurs = np.zeros(n_nodes)
174
+
175
+ for rank, f in enumerate(top_facts):
176
+ subj, obj = f[0].lower(), f[2].lower()
177
+ fs = fact_scores[top_indices[rank]] if fact_scores.ndim > 0 else float(fact_scores)
178
+ for phrase in [subj, obj]:
179
+ pk = compute_mdhash_id(phrase, prefix="entity-")
180
+ pid = retriever.node_name_to_vertex_idx.get(pk)
181
+ if pid is not None:
182
+ wfs = fs
183
+ nc = len(retriever.ent_node_to_chunk_ids.get(pk, set()))
184
+ if nc > 0:
185
+ wfs /= nc
186
+ phrase_weights[pid] += wfs
187
+ number_of_occurs[pid] += 1
188
+
189
+ nonzero = number_of_occurs > 0
190
+ phrase_weights[nonzero] /= number_of_occurs[nonzero]
191
+
192
+ dpr_ids, dpr_scores = retriever._dense_passage_retrieval(query)
193
+ norm_dpr = min_max_normalize(dpr_scores)
194
+ pw = retriever.config.passage_node_weight
195
+ for i, did in enumerate(dpr_ids.tolist()):
196
+ pk = retriever.passage_node_keys[did]
197
+ pid = retriever.node_name_to_vertex_idx.get(pk)
198
+ if pid is not None:
199
+ passage_weights[pid] = norm_dpr[i] * pw
200
+
201
+ node_weights = phrase_weights + passage_weights
202
+ if np.sum(node_weights) == 0:
203
+ sorted_ids, sorted_scores = dpr_ids, dpr_scores
204
+ top_docs = [
205
+ retriever.chunk_store.get_row(retriever.passage_node_keys[idx])["content"]
206
+ for idx in sorted_ids[:200]
207
+ ]
208
+ return {
209
+ "method": "DPR_fallback_zero_seeds",
210
+ "top_docs": top_docs,
211
+ "subgraph_size": 0,
212
+ "leakage_ratio": 1.0,
213
+ "convergence_iters": -1,
214
+ "qafd_time": 0.0,
215
+ "flow_at_relevant": 0.0,
216
+ "flow_at_irrelevant": 0.0,
217
+ }
218
+
219
+ query_emb = retriever._query_emb_fact.get(query) if query_aware else None
220
+ node_embs = retriever._node_emb_dict if query_aware else {}
221
+
222
+ # Run QAFD directly to get raw node scores
223
+ t0 = time.time()
224
+ qafd = IGraphQAFD(
225
+ graph=retriever.graph,
226
+ node_name_to_idx=retriever.node_name_to_vertex_idx,
227
+ source_weights=node_weights,
228
+ node_embeddings=node_embs,
229
+ query_embedding=query_emb,
230
+ alpha=alpha,
231
+ epsilon=retriever.config.qafd_epsilon,
232
+ max_iterations=retriever.config.qafd_max_iterations,
233
+ step_size=retriever.config.qafd_step_size,
234
+ weight_scheme=weight_scheme,
235
+ hybrid_a=hybrid_a,
236
+ hybrid_b=hybrid_b,
237
+ use_node_degree=retriever.config.qafd_use_node_degree,
238
+ random_seed=retriever.config.qafd_random_seed,
239
+ )
240
+ raw_scores = qafd.run()
241
+ elapsed = time.time() - t0
242
+
243
+ # Extract passage scores
244
+ doc_scores = np.array([raw_scores[idx] for idx in retriever.passage_node_idxs])
245
+ total = np.sum(doc_scores)
246
+ if total > 0:
247
+ doc_scores_norm = doc_scores / total
248
+ else:
249
+ doc_scores_norm = np.ones(len(doc_scores)) / max(len(doc_scores), 1)
250
+
251
+ sorted_ids = np.argsort(doc_scores_norm)[::-1]
252
+ sorted_scores = doc_scores_norm[sorted_ids]
253
+
254
+ top_docs = [
255
+ retriever.chunk_store.get_row(retriever.passage_node_keys[idx])["content"]
256
+ for idx in sorted_ids[:200]
257
+ ]
258
+
259
+ # --- Compute detailed metrics ---
260
+
261
+ # Subgraph size: nodes with nonzero flow
262
+ subgraph_size = int(np.sum(raw_scores > 1e-10))
263
+
264
+ # Leakage ratio: flow at irrelevant vs relevant passage nodes
265
+ flow_relevant = 0.0
266
+ flow_irrelevant = 0.0
267
+ relevant_count = 0
268
+ irrelevant_count = 0
269
+
270
+ for i, pk in enumerate(retriever.passage_node_keys):
271
+ content = retriever.chunk_store.get_row(pk)["content"]
272
+ score = doc_scores[i]
273
+ is_gold = any(g in content or content in g for g in gold_doc_set)
274
+ if is_gold:
275
+ flow_relevant += score
276
+ relevant_count += 1
277
+ else:
278
+ flow_irrelevant += score
279
+ irrelevant_count += 1
280
+
281
+ total_flow = flow_relevant + flow_irrelevant
282
+ leakage = flow_irrelevant / total_flow if total_flow > 0 else 1.0
283
+
284
+ return {
285
+ "method": "QAFD",
286
+ "top_docs": top_docs,
287
+ "subgraph_size": subgraph_size,
288
+ "leakage_ratio": round(leakage, 4),
289
+ "flow_at_relevant": round(flow_relevant, 6),
290
+ "flow_at_irrelevant": round(flow_irrelevant, 6),
291
+ "relevant_passages": relevant_count,
292
+ "qafd_time": round(elapsed, 4),
293
+ }
294
+
295
+
296
+ # ===========================================================================
297
+ # Run full experiment
298
+ # ===========================================================================
299
+
300
+ def run_experiment(
301
+ retriever, queries, gold_docs, gold_answers,
302
+ alpha: float, hybrid_a: float, hybrid_b: float,
303
+ query_aware: bool, label: str,
304
+ ) -> Dict:
305
+ """Run all queries with given params, return aggregate metrics."""
306
+ all_docs = []
307
+ subgraph_sizes = []
308
+ leakage_ratios = []
309
+ flow_relevants = []
310
+ flow_irrelevants = []
311
+ qafd_times = []
312
+ per_query = []
313
+
314
+ for qi, q in enumerate(queries):
315
+ gold_set = set(gold_docs[qi]) if gold_docs else set()
316
+ r = run_qafd_detailed(
317
+ retriever, q, gold_set,
318
+ alpha=alpha,
319
+ weight_scheme="original",
320
+ hybrid_a=hybrid_a,
321
+ hybrid_b=hybrid_b,
322
+ query_aware=query_aware,
323
+ )
324
+ all_docs.append(r["top_docs"])
325
+ subgraph_sizes.append(r["subgraph_size"])
326
+ leakage_ratios.append(r["leakage_ratio"])
327
+ flow_relevants.append(r["flow_at_relevant"])
328
+ flow_irrelevants.append(r["flow_at_irrelevant"])
329
+ qafd_times.append(r["qafd_time"])
330
+ per_query.append({
331
+ "query": q[:100],
332
+ "method": r["method"],
333
+ "subgraph_size": r["subgraph_size"],
334
+ "leakage_ratio": r["leakage_ratio"],
335
+ "flow_relevant": r["flow_at_relevant"],
336
+ "flow_irrelevant": r["flow_at_irrelevant"],
337
+ })
338
+
339
+ # Recall
340
+ k_list = [1, 2, 5, 10, 20, 50, 100, 200]
341
+ recall_metrics = recall_at_k(gold_docs, all_docs, k_list) if gold_docs else {}
342
+
343
+ return {
344
+ "label": label,
345
+ "alpha": alpha,
346
+ "hybrid_a": hybrid_a,
347
+ "hybrid_b": hybrid_b,
348
+ "query_aware": query_aware,
349
+ "recall": recall_metrics,
350
+ "avg_subgraph_size": round(np.mean(subgraph_sizes), 1),
351
+ "avg_leakage_ratio": round(np.mean(leakage_ratios), 4),
352
+ "avg_flow_relevant": round(np.mean(flow_relevants), 6),
353
+ "avg_flow_irrelevant": round(np.mean(flow_irrelevants), 6),
354
+ "avg_qafd_time": round(np.mean(qafd_times), 4),
355
+ "per_query": per_query,
356
+ }
357
+
358
+
359
+ def main():
360
+ import argparse
361
+ parser = argparse.ArgumentParser()
362
+ parser.add_argument("--dataset", default="musique")
363
+ parser.add_argument("--num_queries", type=int, default=10)
364
+ args = parser.parse_args()
365
+
366
+ print("=" * 70)
367
+ print(" Query-Aware vs Query-Agnostic Ablation Study")
368
+ print(f" Dataset: {args.dataset}, Queries: {args.num_queries}")
369
+ print("=" * 70)
370
+
371
+ config, retriever, queries, gold_answers, gold_docs, llm_func = \
372
+ load_dataset_and_kg(args.dataset, args.num_queries)
373
+
374
+ print(f" Graph: {retriever.graph.vcount()} nodes, {retriever.graph.ecount()} edges")
375
+ print(f" Entities: {len(retriever.entity_node_keys)}, Passages: {len(retriever.passage_node_keys)}")
376
+
377
+ all_results = {}
378
+
379
+ # ── Run across multiple alpha values ─────────────────────────────
380
+ alpha_values = [2.0, 10.0, 50.0]
381
+
382
+ for alpha in alpha_values:
383
+ print(f"\n{'─' * 70}")
384
+ print(f" Alpha = {alpha}")
385
+ print(f"{'─' * 70}")
386
+
387
+ # Condition A: Query-Agnostic (b=0)
388
+ print(f" Running Query-Agnostic (b=0) ...")
389
+ r_agnostic = run_experiment(
390
+ retriever, queries, gold_docs, gold_answers,
391
+ alpha=alpha, hybrid_a=1.0, hybrid_b=0.0,
392
+ query_aware=False, label=f"agnostic_a{alpha}",
393
+ )
394
+ all_results[f"agnostic_a{alpha}"] = r_agnostic
395
+
396
+ # Condition B: Query-Aware (a=1, b=0.25)
397
+ print(f" Running Query-Aware (a=1, b=0.25) ...")
398
+ r_aware_025 = run_experiment(
399
+ retriever, queries, gold_docs, gold_answers,
400
+ alpha=alpha, hybrid_a=1.0, hybrid_b=0.25,
401
+ query_aware=True, label=f"aware025_a{alpha}",
402
+ )
403
+ all_results[f"aware025_a{alpha}"] = r_aware_025
404
+
405
+ # Condition C: Query-Aware (a=1, b=0.5) — current default
406
+ print(f" Running Query-Aware (a=1, b=0.5) ...")
407
+ r_aware_050 = run_experiment(
408
+ retriever, queries, gold_docs, gold_answers,
409
+ alpha=alpha, hybrid_a=1.0, hybrid_b=0.5,
410
+ query_aware=True, label=f"aware050_a{alpha}",
411
+ )
412
+ all_results[f"aware050_a{alpha}"] = r_aware_050
413
+
414
+ # Print comparison table
415
+ print(f"\n {'Metric':<25} {'Agnostic(b=0)':>15} {'Aware(b=0.25)':>15} {'Aware(b=0.5)':>15} {'Δ(0.25 vs 0)':>15}")
416
+ print(f" {'─' * 85}")
417
+
418
+ for k in [10, 50, 100]:
419
+ key = f"Recall@{k}"
420
+ va = r_agnostic['recall'].get(key, 0)
421
+ vb = r_aware_025['recall'].get(key, 0)
422
+ vc = r_aware_050['recall'].get(key, 0)
423
+ d = vb - va
424
+ marker = " ↑" if d > 0 else (" ↓" if d < 0 else "")
425
+ print(f" {key:<25} {va:>15.4f} {vb:>15.4f} {vc:>15.4f} {d:>+14.4f}{marker}")
426
+
427
+ print(f" {'Subgraph size':<25} {r_agnostic['avg_subgraph_size']:>15.1f} {r_aware_025['avg_subgraph_size']:>15.1f} {r_aware_050['avg_subgraph_size']:>15.1f} {r_aware_025['avg_subgraph_size'] - r_agnostic['avg_subgraph_size']:>+14.1f}")
428
+ print(f" {'Leakage ratio':<25} {r_agnostic['avg_leakage_ratio']:>15.4f} {r_aware_025['avg_leakage_ratio']:>15.4f} {r_aware_050['avg_leakage_ratio']:>15.4f} {r_aware_025['avg_leakage_ratio'] - r_agnostic['avg_leakage_ratio']:>+14.4f}")
429
+ print(f" {'Flow@relevant':<25} {r_agnostic['avg_flow_relevant']:>15.6f} {r_aware_025['avg_flow_relevant']:>15.6f} {r_aware_050['avg_flow_relevant']:>15.6f} {r_aware_025['avg_flow_relevant'] - r_agnostic['avg_flow_relevant']:>+14.6f}")
430
+ print(f" {'Flow@irrelevant':<25} {r_agnostic['avg_flow_irrelevant']:>15.6f} {r_aware_025['avg_flow_irrelevant']:>15.6f} {r_aware_050['avg_flow_irrelevant']:>15.6f} {r_aware_025['avg_flow_irrelevant'] - r_agnostic['avg_flow_irrelevant']:>+14.6f}")
431
+ print(f" {'QAFD time (s)':<25} {r_agnostic['avg_qafd_time']:>15.4f} {r_aware_025['avg_qafd_time']:>15.4f} {r_aware_050['avg_qafd_time']:>15.4f}")
432
+
433
+ # ── Per-query diagnostics (alpha=10, first 5 queries) ────────────
434
+ print(f"\n{'=' * 70}")
435
+ print(" Per-Query Diagnostics (alpha=10.0)")
436
+ print(f"{'=' * 70}")
437
+
438
+ r_ag = all_results.get("agnostic_a10.0", {})
439
+ r_aw = all_results.get("aware025_a10.0", {})
440
+
441
+ if r_ag and r_aw:
442
+ ag_pq = r_ag.get("per_query", [])
443
+ aw_pq = r_aw.get("per_query", [])
444
+ for qi in range(min(5, len(ag_pq))):
445
+ ag = ag_pq[qi]
446
+ aw = aw_pq[qi]
447
+ print(f"\n Q{qi}: {ag['query']}")
448
+ print(f" {'':>20} {'Agnostic':>12} {'Aware':>12} {'Delta':>12}")
449
+ print(f" {'Subgraph size':>20} {ag['subgraph_size']:>12} {aw['subgraph_size']:>12} {aw['subgraph_size']-ag['subgraph_size']:>+12}")
450
+ print(f" {'Leakage ratio':>20} {ag['leakage_ratio']:>12.4f} {aw['leakage_ratio']:>12.4f} {aw['leakage_ratio']-ag['leakage_ratio']:>+12.4f}")
451
+ print(f" {'Flow@relevant':>20} {ag['flow_relevant']:>12.6f} {aw['flow_relevant']:>12.6f} {aw['flow_relevant']-ag['flow_relevant']:>+12.6f}")
452
+ print(f" {'Flow@irrelevant':>20} {ag['flow_irrelevant']:>12.6f} {aw['flow_irrelevant']:>12.6f} {aw['flow_irrelevant']-ag['flow_irrelevant']:>+12.6f}")
453
+
454
+ status = "HELPS" if aw['leakage_ratio'] < ag['leakage_ratio'] else (
455
+ "HURTS" if aw['leakage_ratio'] > ag['leakage_ratio'] else "SAME"
456
+ )
457
+ print(f" → Query awareness {status} (leakage {'decreased' if status == 'HELPS' else 'increased' if status == 'HURTS' else 'unchanged'})")
458
+
459
+ # ── Save ─────────────────────────────────────────────────────────
460
+ out_dir = os.path.join(_project_root, "experiments", "results")
461
+ os.makedirs(out_dir, exist_ok=True)
462
+ out_path = os.path.join(out_dir, f"query_aware_ablation_{args.dataset}.json")
463
+
464
+ # Remove top_docs from saved output (too large)
465
+ save_results = {}
466
+ for k, v in all_results.items():
467
+ sv = dict(v)
468
+ sv.pop("per_query", None)
469
+ save_results[k] = sv
470
+
471
+ with open(out_path, "w") as f:
472
+ json.dump(save_results, f, indent=2, default=str)
473
+ print(f"\n Results saved to {out_path}")
474
+ print("=" * 70)
475
+
476
+
477
+ if __name__ == "__main__":
478
+ main()
experiments/run_all_ablations.sh ADDED
@@ -0,0 +1,49 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/bin/bash
2
+ # Run all query-awareness ablation experiments on musique (100 queries)
3
+ # Each run reuses the pre-built KG — only retrieval is re-run.
4
+
5
+ DATASET="musique"
6
+ N=100
7
+ BASE="python src/hipporag_pipeline/benchmark_runner.py --task multihop --dataset $DATASET --num_queries $N --embedding_model openai-small --skip_qa"
8
+
9
+ echo "======================================================================"
10
+ echo " QAFD-RAG Query-Awareness Ablation (${DATASET}, ${N} queries)"
11
+ echo "======================================================================"
12
+
13
+ # Baseline (all flags = 0)
14
+ echo ""
15
+ echo "[1/7] Baseline (original, all QA flags = 0)"
16
+ $BASE --save_dir outputs_ablation/baseline 2>&1 | grep -E "Recall@|QAFD completed|Retrieval done"
17
+
18
+ # Phase 1: Query-aware sink capacity
19
+ echo ""
20
+ echo "[2/7] Phase 1: qa_sink_gamma=0.5"
21
+ $BASE --qa_sink_gamma 0.5 --save_dir outputs_ablation/sink_05 2>&1 | grep -E "Recall@|QAFD completed|Retrieval done"
22
+
23
+ echo ""
24
+ echo "[3/7] Phase 1: qa_sink_gamma=1.0"
25
+ $BASE --qa_sink_gamma 1.0 --save_dir outputs_ablation/sink_10 2>&1 | grep -E "Recall@|QAFD completed|Retrieval done"
26
+
27
+ # Phase 2: Query-aware warm start
28
+ echo ""
29
+ echo "[4/7] Phase 2: qa_warm_delta=0.5"
30
+ $BASE --qa_warm_delta 0.5 --save_dir outputs_ablation/warm_05 2>&1 | grep -E "Recall@|QAFD completed|Retrieval done"
31
+
32
+ # Phase 3: Post-diffusion reranking
33
+ echo ""
34
+ echo "[5/7] Phase 3: qa_post_lambda=0.5"
35
+ $BASE --qa_post_lambda 0.5 --save_dir outputs_ablation/post_05 2>&1 | grep -E "Recall@|QAFD completed|Retrieval done"
36
+
37
+ echo ""
38
+ echo "[6/7] Phase 3: qa_post_lambda=1.0"
39
+ $BASE --qa_post_lambda 1.0 --save_dir outputs_ablation/post_10 2>&1 | grep -E "Recall@|QAFD completed|Retrieval done"
40
+
41
+ # Combined best
42
+ echo ""
43
+ echo "[7/7] Combined: sink=0.5 + warm=0.5 + post=0.5"
44
+ $BASE --qa_sink_gamma 0.5 --qa_warm_delta 0.5 --qa_post_lambda 0.5 --save_dir outputs_ablation/combined 2>&1 | grep -E "Recall@|QAFD completed|Retrieval done"
45
+
46
+ echo ""
47
+ echo "======================================================================"
48
+ echo " ALL EXPERIMENTS COMPLETE"
49
+ echo "======================================================================"
experiments/run_alpha_qa_sweep.sh ADDED
@@ -0,0 +1,29 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/bin/bash
2
+ # Alpha × QA sweep: find the sweet spot where decoupled QA edge weights help
3
+ cd /home/davoud/QAFD-RAG
4
+ N=100
5
+ BASE="python src/hipporag_pipeline/benchmark_runner.py --task multihop --num_queries $N --embedding_model openai-small --skip_qa"
6
+
7
+ echo "======================================================================"
8
+ echo " Alpha x Query-Awareness Sweep (decoupled push, $N queries)"
9
+ echo "======================================================================"
10
+
11
+ for DATASET in musique hotpotqa 2wikimultihopqa; do
12
+ echo ""
13
+ echo "====== Dataset: $DATASET ======"
14
+
15
+ for ALPHA in 2 3 5 8; do
16
+ echo " --- alpha=$ALPHA ---"
17
+
18
+ echo " QA-aware (original, a=1, b=0.5)"
19
+ $BASE --dataset $DATASET --qafd_alpha $ALPHA 2>&1 | grep "Recall@10:\|Recall@100:"
20
+
21
+ echo " QA-agnostic (b=0)"
22
+ $BASE --dataset $DATASET --qafd_alpha $ALPHA --qafd_weight_scheme none 2>&1 | grep "Recall@10:\|Recall@100:"
23
+ done
24
+ done
25
+
26
+ echo ""
27
+ echo "======================================================================"
28
+ echo " ALL EXPERIMENTS COMPLETE"
29
+ echo "======================================================================"
experiments/run_batch_push_ablation.sh ADDED
@@ -0,0 +1,29 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/bin/bash
2
+ # Batch push ablation: does batch push make edge weights matter?
3
+ # Compare: single vs batch, agnostic vs aware
4
+ cd /home/davoud/QAFD-RAG
5
+ N=100
6
+ BASE="python src/hipporag_pipeline/benchmark_runner.py --task multihop --num_queries $N --embedding_model openai-small --skip_qa"
7
+
8
+ echo "======================================================================"
9
+ echo " Batch Push Ablation ($N queries per dataset)"
10
+ echo "======================================================================"
11
+
12
+ for DATASET in musique hotpotqa 2wikimultihopqa; do
13
+ echo ""
14
+ echo "====== Dataset: $DATASET ======"
15
+
16
+ echo " [1] Single push, query-aware (current default)"
17
+ $BASE --dataset $DATASET 2>&1 | grep -E "Recall@|QAFD:"
18
+
19
+ echo " [2] Batch push, query-aware"
20
+ $BASE --dataset $DATASET --batch_push 2>&1 | grep -E "Recall@|QAFD:"
21
+
22
+ echo " [3] Batch push, query-agnostic (b=0)"
23
+ $BASE --dataset $DATASET --batch_push --qafd_weight_scheme none 2>&1 | grep -E "Recall@|QAFD:"
24
+ done
25
+
26
+ echo ""
27
+ echo "======================================================================"
28
+ echo " ALL EXPERIMENTS COMPLETE"
29
+ echo "======================================================================"
experiments/run_batch_push_nvidia.sh ADDED
@@ -0,0 +1,55 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/bin/bash
2
+ # Batch push sweep on nvidia KG (dense, 1.6M edges)
3
+ # Goal: find config where QA-aware beats agnostic AND recall is high
4
+ cd /home/davoud/QAFD-RAG
5
+ N=100
6
+ BASE="python src/hipporag_pipeline/benchmark_runner.py --task multihop --num_queries $N --embedding_model nvidia-nv-embed-v2 --skip_qa"
7
+ LOG=experiments/results/batch_push_nvidia.log
8
+ mkdir -p experiments/results
9
+ exec > >(tee -a $LOG) 2>&1
10
+
11
+ echo "======================================================================"
12
+ echo " Batch Push Sweep on nvidia KG ($N queries)"
13
+ echo " Started: $(date)"
14
+ echo "======================================================================"
15
+
16
+ for DATASET in musique hotpotqa; do
17
+ echo ""
18
+ echo "====== $DATASET ======"
19
+
20
+ # Baseline: single push, alpha=2, eps=0.01
21
+ echo " [1] single, a=2, eps=0.01, QA-aware"
22
+ $BASE --dataset $DATASET --qafd_alpha 2 --qafd_epsilon 0.01 2>&1 | grep "Recall@"
23
+ echo " [2] single, a=2, eps=0.01, agnostic"
24
+ $BASE --dataset $DATASET --qafd_alpha 2 --qafd_epsilon 0.01 --qafd_weight_scheme none 2>&1 | grep "Recall@"
25
+
26
+ # Batch push, alpha=2, eps=0.01
27
+ echo " [3] batch, a=2, eps=0.01, QA-aware"
28
+ $BASE --dataset $DATASET --qafd_alpha 2 --qafd_epsilon 0.01 --batch_push 2>&1 | grep "Recall@"
29
+ echo " [4] batch, a=2, eps=0.01, agnostic"
30
+ $BASE --dataset $DATASET --qafd_alpha 2 --qafd_epsilon 0.01 --batch_push --qafd_weight_scheme none 2>&1 | grep "Recall@"
31
+
32
+ # Batch push, alpha=2, eps=0.005
33
+ echo " [5] batch, a=2, eps=0.005, QA-aware"
34
+ $BASE --dataset $DATASET --qafd_alpha 2 --qafd_epsilon 0.005 --batch_push 2>&1 | grep "Recall@"
35
+ echo " [6] batch, a=2, eps=0.005, agnostic"
36
+ $BASE --dataset $DATASET --qafd_alpha 2 --qafd_epsilon 0.005 --batch_push --qafd_weight_scheme none 2>&1 | grep "Recall@"
37
+
38
+ # Batch push, alpha=2, eps=0.003
39
+ echo " [7] batch, a=2, eps=0.003, QA-aware"
40
+ $BASE --dataset $DATASET --qafd_alpha 2 --qafd_epsilon 0.003 --batch_push 2>&1 | grep "Recall@"
41
+ echo " [8] batch, a=2, eps=0.003, agnostic"
42
+ $BASE --dataset $DATASET --qafd_alpha 2 --qafd_epsilon 0.003 --batch_push --qafd_weight_scheme none 2>&1 | grep "Recall@"
43
+
44
+ # Batch push, alpha=3, eps=0.005
45
+ echo " [9] batch, a=3, eps=0.005, QA-aware"
46
+ $BASE --dataset $DATASET --qafd_alpha 3 --qafd_epsilon 0.005 --batch_push 2>&1 | grep "Recall@"
47
+ echo " [10] batch, a=3, eps=0.005, agnostic"
48
+ $BASE --dataset $DATASET --qafd_alpha 3 --qafd_epsilon 0.005 --batch_push --qafd_weight_scheme none 2>&1 | grep "Recall@"
49
+ done
50
+
51
+ echo ""
52
+ echo "======================================================================"
53
+ echo " ALL EXPERIMENTS COMPLETE"
54
+ echo " Finished: $(date)"
55
+ echo "======================================================================"
experiments/run_decoupled_ablation.sh ADDED
@@ -0,0 +1,32 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/bin/bash
2
+ # Decoupled push ablation: does separating accumulation from routing
3
+ # make query-aware edge weights effective?
4
+ cd /home/davoud/QAFD-RAG
5
+ N=100
6
+ BASE="python src/hipporag_pipeline/benchmark_runner.py --task multihop --num_queries $N --embedding_model openai-small --skip_qa"
7
+
8
+ echo "======================================================================"
9
+ echo " Decoupled Push Ablation ($N queries per dataset)"
10
+ echo "======================================================================"
11
+
12
+ for DATASET in musique hotpotqa 2wikimultihopqa; do
13
+ echo ""
14
+ echo "====== Dataset: $DATASET ======"
15
+
16
+ echo " [1] Single push, query-aware (decoupled)"
17
+ $BASE --dataset $DATASET 2>&1 | grep -E "Recall@|QAFD:"
18
+
19
+ echo " [2] Single push, query-agnostic (decoupled, b=0 baseline)"
20
+ $BASE --dataset $DATASET --qafd_weight_scheme none 2>&1 | grep -E "Recall@|QAFD:"
21
+
22
+ echo " [3] Batch push, query-aware (decoupled)"
23
+ $BASE --dataset $DATASET --batch_push 2>&1 | grep -E "Recall@|QAFD:"
24
+
25
+ echo " [4] Batch push, query-agnostic (decoupled, b=0 baseline)"
26
+ $BASE --dataset $DATASET --batch_push --qafd_weight_scheme none 2>&1 | grep -E "Recall@|QAFD:"
27
+ done
28
+
29
+ echo ""
30
+ echo "======================================================================"
31
+ echo " ALL EXPERIMENTS COMPLETE"
32
+ echo "======================================================================"
experiments/run_nvidia_qa_ablation.sh ADDED
@@ -0,0 +1,43 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/bin/bash
2
+ # Query-awareness ablation on nvidia-nv-embed-v2 KGs (dense, 1.6M edges)
3
+ # Tests: QA-aware vs agnostic, different alpha/epsilon settings
4
+ cd /home/davoud/QAFD-RAG
5
+ N=100
6
+ BASE="python src/hipporag_pipeline/benchmark_runner.py --task multihop --num_queries $N --embedding_model nvidia-nv-embed-v2 --skip_qa"
7
+
8
+ echo "======================================================================"
9
+ echo " Query-Awareness on nvidia KG (dense graph, $N queries)"
10
+ echo " Started: $(date)"
11
+ echo "======================================================================"
12
+
13
+ for DATASET in musique hotpotqa 2wikimultihopqa; do
14
+ echo ""
15
+ echo "====== $DATASET ======"
16
+
17
+ # Alpha=2, epsilon=0.01 (original settings)
18
+ echo " --- alpha=2, eps=0.01 ---"
19
+ echo " [QA-aware]"
20
+ $BASE --dataset $DATASET --qafd_alpha 2.0 --qafd_epsilon 0.01 2>&1 | grep "Recall@10:\|Recall@100:\|QAFD:"
21
+ echo " [QA-agnostic]"
22
+ $BASE --dataset $DATASET --qafd_alpha 2.0 --qafd_epsilon 0.01 --qafd_weight_scheme none 2>&1 | grep "Recall@10:\|Recall@100:\|QAFD:"
23
+
24
+ # Alpha=3, epsilon=0.01
25
+ echo " --- alpha=3, eps=0.01 ---"
26
+ echo " [QA-aware]"
27
+ $BASE --dataset $DATASET --qafd_alpha 3.0 --qafd_epsilon 0.01 2>&1 | grep "Recall@10:\|Recall@100:\|QAFD:"
28
+ echo " [QA-agnostic]"
29
+ $BASE --dataset $DATASET --qafd_alpha 3.0 --qafd_epsilon 0.01 --qafd_weight_scheme none 2>&1 | grep "Recall@10:\|Recall@100:\|QAFD:"
30
+
31
+ # Alpha=3, epsilon=0.001 (tighter convergence)
32
+ echo " --- alpha=3, eps=0.001 ---"
33
+ echo " [QA-aware]"
34
+ $BASE --dataset $DATASET --qafd_alpha 3.0 --qafd_epsilon 0.001 2>&1 | grep "Recall@10:\|Recall@100:\|QAFD:"
35
+ echo " [QA-agnostic]"
36
+ $BASE --dataset $DATASET --qafd_alpha 3.0 --qafd_epsilon 0.001 --qafd_weight_scheme none 2>&1 | grep "Recall@10:\|Recall@100:\|QAFD:"
37
+ done
38
+
39
+ echo ""
40
+ echo "======================================================================"
41
+ echo " ALL EXPERIMENTS COMPLETE"
42
+ echo " Finished: $(date)"
43
+ echo "======================================================================"
experiments/run_overnight_qa_awareness.sh ADDED
@@ -0,0 +1,136 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/bin/bash
2
+ # Overnight query-awareness experiments
3
+ # Tests decoupled push with alpha=3: QA-aware vs agnostic
4
+ # Across tasks, graph types, and embedding models
5
+ cd /home/davoud/QAFD-RAG
6
+ LOG=/home/davoud/QAFD-RAG/experiments/results/overnight_qa_awareness.log
7
+ mkdir -p experiments/results
8
+ exec > >(tee -a $LOG) 2>&1
9
+ echo "======================================================================"
10
+ echo " Overnight Query-Awareness Experiments"
11
+ echo " Started: $(date)"
12
+ echo "======================================================================"
13
+
14
+ BASE_PE="python src/hipporag_pipeline/benchmark_runner.py"
15
+
16
+ # =====================================================================
17
+ # 1. PASSAGE-ENTITY: Multihop with OpenAI embedding (100q, alpha=3)
18
+ # =====================================================================
19
+ echo ""
20
+ echo "############ PASSAGE-ENTITY: MULTIHOP (openai-small) ############"
21
+
22
+ for DATASET in musique hotpotqa 2wikimultihopqa; do
23
+ echo ""
24
+ echo "====== PE / $DATASET / openai-small ======"
25
+
26
+ echo " [QA-aware] alpha=3"
27
+ $BASE_PE --task multihop --dataset $DATASET --num_queries 100 --skip_qa \
28
+ --embedding_model openai-small --qafd_alpha 3.0 2>&1 | grep "Recall@"
29
+
30
+ echo " [QA-agnostic] alpha=3"
31
+ $BASE_PE --task multihop --dataset $DATASET --num_queries 100 --skip_qa \
32
+ --embedding_model openai-small --qafd_alpha 3.0 --qafd_weight_scheme none 2>&1 | grep "Recall@"
33
+ done
34
+
35
+ # =====================================================================
36
+ # 2. PASSAGE-ENTITY: Multihop with LOCAL embeddings (50q, alpha=3)
37
+ # =====================================================================
38
+ echo ""
39
+ echo "############ PASSAGE-ENTITY: MULTIHOP (local embeddings) ############"
40
+
41
+ echo ""
42
+ echo "====== PE / musique / jina-v3 ======"
43
+ echo " [QA-aware]"
44
+ $BASE_PE --task multihop --dataset musique --num_queries 50 --skip_qa \
45
+ --embedding_model jina-v3 --qafd_alpha 3.0 2>&1 | grep "Recall@"
46
+ echo " [QA-agnostic]"
47
+ $BASE_PE --task multihop --dataset musique --num_queries 50 --skip_qa \
48
+ --embedding_model jina-v3 --qafd_alpha 3.0 --qafd_weight_scheme none 2>&1 | grep "Recall@"
49
+
50
+ echo ""
51
+ echo "====== PE / musique / nvidia-nv-embed-v2 ======"
52
+ echo " [QA-aware]"
53
+ $BASE_PE --task multihop --dataset musique --num_queries 50 --skip_qa \
54
+ --embedding_model nvidia-nv-embed-v2 --qafd_alpha 3.0 2>&1 | grep "Recall@"
55
+ echo " [QA-agnostic]"
56
+ $BASE_PE --task multihop --dataset musique --num_queries 50 --skip_qa \
57
+ --embedding_model nvidia-nv-embed-v2 --qafd_alpha 3.0 --qafd_weight_scheme none 2>&1 | grep "Recall@"
58
+
59
+ echo ""
60
+ echo "====== PE / hotpotqa / gritlm ======"
61
+ echo " [QA-aware]"
62
+ $BASE_PE --task multihop --dataset hotpotqa --num_queries 50 --skip_qa \
63
+ --embedding_model gritlm --qafd_alpha 3.0 2>&1 | grep "Recall@"
64
+ echo " [QA-agnostic]"
65
+ $BASE_PE --task multihop --dataset hotpotqa --num_queries 50 --skip_qa \
66
+ --embedding_model gritlm --qafd_alpha 3.0 --qafd_weight_scheme none 2>&1 | grep "Recall@"
67
+
68
+ # =====================================================================
69
+ # 3. PASSAGE-ENTITY: UltraDomain (10q, quality scores)
70
+ # =====================================================================
71
+ echo ""
72
+ echo "############ PASSAGE-ENTITY: ULTRADOMAIN ############"
73
+
74
+ for DATASET in mix agriculture cs; do
75
+ echo ""
76
+ echo "====== PE / ultradomain / $DATASET / openai-small ======"
77
+
78
+ echo " [QA-aware]"
79
+ $BASE_PE --task ultradomain --dataset $DATASET --num_queries 10 \
80
+ --embedding_model openai-small --qafd_alpha 3.0 2>&1 | grep -E "comprehensiveness|diversity|logicality|relevance|coherence|Overall"
81
+
82
+ echo " [QA-agnostic]"
83
+ $BASE_PE --task ultradomain --dataset $DATASET --num_queries 10 \
84
+ --embedding_model openai-small --qafd_alpha 3.0 --qafd_weight_scheme none 2>&1 | grep -E "comprehensiveness|diversity|logicality|relevance|coherence|Overall"
85
+ done
86
+
87
+ echo ""
88
+ echo "====== PE / ultradomain / mix / jina-v3 ======"
89
+ echo " [QA-aware]"
90
+ $BASE_PE --task ultradomain --dataset mix --num_queries 10 \
91
+ --embedding_model jina-v3 --qafd_alpha 3.0 2>&1 | grep -E "comprehensiveness|diversity|logicality|relevance|coherence|Overall"
92
+ echo " [QA-agnostic]"
93
+ $BASE_PE --task ultradomain --dataset mix --num_queries 10 \
94
+ --embedding_model jina-v3 --qafd_alpha 3.0 --qafd_weight_scheme none 2>&1 | grep -E "comprehensiveness|diversity|logicality|relevance|coherence|Overall"
95
+
96
+ # =====================================================================
97
+ # 4. ENTITY GRAPH: Multihop (100q, alpha=3 vs alpha=2)
98
+ # =====================================================================
99
+ echo ""
100
+ echo "############ ENTITY GRAPH: MULTIHOP ############"
101
+
102
+ for DATASET in musique hotpotqa; do
103
+ echo ""
104
+ echo "====== Entity / $DATASET ======"
105
+
106
+ echo " [alpha=3]"
107
+ QAFD_CONDA_ENV=qafd-rag ./run.sh multihop --dataset $DATASET --questions 100 --alpha 3.0 2>&1 | grep -E "F1 Score|Exact Match"
108
+
109
+ echo " [alpha=2]"
110
+ QAFD_CONDA_ENV=qafd-rag ./run.sh multihop --dataset $DATASET --questions 100 --alpha 2.0 2>&1 | grep -E "F1 Score|Exact Match"
111
+ done
112
+
113
+ # =====================================================================
114
+ # 5. BATCH PUSH + ALPHA=3 (passage-entity, 100q)
115
+ # =====================================================================
116
+ echo ""
117
+ echo "############ BATCH PUSH + ALPHA=3 ############"
118
+
119
+ for DATASET in musique hotpotqa; do
120
+ echo ""
121
+ echo "====== Batch / $DATASET / openai-small ======"
122
+
123
+ echo " [QA-aware, batch]"
124
+ $BASE_PE --task multihop --dataset $DATASET --num_queries 100 --skip_qa \
125
+ --embedding_model openai-small --qafd_alpha 3.0 --batch_push 2>&1 | grep "Recall@"
126
+
127
+ echo " [QA-agnostic, batch]"
128
+ $BASE_PE --task multihop --dataset $DATASET --num_queries 100 --skip_qa \
129
+ --embedding_model openai-small --qafd_alpha 3.0 --batch_push --qafd_weight_scheme none 2>&1 | grep "Recall@"
130
+ done
131
+
132
+ echo ""
133
+ echo "======================================================================"
134
+ echo " ALL EXPERIMENTS COMPLETE"
135
+ echo " Finished: $(date)"
136
+ echo "======================================================================"
experiments/run_qa_awareness_v2.sh ADDED
@@ -0,0 +1,34 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/bin/bash
2
+ # Query-Awareness Ablation V2: warm walk, accumulation, more steps
3
+ cd /home/davoud/QAFD-RAG
4
+ N=100
5
+ BASE="python src/hipporag_pipeline/benchmark_runner.py --task multihop --num_queries $N --embedding_model openai-small --skip_qa"
6
+
7
+ echo "======================================================================"
8
+ echo " Query-Awareness V2 Ablation ($N queries per dataset)"
9
+ echo "======================================================================"
10
+
11
+ for DATASET in musique hotpotqa 2wikimultihopqa; do
12
+ echo ""
13
+ echo "====== Dataset: $DATASET ======"
14
+
15
+ echo " [1] Baseline (original)"
16
+ $BASE --dataset $DATASET 2>&1 | grep -E "Recall@|Retrieval done"
17
+
18
+ echo " [2] QA Warm Walk (edge weights in warm-start)"
19
+ $BASE --dataset $DATASET --qa_warm_walk 2>&1 | grep -E "Recall@|Retrieval done"
20
+
21
+ echo " [3] QA Warm Walk + 5 steps"
22
+ $BASE --dataset $DATASET --qa_warm_walk --qa_warm_steps 5 2>&1 | grep -E "Recall@|Retrieval done"
23
+
24
+ echo " [4] QA Accumulation gamma=0.5"
25
+ $BASE --dataset $DATASET --qa_accum_gamma 0.5 2>&1 | grep -E "Recall@|Retrieval done"
26
+
27
+ echo " [5] QA Warm Walk + Accumulation gamma=0.5"
28
+ $BASE --dataset $DATASET --qa_warm_walk --qa_accum_gamma 0.5 2>&1 | grep -E "Recall@|Retrieval done"
29
+ done
30
+
31
+ echo ""
32
+ echo "======================================================================"
33
+ echo " ALL EXPERIMENTS COMPLETE"
34
+ echo "======================================================================"
experiments/run_sim_mode_ablation.sh ADDED
@@ -0,0 +1,29 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/bin/bash
2
+ # Edge weight ablation: similarity contrast functions across 3 datasets
3
+ # Tests whether sharper similarity contrast makes query-aware edge weights effective
4
+
5
+ cd /home/davoud/QAFD-RAG
6
+ N=100
7
+ BASE="python src/hipporag_pipeline/benchmark_runner.py --task multihop --num_queries $N --embedding_model openai-small --skip_qa"
8
+
9
+ echo "======================================================================"
10
+ echo " Similarity Mode Ablation ($N queries per dataset)"
11
+ echo "======================================================================"
12
+
13
+ for DATASET in musique hotpotqa 2wikimultihopqa; do
14
+ echo ""
15
+ echo "======================================================================"
16
+ echo " Dataset: $DATASET"
17
+ echo "======================================================================"
18
+
19
+ for SIM in normalized relu relu_sq; do
20
+ echo ""
21
+ echo " [$DATASET] sim_mode=$SIM"
22
+ $BASE --dataset $DATASET --sim_mode $SIM 2>&1 | grep -E "Recall@|Retrieval done"
23
+ done
24
+ done
25
+
26
+ echo ""
27
+ echo "======================================================================"
28
+ echo " ALL EXPERIMENTS COMPLETE"
29
+ echo "======================================================================"
requirements.txt ADDED
@@ -0,0 +1,46 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # QAFD-RAG Dependencies
2
+ # Install:
3
+ # conda create -n qafd-rag python=3.10 -y
4
+ # conda activate qafd-rag
5
+ # pip install -r requirements.txt
6
+
7
+ # --- Core ---
8
+ openai>=1.30.0,<2.0
9
+ tenacity>=8.0.0,<9.0
10
+ aiohttp>=3.9.0,<4.0
11
+ pydantic>=2.0.0,<3.0
12
+ numpy>=1.26.0,<2.0
13
+ tiktoken>=0.7.0,<1.0
14
+ networkx>=3.0,<4.0
15
+ graspologic>=3.4.0,<4.0
16
+ nano-vectordb>=0.0.4
17
+ tqdm>=4.60.0
18
+ nest-asyncio>=1.5.0
19
+ python-dotenv>=1.0.0
20
+ httpx>=0.27.0,<1.0
21
+ huggingface_hub>=0.25.0
22
+
23
+ # --- Graph (passage-entity pipeline) ---
24
+ python-igraph>=0.11.0,<1.0
25
+ pyarrow>=15.0.0
26
+ datasets>=3.0.0
27
+
28
+ # --- Embedding Models ---
29
+ torch>=2.0.0
30
+ transformers==4.45.2
31
+ sentence-transformers>=3.0.0,<4.0
32
+ einops>=0.7.0
33
+ gritlm>=1.0.0
34
+ regex>=2024.0.0
35
+
36
+ # --- Data Processing ---
37
+ pandas>=2.0.0
38
+ pyvis>=0.3.0
39
+
40
+ # --- LLM Backends ---
41
+ ollama>=0.5.0
42
+ aioboto3>=11.0.0
43
+ urllib3>=1.26.0,<2.0
44
+
45
+ # --- Benchmark Evaluation ---
46
+ rouge-score>=0.1.2
run.sh ADDED
@@ -0,0 +1,113 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/bin/bash
2
+ # QAFD-RAG Benchmark Runner
3
+ # Unified interface for all benchmarks
4
+
5
+ set -euo pipefail
6
+
7
+ cd "$(dirname "${BASH_SOURCE[0]}")" || { echo "Error: Cannot cd to script directory"; exit 1; }
8
+
9
+ # Conda environment (has torch, transformers, sentence-transformers, etc.)
10
+ CONDA_ENV="${QAFD_CONDA_ENV:-qafd-rag}"
11
+ PYTHON="conda run --live-stream -n ${CONDA_ENV} python"
12
+
13
+ # Set NVIDIA library paths (needed for torch with pip-installed CUDA libs)
14
+ NVIDIA_LIB="$(conda info --base)/envs/${CONDA_ENV}/lib/python3.10/site-packages/nvidia"
15
+ if [ -d "$NVIDIA_LIB" ]; then
16
+ export LD_LIBRARY_PATH="${NVIDIA_LIB}/cudnn/lib:${NVIDIA_LIB}/cublas/lib:${NVIDIA_LIB}/cuda_runtime/lib:${NVIDIA_LIB}/cuda_cupti/lib:${NVIDIA_LIB}/cuda_nvrtc/lib:${NVIDIA_LIB}/cufft/lib:${NVIDIA_LIB}/curand/lib:${NVIDIA_LIB}/cusolver/lib:${NVIDIA_LIB}/cusparse/lib:${NVIDIA_LIB}/nccl/lib:${NVIDIA_LIB}/nvjitlink/lib:${LD_LIBRARY_PATH:-}"
17
+ fi
18
+
19
+ # Load API key from .env if not already set
20
+ if [ -z "${OPENAI_API_KEY:-}" ]; then
21
+ if [ -f .env ]; then
22
+ eval "$(grep -E '^OPENAI_API_KEY=' .env | head -1)"
23
+ export OPENAI_API_KEY
24
+ fi
25
+ if [ -z "${OPENAI_API_KEY:-}" ]; then
26
+ echo "Error: OPENAI_API_KEY not set. Export it or add it to a .env file."
27
+ exit 1
28
+ fi
29
+ fi
30
+
31
+ # Show help
32
+ show_help() {
33
+ echo "Usage: ./run.sh <task> [options]"
34
+ echo ""
35
+ echo "Tasks (graph type):"
36
+ echo " ultradomain QA benchmark on UltraDomain dataset [entity graph]"
37
+ echo " text2sql Text-to-SQL benchmark [entity graph]"
38
+ echo " summarization Document summarization QA benchmark [entity graph]"
39
+ echo " multihop Multi-hop QA benchmark (MuSiQue, etc.) [entity graph]"
40
+ echo ""
41
+ echo "Note: run.sh uses the entity graph pipeline. For the passage-entity graph"
42
+ echo " (recommended for multihop), use the unified runner instead:"
43
+ echo " python benchmarks/run.py --task multihop --dataset musique --questions 100"
44
+ echo ""
45
+ echo "Common Options (all tasks):"
46
+ echo " --questions N Number of questions to benchmark"
47
+ echo " --build Build KG only, don't run benchmark"
48
+ echo " --force-build Force rebuild KG even if exists"
49
+ echo " --max-documents N Max documents for KG building"
50
+ echo " --embedding MODEL Embedding model (openai-small, openai-large, jina-v3)"
51
+ echo " --llm MODEL LLM model (gpt-4o-mini, gpt-4o, gpt-5-nano, gpt-5-mini, gpt-5, gpt-oss-120b)"
52
+ echo ""
53
+ echo "Summarization-specific Options:"
54
+ echo " --dataset NAME Dataset: squality (default: squality)"
55
+ echo ""
56
+ echo "Multihop-specific Options:"
57
+ echo " --dataset NAME Dataset: musique, hotpotqa, 2wikimultihopqa"
58
+ echo ""
59
+ echo "Text2SQL-specific Options:"
60
+ echo " --benchmark NAME Benchmark: spider2-lite (default) or bird"
61
+ echo " --db NAME Database name (e.g., Pagila, superhero)"
62
+ echo ""
63
+ echo "Text2SQL End-to-End Pipeline:"
64
+ echo " Place your .sqlite file in data/text2sql/<benchmark>/databases/<DB>/<DB>.sqlite"
65
+ echo " The DB summary (JSON) is auto-generated on first run."
66
+ echo ""
67
+ echo "Examples:"
68
+ echo " # Build KG"
69
+ echo " ./run.sh ultradomain --build --max-documents 100"
70
+ echo " ./run.sh summarization --build --max-documents 10"
71
+ echo " ./run.sh multihop --dataset musique --build --max-documents 500"
72
+ echo " ./run.sh multihop --dataset hotpotqa --build --max-documents 500"
73
+ echo ""
74
+ echo " # Run benchmark"
75
+ echo " ./run.sh ultradomain --questions 10"
76
+ echo " ./run.sh summarization --questions 50"
77
+ echo " ./run.sh multihop --dataset musique --questions 100"
78
+ echo " ./run.sh multihop --dataset hotpotqa --questions 100"
79
+ echo " ./run.sh multihop --dataset 2wikimultihopqa --questions 100"
80
+ echo " ./run.sh text2sql --questions 5 --db Pagila"
81
+ echo " ./run.sh text2sql --benchmark bird --questions 5 --db superhero"
82
+ echo ""
83
+ echo " # Use a local OpenAI-compatible model"
84
+ echo " export LOCAL_LLM_BASE_URL='http://localhost:8000/v1'"
85
+ echo " ./run.sh ultradomain --questions 10 --llm gpt-oss-120b"
86
+ echo " ./run.sh multihop --dataset musique --questions 10 --llm gpt-oss-120b"
87
+ }
88
+
89
+ case "${1:-help}" in
90
+ ultradomain)
91
+ shift
92
+ $PYTHON benchmarks/ultradomain/benchmark_ultradomain.py "$@"
93
+ ;;
94
+ text2sql)
95
+ shift
96
+ $PYTHON benchmarks/text2sql/benchmark_text2sql.py "$@"
97
+ ;;
98
+ summarization)
99
+ shift
100
+ $PYTHON benchmarks/summarization/benchmark_summarization.py "$@"
101
+ ;;
102
+ multihop)
103
+ shift
104
+ $PYTHON benchmarks/multihop/benchmark_multihop.py "$@"
105
+ ;;
106
+ help|--help|-h)
107
+ show_help
108
+ ;;
109
+ *)
110
+ echo "Error: Unknown task '${1}'. Run './run.sh --help' for usage."
111
+ exit 1
112
+ ;;
113
+ esac
src/QAFD_RAG.py ADDED
@@ -0,0 +1,747 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import asyncio
2
+ import os
3
+ from tqdm.asyncio import tqdm as tqdm_async
4
+ from dataclasses import asdict, dataclass, field
5
+ from datetime import datetime
6
+ from functools import partial
7
+ from typing import Type, cast, List, Dict, Any, Optional
8
+
9
+ from .llm import (
10
+ gpt_4o_mini_complete,
11
+ gpt_oss_120b_complete,
12
+ local_sentence_embedding,
13
+ openai_cloud_embedding,
14
+ is_local_model,
15
+ is_cloud_model,
16
+ )
17
+
18
+ # Try to import the new functions, but don't fail if they don't exist
19
+ try:
20
+ from .llm import get_embedding_func_for_model, EMBEDDING_CONFIGS
21
+ HAS_EMBEDDING_CONFIGS = True
22
+ except ImportError:
23
+ HAS_EMBEDDING_CONFIGS = False
24
+ print("Warning: get_embedding_func_for_model not found in llm.py")
25
+ print("Please update your llm.py file with the new version")
26
+
27
+ from .operate import (
28
+ chunking_by_token_size,
29
+ extract_entities,
30
+ kg_query,
31
+ )
32
+ from .indexing import DatabaseSchemaBuilder
33
+
34
+ from .utils import (
35
+ EmbeddingFunc,
36
+ compute_mdhash_id,
37
+ limit_async_func_call,
38
+ convert_response_to_json,
39
+ logger,
40
+ set_logger,
41
+ )
42
+ from .base import (
43
+ BaseGraphStorage,
44
+ BaseKVStorage,
45
+ BaseVectorStorage,
46
+ StorageNameSpace,
47
+ QueryParam,
48
+ )
49
+
50
+ from .storage import (
51
+ JsonKVStorage,
52
+ NanoVectorDBStorage,
53
+ NetworkXStorage,
54
+ )
55
+
56
+
57
+ async def abuild_from_excel_files(self, excel_paths: List[str]) -> Dict[str, Any]:
58
+ """Build KG from Excel files"""
59
+ from .indexing import ExcelSchemaBuilder
60
+
61
+ builder = ExcelSchemaBuilder(
62
+ graph_storage=self.chunk_entity_relation_graph,
63
+ entities_vdb=self.entities_vdb,
64
+ relationships_vdb=self.relationships_vdb
65
+ )
66
+
67
+ result = await builder.build_from_excel_files(excel_paths)
68
+ await self._insert_done()
69
+
70
+ logger.info(f"Excel KG build completed: {result}")
71
+ return result
72
+
73
+
74
+ def build_from_excel_files(self, excel_paths: List[str]):
75
+ """Sync wrapper"""
76
+ loop = always_get_an_event_loop()
77
+ return loop.run_until_complete(self.abuild_from_excel_files(excel_paths))
78
+
79
+
80
+ def lazy_external_import(module_name: str, class_name: str):
81
+ """Lazily import a class from an external module based on the package of the caller."""
82
+ import inspect
83
+
84
+ caller_frame = inspect.currentframe().f_back
85
+ module = inspect.getmodule(caller_frame)
86
+ package = module.__package__ if module else None
87
+
88
+ def import_class(*args, **kwargs):
89
+ import importlib
90
+ module = importlib.import_module(module_name, package=package)
91
+ cls = getattr(module, class_name)
92
+ return cls(*args, **kwargs)
93
+
94
+ return import_class
95
+
96
+
97
+ Neo4JStorage = lazy_external_import(".kg.neo4j_impl", "Neo4JStorage")
98
+ OracleKVStorage = lazy_external_import(".kg.oracle_impl", "OracleKVStorage")
99
+ OracleGraphStorage = lazy_external_import(".kg.oracle_impl", "OracleGraphStorage")
100
+ OracleVectorDBStorage = lazy_external_import(".kg.oracle_impl", "OracleVectorDBStorage")
101
+ MilvusVectorDBStorge = lazy_external_import(".kg.milvus_impl", "MilvusVectorDBStorge")
102
+ MongoKVStorage = lazy_external_import(".kg.mongo_impl", "MongoKVStorage")
103
+ ChromaVectorDBStorage = lazy_external_import(".kg.chroma_impl", "ChromaVectorDBStorage")
104
+ TiDBKVStorage = lazy_external_import(".kg.tidb_impl", "TiDBKVStorage")
105
+ TiDBVectorDBStorage = lazy_external_import(".kg.tidb_impl", "TiDBVectorDBStorage")
106
+ AGEStorage = lazy_external_import(".kg.age_impl", "AGEStorage")
107
+
108
+
109
+ def always_get_an_event_loop() -> asyncio.AbstractEventLoop:
110
+ """
111
+ Ensure that there is always an event loop available.
112
+
113
+ This function tries to get the current event loop. If the current event loop is closed or does not exist,
114
+ it creates a new event loop and sets it as the current event loop.
115
+
116
+ Returns:
117
+ asyncio.AbstractEventLoop: The current or newly created event loop.
118
+ """
119
+ try:
120
+ current_loop = asyncio.get_event_loop()
121
+ if current_loop.is_closed():
122
+ raise RuntimeError("Event loop is closed.")
123
+ return current_loop
124
+ except RuntimeError:
125
+ logger.info("Creating a new event loop in main thread.")
126
+ new_loop = asyncio.new_event_loop()
127
+ asyncio.set_event_loop(new_loop)
128
+ return new_loop
129
+
130
+
131
+ @dataclass
132
+ class QAFD_RAG:
133
+ working_dir: str = field(
134
+ default_factory=lambda: f"./QAFD_RAG_cache_{datetime.now().strftime('%Y-%m-%d-%H:%M:%S')}"
135
+ )
136
+
137
+ embedding_cache_config: dict = field(
138
+ default_factory=lambda: {
139
+ "enabled": False,
140
+ "similarity_threshold": 0.95,
141
+ "use_llm_check": False,
142
+ }
143
+ )
144
+ kv_storage: str = field(default="JsonKVStorage")
145
+ vector_storage: str = field(default="NanoVectorDBStorage")
146
+ graph_storage: str = field(default="NetworkXStorage")
147
+
148
+ current_log_level = logger.level
149
+ log_level: str = field(default=current_log_level)
150
+
151
+ # Chunking parameters
152
+ chunk_token_size: int = 1200
153
+ chunk_overlap_token_size: int = 100
154
+ tiktoken_model_name: str = "gpt-4o-mini"
155
+
156
+ # Entity extraction parameters
157
+ entity_extract_max_gleaning: int = 1
158
+ entity_summary_to_max_tokens: int = 5000
159
+
160
+ # Node embedding algorithm
161
+ node_embedding_algorithm: str = "node2vec"
162
+ node2vec_params: dict = field(
163
+ default_factory=lambda: {
164
+ "dimensions": 1536,
165
+ "num_walks": 10,
166
+ "walk_length": 40,
167
+ "window_size": 2,
168
+ "iterations": 3,
169
+ "random_seed": 3,
170
+ }
171
+ )
172
+
173
+ # ============================================================================
174
+ # EMBEDDING CONFIGURATION (NEW: Configurable embedding models)
175
+ # ============================================================================
176
+
177
+ # Embedding model key (from EMBEDDING_CONFIGS in llm.py)
178
+ embedding_model_key: Optional[str] = None
179
+
180
+ # Embedding function and dimensions (auto-configured from embedding_model_key)
181
+ embedding_func: Optional[EmbeddingFunc] = None
182
+ embedding_dim: Optional[int] = None
183
+
184
+ # Embedding parameters
185
+ embedding_batch_num: int = 32
186
+ embedding_func_max_async: int = 16
187
+ max_embed_tokens: int = 8192 # Maximum tokens per embedding request
188
+
189
+ # ============================================================================
190
+ # LLM CONFIGURATION
191
+ # ============================================================================
192
+
193
+ llm_model_func: callable = gpt_4o_mini_complete
194
+ llm_model_name: str = "gpt-4o-mini"
195
+ llm_model_max_token_size: int = 32768
196
+ llm_model_max_async: int = 16
197
+ llm_model_kwargs: dict = field(default_factory=dict)
198
+
199
+ # Vector DB storage parameters
200
+ vector_db_storage_cls_kwargs: dict = field(default_factory=dict)
201
+
202
+ enable_llm_cache: bool = True
203
+
204
+ # Additional parameters
205
+ addon_params: dict = field(default_factory=dict)
206
+ convert_response_to_json_func: callable = convert_response_to_json
207
+
208
+ def __post_init__(self):
209
+ log_file = os.path.join("QAFD_RAG.log")
210
+ set_logger(log_file)
211
+ logger.setLevel(self.log_level)
212
+
213
+ logger.info(f"Logger initialized for working directory: {self.working_dir}")
214
+
215
+ # ============================================================================
216
+ # EMBEDDING MODEL CONFIGURATION LOGIC (NEW)
217
+ # ============================================================================
218
+
219
+ if HAS_EMBEDDING_CONFIGS:
220
+ # Priority 1: Use explicitly passed embedding_model_key
221
+ if self.embedding_model_key:
222
+ logger.info(f"[Embedding Config] Using explicit embedding_model_key: {self.embedding_model_key}")
223
+ embedding_func, embedding_dim, emb_config = get_embedding_func_for_model(self.embedding_model_key)
224
+ self.embedding_func = embedding_func
225
+ self.embedding_dim = embedding_dim
226
+ logger.info(f"[Embedding Config] {emb_config['description']}")
227
+ logger.info(f"[Embedding Config] Dimensions: {embedding_dim}, Max tokens: {emb_config['max_tokens']}")
228
+
229
+ # Priority 2: Check environment variable EMBEDDING_MODEL_KEY
230
+ elif os.environ.get("EMBEDDING_MODEL_KEY"):
231
+ embedding_key = os.environ.get("EMBEDDING_MODEL_KEY")
232
+ logger.info(f"[Embedding Config] Using env EMBEDDING_MODEL_KEY: {embedding_key}")
233
+ embedding_func, embedding_dim, emb_config = get_embedding_func_for_model(embedding_key)
234
+ self.embedding_func = embedding_func
235
+ self.embedding_dim = embedding_dim
236
+ self.embedding_model_key = embedding_key
237
+ logger.info(f"[Embedding Config] {emb_config['description']}")
238
+
239
+ # Priority 3: Check environment variable USE_OPENAI_EMBEDDINGS (legacy)
240
+ elif os.environ.get("USE_OPENAI_EMBEDDINGS") == "1":
241
+ logger.info(f"[Embedding Config] Using legacy USE_OPENAI_EMBEDDINGS=1")
242
+ self.embedding_func = openai_cloud_embedding
243
+ self.embedding_dim = 1024
244
+ self.embedding_model_key = "openai-large"
245
+ logger.info(f"[Embedding Config] OpenAI cloud embeddings (1024-dim)")
246
+
247
+ elif os.environ.get("USE_OPENAI_EMBEDDINGS") == "0":
248
+ logger.info(f"[Embedding Config] Using legacy USE_OPENAI_EMBEDDINGS=0")
249
+ self.embedding_func = local_sentence_embedding
250
+ self.embedding_dim = 1024
251
+ self.embedding_model_key = "jina-v3"
252
+ logger.info(f"[Embedding Config] Local Jina v3 embeddings (1024-dim)")
253
+
254
+ # Priority 4: Auto-detect based on LLM model name
255
+ elif self.llm_model_name:
256
+ model_name = self.llm_model_name.lower()
257
+ if is_local_model(model_name):
258
+ logger.info(f"[Embedding Config] Local LLM detected ({model_name}) → using local embeddings")
259
+ self.embedding_func = local_sentence_embedding
260
+ self.embedding_dim = 1024
261
+ self.embedding_model_key = "jina-v3"
262
+ else:
263
+ logger.info(f"[Embedding Config] Cloud LLM detected ({model_name}) → using OpenAI embeddings")
264
+ self.embedding_func = openai_cloud_embedding
265
+ self.embedding_dim = 1024
266
+ self.embedding_model_key = "openai-large"
267
+
268
+ # Priority 5: Default to local Jina v3
269
+ else:
270
+ logger.info(f"[Embedding Config] No configuration found → defaulting to Jina v3 (local)")
271
+ self.embedding_func = local_sentence_embedding
272
+ self.embedding_dim = 1024
273
+ self.embedding_model_key = "jina-v3"
274
+ else:
275
+ # Fallback to old behavior if new functions not available
276
+ logger.warning("[Embedding Config] Using legacy embedding configuration")
277
+ env_embedding_setting = os.environ.get("USE_OPENAI_EMBEDDINGS")
278
+
279
+ if env_embedding_setting == "0":
280
+ self.embedding_func = local_sentence_embedding
281
+ self.embedding_dim = 1024
282
+ logger.info(f"[Embedding Override] Using local embeddings (1024-dim) - forced by USE_OPENAI_EMBEDDINGS=0")
283
+ elif env_embedding_setting == "1":
284
+ self.embedding_func = openai_cloud_embedding
285
+ self.embedding_dim = 1024
286
+ logger.info(f"[Embedding Override] Using OpenAI embeddings (1024-dim) - forced by USE_OPENAI_EMBEDDINGS=1")
287
+ elif is_local_model(self.llm_model_name.lower() if self.llm_model_name else ""):
288
+ self.embedding_func = local_sentence_embedding
289
+ self.embedding_dim = 1024
290
+ logger.info(f"[Embedding] Local model detected → using local embeddings (1024-dim)")
291
+ else:
292
+ self.embedding_func = openai_cloud_embedding
293
+ self.embedding_dim = 1024
294
+ logger.info(f"[Embedding] Cloud model detected → using OpenAI embeddings (1024-dim)")
295
+
296
+ # Validate embedding function was set
297
+ if self.embedding_func is None:
298
+ logger.error("[Embedding Config] Failed to configure embedding function!")
299
+ raise ValueError("Embedding function not configured")
300
+
301
+ if self.embedding_dim is None:
302
+ self.embedding_dim = 1024 # Default
303
+ logger.warning(f"[Embedding Config] embedding_dim not set, defaulting to 1024")
304
+
305
+ logger.info(f"[Embedding Config] ✅ Final: {self.embedding_model_key if self.embedding_model_key else 'auto'} ({self.embedding_dim}-dim)")
306
+
307
+ # ============================================================================
308
+ # STORAGE INITIALIZATION
309
+ # ============================================================================
310
+
311
+ self.key_string_value_json_storage_cls: Type[BaseKVStorage] = (
312
+ self._get_storage_class()[self.kv_storage]
313
+ )
314
+ self.vector_db_storage_cls: Type[BaseVectorStorage] = self._get_storage_class()[
315
+ self.vector_storage
316
+ ]
317
+ self.graph_storage_cls: Type[BaseGraphStorage] = self._get_storage_class()[
318
+ self.graph_storage
319
+ ]
320
+
321
+ if not os.path.exists(self.working_dir):
322
+ logger.info(f"Creating working directory {self.working_dir}")
323
+ os.makedirs(self.working_dir)
324
+
325
+ self.llm_response_cache = (
326
+ self.key_string_value_json_storage_cls(
327
+ namespace="llm_response_cache",
328
+ global_config=asdict(self),
329
+ embedding_func=None,
330
+ )
331
+ if self.enable_llm_cache
332
+ else None
333
+ )
334
+
335
+ # Limit async calls for embedding function
336
+ self.embedding_func = limit_async_func_call(self.embedding_func_max_async)(
337
+ self.embedding_func
338
+ )
339
+
340
+ # Initialize storage components with embedding function
341
+ self.full_docs = self.key_string_value_json_storage_cls(
342
+ namespace="full_docs",
343
+ global_config=asdict(self),
344
+ embedding_func=self.embedding_func,
345
+ )
346
+ self.text_chunks = self.key_string_value_json_storage_cls(
347
+ namespace="text_chunks",
348
+ global_config=asdict(self),
349
+ embedding_func=self.embedding_func,
350
+ )
351
+ self.chunk_entity_relation_graph = self.graph_storage_cls(
352
+ namespace="chunk_entity_relation",
353
+ global_config=asdict(self),
354
+ embedding_func=self.embedding_func,
355
+ )
356
+
357
+ # Vector databases for entities, relationships, and chunks
358
+ self.entities_vdb = self.vector_db_storage_cls(
359
+ namespace="entities",
360
+ global_config=asdict(self),
361
+ embedding_func=self.embedding_func,
362
+ meta_fields={"entity_name"},
363
+ )
364
+ self.relationships_vdb = self.vector_db_storage_cls(
365
+ namespace="relationships",
366
+ global_config=asdict(self),
367
+ embedding_func=self.embedding_func,
368
+ meta_fields={"src_id", "tgt_id"},
369
+ )
370
+ self.chunks_vdb = self.vector_db_storage_cls(
371
+ namespace="chunks",
372
+ global_config=asdict(self),
373
+ embedding_func=self.embedding_func,
374
+ )
375
+
376
+ # Configure LLM function
377
+ self.llm_model_func = limit_async_func_call(self.llm_model_max_async)(
378
+ partial(
379
+ self.llm_model_func,
380
+ hashing_kv=self.llm_response_cache
381
+ if self.llm_response_cache
382
+ and hasattr(self.llm_response_cache, "global_config")
383
+ else self.key_string_value_json_storage_cls(
384
+ global_config=asdict(self),
385
+ ),
386
+ **self.llm_model_kwargs,
387
+ )
388
+ )
389
+
390
+ # Initialize database schema builder
391
+ self.schema_builder = DatabaseSchemaBuilder(
392
+ graph_storage=self.chunk_entity_relation_graph,
393
+ entities_vdb=self.entities_vdb,
394
+ relationships_vdb=self.relationships_vdb,
395
+ llm_model_func=self.llm_model_func
396
+ )
397
+
398
+ def _get_storage_class(self) -> dict[str, Type]:
399
+ return {
400
+ # Key-Value Storage
401
+ "JsonKVStorage": JsonKVStorage,
402
+ "OracleKVStorage": OracleKVStorage,
403
+ "MongoKVStorage": MongoKVStorage,
404
+ "TiDBKVStorage": TiDBKVStorage,
405
+ # Vector Storage
406
+ "NanoVectorDBStorage": NanoVectorDBStorage,
407
+ "OracleVectorDBStorage": OracleVectorDBStorage,
408
+ "MilvusVectorDBStorge": MilvusVectorDBStorge,
409
+ "ChromaVectorDBStorage": ChromaVectorDBStorage,
410
+ "TiDBVectorDBStorage": TiDBVectorDBStorage,
411
+ # Graph Storage
412
+ "NetworkXStorage": NetworkXStorage,
413
+ "Neo4JStorage": Neo4JStorage,
414
+ "OracleGraphStorage": OracleGraphStorage,
415
+ "AGEStorage": AGEStorage,
416
+ }
417
+
418
+ def insert(self, string_or_strings, addon_params=None):
419
+ loop = always_get_an_event_loop()
420
+ return loop.run_until_complete(self.ainsert(string_or_strings, addon_params))
421
+
422
+ async def ainsert(self, string_or_strings, addon_params=None):
423
+ update_storage = False
424
+ try:
425
+ if isinstance(string_or_strings, str):
426
+ string_or_strings = [string_or_strings]
427
+
428
+ new_docs = {
429
+ compute_mdhash_id(c.strip(), prefix="doc-"): {"content": c.strip()}
430
+ for c in string_or_strings
431
+ }
432
+ _add_doc_keys = await self.full_docs.filter_keys(list(new_docs.keys()))
433
+ new_docs = {k: v for k, v in new_docs.items() if k in _add_doc_keys}
434
+ if not len(new_docs):
435
+ logger.warning("All docs are already in the storage")
436
+ return
437
+ update_storage = True
438
+ logger.info(f"[New Docs] inserting {len(new_docs)} docs")
439
+
440
+ inserting_chunks = {}
441
+ for doc_key, doc in tqdm_async(
442
+ new_docs.items(), desc="Chunking documents", unit="doc"
443
+ ):
444
+ chunks = {
445
+ compute_mdhash_id(dp["content"], prefix="chunk-"): {
446
+ **dp,
447
+ "full_doc_id": doc_key,
448
+ }
449
+ for dp in chunking_by_token_size(
450
+ doc["content"],
451
+ overlap_token_size=self.chunk_overlap_token_size,
452
+ max_token_size=self.chunk_token_size,
453
+ tiktoken_model=self.tiktoken_model_name,
454
+ )
455
+ }
456
+ inserting_chunks.update(chunks)
457
+ _add_chunk_keys = await self.text_chunks.filter_keys(
458
+ list(inserting_chunks.keys())
459
+ )
460
+ inserting_chunks = {
461
+ k: v for k, v in inserting_chunks.items() if k in _add_chunk_keys
462
+ }
463
+ if not len(inserting_chunks):
464
+ logger.warning("All chunks are already in the storage")
465
+ return
466
+ logger.info(f"[New Chunks] inserting {len(inserting_chunks)} chunks")
467
+
468
+ await self.chunks_vdb.upsert(inserting_chunks)
469
+
470
+ logger.info("[Entity Extraction]...")
471
+
472
+ # Create a temporary config with custom addon_params if provided
473
+ temp_config = asdict(self)
474
+ if addon_params is not None:
475
+ temp_config["addon_params"] = addon_params
476
+
477
+ maybe_new_kg = await extract_entities(
478
+ inserting_chunks,
479
+ knowledge_graph_inst=self.chunk_entity_relation_graph,
480
+ entity_vdb=self.entities_vdb,
481
+ relationships_vdb=self.relationships_vdb,
482
+ global_config=temp_config,
483
+ )
484
+ if maybe_new_kg is None:
485
+ logger.warning("No new entities and relationships found")
486
+ return
487
+ self.chunk_entity_relation_graph = maybe_new_kg
488
+
489
+ await self.full_docs.upsert(new_docs)
490
+ await self.text_chunks.upsert(inserting_chunks)
491
+ finally:
492
+ if update_storage:
493
+ await self._insert_done()
494
+
495
+ async def _insert_done(self):
496
+ tasks = []
497
+ for storage_inst in [
498
+ self.full_docs,
499
+ self.text_chunks,
500
+ self.llm_response_cache,
501
+ self.entities_vdb,
502
+ self.relationships_vdb,
503
+ self.chunks_vdb,
504
+ self.chunk_entity_relation_graph,
505
+ ]:
506
+ if storage_inst is None:
507
+ continue
508
+ tasks.append(cast(StorageNameSpace, storage_inst).index_done_callback())
509
+ await asyncio.gather(*tasks)
510
+
511
+ def insert_custom_kg(self, custom_kg: dict):
512
+ loop = always_get_an_event_loop()
513
+ return loop.run_until_complete(self.ainsert_custom_kg(custom_kg))
514
+
515
+ async def ainsert_custom_kg(self, custom_kg: dict):
516
+ update_storage = False
517
+ try:
518
+ all_chunks_data = {}
519
+ chunk_to_source_map = {}
520
+ for chunk_data in custom_kg.get("chunks", []):
521
+ chunk_content = chunk_data["content"]
522
+ source_id = chunk_data["source_id"]
523
+ chunk_id = compute_mdhash_id(chunk_content.strip(), prefix="chunk-")
524
+
525
+ chunk_entry = {"content": chunk_content.strip(), "source_id": source_id}
526
+ all_chunks_data[chunk_id] = chunk_entry
527
+ chunk_to_source_map[source_id] = chunk_id
528
+ update_storage = True
529
+
530
+ if self.chunks_vdb is not None and all_chunks_data:
531
+ await self.chunks_vdb.upsert(all_chunks_data)
532
+ if self.text_chunks is not None and all_chunks_data:
533
+ await self.text_chunks.upsert(all_chunks_data)
534
+
535
+ all_entities_data = []
536
+ for entity_data in custom_kg.get("entities", []):
537
+ entity_name = f'"{entity_data["entity_name"].lower()}"'
538
+ entity_type = entity_data.get("entity_type", "UNKNOWN")
539
+ description = entity_data.get("description", "No description provided")
540
+
541
+ source_chunk_id = entity_data.get("source_id", "UNKNOWN")
542
+ source_id = chunk_to_source_map.get(source_chunk_id, "UNKNOWN")
543
+
544
+ if source_id == "UNKNOWN":
545
+ logger.warning(
546
+ f"Entity '{entity_name}' has an UNKNOWN source_id. Please check the source mapping."
547
+ )
548
+
549
+ node_data = {
550
+ "entity_type": entity_type,
551
+ "description": description,
552
+ "source_id": source_id,
553
+ }
554
+
555
+ await self.chunk_entity_relation_graph.upsert_node(
556
+ entity_name, node_data=node_data
557
+ )
558
+ node_data["entity_name"] = entity_name
559
+ all_entities_data.append(node_data)
560
+ update_storage = True
561
+
562
+ all_relationships_data = []
563
+ for relationship_data in custom_kg.get("relationships", []):
564
+ src_id = f'"{relationship_data["src_id"].lower()}"'
565
+ tgt_id = f'"{relationship_data["tgt_id"].lower()}"'
566
+ description = relationship_data["description"]
567
+ keywords = relationship_data["keywords"]
568
+ weight = relationship_data.get("weight", 1.0)
569
+
570
+ source_chunk_id = relationship_data.get("source_id", "UNKNOWN")
571
+ source_id = chunk_to_source_map.get(source_chunk_id, "UNKNOWN")
572
+
573
+ if source_id == "UNKNOWN":
574
+ logger.warning(
575
+ f"Relationship from '{src_id}' to '{tgt_id}' has an UNKNOWN source_id. Please check the source mapping."
576
+ )
577
+
578
+ for need_insert_id in [src_id, tgt_id]:
579
+ if not (
580
+ await self.chunk_entity_relation_graph.has_node(need_insert_id)
581
+ ):
582
+ await self.chunk_entity_relation_graph.upsert_node(
583
+ need_insert_id,
584
+ node_data={
585
+ "source_id": source_id,
586
+ "description": "UNKNOWN",
587
+ "entity_type": "UNKNOWN",
588
+ },
589
+ )
590
+
591
+ await self.chunk_entity_relation_graph.upsert_edge(
592
+ src_id,
593
+ tgt_id,
594
+ edge_data={
595
+ "weight": weight,
596
+ "description": description,
597
+ "keywords": keywords,
598
+ "source_id": source_id,
599
+ },
600
+ )
601
+ edge_data = {
602
+ "src_id": src_id,
603
+ "tgt_id": tgt_id,
604
+ "description": description,
605
+ "keywords": keywords,
606
+ }
607
+ all_relationships_data.append(edge_data)
608
+ update_storage = True
609
+
610
+ if self.entities_vdb is not None:
611
+ data_for_vdb = {
612
+ compute_mdhash_id(dp["entity_name"], prefix="ent-"): {
613
+ "content": dp["entity_name"] + dp["description"],
614
+ "entity_name": dp["entity_name"],
615
+ }
616
+ for dp in all_entities_data
617
+ }
618
+ await self.entities_vdb.upsert(data_for_vdb)
619
+
620
+ if self.relationships_vdb is not None:
621
+ data_for_vdb = {
622
+ compute_mdhash_id(dp["src_id"] + dp["tgt_id"], prefix="rel-"): {
623
+ "src_id": dp["src_id"],
624
+ "tgt_id": dp["tgt_id"],
625
+ "content": dp["keywords"]
626
+ + dp["src_id"]
627
+ + dp["tgt_id"]
628
+ + dp["description"],
629
+ }
630
+ for dp in all_relationships_data
631
+ }
632
+ await self.relationships_vdb.upsert(data_for_vdb)
633
+ finally:
634
+ if update_storage:
635
+ await self._insert_done()
636
+
637
+ def query(self, query: str, param: QueryParam = QueryParam()):
638
+ loop = always_get_an_event_loop()
639
+ return loop.run_until_complete(self.aquery(query, param))
640
+
641
+ async def aquery(self, query: str, param: QueryParam = QueryParam()):
642
+ if param.mode in ["local", "global", "hybrid"]:
643
+ response = await kg_query(
644
+ query,
645
+ self.chunk_entity_relation_graph,
646
+ self.entities_vdb,
647
+ self.relationships_vdb,
648
+ self.text_chunks,
649
+ param,
650
+ asdict(self),
651
+ hashing_kv=self.llm_response_cache
652
+ if self.llm_response_cache
653
+ and hasattr(self.llm_response_cache, "global_config")
654
+ else self.key_string_value_json_storage_cls(
655
+ global_config=asdict(self),
656
+ ),
657
+ )
658
+ else:
659
+ raise ValueError(f"Unknown mode {param.mode}")
660
+ await self._query_done()
661
+ return response
662
+
663
+ async def _query_done(self):
664
+ tasks = []
665
+ for storage_inst in [self.llm_response_cache]:
666
+ if storage_inst is None:
667
+ continue
668
+ tasks.append(cast(StorageNameSpace, storage_inst).index_done_callback())
669
+ await asyncio.gather(*tasks)
670
+
671
+ def delete_by_entity(self, entity_name: str):
672
+ loop = always_get_an_event_loop()
673
+ return loop.run_until_complete(self.adelete_by_entity(entity_name))
674
+
675
+ async def adelete_by_entity(self, entity_name: str):
676
+ entity_name = f'"{entity_name.lower()}"'
677
+
678
+ try:
679
+ await self.entities_vdb.delete_entity(entity_name)
680
+ await self.relationships_vdb.delete_relation(entity_name)
681
+ await self.chunk_entity_relation_graph.delete_node(entity_name)
682
+
683
+ logger.info(
684
+ f"Entity '{entity_name}' and its relationships have been deleted."
685
+ )
686
+ await self._delete_by_entity_done()
687
+ except Exception as e:
688
+ logger.error(f"Error while deleting entity '{entity_name}': {e}")
689
+
690
+ async def _delete_by_entity_done(self):
691
+ tasks = []
692
+ for storage_inst in [
693
+ self.entities_vdb,
694
+ self.relationships_vdb,
695
+ self.chunk_entity_relation_graph,
696
+ ]:
697
+ if storage_inst is None:
698
+ continue
699
+ tasks.append(cast(StorageNameSpace, storage_inst).index_done_callback())
700
+ await asyncio.gather(*tasks)
701
+
702
+ def build_from_database_schema(self,
703
+ schema_file_path: str,
704
+ metadata_file_path: str = None,
705
+ language: str = "English"):
706
+ """
707
+ Build knowledge graph from database schema JSON file
708
+
709
+ This method manually constructs the knowledge graph from a JSON schema file,
710
+ avoiding the chunking issues that can cause LLM errors. It follows the approach
711
+ used in CoFD for database schema processing.
712
+
713
+ Args:
714
+ schema_file_path: Path to the JSON schema file
715
+ metadata_file_path: Optional path to metadata file
716
+ language: Output language for descriptions
717
+
718
+ Returns:
719
+ Dictionary containing build statistics
720
+ """
721
+ loop = always_get_an_event_loop()
722
+ return loop.run_until_complete(self.abuild_from_database_schema(
723
+ schema_file_path, metadata_file_path, language
724
+ ))
725
+
726
+ async def abuild_from_database_schema(self,
727
+ schema_file_path: str,
728
+ metadata_file_path: str = None,
729
+ language: str = "English"):
730
+ """
731
+ Async version of build_from_database_schema
732
+ """
733
+ try:
734
+ # Use the schema builder to construct the knowledge graph
735
+ result = await self.schema_builder.build_from_json_schema(
736
+ schema_file_path, metadata_file_path, language
737
+ )
738
+
739
+ # Update storage after building
740
+ await self._insert_done()
741
+
742
+ logger.info(f"Database schema build completed: {result}")
743
+ return result
744
+
745
+ except Exception as e:
746
+ logger.error(f"Error building from database schema: {e}")
747
+ raise
src/__init__.py ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ from .QAFD_RAG import QAFD_RAG as QAFD_RAG, QueryParam as QueryParam
2
+
3
+
4
+ from .QAFD_RAG import QAFD_RAG, QueryParam
src/answering/__init__.py ADDED
@@ -0,0 +1,42 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Answering module for QAFD-RAG.
3
+
4
+ This module handles the query-time operations: from user question to final answer.
5
+ It coordinates keyword extraction, entity retrieval, graph traversal, context assembly,
6
+ and response generation.
7
+
8
+ Submodules:
9
+ -----------
10
+ - handler: Main query entry point (kg_query)
11
+ - context: Context building based on query modes
12
+ - clusters: Flow diffusion cluster operations
13
+ - text_units: Text chunk retrieval
14
+
15
+ Note: Evaluation functions have been moved to src/evaluation.py
16
+ """
17
+
18
+ from .handler import kg_query
19
+ from .context import build_query_context
20
+ from .clusters import (
21
+ get_embeddings_for_flow_diffusion,
22
+ convert_subgraph_to_json,
23
+ find_flow_diffusion_clusters_and_summarize,
24
+ summarize_clusters_batch_with_llm,
25
+ summarize_cluster_with_llm,
26
+ )
27
+ from .text_units import find_most_related_text_unit_from_entities
28
+
29
+ __all__ = [
30
+ # Handler
31
+ "kg_query",
32
+ # Context
33
+ "build_query_context",
34
+ # Clusters
35
+ "get_embeddings_for_flow_diffusion",
36
+ "convert_subgraph_to_json",
37
+ "find_flow_diffusion_clusters_and_summarize",
38
+ "summarize_clusters_batch_with_llm",
39
+ "summarize_cluster_with_llm",
40
+ # Text units
41
+ "find_most_related_text_unit_from_entities",
42
+ ]
src/answering/clusters.py ADDED
@@ -0,0 +1,739 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Flow diffusion cluster operations for knowledge graph traversal.
3
+
4
+ This module provides functions to find clusters using flow diffusion,
5
+ convert subgraphs to JSON format and summarize clusters using LLM.
6
+ """
7
+
8
+ import asyncio
9
+ import tiktoken
10
+ import networkx as nx
11
+ from ..base import (
12
+ BaseGraphStorage,
13
+ QueryParam,
14
+ )
15
+ from ..retrievers import QueryAwareWeightedFlowDiffusion
16
+ from ..utils import (
17
+ logger,
18
+ truncate_list_by_token_size,
19
+ )
20
+
21
+
22
+ async def get_embeddings_for_flow_diffusion(
23
+ graph: nx.Graph,
24
+ query: str,
25
+ knowledge_graph_inst: BaseGraphStorage,
26
+ global_config: dict,
27
+ query_param: QueryParam = None,
28
+ ) -> tuple[dict, list]:
29
+ """
30
+ Get embeddings for nodes and query for query-aware flow diffusion.
31
+
32
+ Returns:
33
+ --------
34
+ tuple
35
+ (node_embeddings, subquery_embedding)
36
+ """
37
+ node_embeddings = {}
38
+ subquery_embedding = None
39
+
40
+ # Check if query-aware flow diffusion is enabled
41
+ if query_param and not query_param.enable_query_aware_flow_diffusion:
42
+ logger.info("Query-aware flow diffusion is disabled, skipping embedding calculation")
43
+ return {}, None
44
+
45
+ # Try to get embeddings from global config if available
46
+ if "embedding_func" in global_config and global_config["embedding_func"]:
47
+ try:
48
+ # Prioritize using cached node embeddings
49
+ if hasattr(knowledge_graph_inst, 'get_cached_node_embeddings'):
50
+ logger.info("Attempting to use cached node embeddings...")
51
+ node_embeddings = await knowledge_graph_inst.get_cached_node_embeddings(global_config)
52
+ if node_embeddings:
53
+ logger.info(f"Successfully using cached embeddings for {len(node_embeddings)} nodes")
54
+ else:
55
+ logger.info("No cached embeddings found, will compute new ones")
56
+ else:
57
+ logger.info("Storage class does not support cached embeddings, computing new ones...")
58
+ # Get embeddings for all nodes in the graph with batching
59
+ all_nodes = list(graph.nodes())
60
+ if all_nodes:
61
+ # Get batch size and token limit from config
62
+ batch_size = global_config.get("embedding_batch_num", 32)
63
+ max_tokens_per_request = global_config.get("max_embed_tokens", 8192)
64
+
65
+ # Prepare node texts
66
+ node_texts = []
67
+ for node in all_nodes:
68
+ node_info = await knowledge_graph_inst.get_node(node)
69
+ if node_info and "description" in node_info:
70
+ node_texts.append(f"{node} {node_info['description']}")
71
+ else:
72
+ node_texts.append(node)
73
+
74
+ # Process in batches to avoid token limit
75
+ logger.info(f"Computing embeddings for {len(node_texts)} node texts in batches of {batch_size}...")
76
+
77
+ # Import tiktoken for token counting
78
+ tiktoken_model = global_config.get("tiktoken_model_name", "gpt-4o-mini")
79
+ try:
80
+ encoding = tiktoken.encoding_for_model(tiktoken_model)
81
+ except:
82
+ encoding = tiktoken.get_encoding("cl100k_base") # fallback
83
+
84
+ # Process batches
85
+ all_embeddings = []
86
+ for i in range(0, len(node_texts), batch_size):
87
+ batch_texts = node_texts[i:i + batch_size]
88
+ batch_nodes = all_nodes[i:i + batch_size]
89
+
90
+ # Check token count for this batch
91
+ total_tokens = sum(len(encoding.encode(text)) for text in batch_texts)
92
+
93
+ if total_tokens > max_tokens_per_request:
94
+ logger.warning(f"Batch {i//batch_size + 1} exceeds token limit ({total_tokens} > {max_tokens_per_request}), reducing batch size...")
95
+ # Process this batch with smaller chunks
96
+ sub_batch_size = max(1, batch_size // 2)
97
+ for j in range(0, len(batch_texts), sub_batch_size):
98
+ sub_batch_texts = batch_texts[j:j + sub_batch_size]
99
+ sub_batch_nodes = batch_nodes[j:j + sub_batch_size]
100
+
101
+ # Check sub-batch token count
102
+ sub_total_tokens = sum(len(encoding.encode(text)) for text in sub_batch_texts)
103
+ if sub_total_tokens > max_tokens_per_request:
104
+ logger.warning(f"Sub-batch still exceeds token limit ({sub_total_tokens} > {max_tokens_per_request}), processing one by one...")
105
+ # Process one by one
106
+ for k, (text, node) in enumerate(zip(sub_batch_texts, sub_batch_nodes)):
107
+ try:
108
+ embedding_array = await global_config["embedding_func"]([text])
109
+ all_embeddings.append(embedding_array[0])
110
+ logger.debug(f"Processed node {k+1}/{len(sub_batch_texts)} in sub-batch")
111
+ except Exception as e:
112
+ logger.error(f"Failed to process node {node}: {e}")
113
+ # Add zero embedding as fallback
114
+ embedding_dim = 1536 # default dimension
115
+ all_embeddings.append([0.0] * embedding_dim)
116
+ else:
117
+ try:
118
+ embedding_array = await global_config["embedding_func"](sub_batch_texts)
119
+ all_embeddings.extend(embedding_array)
120
+ logger.debug(f"Processed sub-batch {j//sub_batch_size + 1} with {len(sub_batch_texts)} nodes")
121
+ except Exception as e:
122
+ logger.error(f"Failed to process sub-batch: {e}")
123
+ # Add zero embeddings as fallback
124
+ embedding_dim = 1536 # default dimension
125
+ for _ in sub_batch_texts:
126
+ all_embeddings.append([0.0] * embedding_dim)
127
+ else:
128
+ try:
129
+ embedding_array = await global_config["embedding_func"](batch_texts)
130
+ all_embeddings.extend(embedding_array)
131
+ logger.debug(f"Processed batch {i//batch_size + 1} with {len(batch_texts)} nodes")
132
+ except Exception as e:
133
+ logger.error(f"Failed to process batch: {e}")
134
+ # Add zero embeddings as fallback
135
+ embedding_dim = 1536 # default dimension
136
+ for _ in batch_texts:
137
+ all_embeddings.append([0.0] * embedding_dim)
138
+
139
+ # Store embeddings
140
+ for i, node in enumerate(all_nodes):
141
+ if i < len(all_embeddings):
142
+ node_embeddings[node] = all_embeddings[i].tolist()
143
+ else:
144
+ logger.warning(f"Missing embedding for node {node}")
145
+
146
+ logger.info(f"Computed embeddings for {len(node_embeddings)} nodes")
147
+
148
+ # Get query embedding
149
+ if query:
150
+ # Prioritize using cached query embedding
151
+ if hasattr(knowledge_graph_inst, 'get_cached_query_embedding'):
152
+ logger.info("Attempting to use cached query embedding...")
153
+ subquery_embedding = await knowledge_graph_inst.get_cached_query_embedding(query, global_config)
154
+ if subquery_embedding is None:
155
+ logger.warning("Failed to get cached query embedding, computing new one...")
156
+ query_embedding_array = await global_config["embedding_func"]([query])
157
+ subquery_embedding = query_embedding_array[0].tolist()
158
+ logger.info("Query embedding computed successfully")
159
+ else:
160
+ logger.info("Storage class does not support cached query embeddings, computing new one...")
161
+ query_embedding_array = await global_config["embedding_func"]([query])
162
+ subquery_embedding = query_embedding_array[0].tolist()
163
+ logger.info("Query embedding computed successfully")
164
+
165
+ except Exception as e:
166
+ logger.warning(f"Failed to get embeddings for query-aware flow diffusion: {e}")
167
+ # Fallback to non-query-aware mode
168
+ node_embeddings = {}
169
+ subquery_embedding = None
170
+ else:
171
+ logger.warning("No embedding function available in global config")
172
+
173
+ return node_embeddings, subquery_embedding
174
+
175
+
176
+ def convert_subgraph_to_json(G: nx.Graph, cluster_nodes: list, diffused_nodes: dict, source_node: str) -> dict:
177
+ """
178
+ Convert a subgraph to JSON format with nodes and edges information.
179
+
180
+ Parameters:
181
+ -----------
182
+ G : nx.Graph
183
+ The original graph
184
+ cluster_nodes : list
185
+ List of nodes in the cluster
186
+ diffused_nodes : dict
187
+ Dictionary mapping nodes to their flow values
188
+ source_node : str
189
+ The source node for this cluster
190
+
191
+ Returns:
192
+ --------
193
+ dict
194
+ JSON representation of the subgraph
195
+ """
196
+ # Create subgraph from cluster nodes
197
+ support_nodes = set(cluster_nodes)
198
+ subgraph = G.subgraph(support_nodes)
199
+
200
+ # Convert nodes to JSON format
201
+ nodes_json = []
202
+ # Build a local index map so node "id" matches the CSV-style index
203
+ node_index_map = {node: idx for idx, node in enumerate(cluster_nodes)}
204
+ for node in cluster_nodes:
205
+ node_attrs = G.nodes[node] if node in G.nodes else {}
206
+ node_degree = subgraph.degree(node) if node in subgraph else 0
207
+ nodes_json.append({
208
+ "id": node_index_map[node],
209
+ "entity": node,
210
+ "entity_type": node_attrs.get("entity_type", "UNKNOWN"),
211
+ "description": node_attrs.get("description", "UNKNOWN"),
212
+ "rank": node_degree,
213
+ })
214
+
215
+ # Convert edges to JSON format
216
+ edges_json = []
217
+ for u, v, data in subgraph.edges(data=True):
218
+ edges_json.append({
219
+ "source": u,
220
+ "target": v,
221
+ "source_id": node_index_map.get(u),
222
+ "target_id": node_index_map.get(v),
223
+ "weight": data.get('weight', 1.0)
224
+ })
225
+
226
+ return {
227
+ "source_node": source_node,
228
+ "cluster_size": len(cluster_nodes),
229
+ "max_flow_value": max(diffused_nodes.values()) if diffused_nodes else 0.0,
230
+ "nodes": nodes_json,
231
+ "edges": edges_json,
232
+ "total_edges": len(edges_json)
233
+ }
234
+
235
+
236
+ async def find_flow_diffusion_clusters_and_summarize(
237
+ node_datas: list[dict],
238
+ query: str,
239
+ query_param: QueryParam,
240
+ knowledge_graph_inst: BaseGraphStorage,
241
+ global_config: dict,
242
+ ):
243
+ """
244
+ Apply flow diffusion to find clusters and summarize them using LLM.
245
+
246
+ Parameters:
247
+ -----------
248
+ node_datas : list[dict]
249
+ List of node data dictionaries
250
+ query : str
251
+ The original query that provides context for the relationship analysis
252
+ query_param : QueryParam
253
+ Query parameters (includes flow diffusion configuration)
254
+ knowledge_graph_inst : BaseGraphStorage
255
+ Knowledge graph storage instance
256
+ global_config : dict
257
+ Global configuration
258
+
259
+ Returns:
260
+ --------
261
+ list
262
+ If return_raw_clusters=True: List of cluster JSON objects with subgraph data
263
+ If return_raw_clusters=False: List of summarized cluster relationships
264
+ """
265
+ # Use cached NetworkX graph to avoid repeated construction
266
+ if hasattr(knowledge_graph_inst, 'get_cached_nx_graph'):
267
+ logger.info("Attempting to use cached NetworkX graph...")
268
+ G = await knowledge_graph_inst.get_cached_nx_graph()
269
+ logger.info(f"Successfully obtained NetworkX graph with {G.number_of_nodes()} nodes and {G.number_of_edges()} edges")
270
+ else:
271
+ logger.info("Storage class does not support cached graphs, building new one...")
272
+ # Compatibility handling: if no cache method, use original approach
273
+ G = nx.Graph()
274
+ edges = await knowledge_graph_inst.edges()
275
+ nodes = await knowledge_graph_inst.nodes()
276
+
277
+ # Add edges with weights
278
+ for u, v in edges:
279
+ edge_data = await knowledge_graph_inst.get_edge(u, v)
280
+ if edge_data and 'weight' in edge_data:
281
+ G.add_edge(u, v, weight=edge_data['weight'])
282
+ else:
283
+ G.add_edge(u, v, weight=1.0) # Default weight: 1.0 if not specified
284
+
285
+ G.add_nodes_from(nodes)
286
+ logger.info(f"Built new NetworkX graph with {G.number_of_nodes()} nodes and {G.number_of_edges()} edges")
287
+
288
+ # Get source nodes from node_datas (limit to configured maximum)
289
+ source_nodes = [dp["entity_name"] for dp in node_datas[:query_param.max_source_nodes]]
290
+
291
+ # Apply flow diffusion from each source node independently
292
+ all_clusters = []
293
+ use_llm_func = global_config["llm_model_func"]
294
+
295
+ logger.info(f"Starting flow diffusion from {len(source_nodes)} source nodes independently")
296
+
297
+ # Pre-compute node embeddings and query embeddings to avoid repeated computation in loops
298
+ logger.info("Pre-computing embeddings for flow diffusion...")
299
+ node_embeddings, subquery_embedding = await get_embeddings_for_flow_diffusion(
300
+ G, query, knowledge_graph_inst, global_config, query_param
301
+ )
302
+ logger.info(f"Pre-computed embeddings: {len(node_embeddings)} nodes, query embedding: {'Yes' if subquery_embedding else 'No'}")
303
+
304
+ # Initialize cluster processing based on configuration
305
+ all_clusters = []
306
+
307
+ if query_param.use_batch_cluster_summarization:
308
+ logger.info("Using batch cluster summarization mode (more efficient)")
309
+ else:
310
+ logger.info("Using individual cluster summarization mode (original approach)")
311
+
312
+ if query_param.use_batch_cluster_summarization:
313
+ # Collect all clusters first for batch processing
314
+ clusters_to_summarize = []
315
+ raw_clusters = []
316
+
317
+ # Run flow diffusion from each source node
318
+ for source_node in source_nodes:
319
+
320
+ if not G.has_node(source_node):
321
+ continue
322
+
323
+ # Get source node information
324
+ source_node_info = await knowledge_graph_inst.get_node(source_node)
325
+ if not source_node_info:
326
+ source_node_info = {"entity_type": "UNKNOWN", "description": "No description available"}
327
+
328
+ # Calculate confidence based on source node's relevance to the query
329
+ confidence = 0.7 # Default confidence
330
+
331
+ # Apply flow diffusion from this source node
332
+ wfd = QueryAwareWeightedFlowDiffusion(
333
+ G, source_node, source_node, confidence,
334
+ node_embeddings=node_embeddings,
335
+ subquery_embedding=subquery_embedding,
336
+ weight_func=query_param.weight_func
337
+ )
338
+ wfd.initialize(alpha=query_param.alpha)
339
+ diffused_nodes = wfd.flow_diffusion()
340
+
341
+ if len(diffused_nodes) > 1: # Multi-node cluster
342
+ # Get cluster nodes and their flow values
343
+ cluster_nodes = list(diffused_nodes.keys())
344
+ cluster_flow_values = list(diffused_nodes.values())
345
+
346
+ # Only process if cluster has significant flow (using configured threshold)
347
+ if cluster_flow_values:
348
+ max_flow = max(cluster_flow_values)
349
+ if max_flow < query_param.min_flow_threshold:
350
+ continue
351
+ else:
352
+ continue # Skip if no flow values
353
+
354
+ # Get node information for the cluster
355
+ cluster_node_data = []
356
+ for node in cluster_nodes:
357
+ node_info = await knowledge_graph_inst.get_node(node)
358
+ if node_info:
359
+ cluster_node_data.append({
360
+ 'name': node,
361
+ 'entity_type': node_info.get('entity_type', 'UNKNOWN'),
362
+ 'description': node_info.get('description', 'UNKNOWN'),
363
+ 'flow_value': diffused_nodes.get(node, 0.0)
364
+ })
365
+
366
+ # Sort by flow value
367
+ cluster_node_data.sort(key=lambda x: x['flow_value'], reverse=True)
368
+
369
+ if query_param.return_raw_clusters:
370
+ # Convert subgraph to JSON format
371
+ cluster_json = convert_subgraph_to_json(G, cluster_nodes, diffused_nodes, source_node)
372
+ # Add node details to the JSON
373
+ cluster_json["node_details"] = cluster_node_data
374
+ # Enrich source_node with type + description
375
+ for nd in cluster_node_data:
376
+ if nd["name"] == source_node:
377
+ cluster_json["source_node"] = {
378
+ "entity": source_node,
379
+ "entity_type": nd.get("entity_type", "UNKNOWN"),
380
+ "description": nd.get("description", "UNKNOWN")
381
+ }
382
+ break
383
+ # Remove edges before appending
384
+ cluster_json.pop("edges", None)
385
+ cluster_json.pop("total_edges", None)
386
+ cluster_json.pop("node_details", None)
387
+ raw_clusters.append(cluster_json)
388
+ else:
389
+ # Collect cluster data for batch processing
390
+ clusters_to_summarize.append((cluster_node_data, source_node))
391
+
392
+ else: # --- NEW FALLBACK: single-node cluster ---
393
+ node_info = await knowledge_graph_inst.get_node(source_node)
394
+ cluster_node_data = [{
395
+ 'name': source_node,
396
+ 'entity_type': node_info.get('entity_type', 'UNKNOWN') if node_info else 'UNKNOWN',
397
+ 'description': node_info.get('description', 'UNKNOWN') if node_info else 'UNKNOWN',
398
+ 'flow_value': diffused_nodes.get(source_node, 0.0) if diffused_nodes else 0.0
399
+ }]
400
+
401
+ if query_param.return_raw_clusters:
402
+ cluster_json = {
403
+ "source_node": {
404
+ "entity": source_node,
405
+ "entity_type": cluster_node_data[0]["entity_type"],
406
+ "description": cluster_node_data[0]["description"]
407
+ },
408
+ "cluster_size": 1,
409
+ "max_flow_value": 0.0,
410
+ "nodes": [{
411
+ "id": 0,
412
+ "entity": source_node,
413
+ "entity_type": cluster_node_data[0]["entity_type"],
414
+ "description": cluster_node_data[0]["description"],
415
+ "rank": 0
416
+ }],
417
+ "edges": [],
418
+ "total_edges": 0,
419
+ "node_details": cluster_node_data
420
+ }
421
+ # Remove edges before appending
422
+ cluster_json.pop("edges", None)
423
+ cluster_json.pop("total_edges", None)
424
+ cluster_json.pop("node_details", None)
425
+ raw_clusters.append(cluster_json)
426
+ else:
427
+ clusters_to_summarize.append((cluster_node_data, source_node))
428
+
429
+ # Process clusters based on return type
430
+ if query_param.return_raw_clusters:
431
+ all_clusters = raw_clusters
432
+ else:
433
+ # Batch process clusters in chunks to avoid token limits
434
+ if clusters_to_summarize:
435
+ all_clusters = []
436
+ batch_size = query_param.batch_cluster_size
437
+
438
+ # Process clusters in batches
439
+ for i in range(0, len(clusters_to_summarize), batch_size):
440
+ batch_clusters = clusters_to_summarize[i:i + batch_size]
441
+ logger.info(f"Processing batch {i//batch_size + 1}/{(len(clusters_to_summarize) + batch_size - 1)//batch_size} with {len(batch_clusters)} clusters")
442
+
443
+ cluster_summaries = await summarize_clusters_batch_with_llm(
444
+ batch_clusters, use_llm_func, global_config
445
+ )
446
+ all_clusters.extend([summary for summary in cluster_summaries if summary])
447
+ else:
448
+ all_clusters = []
449
+
450
+ else:
451
+ # Process clusters individually (original approach)
452
+ for source_node in source_nodes:
453
+
454
+ if not G.has_node(source_node):
455
+ continue
456
+
457
+ if G.nodes[source_node].get("entity_type", "").lower() == "complete_table":
458
+ continue
459
+
460
+ # Get source node information
461
+ source_node_info = await knowledge_graph_inst.get_node(source_node)
462
+ if not source_node_info:
463
+ source_node_info = {"entity_type": "UNKNOWN", "description": "No description available"}
464
+
465
+ # Calculate confidence based on source node's relevance to the query
466
+ confidence = 0.7 # Default confidence
467
+
468
+ # Apply flow diffusion from this source node
469
+ wfd = QueryAwareWeightedFlowDiffusion(
470
+ G, source_node, source_node, confidence,
471
+ node_embeddings=node_embeddings,
472
+ subquery_embedding=subquery_embedding,
473
+ weight_func=query_param.weight_func
474
+ )
475
+ wfd.initialize(alpha=query_param.alpha)
476
+ diffused_nodes = wfd.flow_diffusion()
477
+
478
+ if len(diffused_nodes) > 1: # Only consider clusters with multiple nodes
479
+ # Get cluster nodes and their flow values
480
+ cluster_nodes = list(diffused_nodes.keys())
481
+ cluster_flow_values = list(diffused_nodes.values())
482
+
483
+ # Only process if cluster has significant flow (using configured threshold)
484
+ if cluster_flow_values:
485
+ max_flow = max(cluster_flow_values)
486
+ if max_flow < query_param.min_flow_threshold:
487
+ continue
488
+ else:
489
+ continue # Skip if no flow values
490
+
491
+ # Get node information for the cluster
492
+ cluster_node_data = []
493
+ for node in cluster_nodes:
494
+ node_info = await knowledge_graph_inst.get_node(node)
495
+ if node_info:
496
+ cluster_node_data.append({
497
+ 'name': node,
498
+ 'entity_type': node_info.get('entity_type', 'UNKNOWN'),
499
+ 'description': node_info.get('description', 'UNKNOWN'),
500
+ 'flow_value': diffused_nodes.get(node, 0.0)
501
+ })
502
+
503
+ # Sort by flow value
504
+ cluster_node_data.sort(key=lambda x: x['flow_value'], reverse=True)
505
+
506
+ # Check if we should return raw cluster data instead of LLM summaries
507
+ if query_param.return_raw_clusters:
508
+ # Convert subgraph to JSON format
509
+ cluster_json = convert_subgraph_to_json(G, cluster_nodes, diffused_nodes, source_node)
510
+ # Add node details to the JSON
511
+ cluster_json["node_details"] = cluster_node_data
512
+ all_clusters.append(cluster_json)
513
+ else:
514
+ # Create cluster summary using LLM (individual processing)
515
+ cluster_summary = await summarize_cluster_with_llm(
516
+ cluster_node_data, source_node, use_llm_func, global_config
517
+ )
518
+
519
+ if cluster_summary:
520
+ all_clusters.append(cluster_summary)
521
+
522
+ logger.info(f"Flow diffusion completed, {len(all_clusters)} clusters found")
523
+
524
+ # Only apply token constraints if we're returning LLM summaries
525
+ if not query_param.return_raw_clusters:
526
+ # Limit the number of clusters based on token constraints
527
+ original_cluster_count = len(all_clusters)
528
+ all_clusters = truncate_list_by_token_size(
529
+ all_clusters,
530
+ key=lambda x: x,
531
+ max_token_size=query_param.max_token_for_local_context,
532
+ )
533
+
534
+ if original_cluster_count != len(all_clusters):
535
+ logger.info(f"Clusters truncated from {original_cluster_count} to {len(all_clusters)} due to token limit")
536
+
537
+ return all_clusters
538
+
539
+
540
+ async def summarize_clusters_batch_with_llm(
541
+ clusters_data: list[tuple[list[dict], str]],
542
+ use_llm_func: callable,
543
+ global_config: dict = None,
544
+ ) -> list[str]:
545
+ """
546
+ Summarize multiple clusters using a single LLM call.
547
+
548
+ Parameters:
549
+ -----------
550
+ clusters_data : list[tuple[list[dict], str]]
551
+ List of tuples containing (cluster_node_data, source_node) for each cluster
552
+ use_llm_func : callable
553
+ LLM function to use for summarization
554
+ global_config : dict
555
+ Global configuration (optional)
556
+
557
+ Returns:
558
+ --------
559
+ list[str]
560
+ List of summarized cluster relationships
561
+ """
562
+ if not clusters_data:
563
+ return []
564
+
565
+ # Filter out clusters with less than 2 nodes
566
+ valid_clusters = []
567
+ for cluster_node_data, source_node in clusters_data:
568
+ if len(cluster_node_data) >= 2:
569
+ valid_clusters.append((cluster_node_data, source_node))
570
+
571
+ if not valid_clusters:
572
+ return []
573
+
574
+ # Create batch prompt for all clusters
575
+ clusters_info = []
576
+ for i, (cluster_node_data, source_node) in enumerate(valid_clusters, 1):
577
+ nodes_info = []
578
+ for node_data in cluster_node_data:
579
+ nodes_info.append(
580
+ f"- {node_data['name']} ({node_data['entity_type']}): {node_data['description']} "
581
+ f"(flow strength: {node_data['flow_value']:.3f})"
582
+ )
583
+
584
+ nodes_text = "\n".join(nodes_info)
585
+ clusters_info.append(f"""Cluster {i}:
586
+ Source Node: {source_node}
587
+ Cluster Nodes:
588
+ {nodes_text}""")
589
+
590
+ all_clusters_text = "\n\n".join(clusters_info)
591
+
592
+ prompt = f"""Please analyze the following clusters of entities and their relationships, then provide a concise summary for each cluster.
593
+
594
+ Clusters Information:
595
+ {all_clusters_text}
596
+
597
+ For each cluster, please provide a summary that explains:
598
+ 1. How these entities are related to each other
599
+ 2. What concept or theme this cluster represents
600
+ 3. The significance of the connections between these entities
601
+ 4. How the source node connects to the other entities in this cluster
602
+
603
+ Please format your response as follows:
604
+ Cluster 1: [summary for cluster 1]
605
+ Cluster 2: [summary for cluster 2]
606
+ ...
607
+ Cluster N: [summary for cluster N]
608
+
609
+ Keep each summary concise but informative, focusing on the relationships and thematic connections."""
610
+
611
+ # Use configuration-based max_tokens if available, otherwise use a reasonable default
612
+ max_tokens = 2000 # Increased for batch processing
613
+ if global_config and "entity_summary_to_max_tokens" in global_config:
614
+ max_tokens = min(global_config["entity_summary_to_max_tokens"] * len(valid_clusters), 16384)
615
+
616
+ logger.info(f"Processing batch of {len(valid_clusters)} clusters with max_tokens={max_tokens}")
617
+
618
+ try:
619
+ batch_summary = await use_llm_func(prompt, max_tokens=max_tokens)
620
+ batch_summary = batch_summary.strip()
621
+
622
+ # Parse the response to extract individual cluster summaries
623
+ summaries = []
624
+ lines = batch_summary.split('\n')
625
+ current_summary = ""
626
+
627
+ for line in lines:
628
+ line = line.strip()
629
+ if line.startswith('Cluster ') and ':' in line:
630
+ # Save previous summary if exists
631
+ if current_summary:
632
+ summaries.append(current_summary.strip())
633
+ # Start new summary
634
+ current_summary = line.split(':', 1)[1].strip()
635
+ elif line and current_summary:
636
+ current_summary += " " + line
637
+
638
+ # Add the last summary
639
+ if current_summary:
640
+ summaries.append(current_summary.strip())
641
+
642
+ # Ensure we have the right number of summaries
643
+ while len(summaries) < len(valid_clusters):
644
+ summaries.append("Cluster summary not available")
645
+
646
+ return summaries[:len(valid_clusters)]
647
+
648
+ except Exception as e:
649
+ logger.error(f"Error summarizing clusters batch: {e}")
650
+ # Fallback: create individual summaries for each cluster
651
+ fallback_summaries = []
652
+ for cluster_node_data, source_node in valid_clusters:
653
+ if len(cluster_node_data) > 0:
654
+ top_nodes = cluster_node_data[:3] # Get top 3 nodes by flow value
655
+ node_names = [node['name'] for node in top_nodes]
656
+ if len(cluster_node_data) > 1:
657
+ fallback_summary = f"Cluster centered around {source_node} connecting to {', '.join(node_names)} and {len(cluster_node_data)-1} other entities with flow values ranging from {cluster_node_data[0]['flow_value']:.3f} to {cluster_node_data[-1]['flow_value']:.3f}"
658
+ else:
659
+ fallback_summary = f"Cluster centered around {source_node} connecting to {', '.join(node_names)}"
660
+ fallback_summaries.append(fallback_summary)
661
+ else:
662
+ fallback_summaries.append(f"Cluster connecting {source_node} with related entities")
663
+
664
+ return fallback_summaries
665
+
666
+
667
+ async def summarize_cluster_with_llm(
668
+ cluster_node_data: list[dict],
669
+ source_node: str,
670
+ use_llm_func: callable,
671
+ global_config: dict = None,
672
+ ) -> str:
673
+ """
674
+ Summarize a cluster of nodes using LLM.
675
+
676
+ Parameters:
677
+ -----------
678
+ cluster_node_data : list[dict]
679
+ List of node data in the cluster
680
+ source_node : str
681
+ Source node name
682
+ use_llm_func : callable
683
+ LLM function to use for summarization
684
+ global_config : dict
685
+ Global configuration (optional)
686
+
687
+ Returns:
688
+ --------
689
+ str
690
+ Summarized cluster relationship
691
+ """
692
+ if len(cluster_node_data) < 2:
693
+ return None
694
+
695
+ # Create prompt for cluster summarization
696
+ nodes_info = []
697
+ for node_data in cluster_node_data:
698
+ nodes_info.append(
699
+ f"- {node_data['name']} ({node_data['entity_type']}): {node_data['description']} "
700
+ f"(flow strength: {node_data['flow_value']:.3f})"
701
+ )
702
+
703
+ nodes_text = "\n".join(nodes_info)
704
+
705
+ prompt = f"""Please analyze the following cluster of entities and their relationships, then provide a concise summary of how they are connected and what this cluster represents.
706
+
707
+ Cluster Information:
708
+ Source Node: {source_node}
709
+ Cluster Nodes:
710
+ {nodes_text}
711
+
712
+ Please provide a summary that explains:
713
+ 1. How these entities are related to each other
714
+ 2. What concept or theme this cluster represents
715
+ 3. The significance of the connections between these entities
716
+ 4. How the source node connects to the other entities in this cluster
717
+
718
+ Keep the summary concise but informative, focusing on the relationships and thematic connections."""
719
+
720
+ # Use configuration-based max_tokens if available, otherwise use a reasonable default
721
+ max_tokens = 1000 # Default value
722
+ if global_config and "entity_summary_to_max_tokens" in global_config:
723
+ max_tokens = global_config["entity_summary_to_max_tokens"]
724
+
725
+ try:
726
+ summary = await use_llm_func(prompt, max_tokens=max_tokens)
727
+ return summary.strip()
728
+ except Exception as e:
729
+ logger.error(f"Error summarizing cluster: {e}")
730
+ # Fallback to more informative description
731
+ if len(cluster_node_data) > 0:
732
+ top_nodes = cluster_node_data[:3] # Get top 3 nodes by flow value
733
+ node_names = [node['name'] for node in top_nodes]
734
+ if len(cluster_node_data) > 1:
735
+ return f"Cluster centered around {source_node} connecting to {', '.join(node_names)} and {len(cluster_node_data)-1} other entities with flow values ranging from {cluster_node_data[0]['flow_value']:.3f} to {cluster_node_data[-1]['flow_value']:.3f}"
736
+ else:
737
+ return f"Cluster centered around {source_node} connecting to {', '.join(node_names)}"
738
+ else:
739
+ return f"Cluster connecting {source_node} with related entities"
src/answering/context.py ADDED
@@ -0,0 +1,322 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Context building for RAG query processing.
3
+
4
+ This module provides functions to build query context from the knowledge graph
5
+ based on different query modes (local, global, hybrid).
6
+ """
7
+
8
+ import asyncio
9
+ from ..base import (
10
+ BaseGraphStorage,
11
+ BaseKVStorage,
12
+ BaseVectorStorage,
13
+ TextChunkSchema,
14
+ QueryParam,
15
+ )
16
+ from ..utils import (
17
+ logger,
18
+ list_of_list_to_csv,
19
+ csv_string_to_list,
20
+ )
21
+ from .clusters import find_flow_diffusion_clusters_and_summarize
22
+ from .text_units import find_most_related_text_unit_from_entities
23
+
24
+
25
+ async def build_query_context(
26
+ query: list,
27
+ knowledge_graph_inst: BaseGraphStorage,
28
+ entities_vdb: BaseVectorStorage,
29
+ relationships_vdb: BaseVectorStorage,
30
+ text_chunks_db: BaseKVStorage[TextChunkSchema],
31
+ query_param: QueryParam,
32
+ global_config: dict,
33
+ ):
34
+ """
35
+ Build query context based on extracted keywords and query mode.
36
+
37
+ Parameters:
38
+ -----------
39
+ query : list
40
+ List containing [ll_keywords, hl_keywords]
41
+ knowledge_graph_inst : BaseGraphStorage
42
+ Knowledge graph storage instance
43
+ entities_vdb : BaseVectorStorage
44
+ Entity vector database
45
+ relationships_vdb : BaseVectorStorage
46
+ Relationships vector database
47
+ text_chunks_db : BaseKVStorage[TextChunkSchema]
48
+ Text chunks database
49
+ query_param : QueryParam
50
+ Query parameters including mode (local/global/hybrid)
51
+ global_config : dict
52
+ Global configuration
53
+
54
+ Returns:
55
+ --------
56
+ str
57
+ Formatted context string for LLM response generation
58
+ """
59
+ ll_keywords, hl_keywords = query[0], query[1]
60
+
61
+ # Initialize context variables
62
+ entities_context, relations_context, text_units_context = "", "", ""
63
+
64
+ if query_param.mode == "local":
65
+ # Local mode: use ll_keywords only
66
+ if ll_keywords == "":
67
+ logger.warning("Low level keywords is empty for local mode")
68
+ return "", "", ""
69
+
70
+ (
71
+ entities_context,
72
+ relations_context,
73
+ text_units_context,
74
+ ) = await _get_node_data_with_flow_diffusion(
75
+ ll_keywords,
76
+ knowledge_graph_inst,
77
+ entities_vdb,
78
+ text_chunks_db,
79
+ query_param,
80
+ global_config,
81
+ )
82
+
83
+ elif query_param.mode == "global":
84
+ # Global mode: use hl_keywords only
85
+ if hl_keywords == "":
86
+ logger.warning("High level keywords is empty for global mode")
87
+ return "", "", ""
88
+
89
+ (
90
+ entities_context,
91
+ relations_context,
92
+ text_units_context,
93
+ ) = await _get_node_data_with_flow_diffusion(
94
+ hl_keywords,
95
+ knowledge_graph_inst,
96
+ entities_vdb,
97
+ text_chunks_db,
98
+ query_param,
99
+ global_config,
100
+ )
101
+
102
+ elif query_param.mode == "hybrid":
103
+ # Hybrid mode: use combined keywords (both ll_keywords and hl_keywords)
104
+ if ll_keywords == "" and hl_keywords == "":
105
+ logger.warning("Both Low Level and High Level keywords are empty for hybrid mode")
106
+ return "", "", ""
107
+
108
+ # Get local information using ll_keywords
109
+ local_entities_context, local_relations_context, local_text_units_context = "", "", ""
110
+ if ll_keywords:
111
+ (
112
+ local_entities_context,
113
+ local_relations_context,
114
+ local_text_units_context,
115
+ ) = await _get_node_data_with_flow_diffusion(
116
+ ll_keywords,
117
+ knowledge_graph_inst,
118
+ entities_vdb,
119
+ text_chunks_db,
120
+ query_param,
121
+ global_config,
122
+ )
123
+
124
+ # Get global information using hl_keywords
125
+ global_entities_context, global_relations_context, global_text_units_context = "", "", ""
126
+ if hl_keywords:
127
+ (
128
+ global_entities_context,
129
+ global_relations_context,
130
+ global_text_units_context,
131
+ ) = await _get_node_data_with_flow_diffusion(
132
+ hl_keywords,
133
+ knowledge_graph_inst,
134
+ entities_vdb,
135
+ text_chunks_db,
136
+ query_param,
137
+ global_config,
138
+ )
139
+
140
+ # Return context based on mode
141
+ if query_param.mode == "local":
142
+ if query_param.return_raw_entities:
143
+ return entities_context
144
+ elif query_param.return_raw_clusters:
145
+ return relations_context
146
+
147
+ return f"""
148
+ -----local-information-----
149
+ -----low-level entity information-----
150
+ ```csv
151
+ {entities_context}
152
+ ```
153
+ -----low-level relationship information-----
154
+ ```csv
155
+ {relations_context}
156
+ ```
157
+ -----Sources-----
158
+ ```csv
159
+ {text_units_context}
160
+ ```
161
+ """
162
+ elif query_param.mode == "global":
163
+ if query_param.return_raw_entities:
164
+ return entities_context
165
+ elif query_param.return_raw_clusters:
166
+ return relations_context
167
+
168
+ return f"""
169
+ -----global-information-----
170
+ -----high-level entity information-----
171
+ ```csv
172
+ {entities_context}
173
+ ```
174
+ -----high-level relationship information-----
175
+ ```csv
176
+ {relations_context}
177
+ ```
178
+ -----Sources-----
179
+ ```csv
180
+ {text_units_context}
181
+ ```
182
+ """
183
+ elif query_param.mode == "hybrid":
184
+ if query_param.return_raw_entities:
185
+ # Merge local + global entities CSV into one CSV and reindex id
186
+ merged_rows = []
187
+ if local_entities_context:
188
+ merged_rows += csv_string_to_list(local_entities_context)[1:]
189
+ if global_entities_context:
190
+ merged_rows += csv_string_to_list(global_entities_context)[1:]
191
+ for idx, row in enumerate(merged_rows):
192
+ if row:
193
+ row[0] = str(idx)
194
+ merged_entities_csv = list_of_list_to_csv(
195
+ [["id", "entity", "entity_type", "description", "rank"]] + merged_rows
196
+ )
197
+ return merged_entities_csv
198
+ elif query_param.return_raw_clusters:
199
+ return local_relations_context + global_relations_context
200
+
201
+ return f"""
202
+ -----hybrid-information-----
203
+ -----local information (from low-level keywords)-----
204
+ -----local entity information-----
205
+ ```csv
206
+ {local_entities_context}
207
+ ```
208
+ -----local relationship information-----
209
+ ```csv
210
+ {local_relations_context}
211
+ ```
212
+ -----local sources-----
213
+ ```csv
214
+ {local_text_units_context}
215
+ ```
216
+ -----global information (from high-level keywords)-----
217
+ -----global entity information-----
218
+ ```csv
219
+ {global_entities_context}
220
+ ```
221
+ -----global relationship information-----
222
+ ```csv
223
+ {global_relations_context}
224
+ ```
225
+ -----global sources-----
226
+ ```csv
227
+ {global_text_units_context}
228
+ ```
229
+ """
230
+ else:
231
+ return ""
232
+
233
+
234
+ async def _get_node_data_with_flow_diffusion(
235
+ query,
236
+ knowledge_graph_inst: BaseGraphStorage,
237
+ entities_vdb: BaseVectorStorage,
238
+ text_chunks_db: BaseKVStorage[TextChunkSchema],
239
+ query_param: QueryParam,
240
+ global_config: dict,
241
+ ):
242
+ """
243
+ Get node data using flow diffusion for finding relationships.
244
+
245
+ Parameters:
246
+ -----------
247
+ query : str
248
+ Query string (can be either ll_keywords or hl_keywords)
249
+ knowledge_graph_inst : BaseGraphStorage
250
+ Knowledge graph storage instance
251
+ entities_vdb : BaseVectorStorage
252
+ Entity vector database
253
+ text_chunks_db : BaseKVStorage[TextChunkSchema]
254
+ Text chunks database
255
+ query_param : QueryParam
256
+ Query parameters
257
+ global_config : dict
258
+ Global configuration
259
+
260
+ Returns:
261
+ --------
262
+ tuple
263
+ (entities_context, relations_context, text_units_context)
264
+ """
265
+ results = await entities_vdb.query(query, top_k=query_param.max_source_nodes)
266
+ if not len(results):
267
+ return "", "", ""
268
+
269
+ node_datas = await asyncio.gather(
270
+ *[knowledge_graph_inst.get_node(r["entity_name"]) for r in results]
271
+ )
272
+ if not all([n is not None for n in node_datas]):
273
+ logger.warning("Some nodes are missing, maybe the storage is damaged")
274
+
275
+ node_degrees = await asyncio.gather(
276
+ *[knowledge_graph_inst.node_degree(r["entity_name"]) for r in results]
277
+ )
278
+ node_datas = [
279
+ {**n, "entity_name": k["entity_name"], "rank": d}
280
+ for k, n, d in zip(results, node_datas, node_degrees)
281
+ if n is not None
282
+ ]
283
+
284
+ use_text_units = await find_most_related_text_unit_from_entities(
285
+ node_datas, query_param, text_chunks_db, knowledge_graph_inst
286
+ )
287
+
288
+ # Use flow diffusion instead of the original relationship finding method
289
+ use_relations = await find_flow_diffusion_clusters_and_summarize(
290
+ node_datas, query, query_param, knowledge_graph_inst, global_config
291
+ )
292
+
293
+ logger.info(
294
+ f"Flow diffusion query uses {len(node_datas)} entities, {len(use_relations)} cluster summaries, {len(use_text_units)} text units"
295
+ )
296
+
297
+ entites_section_list = [["id", "entity", "entity_type", "description", "rank"]]
298
+ for i, n in enumerate(node_datas):
299
+ entites_section_list.append([
300
+ i,
301
+ n["entity_name"],
302
+ n.get("entity_type", "UNKNOWN"),
303
+ n.get("description", "UNKNOWN"),
304
+ n["rank"],
305
+ ])
306
+ entities_context = list_of_list_to_csv(entites_section_list)
307
+
308
+ # Relations context: return JSON clusters when requested; otherwise CSV
309
+ if query_param.return_raw_clusters:
310
+ relations_context = use_relations
311
+ else:
312
+ relations_section_list = [["id", "cluster_summary"]]
313
+ for i, summary in enumerate(use_relations):
314
+ relations_section_list.append([i, summary])
315
+ relations_context = list_of_list_to_csv(relations_section_list)
316
+
317
+ text_units_section_list = [["id", "content"]]
318
+ for i, t in enumerate(use_text_units):
319
+ text_units_section_list.append([i, t["content"]])
320
+ text_units_context = list_of_list_to_csv(text_units_section_list)
321
+
322
+ return entities_context, relations_context, text_units_context