Add code: src, benchmarks, data, requirements, run.sh, README
Browse filesThis view is limited to 50 files because it contains too many changes. See raw diff
- .gitattributes +6 -0
- .gitignore +47 -0
- README.md +345 -0
- benchmarks/multihop/benchmark_multihop.py +531 -0
- benchmarks/run.py +371 -0
- benchmarks/summarization/benchmark_summarization.py +673 -0
- benchmarks/text2sql/benchmark_text2sql.py +385 -0
- benchmarks/text2sql/evaluate_schema.py +213 -0
- benchmarks/ultradomain/benchmark_ultradomain.py +548 -0
- data/multihop/2wikimultihopqa.json +0 -0
- data/multihop/2wikimultihopqa_corpus.json +0 -0
- data/multihop/README.md +23 -0
- data/multihop/hotpotqa.json +0 -0
- data/multihop/hotpotqa_corpus.json +0 -0
- data/multihop/musique.json +3 -0
- data/multihop/musique_corpus.json +0 -0
- data/summarization/README.md +26 -0
- data/text2sql/README.md +70 -0
- data/text2sql/bird/bird.jsonl +129 -0
- data/text2sql/bird/databases/superhero/superhero.sqlite +3 -0
- data/text2sql/bird/databases/superhero/superhero_db_summary.json +0 -0
- data/text2sql/spider2-lite/bigquery/austin/austin_bigquery_summary.json +39 -0
- data/text2sql/spider2-lite/golden_lite_spider_total.json +1443 -0
- data/text2sql/spider2-lite/snowflake/AUSTIN/AUSTIN_db_summary.json +39 -0
- data/text2sql/spider2-lite/spider2-lite.jsonl +0 -0
- data/text2sql/spider2-lite/sqlite/Pagila/Pagila.sqlite +3 -0
- data/text2sql/spider2-lite/sqlite/Pagila/Pagila_db_summary.json +0 -0
- data/ultradomain/README.md +25 -0
- docs/figs/GraphRAG.png +3 -0
- docs/figs/LightRAG.png +3 -0
- docs/figs/QAFD-RAG.png +3 -0
- docs/index.html +452 -0
- experiments/edge_weight_ablation.py +401 -0
- experiments/query_aware_ablation.py +478 -0
- experiments/run_all_ablations.sh +49 -0
- experiments/run_alpha_qa_sweep.sh +29 -0
- experiments/run_batch_push_ablation.sh +29 -0
- experiments/run_batch_push_nvidia.sh +55 -0
- experiments/run_decoupled_ablation.sh +32 -0
- experiments/run_nvidia_qa_ablation.sh +43 -0
- experiments/run_overnight_qa_awareness.sh +136 -0
- experiments/run_qa_awareness_v2.sh +34 -0
- experiments/run_sim_mode_ablation.sh +29 -0
- requirements.txt +46 -0
- run.sh +113 -0
- src/QAFD_RAG.py +747 -0
- src/__init__.py +4 -0
- src/answering/__init__.py +42 -0
- src/answering/clusters.py +739 -0
- 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
|
docs/figs/LightRAG.png
ADDED
|
Git LFS Details
|
docs/figs/QAFD-RAG.png
ADDED
|
Git LFS Details
|
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>—community detection (GraphRAG) or one-hop entity lookup (LightRAG)—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 — Knowledge Graph Construction</h3><ol><li><strong>Chunking</strong> — Documents split into overlapping token-based chunks.</li><li><strong>Entity Extraction</strong> — LLM identifies entities and relationships.</li><li><strong>Graph Assembly</strong> — Entities become nodes, relationships become weighted edges.</li></ol></div>
|
| 382 |
+
<div class="card"><h3>Stage 2 — Query-Aware Retrieval</h3><ol><li><strong>Entity Matching</strong> — Query keywords matched to KG entities via vector similarity.</li><li><strong>Flow Diffusion</strong> — Mass propagates with query-aware edge re-weighting.</li><li><strong>Context Assembly</strong> — 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
|