Spaces:
Runtime error
Runtime error
Upload 27 files
Browse files- .gitattributes +1 -0
- Dockerfile +16 -0
- app.py +67 -0
- core/pipeline_1/.gitkeep +0 -0
- core/pipeline_1/filter.py +95 -0
- core/pipeline_1/logic.py +116 -0
- core/pipeline_1/main.py +105 -0
- core/pipeline_1/metadata_index.json +0 -0
- core/pipeline_2/.gitkeep +0 -0
- core/pipeline_2/indexer.py +151 -0
- core/pipeline_2/logic.py +499 -0
- core/pipeline_2/qdrant_storage/.lock +1 -0
- core/pipeline_2/qdrant_storage/collection/papers/storage.sqlite +3 -0
- core/pipeline_2/qdrant_storage/meta.json +1 -0
- core/pipeline_2/retriever.py +76 -0
- core/pipeline_3/.gitkeep +0 -0
- core/pipeline_3/checkpoint-2.json +9 -0
- core/pipeline_3/embeddings.npy +3 -0
- core/pipeline_3/logic.py +698 -0
- core/pipeline_3/queries.py +141 -0
- core/pipeline_3/setup.py +459 -0
- evaluation_set.json +98 -0
- requirements.txt +134 -0
- run_benchmarks.py +274 -0
- services/metrics_service.py +111 -0
- services/paper_fetcher.py +216 -0
- utils/download_aiml_data.py +80 -0
- utils/token_count.py +183 -0
.gitattributes
CHANGED
|
@@ -33,3 +33,4 @@ 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 |
+
core/pipeline_2/qdrant_storage/collection/papers/storage.sqlite filter=lfs diff=lfs merge=lfs -text
|
Dockerfile
ADDED
|
@@ -0,0 +1,16 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
FROM python:3.11-slim
|
| 2 |
+
|
| 3 |
+
WORKDIR /app
|
| 4 |
+
|
| 5 |
+
RUN apt-get update && apt-get install -y --no-install-recommends \
|
| 6 |
+
build-essential \
|
| 7 |
+
&& rm -rf /var/lib/apt/lists/*
|
| 8 |
+
|
| 9 |
+
COPY requirements.txt .
|
| 10 |
+
RUN pip install --no-cache-dir -r requirements.txt
|
| 11 |
+
|
| 12 |
+
COPY . .
|
| 13 |
+
|
| 14 |
+
EXPOSE 7860
|
| 15 |
+
|
| 16 |
+
CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]
|
app.py
ADDED
|
@@ -0,0 +1,67 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from fastapi import FastAPI, HTTPException
|
| 2 |
+
from pydantic import BaseModel
|
| 3 |
+
import uvicorn
|
| 4 |
+
import importlib.util
|
| 5 |
+
import sys
|
| 6 |
+
from pathlib import Path
|
| 7 |
+
|
| 8 |
+
from core.pipeline_1.logic import PipelineLLMOnly
|
| 9 |
+
from core.pipeline_2.logic import PipelineRAG
|
| 10 |
+
from core.pipeline_3.logic import PipelineGraphRAG
|
| 11 |
+
|
| 12 |
+
app = FastAPI(
|
| 13 |
+
title="SEC Dataset API",
|
| 14 |
+
description="A simple FastAPI setup to fetch and interact with the PleIAs/SEC dataset.",
|
| 15 |
+
version="1.0.0"
|
| 16 |
+
)
|
| 17 |
+
|
| 18 |
+
# Request model
|
| 19 |
+
class QueryRequest(BaseModel):
|
| 20 |
+
query: str
|
| 21 |
+
|
| 22 |
+
# Initialize Pipeline (Scaling up for full context)
|
| 23 |
+
pipeline_baseline = PipelineLLMOnly(top_n=20, max_full_text=20)
|
| 24 |
+
pipeline_rag = PipelineRAG(retrieval_top_k=50, rerank_top_n=10, max_full_text=3)
|
| 25 |
+
pipeline_graph = PipelineGraphRAG(vector_top_k=50, rerank_top_n=10, max_full_text=3)
|
| 26 |
+
|
| 27 |
+
@app.get("/")
|
| 28 |
+
async def root():
|
| 29 |
+
return {"message": "Welcome to the SEC Dataset API", "status": "running"}
|
| 30 |
+
|
| 31 |
+
@app.post("/query/baseline")
|
| 32 |
+
async def query_baseline(request: QueryRequest):
|
| 33 |
+
"""
|
| 34 |
+
Pipeline 1: LLM-Only Baseline.
|
| 35 |
+
Performs metadata filtering and LLM synthesis.
|
| 36 |
+
"""
|
| 37 |
+
if not request.query:
|
| 38 |
+
raise HTTPException(status_code=400, detail="Query cannot be empty")
|
| 39 |
+
|
| 40 |
+
result = pipeline_baseline.run(request.query)
|
| 41 |
+
return result
|
| 42 |
+
|
| 43 |
+
|
| 44 |
+
@app.post("/query/rag")
|
| 45 |
+
async def query_rag(request: QueryRequest):
|
| 46 |
+
"""
|
| 47 |
+
Pipeline 2: Hybrid RAG.
|
| 48 |
+
Hybrid vector search + cross-encoder rerank + Groq LLM synthesis.
|
| 49 |
+
"""
|
| 50 |
+
if not request.query:
|
| 51 |
+
raise HTTPException(status_code=400, detail="Query cannot be empty")
|
| 52 |
+
|
| 53 |
+
result = pipeline_rag.run(request.query)
|
| 54 |
+
return result
|
| 55 |
+
|
| 56 |
+
|
| 57 |
+
@app.post("/query/graph")
|
| 58 |
+
async def query_graph(request: QueryRequest):
|
| 59 |
+
"""
|
| 60 |
+
Pipeline 3: GraphRAG.
|
| 61 |
+
TigerGraph vector search + topic linking + citation expansion + Groq LLM synthesis.
|
| 62 |
+
"""
|
| 63 |
+
if not request.query:
|
| 64 |
+
raise HTTPException(status_code=400, detail="Query cannot be empty")
|
| 65 |
+
|
| 66 |
+
result = pipeline_graph.run(request.query)
|
| 67 |
+
return result
|
core/pipeline_1/.gitkeep
ADDED
|
File without changes
|
core/pipeline_1/filter.py
ADDED
|
@@ -0,0 +1,95 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import os
|
| 2 |
+
import json
|
| 3 |
+
import glob
|
| 4 |
+
import logging
|
| 5 |
+
import re
|
| 6 |
+
from pathlib import Path
|
| 7 |
+
|
| 8 |
+
logging.basicConfig(level=logging.INFO)
|
| 9 |
+
logger = logging.getLogger(__name__)
|
| 10 |
+
|
| 11 |
+
class SimpleFilter:
|
| 12 |
+
def __init__(self, data_dir="data", index_path="core/pipeline-1/metadata_index.json"):
|
| 13 |
+
self.data_dir = Path(data_dir)
|
| 14 |
+
self.index_path = Path(index_path)
|
| 15 |
+
self.metadata = []
|
| 16 |
+
self._load_index()
|
| 17 |
+
|
| 18 |
+
def _load_index(self):
|
| 19 |
+
"""Load or build a lightweight index of all papers."""
|
| 20 |
+
if self.index_path.exists():
|
| 21 |
+
logger.info(f"Loading metadata index from {self.index_path}...")
|
| 22 |
+
with open(self.index_path, 'r') as f:
|
| 23 |
+
self.metadata = json.load(f)
|
| 24 |
+
else:
|
| 25 |
+
self.build_index()
|
| 26 |
+
|
| 27 |
+
def build_index(self):
|
| 28 |
+
"""Scan data directory and build a flat metadata index."""
|
| 29 |
+
logger.info(f"Building metadata index from {self.data_dir}...")
|
| 30 |
+
files = glob.glob(str(self.data_dir / "*.json"))
|
| 31 |
+
|
| 32 |
+
index_data = []
|
| 33 |
+
for f in files:
|
| 34 |
+
try:
|
| 35 |
+
with open(f, 'r') as fh:
|
| 36 |
+
data = json.load(fh)
|
| 37 |
+
index_data.append({
|
| 38 |
+
"id": data.get("id"),
|
| 39 |
+
"title": data.get("title", "").lower(),
|
| 40 |
+
"cited_by_count": data.get("cited_by_count", 0),
|
| 41 |
+
"topics": [t.get("display_name", "").lower() for t in data.get("topics", [])],
|
| 42 |
+
"publication_year": data.get("publication_year")
|
| 43 |
+
})
|
| 44 |
+
except Exception as e:
|
| 45 |
+
continue
|
| 46 |
+
|
| 47 |
+
self.index_path.parent.mkdir(parents=True, exist_ok=True)
|
| 48 |
+
with open(self.index_path, 'w') as f:
|
| 49 |
+
json.dump(index_data, f)
|
| 50 |
+
|
| 51 |
+
self.metadata = index_data
|
| 52 |
+
logger.info(f"Index built with {len(self.metadata)} records.")
|
| 53 |
+
|
| 54 |
+
def get_keywords(self, prompt):
|
| 55 |
+
stop_words = {'a', 'an', 'the', 'is', 'are', 'in', 'on', 'at', 'for', 'to', 'with', 'what', 'how', 'why'}
|
| 56 |
+
words = re.findall(r'\w+', prompt.lower())
|
| 57 |
+
return [w for w in words if w not in stop_words and len(w) > 2]
|
| 58 |
+
|
| 59 |
+
def filter_papers(self, query, top_n=50):
|
| 60 |
+
keywords = self.get_keywords(query)
|
| 61 |
+
scored_papers = []
|
| 62 |
+
|
| 63 |
+
for paper in self.metadata:
|
| 64 |
+
score = 0
|
| 65 |
+
title = paper["title"]
|
| 66 |
+
topics = paper["topics"]
|
| 67 |
+
|
| 68 |
+
# Approach 2: Keyword Title Match
|
| 69 |
+
title_matches = sum(1 for kw in keywords if kw in title)
|
| 70 |
+
score += title_matches * 10
|
| 71 |
+
|
| 72 |
+
# Approach 3: Topic Match
|
| 73 |
+
topic_matches = sum(1 for kw in keywords if any(kw in t for t in topics))
|
| 74 |
+
score += topic_matches * 5
|
| 75 |
+
|
| 76 |
+
# Citation Boost
|
| 77 |
+
score += (paper["cited_by_count"] / 1000)
|
| 78 |
+
|
| 79 |
+
if score > 0:
|
| 80 |
+
scored_papers.append((score, paper))
|
| 81 |
+
|
| 82 |
+
scored_papers.sort(key=lambda x: x[0], reverse=True)
|
| 83 |
+
results = [p for s, p in scored_papers[:top_n]]
|
| 84 |
+
|
| 85 |
+
if len(results) < top_n:
|
| 86 |
+
remaining = top_n - len(results)
|
| 87 |
+
all_sorted = sorted(self.metadata, key=lambda x: x["cited_by_count"], reverse=True)
|
| 88 |
+
for p in all_sorted:
|
| 89 |
+
if p["id"] not in [r["id"] for r in results]:
|
| 90 |
+
results.append(p)
|
| 91 |
+
remaining -= 1
|
| 92 |
+
if remaining == 0:
|
| 93 |
+
break
|
| 94 |
+
|
| 95 |
+
return results
|
core/pipeline_1/logic.py
ADDED
|
@@ -0,0 +1,116 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import os
|
| 2 |
+
import sys
|
| 3 |
+
import logging
|
| 4 |
+
from pathlib import Path
|
| 5 |
+
from dotenv import load_dotenv
|
| 6 |
+
from google import genai
|
| 7 |
+
|
| 8 |
+
# Add current directory to path to allow importing from 'filter.py' in a hyphenated folder
|
| 9 |
+
sys.path.append(str(Path(__file__).parent))
|
| 10 |
+
from filter import SimpleFilter
|
| 11 |
+
from services.paper_fetcher import PaperFetcher
|
| 12 |
+
from services.metrics_service import MetricsService
|
| 13 |
+
import time
|
| 14 |
+
|
| 15 |
+
# Setup logging
|
| 16 |
+
logging.basicConfig(level=logging.INFO)
|
| 17 |
+
logger = logging.getLogger(__name__)
|
| 18 |
+
|
| 19 |
+
load_dotenv()
|
| 20 |
+
|
| 21 |
+
class PipelineLLMOnly:
|
| 22 |
+
"""
|
| 23 |
+
Pipeline 1: LLM-Only (Baseline)
|
| 24 |
+
Logic: Metadata Filter -> Full Text Fetch -> LLM Context
|
| 25 |
+
"""
|
| 26 |
+
def __init__(self, top_n=40, max_full_text=40):
|
| 27 |
+
self.top_n = top_n
|
| 28 |
+
self.max_full_text = max_full_text
|
| 29 |
+
self.filter = SimpleFilter()
|
| 30 |
+
self.fetcher = PaperFetcher()
|
| 31 |
+
|
| 32 |
+
self.api_key = os.environ.get("GEMINI_API_KEY")
|
| 33 |
+
self.model_name = os.environ.get("GEMINI_MODEL", "gemini-1.5-flash")
|
| 34 |
+
self.metrics = MetricsService(self.model_name)
|
| 35 |
+
|
| 36 |
+
if self.api_key:
|
| 37 |
+
self.client = genai.Client(api_key=self.api_key)
|
| 38 |
+
else:
|
| 39 |
+
self.client = None
|
| 40 |
+
logger.warning("GEMINI_API_KEY not found. LLM calls will fail.")
|
| 41 |
+
|
| 42 |
+
def run(self, query, ground_truth=None):
|
| 43 |
+
start_time = time.time()
|
| 44 |
+
logger.info(f"--- Pipeline Run Started ---")
|
| 45 |
+
logger.info(f"Query: {query}")
|
| 46 |
+
logger.info(f"Settings: top_n={self.top_n}, max_full_text={self.max_full_text}")
|
| 47 |
+
|
| 48 |
+
# 1. Filter
|
| 49 |
+
relevant_metadata = self.filter.filter_papers(query, top_n=self.top_n)
|
| 50 |
+
logger.info(f"Filter found {len(relevant_metadata)} papers.")
|
| 51 |
+
|
| 52 |
+
# 2. Fetch Context
|
| 53 |
+
context_parts = []
|
| 54 |
+
abstracts = []
|
| 55 |
+
for i, paper in enumerate(relevant_metadata):
|
| 56 |
+
content = ""
|
| 57 |
+
if i < self.max_full_text:
|
| 58 |
+
logger.info(f"Fetching full text for Paper {i+1}: {paper.get('title', 'Unknown')}")
|
| 59 |
+
full_text = self.fetcher.fetch_full_text(paper['id'])
|
| 60 |
+
if full_text:
|
| 61 |
+
content = full_text
|
| 62 |
+
|
| 63 |
+
meta = self.fetcher.get_work_metadata(paper['id'])
|
| 64 |
+
if meta:
|
| 65 |
+
abstract = self.fetcher.get_abstract(meta)
|
| 66 |
+
abstracts.append(abstract)
|
| 67 |
+
if not content:
|
| 68 |
+
content = f"Title: {paper['title']}\nAbstract: {abstract}"
|
| 69 |
+
elif not content:
|
| 70 |
+
content = f"Title: {paper['title']}"
|
| 71 |
+
|
| 72 |
+
context_parts.append(f"--- PAPER {i+1} ---\n{content}\n")
|
| 73 |
+
|
| 74 |
+
full_context = "\n".join(context_parts)
|
| 75 |
+
reference_text = "\n".join(abstracts) # Use abstracts as BERTScore reference
|
| 76 |
+
|
| 77 |
+
# 3. Prompt
|
| 78 |
+
system_prompt = (
|
| 79 |
+
"You are an AI research assistant. Answer the user query based ONLY on the provided papers. "
|
| 80 |
+
"If the answer isn't there, say so. Cite sources like [Paper 1].\n\n"
|
| 81 |
+
)
|
| 82 |
+
final_prompt = f"{system_prompt}\nQUERY: {query}\n\nCONTEXT:\n{full_context}"
|
| 83 |
+
|
| 84 |
+
# 4. Generate
|
| 85 |
+
if not self.client:
|
| 86 |
+
return {"answer": "LLM Error: GEMINI_API_KEY missing.", "sources": []}
|
| 87 |
+
|
| 88 |
+
try:
|
| 89 |
+
response = self.client.models.generate_content(
|
| 90 |
+
model=self.model_name,
|
| 91 |
+
contents=final_prompt
|
| 92 |
+
)
|
| 93 |
+
|
| 94 |
+
answer = response.text
|
| 95 |
+
usage = response.usage_metadata
|
| 96 |
+
|
| 97 |
+
# 5. Process Metrics
|
| 98 |
+
stats = self.metrics.process_metrics(
|
| 99 |
+
client=self.client,
|
| 100 |
+
query=query,
|
| 101 |
+
answer=answer,
|
| 102 |
+
context=full_context,
|
| 103 |
+
usage_metadata=usage,
|
| 104 |
+
start_time=start_time,
|
| 105 |
+
abstracts_list=abstracts,
|
| 106 |
+
ground_truth=ground_truth
|
| 107 |
+
)
|
| 108 |
+
|
| 109 |
+
return {
|
| 110 |
+
"answer": answer,
|
| 111 |
+
"sources": [p['title'] for p in relevant_metadata],
|
| 112 |
+
"metrics": stats
|
| 113 |
+
}
|
| 114 |
+
except Exception as e:
|
| 115 |
+
logger.error(f"Pipeline error: {e}")
|
| 116 |
+
return {"error": str(e), "answer": "Error generating response."}
|
core/pipeline_1/main.py
ADDED
|
@@ -0,0 +1,105 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import os
|
| 2 |
+
import logging
|
| 3 |
+
from pathlib import Path
|
| 4 |
+
from dotenv import load_dotenv
|
| 5 |
+
import google.generativeai as genai
|
| 6 |
+
|
| 7 |
+
from core.pipeline_1.filter import SimpleFilter
|
| 8 |
+
from services.paper_fetcher import PaperFetcher
|
| 9 |
+
|
| 10 |
+
# Setup logging
|
| 11 |
+
logging.basicConfig(level=logging.INFO)
|
| 12 |
+
logger = logging.getLogger(__name__)
|
| 13 |
+
|
| 14 |
+
# Load environment variables
|
| 15 |
+
load_dotenv()
|
| 16 |
+
|
| 17 |
+
class Pipeline1:
|
| 18 |
+
"""
|
| 19 |
+
LLM-Only Pipeline (Worst-Case Baseline)
|
| 20 |
+
Strategy: Metadata Filter (Approach 2+3) -> Full Text Extraction -> LLM Context
|
| 21 |
+
"""
|
| 22 |
+
def __init__(self, top_n=20, max_full_text=5):
|
| 23 |
+
self.top_n = top_n
|
| 24 |
+
self.max_full_text = max_full_text
|
| 25 |
+
self.filter = SimpleFilter()
|
| 26 |
+
self.fetcher = PaperFetcher()
|
| 27 |
+
|
| 28 |
+
# Configure Gemini
|
| 29 |
+
api_key = os.environ.get("GEMINI_API_KEY")
|
| 30 |
+
if api_key:
|
| 31 |
+
genai.configure(api_key=api_key)
|
| 32 |
+
self.model = genai.GenerativeModel('gemini-1.5-flash')
|
| 33 |
+
else:
|
| 34 |
+
logger.warning("GEMINI_API_KEY not found. LLM calls will fail.")
|
| 35 |
+
self.model = None
|
| 36 |
+
|
| 37 |
+
def run(self, query):
|
| 38 |
+
logger.info(f"Running Pipeline-1 for query: {query}")
|
| 39 |
+
|
| 40 |
+
# 1. Filter papers based on metadata (Approaches 2+3)
|
| 41 |
+
relevant_metadata = self.filter.filter_papers(query, top_n=self.top_n)
|
| 42 |
+
|
| 43 |
+
# 2. Fetch actual content for the top candidates
|
| 44 |
+
# We limit full text fetching because it takes time (Baseline is slow but foolproof)
|
| 45 |
+
context_parts = []
|
| 46 |
+
|
| 47 |
+
for i, paper in enumerate(relevant_metadata):
|
| 48 |
+
logger.info(f"Processing candidate {i+1}/{self.top_n}: {paper['title']}")
|
| 49 |
+
|
| 50 |
+
content = ""
|
| 51 |
+
# Fetch full text for the very top candidates
|
| 52 |
+
if i < self.max_full_text:
|
| 53 |
+
full_text = self.fetcher.fetch_full_text(paper['id'])
|
| 54 |
+
if full_text:
|
| 55 |
+
content = full_text
|
| 56 |
+
|
| 57 |
+
# If full text fetch failed or we are past max_full_text, use title/abstract
|
| 58 |
+
if not content:
|
| 59 |
+
# We try to get metadata for abstract
|
| 60 |
+
meta = self.fetcher.get_work_metadata(paper['id'])
|
| 61 |
+
if meta:
|
| 62 |
+
abstract = self.fetcher.get_abstract(meta)
|
| 63 |
+
content = f"Title: {paper['title']}\nAbstract: {abstract}"
|
| 64 |
+
else:
|
| 65 |
+
content = f"Title: {paper['title']} (No abstract available)"
|
| 66 |
+
|
| 67 |
+
context_parts.append(f"--- PAPER {i+1} ---\n{content}\n")
|
| 68 |
+
|
| 69 |
+
full_context = "\n".join(context_parts)
|
| 70 |
+
|
| 71 |
+
# 3. Construct the Prompt
|
| 72 |
+
system_prompt = (
|
| 73 |
+
"You are an AI research assistant. Below is a collection of research paper excerpts "
|
| 74 |
+
"relevant to a user's query. Answer the query based ONLY on the provided information. "
|
| 75 |
+
"If the information is not sufficient, say so. Provide citations like [Paper 1], [Paper 2].\n\n"
|
| 76 |
+
)
|
| 77 |
+
|
| 78 |
+
final_prompt = f"{system_prompt}\nUSER QUERY: {query}\n\nRESEARCH CONTEXT:\n{full_context}"
|
| 79 |
+
|
| 80 |
+
# 4. Call LLM
|
| 81 |
+
if not self.model:
|
| 82 |
+
return {"error": "LLM not configured", "context_preview": full_context[:500]}
|
| 83 |
+
|
| 84 |
+
try:
|
| 85 |
+
logger.info("Calling Gemini...")
|
| 86 |
+
response = self.model.generate_content(final_prompt)
|
| 87 |
+
return {
|
| 88 |
+
"answer": response.text,
|
| 89 |
+
"sources_used": [p['title'] for p in relevant_metadata],
|
| 90 |
+
"context_length_chars": len(full_context)
|
| 91 |
+
}
|
| 92 |
+
except Exception as e:
|
| 93 |
+
logger.error(f"LLM call failed: {e}")
|
| 94 |
+
return {"error": str(e), "context_preview": full_context[:500]}
|
| 95 |
+
|
| 96 |
+
if __name__ == "__main__":
|
| 97 |
+
pipeline = Pipeline1(top_n=5, max_full_text=2)
|
| 98 |
+
query = "How is reinforcement learning used in large language models?"
|
| 99 |
+
result = pipeline.run(query)
|
| 100 |
+
|
| 101 |
+
if "answer" in result:
|
| 102 |
+
print("\n--- Pipeline 1 Response ---")
|
| 103 |
+
print(result["answer"])
|
| 104 |
+
else:
|
| 105 |
+
print(f"\nPipeline failed: {result.get('error')}")
|
core/pipeline_1/metadata_index.json
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
core/pipeline_2/.gitkeep
ADDED
|
File without changes
|
core/pipeline_2/indexer.py
ADDED
|
@@ -0,0 +1,151 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""One-time indexing: embed 15k papers and store in local Qdrant collection."""
|
| 2 |
+
import glob
|
| 3 |
+
import json
|
| 4 |
+
import logging
|
| 5 |
+
from pathlib import Path
|
| 6 |
+
|
| 7 |
+
from dotenv import load_dotenv
|
| 8 |
+
from fastembed import SparseTextEmbedding
|
| 9 |
+
from qdrant_client import QdrantClient, models
|
| 10 |
+
from sentence_transformers import SentenceTransformer
|
| 11 |
+
|
| 12 |
+
load_dotenv()
|
| 13 |
+
|
| 14 |
+
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
| 15 |
+
logger = logging.getLogger(__name__)
|
| 16 |
+
|
| 17 |
+
_PROJECT_ROOT = Path(__file__).parents[2]
|
| 18 |
+
_DEFAULT_DATA_DIR = str(_PROJECT_ROOT / "data")
|
| 19 |
+
_DEFAULT_STORAGE_PATH = str(Path(__file__).parent / "qdrant_storage")
|
| 20 |
+
|
| 21 |
+
COLLECTION = "papers"
|
| 22 |
+
DENSE_MODEL = "BAAI/bge-large-en-v1.5"
|
| 23 |
+
SPARSE_MODEL = "Qdrant/bm25"
|
| 24 |
+
DENSE_DIM = 1024
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
def _reconstruct_abstract(inverted: dict) -> str:
|
| 28 |
+
if not inverted:
|
| 29 |
+
return ""
|
| 30 |
+
max_idx = max((max(v) for v in inverted.values() if v), default=0)
|
| 31 |
+
words = [""] * (max_idx + 1)
|
| 32 |
+
for word, idxs in inverted.items():
|
| 33 |
+
for i in idxs:
|
| 34 |
+
words[i] = word
|
| 35 |
+
return " ".join(words).strip()
|
| 36 |
+
|
| 37 |
+
|
| 38 |
+
def _doc_text(paper: dict) -> str:
|
| 39 |
+
title = paper.get("title") or ""
|
| 40 |
+
abstract = _reconstruct_abstract(paper.get("abstract_inverted_index") or {})
|
| 41 |
+
return f"Title: {title}\n\nAbstract: {abstract}"
|
| 42 |
+
|
| 43 |
+
|
| 44 |
+
def build_index(
|
| 45 |
+
data_dir: str = _DEFAULT_DATA_DIR,
|
| 46 |
+
storage_path: str = _DEFAULT_STORAGE_PATH,
|
| 47 |
+
batch_size: int = 128,
|
| 48 |
+
) -> None:
|
| 49 |
+
client = QdrantClient(path=storage_path)
|
| 50 |
+
|
| 51 |
+
existing = {c.name for c in client.get_collections().collections}
|
| 52 |
+
if COLLECTION in existing:
|
| 53 |
+
n = client.count(COLLECTION).count
|
| 54 |
+
logger.info(f"Collection '{COLLECTION}' exists with {n} vectors β skipping model load and indexing.")
|
| 55 |
+
return
|
| 56 |
+
|
| 57 |
+
logger.info(f"Loading dense model: {DENSE_MODEL}")
|
| 58 |
+
dense_model = SentenceTransformer(DENSE_MODEL)
|
| 59 |
+
|
| 60 |
+
logger.info(f"Loading sparse model: {SPARSE_MODEL}...")
|
| 61 |
+
sparse_model = SparseTextEmbedding(SPARSE_MODEL, providers=["CPUExecutionProvider"])
|
| 62 |
+
logger.info("Sparse model loaded.")
|
| 63 |
+
|
| 64 |
+
logger.info(f"Creating collection '{COLLECTION}' in Qdrant...")
|
| 65 |
+
client.create_collection(
|
| 66 |
+
collection_name=COLLECTION,
|
| 67 |
+
vectors_config={
|
| 68 |
+
"dense": models.VectorParams(size=DENSE_DIM, distance=models.Distance.COSINE)
|
| 69 |
+
},
|
| 70 |
+
sparse_vectors_config={
|
| 71 |
+
"sparse": models.SparseVectorParams(modifier=models.Modifier.IDF)
|
| 72 |
+
},
|
| 73 |
+
)
|
| 74 |
+
logger.info(f"Collection '{COLLECTION}' created.")
|
| 75 |
+
|
| 76 |
+
files = sorted(glob.glob(str(Path(data_dir) / "*.json")))
|
| 77 |
+
logger.info(f"Found {len(files)} papers in {data_dir}. Starting indexing...")
|
| 78 |
+
|
| 79 |
+
for start in range(0, len(files), batch_size):
|
| 80 |
+
end = min(start + batch_size, len(files))
|
| 81 |
+
logger.info(f"Processing batch: {start} to {end}...")
|
| 82 |
+
batch = files[start : end]
|
| 83 |
+
papers, texts = [], []
|
| 84 |
+
|
| 85 |
+
for fpath in batch:
|
| 86 |
+
try:
|
| 87 |
+
with open(fpath) as f:
|
| 88 |
+
data = json.load(f)
|
| 89 |
+
papers.append(data)
|
| 90 |
+
texts.append(_doc_text(data))
|
| 91 |
+
except Exception as e:
|
| 92 |
+
logger.error(f"Error reading {fpath}: {e}")
|
| 93 |
+
continue
|
| 94 |
+
|
| 95 |
+
if not texts:
|
| 96 |
+
continue
|
| 97 |
+
|
| 98 |
+
logger.info(f" Encoding {len(texts)} dense vectors...")
|
| 99 |
+
dense_vecs = dense_model.encode(
|
| 100 |
+
texts, batch_size=64, normalize_embeddings=True, show_progress_bar=False
|
| 101 |
+
)
|
| 102 |
+
|
| 103 |
+
logger.info(f" Encoding {len(texts)} sparse vectors...")
|
| 104 |
+
sparse_vecs = list(sparse_model.embed(texts))
|
| 105 |
+
|
| 106 |
+
logger.info(f" Upserting {len(texts)} points to Qdrant...")
|
| 107 |
+
points = []
|
| 108 |
+
for i, (paper, dv, sv) in enumerate(zip(papers, dense_vecs, sparse_vecs)):
|
| 109 |
+
abstract = _reconstruct_abstract(paper.get("abstract_inverted_index") or {})
|
| 110 |
+
topics = [t.get("display_name", "") for t in (paper.get("topics") or [])]
|
| 111 |
+
keywords = [k.get("display_name", "") for k in (paper.get("keywords") or [])]
|
| 112 |
+
|
| 113 |
+
# Extract full-text access URLs from OpenAlex metadata
|
| 114 |
+
best_oa = paper.get("best_oa_location") or {}
|
| 115 |
+
pdf_url = best_oa.get("pdf_url")
|
| 116 |
+
landing_page_url = best_oa.get("landing_page_url")
|
| 117 |
+
|
| 118 |
+
points.append(
|
| 119 |
+
models.PointStruct(
|
| 120 |
+
id=start + i,
|
| 121 |
+
vector={
|
| 122 |
+
"dense": dv.tolist(),
|
| 123 |
+
"sparse": models.SparseVector(
|
| 124 |
+
indices=sv.indices.tolist(),
|
| 125 |
+
values=sv.values.tolist(),
|
| 126 |
+
),
|
| 127 |
+
},
|
| 128 |
+
payload={
|
| 129 |
+
"paper_id": paper.get("id"),
|
| 130 |
+
"title": paper.get("title", ""),
|
| 131 |
+
"abstract": abstract, # Full abstract β NO truncation
|
| 132 |
+
"year": paper.get("publication_year"),
|
| 133 |
+
"cited_by_count": paper.get("cited_by_count", 0),
|
| 134 |
+
"topics": topics[:10], # More topic coverage
|
| 135 |
+
"keywords": keywords[:10], # New: keyword metadata
|
| 136 |
+
"doi": paper.get("doi"),
|
| 137 |
+
"pdf_url": pdf_url, # New: full-text access URL
|
| 138 |
+
"landing_page_url": landing_page_url, # New: landing page URL
|
| 139 |
+
},
|
| 140 |
+
)
|
| 141 |
+
)
|
| 142 |
+
|
| 143 |
+
client.upsert(collection_name=COLLECTION, points=points)
|
| 144 |
+
logger.info(f" DONE: Indexed {start + len(points)}/{len(files)}")
|
| 145 |
+
|
| 146 |
+
n = client.count(COLLECTION).count
|
| 147 |
+
logger.info(f"Indexing complete β {n} vectors in '{COLLECTION}'.")
|
| 148 |
+
|
| 149 |
+
|
| 150 |
+
if __name__ == "__main__":
|
| 151 |
+
build_index()
|
core/pipeline_2/logic.py
ADDED
|
@@ -0,0 +1,499 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Pipeline 2: Hybrid RAG β vector search + reranker + Groq LLM.
|
| 2 |
+
|
| 3 |
+
Architectural improvements over baseline:
|
| 4 |
+
- Query decomposition for multi-hop questions (LLM-driven)
|
| 5 |
+
- Multi-paper full-text fetch (up to 3 papers)
|
| 6 |
+
- Fallback: if top papers lack full text, find the highest-scoring paper that does
|
| 7 |
+
- Structure-aware sentence-level passage extraction with contextual prepend
|
| 8 |
+
- Improved synthesis prompt for completeness and specificity
|
| 9 |
+
"""
|
| 10 |
+
import json
|
| 11 |
+
import logging
|
| 12 |
+
import os
|
| 13 |
+
import re
|
| 14 |
+
import sys
|
| 15 |
+
import time
|
| 16 |
+
from dataclasses import dataclass
|
| 17 |
+
from pathlib import Path
|
| 18 |
+
from concurrent.futures import ThreadPoolExecutor, as_completed
|
| 19 |
+
|
| 20 |
+
from dotenv import load_dotenv
|
| 21 |
+
from groq import Groq, RateLimitError as GroqRateLimitError
|
| 22 |
+
|
| 23 |
+
# Project root on path so services/ imports resolve
|
| 24 |
+
sys.path.append(str(Path(__file__).parents[2]))
|
| 25 |
+
# pipeline-2 dir on path so indexer/retriever imports resolve
|
| 26 |
+
sys.path.append(str(Path(__file__).parent))
|
| 27 |
+
|
| 28 |
+
from indexer import build_index
|
| 29 |
+
from retriever import HybridRetriever
|
| 30 |
+
from services.metrics_service import MetricsService
|
| 31 |
+
from services.paper_fetcher import PaperFetcher
|
| 32 |
+
|
| 33 |
+
load_dotenv()
|
| 34 |
+
|
| 35 |
+
logging.basicConfig(level=logging.INFO)
|
| 36 |
+
logger = logging.getLogger(__name__)
|
| 37 |
+
|
| 38 |
+
_STORAGE_PATH = str(Path(__file__).parent / "qdrant_storage")
|
| 39 |
+
_DEFAULT_MODEL = "llama-3.3-70b-versatile"
|
| 40 |
+
_DECOMPOSE_MODEL = "llama-3.1-8b-instant" # cheap classifier β no need for 70B
|
| 41 |
+
|
| 42 |
+
|
| 43 |
+
@dataclass
|
| 44 |
+
class _GroqUsage:
|
| 45 |
+
"""Maps Groq usage fields to the shape MetricsService expects."""
|
| 46 |
+
prompt_token_count: int
|
| 47 |
+
candidates_token_count: int
|
| 48 |
+
|
| 49 |
+
|
| 50 |
+
class PipelineRAG:
|
| 51 |
+
"""
|
| 52 |
+
Pipeline 2: Hybrid RAG.
|
| 53 |
+
Flow: query decomposition β hybrid search (RRF) β cross-encoder rerank
|
| 54 |
+
β multi-paper full-text fetch β structure-aware chunking β Groq synthesis.
|
| 55 |
+
|
| 56 |
+
Run the indexer once before starting the server:
|
| 57 |
+
.venv/bin/python core/pipeline-2/indexer.py
|
| 58 |
+
"""
|
| 59 |
+
|
| 60 |
+
def __init__(
|
| 61 |
+
self,
|
| 62 |
+
retrieval_top_k: int = 50,
|
| 63 |
+
rerank_top_n: int = 10,
|
| 64 |
+
max_full_text: int = 3,
|
| 65 |
+
):
|
| 66 |
+
self.retrieval_top_k = retrieval_top_k
|
| 67 |
+
self.rerank_top_n = rerank_top_n
|
| 68 |
+
self.max_full_text = max_full_text
|
| 69 |
+
|
| 70 |
+
self.model_name = os.environ.get("GROQ_MODEL", _DEFAULT_MODEL)
|
| 71 |
+
_keys = [os.environ.get(f"GROQ_API_KEY_{i}") for i in range(1, 4)]
|
| 72 |
+
_keys = [k for k in _keys if k]
|
| 73 |
+
if not _keys:
|
| 74 |
+
_keys = [os.environ["GROQ_API_KEY"]]
|
| 75 |
+
self._groq_clients = [Groq(api_key=k) for k in _keys]
|
| 76 |
+
self._groq_key_idx = 0
|
| 77 |
+
logger.info(f"Loaded {len(self._groq_clients)} Groq API key(s).")
|
| 78 |
+
self.fetcher = PaperFetcher()
|
| 79 |
+
self.metrics = MetricsService(self.model_name)
|
| 80 |
+
|
| 81 |
+
logger.info("Ensuring Qdrant index exists...")
|
| 82 |
+
build_index(storage_path=_STORAGE_PATH)
|
| 83 |
+
|
| 84 |
+
logger.info("Initializing retriever (loading models)...")
|
| 85 |
+
self.retriever = HybridRetriever(storage_path=_STORAGE_PATH)
|
| 86 |
+
|
| 87 |
+
# ------------------------------------------------------------------
|
| 88 |
+
# Groq API call with key rotation on rate limit
|
| 89 |
+
# ------------------------------------------------------------------
|
| 90 |
+
def _groq_call(self, **kwargs):
|
| 91 |
+
for _ in range(len(self._groq_clients)):
|
| 92 |
+
try:
|
| 93 |
+
return self._groq_clients[self._groq_key_idx].chat.completions.create(**kwargs)
|
| 94 |
+
except GroqRateLimitError:
|
| 95 |
+
next_idx = (self._groq_key_idx + 1) % len(self._groq_clients)
|
| 96 |
+
if next_idx == self._groq_key_idx:
|
| 97 |
+
raise
|
| 98 |
+
logger.warning(
|
| 99 |
+
f"Rate limit on Groq key {self._groq_key_idx + 1}, "
|
| 100 |
+
f"rotating to key {next_idx + 1}..."
|
| 101 |
+
)
|
| 102 |
+
self._groq_key_idx = next_idx
|
| 103 |
+
raise GroqRateLimitError("All Groq API keys exhausted.")
|
| 104 |
+
|
| 105 |
+
# ------------------------------------------------------------------
|
| 106 |
+
# Query Decomposition
|
| 107 |
+
# ------------------------------------------------------------------
|
| 108 |
+
def _decompose_query(self, query: str) -> dict:
|
| 109 |
+
"""Use a lightweight LLM to detect true multi-paper questions and decompose them.
|
| 110 |
+
|
| 111 |
+
MULTI-HOP = the question explicitly references 2+ DISTINCT papers/methods
|
| 112 |
+
that must each be independently retrieved (e.g. "Both CGCNN and DPMD...").
|
| 113 |
+
Multi-part questions about ONE paper are NOT multi-hop.
|
| 114 |
+
|
| 115 |
+
Returns {"is_multi_hop": bool, "sub_queries": [str, ...]}.
|
| 116 |
+
For single-hop questions, sub_queries contains only the original query.
|
| 117 |
+
"""
|
| 118 |
+
try:
|
| 119 |
+
response = self._groq_call(
|
| 120 |
+
model=_DECOMPOSE_MODEL,
|
| 121 |
+
messages=[
|
| 122 |
+
{
|
| 123 |
+
"role": "system",
|
| 124 |
+
"content": (
|
| 125 |
+
"You classify scientific search queries as single-hop or multi-hop.\n\n"
|
| 126 |
+
"MULTI-HOP = the question EXPLICITLY references 2+ DISTINCT papers or methods "
|
| 127 |
+
"that must each be independently retrieved. "
|
| 128 |
+
"Look for: 'Both X and Y', 'X paper and Y paper', 'X review and Y review', "
|
| 129 |
+
"or a comparison between two named systems from different papers.\n\n"
|
| 130 |
+
"SINGLE-HOP = anything else, INCLUDING multi-part questions about ONE topic.\n\n"
|
| 131 |
+
"SINGLE-HOP examples (is_multi_hop=false):\n"
|
| 132 |
+
" 'What three components does X have and what does each do?' β one paper\n"
|
| 133 |
+
" 'What does ADMET stand for and why does ADMETlab 2.0 address it?' β one paper\n"
|
| 134 |
+
" 'What makes DPMD first-principles based and what symmetries does it preserve?' β one paper\n"
|
| 135 |
+
" 'What two factors drove NLP progress and what is the library goal?' β one paper\n\n"
|
| 136 |
+
"MULTI-HOP examples (is_multi_hop=true, max 2 sub-queries):\n"
|
| 137 |
+
" 'Both CGCNN and DPMD apply neural networks... what does each predict?' β two papers\n"
|
| 138 |
+
" 'The HuggingFace paper and the molecular representations paper both...' β two papers\n"
|
| 139 |
+
" 'The ML for fluid mechanics review and the ML in materials science review...' β two papers\n\n"
|
| 140 |
+
"Return ONLY valid JSON, no markdown fences:\n"
|
| 141 |
+
'{"is_multi_hop": false, "sub_queries": ["original query"]}\n'
|
| 142 |
+
"or\n"
|
| 143 |
+
'{"is_multi_hop": true, "sub_queries": ["focused query about paper A", "focused query about paper B"]}'
|
| 144 |
+
),
|
| 145 |
+
},
|
| 146 |
+
{"role": "user", "content": query},
|
| 147 |
+
],
|
| 148 |
+
temperature=0.0,
|
| 149 |
+
max_tokens=256,
|
| 150 |
+
)
|
| 151 |
+
raw = response.choices[0].message.content.strip()
|
| 152 |
+
# Strip markdown code fences if the LLM wraps the JSON
|
| 153 |
+
raw = re.sub(r"^```(?:json)?\s*", "", raw)
|
| 154 |
+
raw = re.sub(r"\s*```$", "", raw)
|
| 155 |
+
result = json.loads(raw)
|
| 156 |
+
if not isinstance(result.get("sub_queries"), list) or not result["sub_queries"]:
|
| 157 |
+
return {"is_multi_hop": False, "sub_queries": [query]}
|
| 158 |
+
|
| 159 |
+
# Post-LLM sanity checks: override false multi-hop classifications
|
| 160 |
+
if result.get("is_multi_hop") and len(result["sub_queries"]) >= 2:
|
| 161 |
+
subs = result["sub_queries"]
|
| 162 |
+
override = False
|
| 163 |
+
|
| 164 |
+
# Check 1: both sub-queries share a capitalized word-pair β same paper
|
| 165 |
+
def _cap_pairs(text):
|
| 166 |
+
tokens = text.split()
|
| 167 |
+
return {
|
| 168 |
+
(tokens[i].lower(), tokens[i + 1].lower())
|
| 169 |
+
for i in range(len(tokens) - 1)
|
| 170 |
+
if tokens[i][0].isupper() and tokens[i + 1][0].isupper()
|
| 171 |
+
and len(tokens[i]) > 1 and len(tokens[i + 1]) > 1
|
| 172 |
+
}
|
| 173 |
+
if _cap_pairs(subs[0]) & _cap_pairs(subs[1]):
|
| 174 |
+
override = True
|
| 175 |
+
|
| 176 |
+
# Check 2: second sub-query uses "it" as a pronoun and neither sub-query
|
| 177 |
+
# mentions a specific named system (ALL-CAPS or CamelCase) β same topic
|
| 178 |
+
if not override and re.search(r'\bit\b', subs[1], re.IGNORECASE):
|
| 179 |
+
has_named_system = re.search(
|
| 180 |
+
r'\b([A-Z]{2,}|[A-Z][a-z]+[A-Z][a-zA-Z]*)\b',
|
| 181 |
+
subs[0] + " " + subs[1]
|
| 182 |
+
)
|
| 183 |
+
if not has_named_system:
|
| 184 |
+
override = True
|
| 185 |
+
|
| 186 |
+
if override:
|
| 187 |
+
logger.info("Multi-hop override: sub-queries target the same paper β single-hop.")
|
| 188 |
+
result = {"is_multi_hop": False, "sub_queries": [query]}
|
| 189 |
+
|
| 190 |
+
logger.info(f"Query decomposition: multi_hop={result['is_multi_hop']}, "
|
| 191 |
+
f"sub_queries={result['sub_queries']}")
|
| 192 |
+
return result
|
| 193 |
+
except Exception as e:
|
| 194 |
+
logger.warning(f"Query decomposition failed ({e}), using original query.")
|
| 195 |
+
return {"is_multi_hop": False, "sub_queries": [query]}
|
| 196 |
+
|
| 197 |
+
# ------------------------------------------------------------------
|
| 198 |
+
# Structure-aware passage extraction
|
| 199 |
+
# ------------------------------------------------------------------
|
| 200 |
+
def _extract_relevant_passages(
|
| 201 |
+
self, full_text: str, query: str, paper_title: str = "",
|
| 202 |
+
top_k: int = 8, target_tokens: int = 600,
|
| 203 |
+
) -> str:
|
| 204 |
+
"""Sentence-aware passage extraction with contextual prepend.
|
| 205 |
+
|
| 206 |
+
1. Split full text on sentence boundaries (preserving scientific structure).
|
| 207 |
+
2. Group sentences into coherent passages (~3-5 sentences each).
|
| 208 |
+
3. Embed passages and rank by cosine similarity to the query.
|
| 209 |
+
4. Prepend paper title to each selected passage for context.
|
| 210 |
+
"""
|
| 211 |
+
# Split on sentence boundaries β handles abbreviations reasonably
|
| 212 |
+
sentences = re.split(r'(?<=[.!?])\s+(?=[A-Z])', full_text)
|
| 213 |
+
if not sentences:
|
| 214 |
+
return full_text
|
| 215 |
+
|
| 216 |
+
# Group sentences into passages of roughly target_tokens characters
|
| 217 |
+
passages: list[str] = []
|
| 218 |
+
current: list[str] = []
|
| 219 |
+
current_len = 0
|
| 220 |
+
for sent in sentences:
|
| 221 |
+
current.append(sent)
|
| 222 |
+
current_len += len(sent)
|
| 223 |
+
if current_len >= target_tokens:
|
| 224 |
+
passages.append(" ".join(current))
|
| 225 |
+
# Carry the last sentence over for overlap context
|
| 226 |
+
current = [current[-1]] if len(current) > 1 else []
|
| 227 |
+
current_len = len(current[0]) if current else 0
|
| 228 |
+
if current:
|
| 229 |
+
passages.append(" ".join(current))
|
| 230 |
+
|
| 231 |
+
if len(passages) <= top_k:
|
| 232 |
+
return full_text
|
| 233 |
+
|
| 234 |
+
# Stage 1: bi-encoder to narrow to top-30 candidates (fast)
|
| 235 |
+
texts = [query] + passages
|
| 236 |
+
embeddings = self.retriever.dense.encode(
|
| 237 |
+
texts, normalize_embeddings=True, show_progress_bar=False, batch_size=64
|
| 238 |
+
)
|
| 239 |
+
dense_scores = embeddings[1:] @ embeddings[0]
|
| 240 |
+
stage1_k = min(30, len(passages))
|
| 241 |
+
stage1_indices = dense_scores.argsort()[-stage1_k:][::-1].tolist()
|
| 242 |
+
|
| 243 |
+
# Always include passage[0] (intro/abstract) β key summary statements live there
|
| 244 |
+
mandatory = {0}
|
| 245 |
+
stage1_pool = sorted(set(stage1_indices) | mandatory)
|
| 246 |
+
|
| 247 |
+
# Stage 2: cross-encoder rerank for accurate passage selection
|
| 248 |
+
pairs = [[query, passages[i]] for i in stage1_pool]
|
| 249 |
+
rerank_scores = self.retriever.reranker.predict(pairs, batch_size=32)
|
| 250 |
+
scored = sorted(zip(rerank_scores, stage1_pool), key=lambda x: x[0], reverse=True)
|
| 251 |
+
|
| 252 |
+
# Take top_k-1 by reranker score + always keep passage[0]
|
| 253 |
+
top_by_reranker = [idx for _, idx in scored if idx not in mandatory][: top_k - 1]
|
| 254 |
+
top_indices = sorted(mandatory | set(top_by_reranker))
|
| 255 |
+
|
| 256 |
+
# Contextual prepend: attach paper title to each passage
|
| 257 |
+
prefix = f"[Source: {paper_title}]" if paper_title else ""
|
| 258 |
+
selected = []
|
| 259 |
+
for idx in top_indices:
|
| 260 |
+
passage_text = passages[idx].strip()
|
| 261 |
+
if prefix:
|
| 262 |
+
selected.append(f"{prefix}\n{passage_text}")
|
| 263 |
+
else:
|
| 264 |
+
selected.append(passage_text)
|
| 265 |
+
|
| 266 |
+
return "\n\n[...]\n\n".join(selected)
|
| 267 |
+
|
| 268 |
+
# ------------------------------------------------------------------
|
| 269 |
+
# Retrieval with dedup support for multi-hop
|
| 270 |
+
# ------------------------------------------------------------------
|
| 271 |
+
def _deduplicate_papers(self, papers: list[dict], top_n: int) -> list[dict]:
|
| 272 |
+
"""Deduplicate papers by paper_id, keeping the first (highest-ranked) occurrence."""
|
| 273 |
+
seen = set()
|
| 274 |
+
deduped = []
|
| 275 |
+
for p in papers:
|
| 276 |
+
pid = p.get("paper_id")
|
| 277 |
+
if pid and pid not in seen:
|
| 278 |
+
seen.add(pid)
|
| 279 |
+
deduped.append(p)
|
| 280 |
+
if len(deduped) >= top_n:
|
| 281 |
+
break
|
| 282 |
+
return deduped
|
| 283 |
+
|
| 284 |
+
# ------------------------------------------------------------------
|
| 285 |
+
# Multi-hop coverage enforcement
|
| 286 |
+
# ------------------------------------------------------------------
|
| 287 |
+
def _ensure_sub_query_coverage(
|
| 288 |
+
self, top_papers: list[dict], sub_queries: list[str], top_n: int
|
| 289 |
+
) -> list[dict]:
|
| 290 |
+
"""Guarantee that each sub-query is represented by at least one paper.
|
| 291 |
+
|
| 292 |
+
For each sub-query, score every paper in top_papers with the cross-encoder.
|
| 293 |
+
If no paper scores above a minimal threshold, run a targeted search for that
|
| 294 |
+
sub-query and inject the best hit at the end of the list.
|
| 295 |
+
"""
|
| 296 |
+
# bge-reranker-base returns sigmoid-activated [0,1] scores; 0.1 means "weak relevance"
|
| 297 |
+
COVERAGE_THRESHOLD = 0.1
|
| 298 |
+
|
| 299 |
+
for sub_q in sub_queries:
|
| 300 |
+
pairs = [[sub_q, f"{p.get('title', '')}\n{p.get('abstract', '')}"]
|
| 301 |
+
for p in top_papers]
|
| 302 |
+
scores = self.retriever.reranker.predict(pairs)
|
| 303 |
+
if float(max(scores)) > COVERAGE_THRESHOLD:
|
| 304 |
+
continue # sub-query already covered
|
| 305 |
+
|
| 306 |
+
# No paper covers this sub-query β fetch a targeted result and append
|
| 307 |
+
logger.info(f"Coverage gap detected for sub-query: {sub_q[:60]}... Fetching targeted paper.")
|
| 308 |
+
extra = self.retriever.search(sub_q, top_k=5)
|
| 309 |
+
reranked_extra = self.retriever.rerank(sub_q, extra, top_n=1)
|
| 310 |
+
if reranked_extra:
|
| 311 |
+
pid = reranked_extra[0].get("paper_id")
|
| 312 |
+
if pid not in {p.get("paper_id") for p in top_papers}:
|
| 313 |
+
top_papers.append(reranked_extra[0])
|
| 314 |
+
|
| 315 |
+
return self._deduplicate_papers(top_papers, top_n)
|
| 316 |
+
|
| 317 |
+
# ------------------------------------------------------------------
|
| 318 |
+
# Context assembly with multi-paper full-text + fallback
|
| 319 |
+
# ------------------------------------------------------------------
|
| 320 |
+
def _build_context(self, top_papers: list[dict], query: str, max_full_text: int = None) -> tuple[str, list[str], list[str]]:
|
| 321 |
+
"""Build the context string for LLM synthesis.
|
| 322 |
+
|
| 323 |
+
Fetches full text for the top-N papers IN PARALLEL to avoid sequential
|
| 324 |
+
30s-per-paper timeouts stacking up. Falls back to abstract when fetch fails.
|
| 325 |
+
|
| 326 |
+
Returns (full_context, abstracts_list, full_text_paper_titles).
|
| 327 |
+
"""
|
| 328 |
+
if max_full_text is None:
|
| 329 |
+
max_full_text = self.max_full_text
|
| 330 |
+
papers_to_fetch = top_papers[:max_full_text]
|
| 331 |
+
remaining = top_papers[max_full_text:]
|
| 332 |
+
|
| 333 |
+
def _fetch(args: tuple[int, dict]) -> tuple[int, dict, str | None]:
|
| 334 |
+
i, paper = args
|
| 335 |
+
title = paper.get("title", "")
|
| 336 |
+
logger.info(f"Attempting full text fetch for: {title[:60]}")
|
| 337 |
+
text = self.fetcher.fetch_full_text(paper["paper_id"])
|
| 338 |
+
return i, paper, text
|
| 339 |
+
|
| 340 |
+
# Parallel fetch β wall time is max(individual fetches) instead of sum
|
| 341 |
+
fetch_results: dict[int, tuple[dict, str | None]] = {}
|
| 342 |
+
with ThreadPoolExecutor(max_workers=self.max_full_text) as pool:
|
| 343 |
+
for i, paper, text in pool.map(_fetch, enumerate(papers_to_fetch)):
|
| 344 |
+
fetch_results[i] = (paper, text)
|
| 345 |
+
|
| 346 |
+
context_parts: list[str] = []
|
| 347 |
+
abstracts: list[str] = []
|
| 348 |
+
full_text_titles: list[str] = []
|
| 349 |
+
|
| 350 |
+
# Adaptive top_k: distribute passage budget evenly across however many
|
| 351 |
+
# papers actually returned full text, capped at 8 for a single paper.
|
| 352 |
+
full_text_count = sum(
|
| 353 |
+
1 for i in range(len(papers_to_fetch))
|
| 354 |
+
if fetch_results[i][1] and len(fetch_results[i][1]) > 1000
|
| 355 |
+
)
|
| 356 |
+
adaptive_top_k = max(6, 12 // max(1, full_text_count))
|
| 357 |
+
|
| 358 |
+
for i in range(len(papers_to_fetch)):
|
| 359 |
+
paper, full_text = fetch_results[i]
|
| 360 |
+
abstract = paper.get("abstract", "")
|
| 361 |
+
title = paper.get("title", "")
|
| 362 |
+
abstracts.append(abstract)
|
| 363 |
+
|
| 364 |
+
if full_text and len(full_text) > 1000:
|
| 365 |
+
content = self._extract_relevant_passages(
|
| 366 |
+
full_text, query, paper_title=title, top_k=adaptive_top_k
|
| 367 |
+
)
|
| 368 |
+
context_parts.append(
|
| 369 |
+
f"--- PAPER {i} [{title}] (FULL TEXT, RELEVANT PASSAGES) ---\n{content}\n"
|
| 370 |
+
)
|
| 371 |
+
full_text_titles.append(title)
|
| 372 |
+
else:
|
| 373 |
+
context_parts.append(f"--- PAPER {i} [{title}] (ABSTRACT) ---\n{abstract}\n")
|
| 374 |
+
|
| 375 |
+
# Remaining papers (beyond max_full_text) always use abstract
|
| 376 |
+
for j, paper in enumerate(remaining, start=len(papers_to_fetch)):
|
| 377 |
+
abstract = paper.get("abstract", "")
|
| 378 |
+
title = paper.get("title", "")
|
| 379 |
+
abstracts.append(abstract)
|
| 380 |
+
context_parts.append(f"--- PAPER {j} [{title}] (ABSTRACT) ---\n{abstract}\n")
|
| 381 |
+
|
| 382 |
+
# If no full text was fetched at all, scan remaining for any accessible paper
|
| 383 |
+
if not full_text_titles:
|
| 384 |
+
logger.info("No full text found in top papers β scanning remaining for fallback...")
|
| 385 |
+
for paper in remaining:
|
| 386 |
+
title = paper.get("title", "")
|
| 387 |
+
full_text = self.fetcher.fetch_full_text(paper["paper_id"])
|
| 388 |
+
if full_text and len(full_text) > 1000:
|
| 389 |
+
content = self._extract_relevant_passages(
|
| 390 |
+
full_text, query, paper_title=title, top_k=8
|
| 391 |
+
)
|
| 392 |
+
context_parts.append(
|
| 393 |
+
f"--- FALLBACK PAPER [{title}] (FULL TEXT, RELEVANT PASSAGES) ---\n{content}\n"
|
| 394 |
+
)
|
| 395 |
+
full_text_titles.append(title)
|
| 396 |
+
break
|
| 397 |
+
|
| 398 |
+
full_context = "\n".join(context_parts)
|
| 399 |
+
return full_context, abstracts, full_text_titles
|
| 400 |
+
|
| 401 |
+
# ------------------------------------------------------------------
|
| 402 |
+
# Main entry point
|
| 403 |
+
# ------------------------------------------------------------------
|
| 404 |
+
def run(self, query: str, ground_truth: str = None) -> dict:
|
| 405 |
+
start = time.time()
|
| 406 |
+
logger.info(f"Pipeline 2 query: {query}")
|
| 407 |
+
|
| 408 |
+
# Step 1: Query decomposition
|
| 409 |
+
decomposition = self._decompose_query(query)
|
| 410 |
+
|
| 411 |
+
# Step 2: Retrieval + Reranking
|
| 412 |
+
if decomposition["is_multi_hop"]:
|
| 413 |
+
# Parallel sub-query retrieval: each sub-query targets one required paper
|
| 414 |
+
def _search_sub_query(sub_q: str) -> list[dict]:
|
| 415 |
+
return self.retriever.search(sub_q, top_k=20)
|
| 416 |
+
|
| 417 |
+
all_candidates: list[dict] = []
|
| 418 |
+
with ThreadPoolExecutor(max_workers=len(decomposition["sub_queries"])) as pool:
|
| 419 |
+
futures = [pool.submit(_search_sub_query, sq)
|
| 420 |
+
for sq in decomposition["sub_queries"]]
|
| 421 |
+
for fut in as_completed(futures):
|
| 422 |
+
all_candidates.extend(fut.result())
|
| 423 |
+
|
| 424 |
+
# Deduplicate then do ONE final rerank against the ORIGINAL question (Fix 3)
|
| 425 |
+
unique_candidates = self._deduplicate_papers(all_candidates, top_n=self.retrieval_top_k)
|
| 426 |
+
top_papers = self.retriever.rerank(query, unique_candidates, top_n=self.rerank_top_n)
|
| 427 |
+
|
| 428 |
+
# Fix 4: Mandatory coverage check β ensure each sub-query has a representative paper
|
| 429 |
+
top_papers = self._ensure_sub_query_coverage(
|
| 430 |
+
top_papers, decomposition["sub_queries"], top_n=self.rerank_top_n
|
| 431 |
+
)
|
| 432 |
+
else:
|
| 433 |
+
candidates = self.retriever.search(query, top_k=self.retrieval_top_k)
|
| 434 |
+
top_papers = self.retriever.rerank(query, candidates, top_n=self.rerank_top_n)
|
| 435 |
+
|
| 436 |
+
# Step 3: Context assembly β single-hop gets 2 full-text papers (less noise),
|
| 437 |
+
# multi-hop gets 3 (needs evidence from multiple sources)
|
| 438 |
+
max_ft = self.max_full_text if decomposition["is_multi_hop"] else min(2, self.max_full_text)
|
| 439 |
+
full_context, abstracts, full_text_titles = self._build_context(top_papers, query, max_full_text=max_ft)
|
| 440 |
+
|
| 441 |
+
# Step 4: LLM Synthesis
|
| 442 |
+
system_msg = (
|
| 443 |
+
"You are an expert AI research assistant specializing in scientific literature synthesis. "
|
| 444 |
+
"Answer the query using ONLY the provided paper context. Speak with absolute certainty.\n\n"
|
| 445 |
+
"Rules:\n"
|
| 446 |
+
"1. Answer directly and immediately. Do NOT use any preamble like 'Based on the provided papers', 'According to the context', or 'It can be inferred'. Start the response with the subject of the question.\n"
|
| 447 |
+
"2. Quote technical terms EXACTLY as they appear in the papers. "
|
| 448 |
+
"For any list of methods, data sources, properties, or named items: copy the exact words. "
|
| 449 |
+
"Do NOT substitute synonyms or rephrase β "
|
| 450 |
+
"e.g. 'large-scale simulations' must not become 'multi-scale simulations'; "
|
| 451 |
+
"'field measurements' must not become 'real data from the field'.\n"
|
| 452 |
+
"3. Be complete β address EVERY part of the question. For list questions, include "
|
| 453 |
+
"every item. For comparison questions, cover EACH paper explicitly.\n"
|
| 454 |
+
"4. Never add disclaimers, caveats, hedges, or uncertainty statements. "
|
| 455 |
+
"Do not state that information is missing or unclear. If a detail is missing, omit it silently.\n"
|
| 456 |
+
"5. Do not use square brackets or cite PAPER numbers in the final text. "
|
| 457 |
+
"6. Prefer concise, precise, and highly technical language."
|
| 458 |
+
)
|
| 459 |
+
user_msg = f"QUERY: {query}\n\nCONTEXT:\n{full_context}"
|
| 460 |
+
|
| 461 |
+
try:
|
| 462 |
+
response = self._groq_call(
|
| 463 |
+
model=self.model_name,
|
| 464 |
+
messages=[
|
| 465 |
+
{"role": "system", "content": system_msg},
|
| 466 |
+
{"role": "user", "content": user_msg},
|
| 467 |
+
],
|
| 468 |
+
temperature=0.1,
|
| 469 |
+
max_tokens=1024,
|
| 470 |
+
)
|
| 471 |
+
|
| 472 |
+
answer = response.choices[0].message.content
|
| 473 |
+
usage = _GroqUsage(
|
| 474 |
+
prompt_token_count=response.usage.prompt_tokens,
|
| 475 |
+
candidates_token_count=response.usage.completion_tokens,
|
| 476 |
+
)
|
| 477 |
+
|
| 478 |
+
stats = self.metrics.process_metrics(
|
| 479 |
+
client=None,
|
| 480 |
+
query=query,
|
| 481 |
+
answer=answer,
|
| 482 |
+
context=full_context,
|
| 483 |
+
usage_metadata=usage,
|
| 484 |
+
start_time=start,
|
| 485 |
+
abstracts_list=abstracts,
|
| 486 |
+
ground_truth=ground_truth
|
| 487 |
+
)
|
| 488 |
+
|
| 489 |
+
return {
|
| 490 |
+
"answer": answer,
|
| 491 |
+
"sources": [p.get("title", "") for p in top_papers],
|
| 492 |
+
"metrics": stats,
|
| 493 |
+
"selected_papers": full_text_titles,
|
| 494 |
+
"query_decomposition": decomposition,
|
| 495 |
+
}
|
| 496 |
+
|
| 497 |
+
except Exception as e:
|
| 498 |
+
logger.error(f"Pipeline 2 error: {e}")
|
| 499 |
+
return {"error": str(e), "answer": "Error generating response."}
|
core/pipeline_2/qdrant_storage/.lock
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
tmp lock file
|
core/pipeline_2/qdrant_storage/collection/papers/storage.sqlite
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:7137fccf970a9b1a45152d91866c45b6d5c8b081217364467723ad06edac3787
|
| 3 |
+
size 182099968
|
core/pipeline_2/qdrant_storage/meta.json
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
{"collections": {"papers": {"vectors": {"dense": {"size": 1024, "distance": "Cosine", "hnsw_config": null, "quantization_config": null, "on_disk": null, "datatype": null, "multivector_config": null}}, "shard_number": null, "sharding_method": null, "replication_factor": null, "write_consistency_factor": null, "on_disk_payload": null, "hnsw_config": null, "wal_config": null, "optimizers_config": null, "quantization_config": null, "sparse_vectors": {"sparse": {"index": null, "modifier": "idf"}}, "strict_mode_config": null, "metadata": null}}, "aliases": {}}
|
core/pipeline_2/retriever.py
ADDED
|
@@ -0,0 +1,76 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Hybrid retrieval (dense + sparse RRF fusion) + cross-encoder reranking."""
|
| 2 |
+
import logging
|
| 3 |
+
from pathlib import Path
|
| 4 |
+
from typing import Any
|
| 5 |
+
|
| 6 |
+
from fastembed import SparseTextEmbedding
|
| 7 |
+
from qdrant_client import QdrantClient, models
|
| 8 |
+
from sentence_transformers import CrossEncoder, SentenceTransformer
|
| 9 |
+
|
| 10 |
+
logger = logging.getLogger(__name__)
|
| 11 |
+
|
| 12 |
+
_DEFAULT_STORAGE_PATH = str(Path(__file__).parent / "qdrant_storage")
|
| 13 |
+
|
| 14 |
+
COLLECTION = "papers"
|
| 15 |
+
DENSE_MODEL = "BAAI/bge-large-en-v1.5"
|
| 16 |
+
SPARSE_MODEL = "Qdrant/bm25"
|
| 17 |
+
RERANKER_MODEL = "BAAI/bge-reranker-v2-m3" # Upgraded from bge-reranker-base
|
| 18 |
+
|
| 19 |
+
|
| 20 |
+
class HybridRetriever:
|
| 21 |
+
def __init__(self, storage_path: str = _DEFAULT_STORAGE_PATH):
|
| 22 |
+
logger.info("Initializing HybridRetriever (loading models at runtime)...")
|
| 23 |
+
self.storage_path = storage_path
|
| 24 |
+
|
| 25 |
+
logger.info("Loading dense model...")
|
| 26 |
+
self.dense = SentenceTransformer(DENSE_MODEL)
|
| 27 |
+
|
| 28 |
+
logger.info("Loading sparse model...")
|
| 29 |
+
self.sparse = SparseTextEmbedding(SPARSE_MODEL, providers=["CPUExecutionProvider"])
|
| 30 |
+
|
| 31 |
+
logger.info("Loading reranker...")
|
| 32 |
+
self.reranker = CrossEncoder(RERANKER_MODEL, max_length=1024)
|
| 33 |
+
|
| 34 |
+
self.client = QdrantClient(path=storage_path)
|
| 35 |
+
logger.info("HybridRetriever ready.")
|
| 36 |
+
|
| 37 |
+
def _dense_vec(self, text: str) -> list[float]:
|
| 38 |
+
return self.dense.encode(text, normalize_embeddings=True).tolist()
|
| 39 |
+
|
| 40 |
+
def _sparse_vec(self, text: str) -> models.SparseVector:
|
| 41 |
+
result = next(self.sparse.embed([text]))
|
| 42 |
+
return models.SparseVector(
|
| 43 |
+
indices=result.indices.tolist(),
|
| 44 |
+
values=result.values.tolist(),
|
| 45 |
+
)
|
| 46 |
+
|
| 47 |
+
def search(self, query: str, top_k: int = 50) -> list[dict[str, Any]]:
|
| 48 |
+
"""Hybrid dense+sparse search with RRF fusion."""
|
| 49 |
+
results = self.client.query_points(
|
| 50 |
+
collection_name=COLLECTION,
|
| 51 |
+
prefetch=[
|
| 52 |
+
models.Prefetch(query=self._dense_vec(query), using="dense", limit=top_k),
|
| 53 |
+
models.Prefetch(query=self._sparse_vec(query), using="sparse", limit=top_k),
|
| 54 |
+
],
|
| 55 |
+
query=models.FusionQuery(fusion=models.Fusion.RRF),
|
| 56 |
+
limit=top_k,
|
| 57 |
+
with_payload=True,
|
| 58 |
+
)
|
| 59 |
+
return [{**pt.payload, "_score": pt.score} for pt in results.points]
|
| 60 |
+
|
| 61 |
+
def rerank(self, query: str, candidates: list[dict], top_n: int = 10) -> list[dict]:
|
| 62 |
+
"""Pure semantic cross-encoder reranking over title + abstract.
|
| 63 |
+
|
| 64 |
+
No artificial boosts β ranking is based entirely on query-document
|
| 65 |
+
relevance as scored by the cross-encoder.
|
| 66 |
+
"""
|
| 67 |
+
if not candidates:
|
| 68 |
+
return []
|
| 69 |
+
pairs = [
|
| 70 |
+
[query, f"{c.get('title', '')}\n{c.get('abstract', '')}"]
|
| 71 |
+
for c in candidates
|
| 72 |
+
]
|
| 73 |
+
scores = self.reranker.predict(pairs, batch_size=16)
|
| 74 |
+
|
| 75 |
+
ranked = sorted(zip(scores, candidates), key=lambda x: x[0], reverse=True)
|
| 76 |
+
return [c for _, c in ranked[:top_n]]
|
core/pipeline_3/.gitkeep
ADDED
|
File without changes
|
core/pipeline_3/checkpoint-2.json
ADDED
|
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"embed_batches": 223,
|
| 3 |
+
"papers": 56,
|
| 4 |
+
"authors": true,
|
| 5 |
+
"topics": true,
|
| 6 |
+
"authored_by": true,
|
| 7 |
+
"has_topic": true,
|
| 8 |
+
"cites": true
|
| 9 |
+
}
|
core/pipeline_3/embeddings.npy
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:b2bcc131125783a5ab4d8a2a75dc487a7b81cc617f95baa7a64a975efd4654d0
|
| 3 |
+
size 58355840
|
core/pipeline_3/logic.py
ADDED
|
@@ -0,0 +1,698 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Pipeline 3: GraphRAG β Qdrant hybrid retrieval + TigerGraph graph traversal + Groq LLM.
|
| 2 |
+
|
| 3 |
+
Architecture:
|
| 4 |
+
Qdrant (dense + BM25) β seed papers [same quality as Pipeline 2]
|
| 5 |
+
TigerGraph β graph expansion [CITES / HAS_TOPIC, only when query warrants]
|
| 6 |
+
RRF merge + cross-encoder rerank β Groq synthesis
|
| 7 |
+
|
| 8 |
+
TigerGraph is used exclusively for what it does uniquely: traversing citation and topic
|
| 9 |
+
edges to surface papers that are graph-adjacent to the retrieval seeds. Everything else
|
| 10 |
+
(hybrid text search, reranking, context assembly) reuses Pipeline 2's stack.
|
| 11 |
+
"""
|
| 12 |
+
import json
|
| 13 |
+
import logging
|
| 14 |
+
import os
|
| 15 |
+
import re
|
| 16 |
+
import sys
|
| 17 |
+
import time
|
| 18 |
+
from dataclasses import dataclass
|
| 19 |
+
from pathlib import Path
|
| 20 |
+
from concurrent.futures import ThreadPoolExecutor
|
| 21 |
+
|
| 22 |
+
import pyTigerGraph as tg
|
| 23 |
+
from dotenv import load_dotenv
|
| 24 |
+
from groq import Groq, RateLimitError as GroqRateLimitError
|
| 25 |
+
from sentence_transformers import SentenceTransformer
|
| 26 |
+
|
| 27 |
+
sys.path.append(str(Path(__file__).parents[2]))
|
| 28 |
+
|
| 29 |
+
from core.pipeline_2.retriever import HybridRetriever
|
| 30 |
+
from services.metrics_service import MetricsService
|
| 31 |
+
from services.paper_fetcher import PaperFetcher
|
| 32 |
+
from core.pipeline_3.queries import (
|
| 33 |
+
install_queries,
|
| 34 |
+
parse_paper_results,
|
| 35 |
+
TG_HOST, TG_PORT, TG_USER, TG_PASS, GRAPH,
|
| 36 |
+
)
|
| 37 |
+
|
| 38 |
+
load_dotenv()
|
| 39 |
+
|
| 40 |
+
logging.basicConfig(level=logging.INFO)
|
| 41 |
+
logger = logging.getLogger(__name__)
|
| 42 |
+
|
| 43 |
+
_DEFAULT_MODEL = "llama-3.3-70b-versatile"
|
| 44 |
+
_DECOMPOSE_MODEL = "llama-3.1-8b-instant"
|
| 45 |
+
|
| 46 |
+
_PIN_K = 3 # top Qdrant seeds always guaranteed a slot in final top_n
|
| 47 |
+
|
| 48 |
+
_STOP_WORDS = {
|
| 49 |
+
"what", "which", "when", "where", "who", "how", "why", "does", "do",
|
| 50 |
+
"is", "are", "was", "were", "the", "this", "that", "these", "those",
|
| 51 |
+
"a", "an", "and", "or", "but", "of", "in", "on", "at", "to", "for",
|
| 52 |
+
"with", "by", "from", "into", "about", "between", "each", "other",
|
| 53 |
+
"have", "has", "had", "their", "they", "both", "used", "uses", "using",
|
| 54 |
+
"also", "can", "its", "such", "more", "than", "some", "any", "only",
|
| 55 |
+
}
|
| 56 |
+
|
| 57 |
+
|
| 58 |
+
@dataclass
|
| 59 |
+
class _GroqUsage:
|
| 60 |
+
prompt_token_count: int
|
| 61 |
+
candidates_token_count: int
|
| 62 |
+
|
| 63 |
+
|
| 64 |
+
class PipelineGraphRAG:
|
| 65 |
+
"""Pipeline 3: Qdrant hybrid retrieval + TigerGraph graph expansion.
|
| 66 |
+
|
| 67 |
+
Differentiator over Pipeline 2: citation (CITES) and topic (HAS_TOPIC)
|
| 68 |
+
graph expansion via TigerGraph, activated only when the query semantically
|
| 69 |
+
justifies it (decided by the decomposition layer).
|
| 70 |
+
"""
|
| 71 |
+
|
| 72 |
+
def __init__(self, rerank_top_n: int = 10, max_full_text: int = 3):
|
| 73 |
+
self.rerank_top_n = rerank_top_n
|
| 74 |
+
self.max_full_text = max_full_text
|
| 75 |
+
|
| 76 |
+
self.model_name = os.environ.get("GROQ_MODEL", _DEFAULT_MODEL)
|
| 77 |
+
_keys = [os.environ.get(f"GROQ_API_KEY_{i}") for i in range(1, 4)]
|
| 78 |
+
_keys = [k for k in _keys if k]
|
| 79 |
+
if not _keys:
|
| 80 |
+
_keys = [os.environ["GROQ_API_KEY"]]
|
| 81 |
+
self._groq_clients = [Groq(api_key=k) for k in _keys]
|
| 82 |
+
self._groq_key_idx = 0
|
| 83 |
+
logger.info(f"Loaded {len(self._groq_clients)} Groq API key(s).")
|
| 84 |
+
|
| 85 |
+
self.fetcher = PaperFetcher()
|
| 86 |
+
self.metrics = MetricsService(self.model_name)
|
| 87 |
+
|
| 88 |
+
logger.info("Initializing Qdrant retriever (dense + BM25)...")
|
| 89 |
+
self.retriever = HybridRetriever()
|
| 90 |
+
# Reuse models already loaded by HybridRetriever β no double loading
|
| 91 |
+
self.dense = self.retriever.dense
|
| 92 |
+
self.reranker = self.retriever.reranker
|
| 93 |
+
|
| 94 |
+
logger.info("Loading fast passage embedding model (bge-small-en-v1.5)...")
|
| 95 |
+
self.fast_dense = SentenceTransformer("BAAI/bge-small-en-v1.5")
|
| 96 |
+
|
| 97 |
+
logger.info("Connecting to TigerGraph (graph traversal only)...")
|
| 98 |
+
_tg_secret = os.environ.get("TG_SECRET", "")
|
| 99 |
+
if _tg_secret:
|
| 100 |
+
self.conn = tg.TigerGraphConnection(
|
| 101 |
+
host=TG_HOST,
|
| 102 |
+
graphname=GRAPH,
|
| 103 |
+
gsqlSecret=_tg_secret,
|
| 104 |
+
restppPort=str(TG_PORT),
|
| 105 |
+
gsPort=str(TG_PORT),
|
| 106 |
+
)
|
| 107 |
+
self.conn.getToken(_tg_secret)
|
| 108 |
+
else:
|
| 109 |
+
self.conn = tg.TigerGraphConnection(
|
| 110 |
+
host=TG_HOST,
|
| 111 |
+
graphname=GRAPH,
|
| 112 |
+
username=TG_USER,
|
| 113 |
+
password=TG_PASS,
|
| 114 |
+
restppPort=str(TG_PORT),
|
| 115 |
+
gsPort=str(TG_PORT),
|
| 116 |
+
)
|
| 117 |
+
|
| 118 |
+
logger.info("Installing/verifying GSQL queries...")
|
| 119 |
+
install_queries()
|
| 120 |
+
|
| 121 |
+
logger.info("PipelineGraphRAG ready.")
|
| 122 |
+
|
| 123 |
+
# ------------------------------------------------------------------
|
| 124 |
+
# Groq call with key rotation on rate limit
|
| 125 |
+
# ------------------------------------------------------------------
|
| 126 |
+
def _groq_call(self, **kwargs):
|
| 127 |
+
for _ in range(len(self._groq_clients)):
|
| 128 |
+
try:
|
| 129 |
+
return self._groq_clients[self._groq_key_idx].chat.completions.create(**kwargs)
|
| 130 |
+
except GroqRateLimitError:
|
| 131 |
+
next_idx = (self._groq_key_idx + 1) % len(self._groq_clients)
|
| 132 |
+
if next_idx == self._groq_key_idx:
|
| 133 |
+
raise
|
| 134 |
+
logger.warning(
|
| 135 |
+
f"Rate limit on Groq key {self._groq_key_idx + 1}, "
|
| 136 |
+
f"rotating to key {next_idx + 1}..."
|
| 137 |
+
)
|
| 138 |
+
self._groq_key_idx = next_idx
|
| 139 |
+
raise GroqRateLimitError("All Groq API keys exhausted.")
|
| 140 |
+
|
| 141 |
+
# ------------------------------------------------------------------
|
| 142 |
+
# Query decomposition + graph signal selection
|
| 143 |
+
# ------------------------------------------------------------------
|
| 144 |
+
def _decompose_query(self, query: str) -> dict:
|
| 145 |
+
"""Classify query and select which graph signals to activate.
|
| 146 |
+
|
| 147 |
+
Returns:
|
| 148 |
+
is_multi_hop β True if 2+ distinct papers must be retrieved
|
| 149 |
+
sub_queries β list of focused retrieval queries
|
| 150 |
+
use_topic β True if HAS_TOPIC expansion should run
|
| 151 |
+
use_cites β True if CITES expansion should run
|
| 152 |
+
"""
|
| 153 |
+
_FALLBACK = {
|
| 154 |
+
"is_multi_hop": False,
|
| 155 |
+
"sub_queries": [query],
|
| 156 |
+
"use_topic": False,
|
| 157 |
+
"use_cites": False,
|
| 158 |
+
}
|
| 159 |
+
try:
|
| 160 |
+
response = self._groq_call(
|
| 161 |
+
model=_DECOMPOSE_MODEL,
|
| 162 |
+
messages=[
|
| 163 |
+
{
|
| 164 |
+
"role": "system",
|
| 165 |
+
"content": (
|
| 166 |
+
"You classify scientific search queries to decide retrieval strategy.\n\n"
|
| 167 |
+
"## HOP CLASSIFICATION\n"
|
| 168 |
+
"MULTI-HOP (is_multi_hop=true): question EXPLICITLY references 2+ DISTINCT "
|
| 169 |
+
"papers/methods that must each be independently retrieved. "
|
| 170 |
+
"Look for: 'Both X and Y', 'X paper and Y paper', comparison between two "
|
| 171 |
+
"named systems from different papers.\n"
|
| 172 |
+
"SINGLE-HOP (is_multi_hop=false): anything else, including multi-part "
|
| 173 |
+
"questions about ONE topic or ONE named paper.\n\n"
|
| 174 |
+
"## GRAPH SIGNAL SELECTION\n"
|
| 175 |
+
"use_topic=true: query is a broad survey of a research area or field "
|
| 176 |
+
"(e.g. 'approaches to X', 'methods for Y', 'overview of Z'). "
|
| 177 |
+
"False for questions about a specific named paper.\n"
|
| 178 |
+
"use_cites=true: query asks about related/influential work, requires "
|
| 179 |
+
"synthesis across papers, or is multi-hop. "
|
| 180 |
+
"False for specific factual questions about one named paper.\n\n"
|
| 181 |
+
"## EXAMPLES\n"
|
| 182 |
+
" 'What two factors drove NLP progress according to HuggingFace Transformers?' "
|
| 183 |
+
"β is_multi_hop=false, use_topic=false, use_cites=false\n"
|
| 184 |
+
" 'What are the main approaches to knowledge graph embeddings?' "
|
| 185 |
+
"β is_multi_hop=false, use_topic=true, use_cites=true\n"
|
| 186 |
+
" 'Both CGCNN and DPMD apply neural networks β what does each predict?' "
|
| 187 |
+
"β is_multi_hop=true, use_topic=false, use_cites=true\n"
|
| 188 |
+
" 'How does ADMETlab relate to de novo drug design?' "
|
| 189 |
+
"β is_multi_hop=true, use_topic=true, use_cites=true\n\n"
|
| 190 |
+
"Return ONLY valid JSON, no markdown:\n"
|
| 191 |
+
'{"is_multi_hop": false, "sub_queries": ["original query"], '
|
| 192 |
+
'"use_topic": false, "use_cites": false}'
|
| 193 |
+
),
|
| 194 |
+
},
|
| 195 |
+
{"role": "user", "content": query},
|
| 196 |
+
],
|
| 197 |
+
temperature=0.0,
|
| 198 |
+
max_tokens=300,
|
| 199 |
+
)
|
| 200 |
+
raw = response.choices[0].message.content.strip()
|
| 201 |
+
raw = re.sub(r"^```(?:json)?\s*", "", raw)
|
| 202 |
+
raw = re.sub(r"\s*```$", "", raw)
|
| 203 |
+
result = json.loads(raw)
|
| 204 |
+
if not isinstance(result.get("sub_queries"), list) or not result["sub_queries"]:
|
| 205 |
+
return _FALLBACK
|
| 206 |
+
|
| 207 |
+
result.setdefault("use_topic", False)
|
| 208 |
+
result.setdefault("use_cites", False)
|
| 209 |
+
|
| 210 |
+
if result.get("is_multi_hop") and len(result["sub_queries"]) >= 2:
|
| 211 |
+
subs = result["sub_queries"]
|
| 212 |
+
override = False
|
| 213 |
+
|
| 214 |
+
def _cap_pairs(text):
|
| 215 |
+
tokens = text.split()
|
| 216 |
+
return {
|
| 217 |
+
(tokens[i].lower(), tokens[i + 1].lower())
|
| 218 |
+
for i in range(len(tokens) - 1)
|
| 219 |
+
if tokens[i][0].isupper() and tokens[i + 1][0].isupper()
|
| 220 |
+
and len(tokens[i]) > 1 and len(tokens[i + 1]) > 1
|
| 221 |
+
}
|
| 222 |
+
if _cap_pairs(subs[0]) & _cap_pairs(subs[1]):
|
| 223 |
+
override = True
|
| 224 |
+
|
| 225 |
+
if not override and re.search(r'\bit\b', subs[1], re.IGNORECASE):
|
| 226 |
+
has_named_system = re.search(
|
| 227 |
+
r'\b([A-Z]{2,}|[A-Z][a-z]+[A-Z][a-zA-Z]*)\b',
|
| 228 |
+
subs[0] + " " + subs[1]
|
| 229 |
+
)
|
| 230 |
+
if not has_named_system:
|
| 231 |
+
override = True
|
| 232 |
+
|
| 233 |
+
if not override and re.search(
|
| 234 |
+
r'\bthe\s+(library|paper|model|system|method|framework|approach|algorithm|tool|platform)\b',
|
| 235 |
+
subs[1], re.IGNORECASE,
|
| 236 |
+
):
|
| 237 |
+
override = True
|
| 238 |
+
|
| 239 |
+
if override:
|
| 240 |
+
logger.info("Multi-hop override: sub-queries target same paper β single-hop.")
|
| 241 |
+
result = {
|
| 242 |
+
"is_multi_hop": False,
|
| 243 |
+
"sub_queries": [query],
|
| 244 |
+
"use_topic": result["use_topic"],
|
| 245 |
+
"use_cites": result["use_cites"],
|
| 246 |
+
}
|
| 247 |
+
|
| 248 |
+
logger.info(
|
| 249 |
+
f"Decomposition: multi_hop={result['is_multi_hop']}, "
|
| 250 |
+
f"use_topic={result['use_topic']}, use_cites={result['use_cites']}, "
|
| 251 |
+
f"subs={result['sub_queries']}"
|
| 252 |
+
)
|
| 253 |
+
return result
|
| 254 |
+
except Exception as e:
|
| 255 |
+
logger.warning(f"Query decomposition failed ({e}), using original query.")
|
| 256 |
+
return _FALLBACK
|
| 257 |
+
|
| 258 |
+
# ------------------------------------------------------------------
|
| 259 |
+
# TigerGraph graph traversal
|
| 260 |
+
# ------------------------------------------------------------------
|
| 261 |
+
def _extract_topic_keywords(self, query: str) -> list[str]:
|
| 262 |
+
words = re.findall(r'\b[a-zA-Z]{4,}\b', query)
|
| 263 |
+
keywords = [w for w in words if w.lower() not in _STOP_WORDS]
|
| 264 |
+
return sorted(set(keywords), key=len, reverse=True)[:3]
|
| 265 |
+
|
| 266 |
+
def _topic_search(self, keywords: list[str], top_k: int = 20) -> list[dict]:
|
| 267 |
+
if not keywords:
|
| 268 |
+
return []
|
| 269 |
+
|
| 270 |
+
def _search_one(kw: str) -> list[dict]:
|
| 271 |
+
try:
|
| 272 |
+
response = self.conn.runInstalledQuery(
|
| 273 |
+
"search_papers_by_topic",
|
| 274 |
+
{"pattern": f"%{kw}%", "top_k": top_k},
|
| 275 |
+
)
|
| 276 |
+
return parse_paper_results(response)
|
| 277 |
+
except Exception as e:
|
| 278 |
+
logger.warning(f"Topic search for '{kw}' failed: {e}")
|
| 279 |
+
return []
|
| 280 |
+
|
| 281 |
+
papers: list[dict] = []
|
| 282 |
+
seen: set[str] = set()
|
| 283 |
+
with ThreadPoolExecutor(max_workers=len(keywords)) as pool:
|
| 284 |
+
futs = [pool.submit(_search_one, kw) for kw in keywords]
|
| 285 |
+
for fut in futs:
|
| 286 |
+
try:
|
| 287 |
+
for p in fut.result(timeout=15):
|
| 288 |
+
pid = p.get("paper_id", "")
|
| 289 |
+
if pid and pid not in seen:
|
| 290 |
+
seen.add(pid)
|
| 291 |
+
papers.append(p)
|
| 292 |
+
except Exception:
|
| 293 |
+
pass
|
| 294 |
+
return papers
|
| 295 |
+
|
| 296 |
+
def _expand_neighborhood(self, paper_ids: list[str]) -> list[dict]:
|
| 297 |
+
if not paper_ids:
|
| 298 |
+
return []
|
| 299 |
+
try:
|
| 300 |
+
# Run GSQL query with timeout and limited seeds to avoid hangs
|
| 301 |
+
response = self.conn.runInstalledQuery(
|
| 302 |
+
"expand_paper_neighborhood",
|
| 303 |
+
{"seed_ids": paper_ids[:10]},
|
| 304 |
+
timeout=15000 # 15s timeout
|
| 305 |
+
)
|
| 306 |
+
return parse_paper_results(response)
|
| 307 |
+
except Exception as e:
|
| 308 |
+
logger.warning(f"Neighborhood expansion failed or timed out: {e}")
|
| 309 |
+
return []
|
| 310 |
+
|
| 311 |
+
# ------------------------------------------------------------------
|
| 312 |
+
# Retrieval
|
| 313 |
+
# ------------------------------------------------------------------
|
| 314 |
+
@staticmethod
|
| 315 |
+
def _normalize_ids(papers: list[dict]) -> list[dict]:
|
| 316 |
+
"""Strip OpenAlex URL prefix from paper_id so Qdrant and TG IDs match.
|
| 317 |
+
|
| 318 |
+
Qdrant stores full URLs (https://openalex.org/W...).
|
| 319 |
+
TigerGraph uses short IDs (W...). Normalize to short form for RRF merge.
|
| 320 |
+
"""
|
| 321 |
+
for p in papers:
|
| 322 |
+
pid = p.get("paper_id", "")
|
| 323 |
+
if "/" in pid:
|
| 324 |
+
p["paper_id"] = pid.rsplit("/", 1)[-1]
|
| 325 |
+
return papers
|
| 326 |
+
|
| 327 |
+
def _rrf_merge(
|
| 328 |
+
self, *lists: list[dict], weights: list[float] = None, k: int = 60
|
| 329 |
+
) -> list[dict]:
|
| 330 |
+
"""Reciprocal Rank Fusion with optional list weighting.
|
| 331 |
+
|
| 332 |
+
Weights allow favoring trust-worthy sources (like Vector Search)
|
| 333 |
+
over noisy expansions (like broad topic search).
|
| 334 |
+
"""
|
| 335 |
+
if not weights:
|
| 336 |
+
weights = [1.0] * len(lists)
|
| 337 |
+
|
| 338 |
+
scores: dict[str, float] = {}
|
| 339 |
+
paper_map: dict[str, dict] = {}
|
| 340 |
+
|
| 341 |
+
for i, result_list in enumerate(lists):
|
| 342 |
+
w = weights[i] if i < len(weights) else 1.0
|
| 343 |
+
for rank, paper in enumerate(result_list, start=1):
|
| 344 |
+
pid = paper.get("paper_id", "")
|
| 345 |
+
if not pid:
|
| 346 |
+
continue
|
| 347 |
+
# RRF Score = weight * (1 / (k + rank))
|
| 348 |
+
scores[pid] = scores.get(pid, 0.0) + w * (1.0 / (k + rank))
|
| 349 |
+
if pid not in paper_map:
|
| 350 |
+
paper_map[pid] = paper
|
| 351 |
+
|
| 352 |
+
sorted_pids = sorted(scores, key=lambda p: scores[p], reverse=True)
|
| 353 |
+
return [paper_map[pid] for pid in sorted_pids]
|
| 354 |
+
|
| 355 |
+
def _retrieve_for_sub_query(
|
| 356 |
+
self, sub_q: str, use_topic: bool = False, use_cites: bool = False
|
| 357 |
+
) -> list[dict]:
|
| 358 |
+
"""Return candidate_pool from hybrid search + optional graph expansion.
|
| 359 |
+
|
| 360 |
+
Qdrant handles dense + BM25 hybrid retrieval (same as Pipeline 2).
|
| 361 |
+
TigerGraph adds citation and/or topic expansion only when the query
|
| 362 |
+
decomposition layer determined those signals are relevant.
|
| 363 |
+
"""
|
| 364 |
+
qdrant_results = self._normalize_ids(self.retriever.search(sub_q, top_k=50))
|
| 365 |
+
|
| 366 |
+
if not use_topic and not use_cites:
|
| 367 |
+
return qdrant_results
|
| 368 |
+
|
| 369 |
+
seed_ids = [p["paper_id"] for p in qdrant_results[:20] if p.get("paper_id")]
|
| 370 |
+
|
| 371 |
+
with ThreadPoolExecutor(max_workers=2) as pool:
|
| 372 |
+
cites_fut = (
|
| 373 |
+
pool.submit(self._expand_neighborhood, seed_ids) if use_cites else None
|
| 374 |
+
)
|
| 375 |
+
topic_fut = (
|
| 376 |
+
pool.submit(self._topic_search, self._extract_topic_keywords(sub_q), 20)
|
| 377 |
+
if use_topic else None
|
| 378 |
+
)
|
| 379 |
+
cites_results = cites_fut.result() if cites_fut else []
|
| 380 |
+
topic_results = topic_fut.result() if topic_fut else []
|
| 381 |
+
|
| 382 |
+
# Prioritize Vector Search (2.0) over Graph Expansions (1.0)
|
| 383 |
+
return self._rrf_merge(
|
| 384 |
+
qdrant_results, topic_results, cites_results, weights=[4.0, 1.0, 1.0]
|
| 385 |
+
)
|
| 386 |
+
|
| 387 |
+
|
| 388 |
+
# ------------------------------------------------------------------
|
| 389 |
+
# Reranking
|
| 390 |
+
# ------------------------------------------------------------------
|
| 391 |
+
def _rerank(self, query: str, candidates: list[dict], top_n: int = 10) -> list[dict]:
|
| 392 |
+
if not candidates:
|
| 393 |
+
return []
|
| 394 |
+
pairs = [
|
| 395 |
+
[query, f"{c.get('title', '')}\n{c.get('abstract', '')}"]
|
| 396 |
+
for c in candidates
|
| 397 |
+
]
|
| 398 |
+
scores = self.reranker.predict(pairs, batch_size=16)
|
| 399 |
+
ranked = sorted(zip(scores, candidates), key=lambda x: x[0], reverse=True)
|
| 400 |
+
return [c for _, c in ranked[:top_n]]
|
| 401 |
+
|
| 402 |
+
# ------------------------------------------------------------------
|
| 403 |
+
# Structure-aware passage extraction
|
| 404 |
+
# ------------------------------------------------------------------
|
| 405 |
+
def _extract_relevant_passages(
|
| 406 |
+
self, full_text: str, query: str, paper_title: str = "",
|
| 407 |
+
top_k: int = 8, target_tokens: int = 600,
|
| 408 |
+
sub_queries: list[str] | None = None,
|
| 409 |
+
) -> str:
|
| 410 |
+
sentences = re.split(r'(?<=[.!?])\s+(?=[A-Z])', full_text)
|
| 411 |
+
if not sentences:
|
| 412 |
+
return full_text
|
| 413 |
+
|
| 414 |
+
passages: list[str] = []
|
| 415 |
+
current: list[str] = []
|
| 416 |
+
current_len = 0
|
| 417 |
+
for sent in sentences:
|
| 418 |
+
current.append(sent)
|
| 419 |
+
current_len += len(sent)
|
| 420 |
+
if current_len >= target_tokens:
|
| 421 |
+
passages.append(" ".join(current))
|
| 422 |
+
current = [current[-1]] if len(current) > 1 else []
|
| 423 |
+
current_len = len(current[0]) if current else 0
|
| 424 |
+
if current:
|
| 425 |
+
passages.append(" ".join(current))
|
| 426 |
+
|
| 427 |
+
if len(passages) <= top_k:
|
| 428 |
+
return full_text
|
| 429 |
+
|
| 430 |
+
# Cap to avoid slow CPU encoding on very long papers
|
| 431 |
+
MAX_PASSAGES = 50
|
| 432 |
+
if len(passages) > MAX_PASSAGES:
|
| 433 |
+
step = len(passages) / MAX_PASSAGES
|
| 434 |
+
passages = [passages[int(i * step)] for i in range(MAX_PASSAGES)]
|
| 435 |
+
|
| 436 |
+
# Multi-query dense selection: each sub-query gets its own top slots.
|
| 437 |
+
# Uses fast bge-small (10x faster than bge-large on CPU, same relative accuracy
|
| 438 |
+
# for within-document passage selection where only relative scores matter).
|
| 439 |
+
search_queries = sub_queries if sub_queries else [query]
|
| 440 |
+
n_q = len(search_queries)
|
| 441 |
+
per_q = min(max(4, top_k // n_q + 2), len(passages))
|
| 442 |
+
|
| 443 |
+
all_texts = search_queries + passages
|
| 444 |
+
embeddings = self.fast_dense.encode(
|
| 445 |
+
all_texts, normalize_embeddings=True, show_progress_bar=False, batch_size=64
|
| 446 |
+
)
|
| 447 |
+
query_vecs = embeddings[:n_q]
|
| 448 |
+
passage_vecs = embeddings[n_q:]
|
| 449 |
+
|
| 450 |
+
best_score: dict[int, float] = {}
|
| 451 |
+
for q_vec in query_vecs:
|
| 452 |
+
scores = passage_vecs @ q_vec
|
| 453 |
+
for idx in scores.argsort()[-per_q:][::-1]:
|
| 454 |
+
best_score[int(idx)] = max(best_score.get(int(idx), -1.0), float(scores[idx]))
|
| 455 |
+
|
| 456 |
+
top_indices = sorted(
|
| 457 |
+
idx for idx, _ in sorted(best_score.items(), key=lambda x: x[1], reverse=True)[:top_k]
|
| 458 |
+
)
|
| 459 |
+
|
| 460 |
+
prefix = f"[Source: {paper_title}]" if paper_title else ""
|
| 461 |
+
selected = []
|
| 462 |
+
for idx in top_indices:
|
| 463 |
+
passage_text = passages[idx].strip()
|
| 464 |
+
selected.append(f"{prefix}\n{passage_text}" if prefix else passage_text)
|
| 465 |
+
|
| 466 |
+
return "\n\n[...]\n\n".join(selected)
|
| 467 |
+
|
| 468 |
+
# ------------------------------------------------------------------
|
| 469 |
+
# Deduplication
|
| 470 |
+
# ------------------------------------------------------------------
|
| 471 |
+
def _deduplicate_papers(self, papers: list[dict], top_n: int) -> list[dict]:
|
| 472 |
+
seen: set[str] = set()
|
| 473 |
+
deduped = []
|
| 474 |
+
for p in papers:
|
| 475 |
+
pid = p.get("paper_id")
|
| 476 |
+
if pid and pid not in seen:
|
| 477 |
+
seen.add(pid)
|
| 478 |
+
deduped.append(p)
|
| 479 |
+
if len(deduped) >= top_n:
|
| 480 |
+
break
|
| 481 |
+
return deduped
|
| 482 |
+
|
| 483 |
+
# ------------------------------------------------------------------
|
| 484 |
+
# Sub-query coverage enforcement
|
| 485 |
+
# ------------------------------------------------------------------
|
| 486 |
+
def _ensure_sub_query_coverage(
|
| 487 |
+
self, top_papers: list[dict], sub_queries: list[str], top_n: int
|
| 488 |
+
) -> list[dict]:
|
| 489 |
+
if not top_papers:
|
| 490 |
+
return []
|
| 491 |
+
COVERAGE_THRESHOLD = 0.1
|
| 492 |
+
for sub_q in sub_queries:
|
| 493 |
+
pairs = [
|
| 494 |
+
[sub_q, f"{p.get('title', '')}\n{p.get('abstract', '')}"]
|
| 495 |
+
for p in top_papers
|
| 496 |
+
]
|
| 497 |
+
scores = self.reranker.predict(pairs)
|
| 498 |
+
if float(max(scores)) > COVERAGE_THRESHOLD:
|
| 499 |
+
continue
|
| 500 |
+
logger.info(f"Coverage gap for: {sub_q[:60]}... fetching targeted paper.")
|
| 501 |
+
extra = self._normalize_ids(self.retriever.search(sub_q, top_k=5))
|
| 502 |
+
reranked_extra = self._rerank(sub_q, extra, top_n=1)
|
| 503 |
+
if reranked_extra:
|
| 504 |
+
pid = reranked_extra[0].get("paper_id")
|
| 505 |
+
if pid not in {p.get("paper_id") for p in top_papers}:
|
| 506 |
+
insert_pos = min(_PIN_K, len(top_papers))
|
| 507 |
+
top_papers.insert(insert_pos, reranked_extra[0])
|
| 508 |
+
return self._deduplicate_papers(top_papers, top_n)
|
| 509 |
+
|
| 510 |
+
# ------------------------------------------------------------------
|
| 511 |
+
# Context assembly with parallel full-text fetch
|
| 512 |
+
# ------------------------------------------------------------------
|
| 513 |
+
def _build_context(
|
| 514 |
+
self, top_papers: list[dict], query: str, max_full_text: int = None,
|
| 515 |
+
sub_queries: list[str] | None = None,
|
| 516 |
+
) -> tuple[str, list[str], list[str]]:
|
| 517 |
+
if max_full_text is None:
|
| 518 |
+
max_full_text = self.max_full_text
|
| 519 |
+
papers_to_fetch = top_papers[:max_full_text]
|
| 520 |
+
remaining = top_papers[max_full_text:]
|
| 521 |
+
|
| 522 |
+
def _fetch(args: tuple[int, dict]) -> tuple[int, dict, str | None]:
|
| 523 |
+
i, paper = args
|
| 524 |
+
logger.info(f"Fetching full text: {paper.get('title', '')[:60]}")
|
| 525 |
+
return i, paper, self.fetcher.fetch_full_text(paper["paper_id"])
|
| 526 |
+
|
| 527 |
+
_FETCH_TIMEOUT = 30
|
| 528 |
+
fetch_results: dict[int, tuple[dict, str | None]] = {}
|
| 529 |
+
with ThreadPoolExecutor(max_workers=max(1, max_full_text)) as pool:
|
| 530 |
+
futures = {pool.submit(_fetch, args): args[0] for args in enumerate(papers_to_fetch)}
|
| 531 |
+
for future, idx in futures.items():
|
| 532 |
+
try:
|
| 533 |
+
i, paper, text = future.result(timeout=_FETCH_TIMEOUT)
|
| 534 |
+
fetch_results[i] = (paper, text)
|
| 535 |
+
except Exception as exc:
|
| 536 |
+
logger.warning(f"Full text fetch timed out or failed for paper {idx}: {exc}")
|
| 537 |
+
fetch_results[idx] = (papers_to_fetch[idx], None)
|
| 538 |
+
|
| 539 |
+
context_parts: list[str] = []
|
| 540 |
+
abstracts: list[str] = []
|
| 541 |
+
full_text_titles: list[str] = []
|
| 542 |
+
|
| 543 |
+
full_text_count = sum(
|
| 544 |
+
1 for i in range(len(papers_to_fetch))
|
| 545 |
+
if fetch_results[i][1] and len(fetch_results[i][1]) > 1000
|
| 546 |
+
)
|
| 547 |
+
# Match Pipeline 2's adaptive context logic
|
| 548 |
+
adaptive_top_k = max(6, 12 // max(1, full_text_count))
|
| 549 |
+
|
| 550 |
+
for i in range(len(papers_to_fetch)):
|
| 551 |
+
paper, full_text = fetch_results[i]
|
| 552 |
+
abstract = paper.get("abstract", "")
|
| 553 |
+
title = paper.get("title", "")
|
| 554 |
+
abstracts.append(abstract)
|
| 555 |
+
if full_text and len(full_text) > 1000:
|
| 556 |
+
# Process sequentially to avoid CPU thrashing on ML models
|
| 557 |
+
content = self._extract_relevant_passages(
|
| 558 |
+
full_text, query, paper_title=title, top_k=adaptive_top_k,
|
| 559 |
+
sub_queries=sub_queries,
|
| 560 |
+
)
|
| 561 |
+
context_parts.append(
|
| 562 |
+
f"--- PAPER {i} [{title}] (FULL TEXT, RELEVANT PASSAGES) ---\n{content}\n"
|
| 563 |
+
)
|
| 564 |
+
full_text_titles.append(title)
|
| 565 |
+
else:
|
| 566 |
+
context_parts.append(f"--- PAPER {i} [{title}] (ABSTRACT) ---\n{abstract}\n")
|
| 567 |
+
|
| 568 |
+
for j, paper in enumerate(remaining, start=len(papers_to_fetch)):
|
| 569 |
+
abstract = paper.get("abstract", "")
|
| 570 |
+
title = paper.get("title", "")
|
| 571 |
+
abstracts.append(abstract)
|
| 572 |
+
context_parts.append(f"--- PAPER {j} [{title}] (ABSTRACT) ---\n{abstract}\n")
|
| 573 |
+
|
| 574 |
+
if not full_text_titles:
|
| 575 |
+
logger.info("No full text in top papers β scanning remaining for fallback (max 3)...")
|
| 576 |
+
for paper in remaining[:3]:
|
| 577 |
+
title = paper.get("title", "")
|
| 578 |
+
try:
|
| 579 |
+
with ThreadPoolExecutor(max_workers=1) as _pool:
|
| 580 |
+
full_text = _pool.submit(
|
| 581 |
+
self.fetcher.fetch_full_text, paper["paper_id"]
|
| 582 |
+
).result(timeout=_FETCH_TIMEOUT)
|
| 583 |
+
except Exception:
|
| 584 |
+
full_text = None
|
| 585 |
+
if full_text and len(full_text) > 1000:
|
| 586 |
+
content = self._extract_relevant_passages(
|
| 587 |
+
full_text, query, paper_title=title, top_k=8,
|
| 588 |
+
sub_queries=sub_queries,
|
| 589 |
+
)
|
| 590 |
+
context_parts.append(
|
| 591 |
+
f"--- FALLBACK PAPER [{title}] (FULL TEXT, RELEVANT PASSAGES) ---\n{content}\n"
|
| 592 |
+
)
|
| 593 |
+
full_text_titles.append(title)
|
| 594 |
+
break
|
| 595 |
+
|
| 596 |
+
return "\n".join(context_parts), abstracts, full_text_titles
|
| 597 |
+
|
| 598 |
+
# ------------------------------------------------------------------
|
| 599 |
+
# Main entry point
|
| 600 |
+
# ------------------------------------------------------------------
|
| 601 |
+
def run(self, query: str, ground_truth: str = None) -> dict:
|
| 602 |
+
start = time.time()
|
| 603 |
+
logger.info(f"Pipeline 3 (GraphRAG) query: {query}")
|
| 604 |
+
|
| 605 |
+
decomposition = self._decompose_query(query)
|
| 606 |
+
use_topic = decomposition["use_topic"]
|
| 607 |
+
use_cites = decomposition["use_cites"]
|
| 608 |
+
|
| 609 |
+
# 1. Retrieval β always use the original query for Qdrant search.
|
| 610 |
+
# Sub-queries (used for passage extraction) are too vague/generic to be
|
| 611 |
+
# reliable Qdrant queries and introduce noise when used independently.
|
| 612 |
+
if use_topic or use_cites:
|
| 613 |
+
unique_candidates = self._retrieve_for_sub_query(query, use_topic, use_cites)
|
| 614 |
+
else:
|
| 615 |
+
unique_candidates = self._normalize_ids(self.retriever.search(query, top_k=50))
|
| 616 |
+
|
| 617 |
+
# 2. Final Single Rerank β cap at 50 to bound cross-encoder CPU time
|
| 618 |
+
top_papers = self._rerank(query, unique_candidates[:50], top_n=self.rerank_top_n)
|
| 619 |
+
|
| 620 |
+
# 3. Entity pinning: if an acronym is explicitly named in the query (CGCNN, DPMD, CFDβ¦),
|
| 621 |
+
# guarantee at least one matching paper in top_papers. The cross-encoder optimizes
|
| 622 |
+
# for global relevance and can demote a paper that is explicitly the subject of the query.
|
| 623 |
+
query_acronyms = list(dict.fromkeys(re.findall(r'\b[A-Z]{2,}\b', query)))
|
| 624 |
+
if query_acronyms:
|
| 625 |
+
top_ids = {p.get("paper_id") for p in top_papers}
|
| 626 |
+
for acronym in query_acronyms[:3]:
|
| 627 |
+
hits = self._normalize_ids(self.retriever.search(acronym, top_k=5))
|
| 628 |
+
if hits and hits[0].get("paper_id") not in top_ids:
|
| 629 |
+
logger.info(f"Pinning '{acronym}': {hits[0].get('title', '')[:60]}")
|
| 630 |
+
top_papers.insert(min(_PIN_K, len(top_papers)), hits[0])
|
| 631 |
+
top_papers = top_papers[:self.rerank_top_n]
|
| 632 |
+
top_ids.add(hits[0].get("paper_id"))
|
| 633 |
+
|
| 634 |
+
max_ft = self.max_full_text if decomposition["is_multi_hop"] else min(2, self.max_full_text)
|
| 635 |
+
full_context, abstracts, full_text_titles = self._build_context(
|
| 636 |
+
top_papers, query, max_full_text=max_ft,
|
| 637 |
+
sub_queries=decomposition["sub_queries"],
|
| 638 |
+
)
|
| 639 |
+
|
| 640 |
+
system_msg = (
|
| 641 |
+
"You are an expert AI research assistant specializing in scientific literature synthesis. "
|
| 642 |
+
"Answer the query using ONLY the provided paper context. Speak with absolute certainty.\n\n"
|
| 643 |
+
"Rules:\n"
|
| 644 |
+
"1. Answer directly and immediately. Do NOT use any preamble like 'Based on the provided papers', "
|
| 645 |
+
"'According to the context', or 'It can be inferred'. Start the response with the subject of the question.\n"
|
| 646 |
+
"2. Quote technical terms EXACTLY as they appear in the papers. "
|
| 647 |
+
"For any list of methods, data sources, properties, or named items: copy the exact words. "
|
| 648 |
+
"Do NOT substitute synonyms or rephrase β "
|
| 649 |
+
"e.g. 'large-scale simulations' must not become 'multi-scale simulations'; "
|
| 650 |
+
"'field measurements' must not become 'real data from the field'.\n"
|
| 651 |
+
"3. Be complete β address EVERY part of the question. For list questions, include "
|
| 652 |
+
"every item. For comparison questions, cover EACH paper explicitly.\n"
|
| 653 |
+
"4. Never add disclaimers, caveats, hedges, or uncertainty statements. "
|
| 654 |
+
"Do not state that information is missing or unclear. If a detail is missing, omit it silently.\n"
|
| 655 |
+
"5. Do not use square brackets or cite PAPER numbers in the final text.\n"
|
| 656 |
+
"6. Prefer concise, precise, and highly technical language."
|
| 657 |
+
)
|
| 658 |
+
user_msg = f"QUERY: {query}\n\nCONTEXT:\n{full_context}"
|
| 659 |
+
|
| 660 |
+
try:
|
| 661 |
+
response = self._groq_call(
|
| 662 |
+
model=self.model_name,
|
| 663 |
+
messages=[
|
| 664 |
+
{"role": "system", "content": system_msg},
|
| 665 |
+
{"role": "user", "content": user_msg},
|
| 666 |
+
],
|
| 667 |
+
temperature=0.1,
|
| 668 |
+
max_tokens=1024,
|
| 669 |
+
)
|
| 670 |
+
|
| 671 |
+
answer = response.choices[0].message.content
|
| 672 |
+
usage = _GroqUsage(
|
| 673 |
+
prompt_token_count=response.usage.prompt_tokens,
|
| 674 |
+
candidates_token_count=response.usage.completion_tokens,
|
| 675 |
+
)
|
| 676 |
+
|
| 677 |
+
stats = self.metrics.process_metrics(
|
| 678 |
+
client=None,
|
| 679 |
+
query=query,
|
| 680 |
+
answer=answer,
|
| 681 |
+
context=full_context,
|
| 682 |
+
usage_metadata=usage,
|
| 683 |
+
start_time=start,
|
| 684 |
+
abstracts_list=abstracts,
|
| 685 |
+
ground_truth=ground_truth,
|
| 686 |
+
)
|
| 687 |
+
|
| 688 |
+
return {
|
| 689 |
+
"answer": answer,
|
| 690 |
+
"sources": [p.get("title", "") for p in top_papers],
|
| 691 |
+
"metrics": stats,
|
| 692 |
+
"selected_papers": full_text_titles,
|
| 693 |
+
"query_decomposition": decomposition,
|
| 694 |
+
}
|
| 695 |
+
|
| 696 |
+
except Exception as e:
|
| 697 |
+
logger.error(f"Pipeline 3 error: {e}")
|
| 698 |
+
return {"error": str(e), "answer": "Error generating response."}
|
core/pipeline_3/queries.py
ADDED
|
@@ -0,0 +1,141 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""GSQL query definitions and installation helpers for PaperGraph."""
|
| 2 |
+
import logging
|
| 3 |
+
import sys
|
| 4 |
+
from pathlib import Path
|
| 5 |
+
|
| 6 |
+
import requests
|
| 7 |
+
|
| 8 |
+
sys.path.append(str(Path(__file__).parents[2]))
|
| 9 |
+
|
| 10 |
+
logger = logging.getLogger(__name__)
|
| 11 |
+
|
| 12 |
+
import os
|
| 13 |
+
|
| 14 |
+
TG_HOST = os.environ.get("TG_HOST", "http://localhost")
|
| 15 |
+
TG_PORT = int(os.environ.get("TG_PORT", "14240"))
|
| 16 |
+
TG_USER = os.environ.get("TG_USER", "tigergraph")
|
| 17 |
+
TG_PASS = os.environ.get("TG_PASS", "tigergraph")
|
| 18 |
+
GRAPH = "PaperGraph"
|
| 19 |
+
|
| 20 |
+
# Q1: Fetch Paper vertices by ID set β used for post-search attribute lookup
|
| 21 |
+
QUERY_GET_PAPERS_BY_IDS = """\
|
| 22 |
+
CREATE OR REPLACE QUERY get_papers_by_ids(
|
| 23 |
+
SET<STRING> paper_ids
|
| 24 |
+
) FOR GRAPH PaperGraph SYNTAX V3 {
|
| 25 |
+
result = SELECT p FROM Paper:p WHERE p.paper_id IN paper_ids;
|
| 26 |
+
PRINT result[result.paper_id, result.title, result.abstract,
|
| 27 |
+
result.year, result.cited_by_count, result.doi, result.pdf_url];
|
| 28 |
+
}
|
| 29 |
+
"""
|
| 30 |
+
|
| 31 |
+
# Q2: 1-hop expansion via CITES (undirected form required for TG CE)
|
| 32 |
+
QUERY_EXPAND_NEIGHBORHOOD = """\
|
| 33 |
+
CREATE OR REPLACE QUERY expand_paper_neighborhood(
|
| 34 |
+
SET<STRING> seed_ids
|
| 35 |
+
) FOR GRAPH PaperGraph SYNTAX V3 {
|
| 36 |
+
seeds = SELECT p FROM Paper:p WHERE p.paper_id IN seed_ids;
|
| 37 |
+
neighbors = SELECT t FROM seeds:s -[e:CITES]- Paper:t;
|
| 38 |
+
result = seeds UNION neighbors;
|
| 39 |
+
PRINT result[result.paper_id, result.title, result.abstract,
|
| 40 |
+
result.year, result.cited_by_count, result.doi, result.pdf_url];
|
| 41 |
+
}
|
| 42 |
+
"""
|
| 43 |
+
|
| 44 |
+
# Q3: Topic entity linking β papers connected via HAS_TOPIC edges (score-weighted)
|
| 45 |
+
QUERY_TOPIC_SEARCH = """\
|
| 46 |
+
CREATE OR REPLACE QUERY search_papers_by_topic(
|
| 47 |
+
STRING pattern, INT top_k = 20
|
| 48 |
+
) FOR GRAPH PaperGraph SYNTAX V3 {
|
| 49 |
+
MaxAccum<FLOAT> @topicScore;
|
| 50 |
+
topics = SELECT t FROM Topic:t
|
| 51 |
+
WHERE t.display_name LIKE pattern
|
| 52 |
+
OR t.subfield LIKE pattern
|
| 53 |
+
OR t.field LIKE pattern;
|
| 54 |
+
papers = SELECT p FROM Paper:p -[ht:HAS_TOPIC]- topics:t
|
| 55 |
+
ACCUM p.@topicScore += ht.score
|
| 56 |
+
ORDER BY p.@topicScore DESC
|
| 57 |
+
LIMIT top_k;
|
| 58 |
+
PRINT papers[papers.paper_id, papers.title, papers.abstract,
|
| 59 |
+
papers.year, papers.cited_by_count, papers.doi, papers.pdf_url];
|
| 60 |
+
}
|
| 61 |
+
"""
|
| 62 |
+
|
| 63 |
+
ALL_QUERIES = {
|
| 64 |
+
"get_papers_by_ids": QUERY_GET_PAPERS_BY_IDS,
|
| 65 |
+
"expand_paper_neighborhood": QUERY_EXPAND_NEIGHBORHOOD,
|
| 66 |
+
"search_papers_by_topic": QUERY_TOPIC_SEARCH,
|
| 67 |
+
}
|
| 68 |
+
|
| 69 |
+
|
| 70 |
+
def _gsql(stmt: str) -> str:
|
| 71 |
+
resp = requests.post(
|
| 72 |
+
f"{TG_HOST}:{TG_PORT}/gsql/v1/statements",
|
| 73 |
+
data=stmt.encode("utf-8"),
|
| 74 |
+
headers={"Content-Type": "text/plain"},
|
| 75 |
+
auth=(TG_USER, TG_PASS),
|
| 76 |
+
timeout=120,
|
| 77 |
+
)
|
| 78 |
+
return resp.text
|
| 79 |
+
|
| 80 |
+
|
| 81 |
+
def _query_is_installed(name: str) -> bool:
|
| 82 |
+
"""Return True if the named query is already installed.
|
| 83 |
+
|
| 84 |
+
Uses GSQL SHOW QUERY β authoritative and fast (~1s vs ~60s for compile).
|
| 85 |
+
"""
|
| 86 |
+
try:
|
| 87 |
+
result = _gsql(f"USE GRAPH {GRAPH}\nSHOW QUERY {name}")
|
| 88 |
+
return '"error":true' not in result and "does not exist" not in result.lower()
|
| 89 |
+
except Exception:
|
| 90 |
+
return False
|
| 91 |
+
|
| 92 |
+
|
| 93 |
+
def install_queries() -> None:
|
| 94 |
+
"""Create and compile all GSQL queries into PaperGraph. Skips already-installed queries."""
|
| 95 |
+
for name, gsql in ALL_QUERIES.items():
|
| 96 |
+
if _query_is_installed(name):
|
| 97 |
+
logger.info(f" {name}: already installed, skipping.")
|
| 98 |
+
continue
|
| 99 |
+
logger.info(f"Installing GSQL query: {name}")
|
| 100 |
+
result = _gsql(f"USE GRAPH {GRAPH}\n{gsql}\nINSTALL QUERY {name}")
|
| 101 |
+
result_lower = result.lower()
|
| 102 |
+
already = "already" in result_lower
|
| 103 |
+
failed = (
|
| 104 |
+
not already and (
|
| 105 |
+
'"error":true' in result
|
| 106 |
+
or "installation failed" in result_lower
|
| 107 |
+
or ("draft" in result_lower and "installed" not in result_lower)
|
| 108 |
+
)
|
| 109 |
+
)
|
| 110 |
+
if failed:
|
| 111 |
+
logger.warning(f" {name}: {result[:300]}")
|
| 112 |
+
else:
|
| 113 |
+
logger.info(f" {name}: OK")
|
| 114 |
+
|
| 115 |
+
|
| 116 |
+
def parse_paper_results(response: list) -> list[dict]:
|
| 117 |
+
"""Extract flat paper attribute dicts from runInstalledQuery response.
|
| 118 |
+
|
| 119 |
+
pyTigerGraph wraps PRINT output as:
|
| 120 |
+
[{"<var_name>": [{"v_id": "...", "v_type": "Paper", "attributes": {...}}]}]
|
| 121 |
+
"""
|
| 122 |
+
papers = []
|
| 123 |
+
for item in response:
|
| 124 |
+
if not isinstance(item, dict):
|
| 125 |
+
continue
|
| 126 |
+
for val in item.values():
|
| 127 |
+
if not isinstance(val, list):
|
| 128 |
+
continue
|
| 129 |
+
for v in val:
|
| 130 |
+
if isinstance(v, dict) and "attributes" in v:
|
| 131 |
+
attrs = dict(v["attributes"])
|
| 132 |
+
if not attrs.get("paper_id") and isinstance(v.get("v_id"), str):
|
| 133 |
+
attrs["paper_id"] = v["v_id"]
|
| 134 |
+
papers.append(attrs)
|
| 135 |
+
return papers
|
| 136 |
+
|
| 137 |
+
|
| 138 |
+
if __name__ == "__main__":
|
| 139 |
+
logging.basicConfig(level=logging.INFO)
|
| 140 |
+
install_queries()
|
| 141 |
+
logger.info("All queries installed.")
|
core/pipeline_3/setup.py
ADDED
|
@@ -0,0 +1,459 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""One-time schema creation + data ingestion for Pipeline 3 (TigerGraph GraphRAG).
|
| 2 |
+
|
| 3 |
+
Run once before starting the server:
|
| 4 |
+
python core/pipeline_3/setup.py
|
| 5 |
+
|
| 6 |
+
Requires TigerGraph community container running on localhost:14240:
|
| 7 |
+
docker run -d --init -p 14240:14240 --name tigergraph tigergraph/community:4.2.2
|
| 8 |
+
|
| 9 |
+
Safe to re-run after interruption β resumes from last completed checkpoint.
|
| 10 |
+
"""
|
| 11 |
+
import json
|
| 12 |
+
import logging
|
| 13 |
+
import os
|
| 14 |
+
import requests
|
| 15 |
+
from pathlib import Path
|
| 16 |
+
from typing import Optional
|
| 17 |
+
|
| 18 |
+
import numpy as np
|
| 19 |
+
import pyTigerGraph as tg
|
| 20 |
+
from dotenv import load_dotenv
|
| 21 |
+
from sentence_transformers import SentenceTransformer
|
| 22 |
+
|
| 23 |
+
load_dotenv()
|
| 24 |
+
|
| 25 |
+
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
| 26 |
+
logger = logging.getLogger(__name__)
|
| 27 |
+
|
| 28 |
+
# ββ Config βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 29 |
+
TG_HOST = os.environ.get("TG_HOST", "http://localhost")
|
| 30 |
+
TG_PORT = int(os.environ.get("TG_PORT", "14240"))
|
| 31 |
+
TG_USER = os.environ.get("TG_USER", "tigergraph")
|
| 32 |
+
TG_PASS = os.environ.get("TG_PASS", "tigergraph")
|
| 33 |
+
TG_SECRET = os.environ.get("TG_SECRET", "")
|
| 34 |
+
GRAPH = "PaperGraph"
|
| 35 |
+
EMBED_MODEL = "BAAI/bge-large-en-v1.5"
|
| 36 |
+
EMBED_DIM = 1024
|
| 37 |
+
EMBED_BATCH = 64
|
| 38 |
+
UPSERT_BATCH = 256
|
| 39 |
+
|
| 40 |
+
_PROJECT_ROOT = Path(__file__).parents[2]
|
| 41 |
+
_DATA_DIR = _PROJECT_ROOT / "data"
|
| 42 |
+
_PIPELINE_DIR = Path(__file__).parent
|
| 43 |
+
_EMBED_CACHE = _PIPELINE_DIR / "embeddings.npy"
|
| 44 |
+
_EMBED_PARTIAL = _PIPELINE_DIR / "embed_partials"
|
| 45 |
+
_CHECKPOINT = _PIPELINE_DIR / "checkpoint.json"
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
# ββ TigerGraph connection ββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 49 |
+
def _get_conn() -> tg.TigerGraphConnection:
|
| 50 |
+
if TG_SECRET:
|
| 51 |
+
conn = tg.TigerGraphConnection(
|
| 52 |
+
host=TG_HOST,
|
| 53 |
+
graphname=GRAPH,
|
| 54 |
+
gsqlSecret=TG_SECRET,
|
| 55 |
+
restppPort=str(TG_PORT),
|
| 56 |
+
gsPort=str(TG_PORT),
|
| 57 |
+
)
|
| 58 |
+
conn.getToken(TG_SECRET)
|
| 59 |
+
else:
|
| 60 |
+
conn = tg.TigerGraphConnection(
|
| 61 |
+
host=TG_HOST,
|
| 62 |
+
graphname=GRAPH,
|
| 63 |
+
username=TG_USER,
|
| 64 |
+
password=TG_PASS,
|
| 65 |
+
restppPort=str(TG_PORT),
|
| 66 |
+
gsPort=str(TG_PORT),
|
| 67 |
+
)
|
| 68 |
+
return conn
|
| 69 |
+
|
| 70 |
+
|
| 71 |
+
# ββ Checkpoint helpers βββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 72 |
+
def _load_checkpoint() -> dict:
|
| 73 |
+
if _CHECKPOINT.exists():
|
| 74 |
+
ckpt = json.loads(_CHECKPOINT.read_text())
|
| 75 |
+
ckpt.setdefault("embed_batches", 0)
|
| 76 |
+
return ckpt
|
| 77 |
+
return {
|
| 78 |
+
"embed_batches": 0,
|
| 79 |
+
"papers": 0,
|
| 80 |
+
"authors": False,
|
| 81 |
+
"topics": False,
|
| 82 |
+
"authored_by": False,
|
| 83 |
+
"has_topic": False,
|
| 84 |
+
"cites": False,
|
| 85 |
+
}
|
| 86 |
+
|
| 87 |
+
|
| 88 |
+
def _save_checkpoint(ckpt: dict) -> None:
|
| 89 |
+
_CHECKPOINT.write_text(json.dumps(ckpt, indent=2))
|
| 90 |
+
|
| 91 |
+
|
| 92 |
+
# ββ GSQL helper (raw REST β confirmed working with TG 4.2.2) ββββββββββββββββββ
|
| 93 |
+
def _gsql(stmt: str) -> str:
|
| 94 |
+
resp = requests.post(
|
| 95 |
+
f"{TG_HOST}:{TG_PORT}/gsql/v1/statements",
|
| 96 |
+
data=stmt.encode("utf-8"),
|
| 97 |
+
headers={"Content-Type": "text/plain"},
|
| 98 |
+
auth=(TG_USER, TG_PASS),
|
| 99 |
+
timeout=120,
|
| 100 |
+
)
|
| 101 |
+
return resp.text
|
| 102 |
+
|
| 103 |
+
|
| 104 |
+
def _graph_ready() -> bool:
|
| 105 |
+
try:
|
| 106 |
+
resp = requests.get(
|
| 107 |
+
f"{TG_HOST}:{TG_PORT}/restpp/graph/{GRAPH}/vertices/Paper",
|
| 108 |
+
params={"limit": 1},
|
| 109 |
+
auth=(TG_USER, TG_PASS),
|
| 110 |
+
timeout=10,
|
| 111 |
+
)
|
| 112 |
+
return resp.status_code == 200
|
| 113 |
+
except Exception:
|
| 114 |
+
return False
|
| 115 |
+
|
| 116 |
+
|
| 117 |
+
# ββ Schema βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 118 |
+
_SCHEMA = f"""
|
| 119 |
+
CREATE GRAPH {GRAPH}()
|
| 120 |
+
USE GRAPH {GRAPH}
|
| 121 |
+
|
| 122 |
+
CREATE SCHEMA_CHANGE JOB init_schema FOR GRAPH {GRAPH} {{
|
| 123 |
+
ADD VERTEX Paper(
|
| 124 |
+
PRIMARY_ID paper_id STRING,
|
| 125 |
+
title STRING DEFAULT "",
|
| 126 |
+
abstract STRING DEFAULT "",
|
| 127 |
+
year INT DEFAULT 0,
|
| 128 |
+
cited_by_count INT DEFAULT 0,
|
| 129 |
+
doi STRING DEFAULT "",
|
| 130 |
+
pdf_url STRING DEFAULT "",
|
| 131 |
+
venue STRING DEFAULT "",
|
| 132 |
+
corresponding_author STRING DEFAULT "",
|
| 133 |
+
authors STRING DEFAULT "",
|
| 134 |
+
funders STRING DEFAULT ""
|
| 135 |
+
) WITH primary_id_as_attribute="true";
|
| 136 |
+
|
| 137 |
+
ADD VERTEX Author(
|
| 138 |
+
PRIMARY_ID author_id STRING,
|
| 139 |
+
display_name STRING DEFAULT ""
|
| 140 |
+
) WITH primary_id_as_attribute="true";
|
| 141 |
+
|
| 142 |
+
ADD VERTEX Topic(
|
| 143 |
+
PRIMARY_ID topic_id STRING,
|
| 144 |
+
display_name STRING DEFAULT "",
|
| 145 |
+
subfield STRING DEFAULT "",
|
| 146 |
+
field STRING DEFAULT "",
|
| 147 |
+
domain STRING DEFAULT ""
|
| 148 |
+
) WITH primary_id_as_attribute="true";
|
| 149 |
+
|
| 150 |
+
ADD DIRECTED EDGE AUTHORED_BY(FROM Paper, TO Author);
|
| 151 |
+
ADD DIRECTED EDGE HAS_TOPIC(FROM Paper, TO Topic, score FLOAT DEFAULT 0.0);
|
| 152 |
+
ADD DIRECTED EDGE CITES(FROM Paper, TO Paper);
|
| 153 |
+
}}
|
| 154 |
+
RUN SCHEMA_CHANGE JOB init_schema
|
| 155 |
+
DROP JOB init_schema
|
| 156 |
+
|
| 157 |
+
CREATE SCHEMA_CHANGE JOB add_embedding FOR GRAPH {GRAPH} {{
|
| 158 |
+
ALTER VERTEX Paper ADD VECTOR ATTRIBUTE embedding(dimension={EMBED_DIM});
|
| 159 |
+
}}
|
| 160 |
+
RUN SCHEMA_CHANGE JOB add_embedding
|
| 161 |
+
DROP JOB add_embedding
|
| 162 |
+
"""
|
| 163 |
+
|
| 164 |
+
|
| 165 |
+
def create_schema() -> None:
|
| 166 |
+
if _graph_ready():
|
| 167 |
+
logger.info("Graph already exists β skipping schema creation.")
|
| 168 |
+
return
|
| 169 |
+
logger.info("Creating graph schemaβ¦")
|
| 170 |
+
out = _gsql(_SCHEMA)
|
| 171 |
+
logger.info(f"Schema output:\n{out[:800]}")
|
| 172 |
+
if "error" in out.lower() or "fail" in out.lower():
|
| 173 |
+
raise RuntimeError(f"Schema creation failed:\n{out}")
|
| 174 |
+
|
| 175 |
+
|
| 176 |
+
# ββ Embedding with per-batch checkpointing βββββββββββββββββββββββββββββββββββββ
|
| 177 |
+
def _embed_papers(papers: list[dict], ckpt: dict) -> np.ndarray:
|
| 178 |
+
if _EMBED_CACHE.exists():
|
| 179 |
+
logger.info(f"Loading full embedding cache from {_EMBED_CACHE}β¦")
|
| 180 |
+
embs = np.load(_EMBED_CACHE)
|
| 181 |
+
logger.info(f"Loaded embeddings shape: {embs.shape}")
|
| 182 |
+
return embs
|
| 183 |
+
|
| 184 |
+
_EMBED_PARTIAL.mkdir(exist_ok=True)
|
| 185 |
+
texts = [f"Title: {p['title']}\n\nAbstract: {p['abstract']}" for p in papers]
|
| 186 |
+
n_batches = (len(texts) + EMBED_BATCH - 1) // EMBED_BATCH
|
| 187 |
+
completed = ckpt.get("embed_batches", 0)
|
| 188 |
+
|
| 189 |
+
logger.info(f"Loading embedding model {EMBED_MODEL}β¦")
|
| 190 |
+
model = SentenceTransformer(EMBED_MODEL)
|
| 191 |
+
logger.info(f"Embedding: {completed}/{n_batches} batches already done, resumingβ¦")
|
| 192 |
+
|
| 193 |
+
for i in range(completed, n_batches):
|
| 194 |
+
start = i * EMBED_BATCH
|
| 195 |
+
batch_embs = model.encode(
|
| 196 |
+
texts[start : start + EMBED_BATCH],
|
| 197 |
+
normalize_embeddings=True,
|
| 198 |
+
)
|
| 199 |
+
np.save(_EMBED_PARTIAL / f"batch_{i:04d}.npy", batch_embs)
|
| 200 |
+
ckpt["embed_batches"] = i + 1
|
| 201 |
+
_save_checkpoint(ckpt)
|
| 202 |
+
if (i + 1) % 10 == 0 or i + 1 == n_batches:
|
| 203 |
+
logger.info(f" Embedded {i + 1}/{n_batches} batches β")
|
| 204 |
+
|
| 205 |
+
# Merge partials into single cache file
|
| 206 |
+
all_embs = [np.load(_EMBED_PARTIAL / f"batch_{i:04d}.npy") for i in range(n_batches)]
|
| 207 |
+
embeddings = np.vstack(all_embs)
|
| 208 |
+
np.save(_EMBED_CACHE, embeddings)
|
| 209 |
+
logger.info(f"Full embeddings saved β {_EMBED_CACHE}")
|
| 210 |
+
|
| 211 |
+
for f in sorted(_EMBED_PARTIAL.glob("*.npy")):
|
| 212 |
+
f.unlink()
|
| 213 |
+
_EMBED_PARTIAL.rmdir()
|
| 214 |
+
return embeddings
|
| 215 |
+
|
| 216 |
+
|
| 217 |
+
# ββ Data parsing βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 218 |
+
def _reconstruct_abstract(inverted: dict) -> str:
|
| 219 |
+
if not inverted:
|
| 220 |
+
return ""
|
| 221 |
+
max_idx = max((max(v) for v in inverted.values() if v), default=0)
|
| 222 |
+
words = [""] * (max_idx + 1)
|
| 223 |
+
for word, positions in inverted.items():
|
| 224 |
+
for i in positions:
|
| 225 |
+
words[i] = word
|
| 226 |
+
return " ".join(words).strip()
|
| 227 |
+
|
| 228 |
+
|
| 229 |
+
def _parse_paper(path: Path) -> Optional[dict]:
|
| 230 |
+
try:
|
| 231 |
+
d = json.loads(path.read_text())
|
| 232 |
+
paper_id = path.stem
|
| 233 |
+
|
| 234 |
+
abstract = _reconstruct_abstract(d.get("abstract_inverted_index") or {})
|
| 235 |
+
|
| 236 |
+
loc = d.get("primary_location") or {}
|
| 237 |
+
source = loc.get("source") or {}
|
| 238 |
+
venue = source.get("display_name", "") or ""
|
| 239 |
+
|
| 240 |
+
authorships = d.get("authorships") or []
|
| 241 |
+
author_names = [
|
| 242 |
+
a["author"]["display_name"]
|
| 243 |
+
for a in authorships[:5]
|
| 244 |
+
if (a.get("author") or {}).get("display_name")
|
| 245 |
+
]
|
| 246 |
+
corresponding_author = next(
|
| 247 |
+
(
|
| 248 |
+
a["author"].get("display_name", "")
|
| 249 |
+
for a in authorships
|
| 250 |
+
if a.get("is_corresponding") and a.get("author")
|
| 251 |
+
),
|
| 252 |
+
author_names[0] if author_names else "",
|
| 253 |
+
)
|
| 254 |
+
|
| 255 |
+
funder_names = list({
|
| 256 |
+
a["funder_display_name"]
|
| 257 |
+
for a in (d.get("awards") or [])
|
| 258 |
+
if a.get("funder_display_name")
|
| 259 |
+
})
|
| 260 |
+
|
| 261 |
+
best_oa = d.get("best_oa_location") or {}
|
| 262 |
+
|
| 263 |
+
topics = [
|
| 264 |
+
{
|
| 265 |
+
"topic_id": t["id"].split("/")[-1],
|
| 266 |
+
"display_name": t.get("display_name", ""),
|
| 267 |
+
"subfield": (t.get("subfield") or {}).get("display_name", ""),
|
| 268 |
+
"field": (t.get("field") or {}).get("display_name", ""),
|
| 269 |
+
"domain": (t.get("domain") or {}).get("display_name", ""),
|
| 270 |
+
"score": float(t.get("score", 0.0)),
|
| 271 |
+
}
|
| 272 |
+
for t in (d.get("topics") or [])[:3]
|
| 273 |
+
if t.get("id")
|
| 274 |
+
]
|
| 275 |
+
|
| 276 |
+
author_vertices = [
|
| 277 |
+
{
|
| 278 |
+
"author_id": a["author"]["id"].split("/")[-1],
|
| 279 |
+
"display_name": a["author"].get("display_name", ""),
|
| 280 |
+
}
|
| 281 |
+
for a in authorships
|
| 282 |
+
if (a.get("author") or {}).get("id")
|
| 283 |
+
]
|
| 284 |
+
|
| 285 |
+
referenced = [r.split("/")[-1] for r in (d.get("referenced_works") or [])]
|
| 286 |
+
|
| 287 |
+
return {
|
| 288 |
+
"paper_id": paper_id,
|
| 289 |
+
"title": d.get("title", "") or "",
|
| 290 |
+
"abstract": abstract,
|
| 291 |
+
"year": d.get("publication_year") or 0,
|
| 292 |
+
"cited_by_count": d.get("cited_by_count", 0) or 0,
|
| 293 |
+
"doi": d.get("doi", "") or "",
|
| 294 |
+
"pdf_url": best_oa.get("pdf_url", "") or "",
|
| 295 |
+
"venue": venue,
|
| 296 |
+
"corresponding_author": corresponding_author,
|
| 297 |
+
"authors": ", ".join(author_names),
|
| 298 |
+
"funders": ", ".join(funder_names),
|
| 299 |
+
"topics": topics,
|
| 300 |
+
"author_vertices": author_vertices,
|
| 301 |
+
"referenced": referenced,
|
| 302 |
+
}
|
| 303 |
+
except Exception as e:
|
| 304 |
+
logger.warning(f"Skipping {path.name}: {e}")
|
| 305 |
+
return None
|
| 306 |
+
|
| 307 |
+
|
| 308 |
+
# ββ Ingestion ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 309 |
+
def ingest(data_dir: Path = _DATA_DIR) -> None:
|
| 310 |
+
ckpt = _load_checkpoint()
|
| 311 |
+
logger.info(f"Resuming from checkpoint: {ckpt}")
|
| 312 |
+
|
| 313 |
+
files = sorted(data_dir.glob("*.json"))
|
| 314 |
+
corpus_ids = {f.stem for f in files}
|
| 315 |
+
logger.info(f"Found {len(files)} paper files.")
|
| 316 |
+
|
| 317 |
+
papers = [p for f in files if (p := _parse_paper(f))]
|
| 318 |
+
logger.info(f"Parsed {len(papers)} papers.")
|
| 319 |
+
|
| 320 |
+
embeddings = _embed_papers(papers, ckpt)
|
| 321 |
+
conn = _get_conn()
|
| 322 |
+
total_batches = (len(papers) + UPSERT_BATCH - 1) // UPSERT_BATCH
|
| 323 |
+
|
| 324 |
+
# Paper vertices βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 325 |
+
completed = ckpt["papers"]
|
| 326 |
+
if completed < total_batches:
|
| 327 |
+
logger.info(f"Upserting Paper vertices β resuming from batch {completed}/{total_batches}β¦")
|
| 328 |
+
for batch_idx in range(completed, total_batches):
|
| 329 |
+
start = batch_idx * UPSERT_BATCH
|
| 330 |
+
chunk = papers[start : start + UPSERT_BATCH]
|
| 331 |
+
embs = embeddings[start : start + UPSERT_BATCH]
|
| 332 |
+
conn.upsertVertices("Paper", [
|
| 333 |
+
(p["paper_id"], {
|
| 334 |
+
"title": p["title"],
|
| 335 |
+
"abstract": p["abstract"],
|
| 336 |
+
"year": p["year"],
|
| 337 |
+
"cited_by_count": p["cited_by_count"],
|
| 338 |
+
"doi": p["doi"],
|
| 339 |
+
"pdf_url": p["pdf_url"],
|
| 340 |
+
"venue": p["venue"],
|
| 341 |
+
"corresponding_author": p["corresponding_author"],
|
| 342 |
+
"authors": p["authors"],
|
| 343 |
+
"funders": p["funders"],
|
| 344 |
+
"embedding": emb.tolist(),
|
| 345 |
+
})
|
| 346 |
+
for p, emb in zip(chunk, embs)
|
| 347 |
+
])
|
| 348 |
+
ckpt["papers"] = batch_idx + 1
|
| 349 |
+
_save_checkpoint(ckpt)
|
| 350 |
+
logger.info(f" Papers batch {batch_idx + 1}/{total_batches} β")
|
| 351 |
+
else:
|
| 352 |
+
logger.info("Paper vertices already complete β skipping.")
|
| 353 |
+
|
| 354 |
+
# Collect Author + Topic data ββββββββββββββββββββββββββββββββββββββββββββββ
|
| 355 |
+
all_authors: dict[str, dict] = {}
|
| 356 |
+
all_topics: dict[str, dict] = {}
|
| 357 |
+
for p in papers:
|
| 358 |
+
for av in p["author_vertices"]:
|
| 359 |
+
all_authors.setdefault(av["author_id"], av)
|
| 360 |
+
for t in p["topics"]:
|
| 361 |
+
all_topics.setdefault(t["topic_id"], t)
|
| 362 |
+
|
| 363 |
+
# Author vertices ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 364 |
+
if not ckpt["authors"]:
|
| 365 |
+
author_list = list(all_authors.values())
|
| 366 |
+
logger.info(f"Upserting {len(author_list)} Author verticesβ¦")
|
| 367 |
+
for start in range(0, len(author_list), UPSERT_BATCH):
|
| 368 |
+
chunk = author_list[start : start + UPSERT_BATCH]
|
| 369 |
+
conn.upsertVertices("Author", [
|
| 370 |
+
(a["author_id"], {"display_name": a["display_name"]})
|
| 371 |
+
for a in chunk
|
| 372 |
+
])
|
| 373 |
+
ckpt["authors"] = True
|
| 374 |
+
_save_checkpoint(ckpt)
|
| 375 |
+
logger.info("Author vertices β")
|
| 376 |
+
else:
|
| 377 |
+
logger.info("Author vertices already complete β skipping.")
|
| 378 |
+
|
| 379 |
+
# Topic vertices βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 380 |
+
if not ckpt["topics"]:
|
| 381 |
+
topic_list = list(all_topics.values())
|
| 382 |
+
logger.info(f"Upserting {len(topic_list)} Topic verticesβ¦")
|
| 383 |
+
for start in range(0, len(topic_list), UPSERT_BATCH):
|
| 384 |
+
chunk = topic_list[start : start + UPSERT_BATCH]
|
| 385 |
+
conn.upsertVertices("Topic", [
|
| 386 |
+
(t["topic_id"], {
|
| 387 |
+
"display_name": t["display_name"],
|
| 388 |
+
"subfield": t["subfield"],
|
| 389 |
+
"field": t["field"],
|
| 390 |
+
"domain": t["domain"],
|
| 391 |
+
})
|
| 392 |
+
for t in chunk
|
| 393 |
+
])
|
| 394 |
+
ckpt["topics"] = True
|
| 395 |
+
_save_checkpoint(ckpt)
|
| 396 |
+
logger.info("Topic vertices β")
|
| 397 |
+
else:
|
| 398 |
+
logger.info("Topic vertices already complete β skipping.")
|
| 399 |
+
|
| 400 |
+
# AUTHORED_BY edges ββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 401 |
+
if not ckpt["authored_by"]:
|
| 402 |
+
authored_buf = [
|
| 403 |
+
(p["paper_id"], av["author_id"], {})
|
| 404 |
+
for p in papers
|
| 405 |
+
for av in p["author_vertices"]
|
| 406 |
+
if av["author_id"]
|
| 407 |
+
]
|
| 408 |
+
logger.info(f"Upserting {len(authored_buf)} AUTHORED_BY edgesβ¦")
|
| 409 |
+
for start in range(0, len(authored_buf), UPSERT_BATCH):
|
| 410 |
+
conn.upsertEdges("Paper", "AUTHORED_BY", "Author",
|
| 411 |
+
authored_buf[start : start + UPSERT_BATCH])
|
| 412 |
+
ckpt["authored_by"] = True
|
| 413 |
+
_save_checkpoint(ckpt)
|
| 414 |
+
logger.info("AUTHORED_BY edges β")
|
| 415 |
+
else:
|
| 416 |
+
logger.info("AUTHORED_BY edges already complete β skipping.")
|
| 417 |
+
|
| 418 |
+
# HAS_TOPIC edges ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 419 |
+
if not ckpt["has_topic"]:
|
| 420 |
+
topic_buf = [
|
| 421 |
+
(p["paper_id"], t["topic_id"], {"score": t["score"]})
|
| 422 |
+
for p in papers
|
| 423 |
+
for t in p["topics"]
|
| 424 |
+
if t["topic_id"]
|
| 425 |
+
]
|
| 426 |
+
logger.info(f"Upserting {len(topic_buf)} HAS_TOPIC edgesβ¦")
|
| 427 |
+
for start in range(0, len(topic_buf), UPSERT_BATCH):
|
| 428 |
+
conn.upsertEdges("Paper", "HAS_TOPIC", "Topic",
|
| 429 |
+
topic_buf[start : start + UPSERT_BATCH])
|
| 430 |
+
ckpt["has_topic"] = True
|
| 431 |
+
_save_checkpoint(ckpt)
|
| 432 |
+
logger.info("HAS_TOPIC edges β")
|
| 433 |
+
else:
|
| 434 |
+
logger.info("HAS_TOPIC edges already complete β skipping.")
|
| 435 |
+
|
| 436 |
+
# CITES edges (in-corpus only) βββββββββββββββββββββββββββββββββββββββββββββ
|
| 437 |
+
if not ckpt["cites"]:
|
| 438 |
+
cites_buf = [
|
| 439 |
+
(p["paper_id"], ref_id, {})
|
| 440 |
+
for p in papers
|
| 441 |
+
for ref_id in p["referenced"]
|
| 442 |
+
if ref_id in corpus_ids
|
| 443 |
+
]
|
| 444 |
+
logger.info(f"Upserting {len(cites_buf)} CITES edges (in-corpus)β¦")
|
| 445 |
+
for start in range(0, len(cites_buf), UPSERT_BATCH):
|
| 446 |
+
conn.upsertEdges("Paper", "CITES", "Paper",
|
| 447 |
+
cites_buf[start : start + UPSERT_BATCH])
|
| 448 |
+
ckpt["cites"] = True
|
| 449 |
+
_save_checkpoint(ckpt)
|
| 450 |
+
logger.info("CITES edges β")
|
| 451 |
+
else:
|
| 452 |
+
logger.info("CITES edges already complete β skipping.")
|
| 453 |
+
|
| 454 |
+
logger.info("Ingestion complete.")
|
| 455 |
+
|
| 456 |
+
|
| 457 |
+
if __name__ == "__main__":
|
| 458 |
+
create_schema()
|
| 459 |
+
ingest()
|
evaluation_set.json
ADDED
|
@@ -0,0 +1,98 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
[
|
| 2 |
+
{
|
| 3 |
+
"work_id": "W2980282514",
|
| 4 |
+
"type": "factual",
|
| 5 |
+
"question": "What two factors have driven recent progress in NLP according to the HuggingFace Transformers paper, and what is the library's stated goal?",
|
| 6 |
+
"correct_answer": "Two factors have driven recent progress in NLP: advances in model architecture and model pretraining. The HuggingFace Transformers library's stated goal is to support industrial-strength implementations of popular model variants that are easy to read, extend, and deploy, and to provide a centralized model hub for the distribution and usage of a wide variety of pretrained models."
|
| 7 |
+
},
|
| 8 |
+
{
|
| 9 |
+
"work_id": "W2529996553",
|
| 10 |
+
"type": "factual",
|
| 11 |
+
"question": "What three coupled functions does the neural network construct in the automatic chemical design paper and what does each one do?",
|
| 12 |
+
"correct_answer": "The neural network constructs three coupled functions: an encoder, a decoder, and a predictor. The encoder converts the discrete representation of a molecule into a real-valued continuous vector (latent space), the decoder converts these continuous vectors back to discrete molecular representations, and the property predictor estimates chemical properties from the latent continuous vector representation of the molecule."
|
| 13 |
+
},
|
| 14 |
+
{
|
| 15 |
+
"work_id": "W2766856748",
|
| 16 |
+
"type": "factual",
|
| 17 |
+
"question": "In the 2018 CGCNN paper by Xie and Grossman, what specific limitation of prior machine learning methods for crystal property prediction required manually constructed feature vectors, and how does their crystal graph framework address it?",
|
| 18 |
+
"correct_answer": "Prior machine learning methods required manually constructed feature vectors or complex transformations of atom coordinates to input the crystal structure, which either constrained the model to certain crystal types or made chemical interpretation difficult. The crystal graph framework addresses this by directly learning material properties from the connection of atoms in the crystal, providing a universal and interpretable representation across different crystal types and chemistries."
|
| 19 |
+
},
|
| 20 |
+
{
|
| 21 |
+
"work_id": "W2887306621",
|
| 22 |
+
"type": "factual",
|
| 23 |
+
"question": "In the 2018 review by Wu and Sun on machine learning in materials science, what four interdisciplinary fields are combined, and what is the primary advantage of ML methods over traditional theoretical simulations?",
|
| 24 |
+
"correct_answer": "Machine learning in materials science is an interdisciplinary field combining computer science, statistics, computational mathematics, and engineering. Its primary advantage is providing faster calculation speeds and higher prediction accuracy compared to traditional theoretical simulations that rely on solving fundamental physical or chemical equations."
|
| 25 |
+
},
|
| 26 |
+
{
|
| 27 |
+
"work_id": "W3154258817",
|
| 28 |
+
"type": "factual",
|
| 29 |
+
"question": "What does ADMET stand for, why are these properties the primary reason for drug development failure, and what does ADMETlab 2.0 provide to address this?",
|
| 30 |
+
"correct_answer": "ADMET stands for Absorption, Distribution, Metabolism, Excretion, and Toxicity. These properties are the primary reason for drug development failure because undesirable pharmacokinetics and toxicity of candidate compounds are the main reasons for failure. ADMETlab 2.0 provides an integrated online platform for accurate and comprehensive predictions of 17 physicochemical properties, 13 medicinal chemistry properties, 23 ADME properties, 27 toxicity endpoints, and 8 toxicophore rules."
|
| 31 |
+
},
|
| 32 |
+
{
|
| 33 |
+
"work_id": "W2742127985",
|
| 34 |
+
"type": "factual",
|
| 35 |
+
"question": "What makes the Deep Potential Molecular Dynamics (DPMD) method first-principles based, and what three natural symmetries does the neural network model explicitly preserve?",
|
| 36 |
+
"correct_answer": "The Deep Potential Molecular Dynamics (DPMD) method is first-principles based because there are no ad hoc components aside from the network model itself. The neural network model explicitly preserves the translational, rotational, and permutational symmetries of the atomic system."
|
| 37 |
+
},
|
| 38 |
+
{
|
| 39 |
+
"work_id": "W2968923792",
|
| 40 |
+
"type": "factual",
|
| 41 |
+
"question": "What collection of statistical methods is described as one of the most exciting new tools in the material science toolbox, and what two types of research has it proved capable of speeding up?",
|
| 42 |
+
"correct_answer": "Machine learning, described as a collection of statistical methods, is identified as one of the most exciting new tools in the material science toolbox. It has proved capable of considerably speeding up both fundamental and applied research in the field."
|
| 43 |
+
},
|
| 44 |
+
{
|
| 45 |
+
"work_id": "W3207969687",
|
| 46 |
+
"type": "factual",
|
| 47 |
+
"question": "According to the 2021 review on de novo drug design by Bai et al., what are the three main types of deep learning architectures used for molecule generation, and how do they benefit molecular dynamics simulations?",
|
| 48 |
+
"correct_answer": "The three main deep learning architectures for molecule generation are Recurrent Neural Networks (RNNs), Variational Autoencoders (VAEs), and Generative Adversarial Networks (GANs). In molecular dynamics (MD), deep learning models provide high-accuracy potential energy surfaces and force fields at a computational cost similar to traditional force fields, improving simulation efficiency."
|
| 49 |
+
},
|
| 50 |
+
{
|
| 51 |
+
"work_id": "W2966357564",
|
| 52 |
+
"type": "factual",
|
| 53 |
+
"question": "What two classes of neural network models have yielded promising results for molecular property prediction, and what does each class use as its input representation?",
|
| 54 |
+
"correct_answer": "Two classes of neural network models have yielded promising results for molecular property prediction: 1. neural networks applied to computed molecular fingerprints or expert-crafted descriptors, 2. graph convolutional neural networks that construct a learned molecular representation by operating on the graph structure of the molecule. The first class uses fingerprints or descriptors as input, while the second class uses the graph structure (atoms and bonds)."
|
| 55 |
+
},
|
| 56 |
+
{
|
| 57 |
+
"work_id": "W2900489253",
|
| 58 |
+
"type": "factual",
|
| 59 |
+
"question": "What is the key architectural feature of CheMixNet that allows it to predict multiple chemical properties simultaneously with high accuracy?",
|
| 60 |
+
"correct_answer": "The key architectural feature of CheMixNet is its ability to learn from a mixture of features learned from two different input representations: a vector input (molecular fingerprints) and a sequence input (SMILES strings), using a variation of multi-input-single-output (MISO) architectures to achieve high accuracy."
|
| 61 |
+
},
|
| 62 |
+
{
|
| 63 |
+
"work_ids": ["W2766856748", "W2742127985"],
|
| 64 |
+
"type": "multi_hop",
|
| 65 |
+
"question": "How does the CGCNN framework represent crystal structures, and what specific physical symmetries do both CGCNN and Deep Potential Molecular Dynamics (DPMD) explicitly preserve?",
|
| 66 |
+
"correct_answer": "CGCNN represents crystals as crystal graphs where nodes are atoms and edges are bonds, using atom/bond feature vectors. Both CGCNN and DPMD explicitly preserve translational, rotational, and permutational symmetries of the atomic system to ensure physical consistency and invariance."
|
| 67 |
+
},
|
| 68 |
+
{
|
| 69 |
+
"work_ids": ["W3154258817", "W3207969687"],
|
| 70 |
+
"type": "multi_hop",
|
| 71 |
+
"question": "The ADMETlab 2.0 paper and the de novo drug design review both describe ML applications to drug development. At what specific stage does ADMETlab 2.0 intervene, and how does this complement generative design described in the review?",
|
| 72 |
+
"correct_answer": "ADMETlab 2.0 intervenes at the stage of evaluating and optimizing absorption, distribution, metabolism, excretion, and toxicity (ADMET) properties in lead compounds. This complements generative design (using RNNs, VAEs, GANs) by providing a crucial evaluation step where generated molecules are assessed for pharmacokinetic and toxicity properties, increasing the likelihood of identifying successful drug candidates."
|
| 73 |
+
},
|
| 74 |
+
{
|
| 75 |
+
"work_ids": ["W2887306621", "W2968923792"],
|
| 76 |
+
"type": "multi_hop",
|
| 77 |
+
"question": "How does the application of ML in fluid dynamics surrogate modeling relate to the broader acceleration of research described in the materials science review?",
|
| 78 |
+
"correct_answer": "Machine learning in fluid dynamics surrogate modeling accelerates research by supporting design exploration and optimization of material properties, such as mechanical properties of alloy materials and energy-related applications. By utilizing multifidelity models, physics-based constraints, and surrogate modeling techniques, ML reduces computational costs and improves simulation accuracy, bridging the gap between fundamental research and applied engineering."
|
| 79 |
+
},
|
| 80 |
+
{
|
| 81 |
+
"work_ids": ["W2900489253", "W2968923792"],
|
| 82 |
+
"type": "multi_hop",
|
| 83 |
+
"question": "The CheMixNet paper and the materials science review both discuss property prediction. How does CheMixNet's mixed architecture address the general need for better 'statistical tools' in the materials science toolbox?",
|
| 84 |
+
"correct_answer": "CheMixNet's mixed architecture addresses the need for advanced statistical tools by combining sequence (SMILES) and fingerprint representations using RNNs and 1-D CNNs. This hybrid approach generalizes across diverse datasets like the Harvard Clean Energy Project and MoleculeNet, outperforming single-representation models (like MLP, CNN, or RNN alone) and providing a robust framework for predicting multiple chemical properties simultaneously."
|
| 85 |
+
},
|
| 86 |
+
{
|
| 87 |
+
"work_ids": ["W3207969687", "W3154258817"],
|
| 88 |
+
"type": "multi_hop",
|
| 89 |
+
"question": "How does the challenge of screening de novo designed molecules relate to the integrated platform provided by ADMETlab 2.0?",
|
| 90 |
+
"correct_answer": "ADMETlab 2.0 provides an integrated platform to screen the vast chemical space generated by de novo design. It offers high-throughput evaluation of over 50 endpoints, including physicochemical properties, medicinal chemistry rules, ADME properties, and toxicity. Using a multi-task graph attention framework and batch computation, it facilitates the rapid prioritization of candidates, significantly reducing the downstream synthesis workload."
|
| 91 |
+
},
|
| 92 |
+
{
|
| 93 |
+
"work_ids": ["W2980282514", "W2887306621"],
|
| 94 |
+
"type": "multi_hop",
|
| 95 |
+
"question": "The HuggingFace Transformers paper describes a unified API for pretrained models. How could this unified approach benefit the development of surrogate models for fluid dynamics?",
|
| 96 |
+
"correct_answer": "A unified API similar to HuggingFace Transformers could revolutionize fluid dynamics by establishing a centralized model hub for distributing and usage of pretrained surrogate models. This would allow users to compare model variants using the same minimal API and facilitate the use of transfer learning, where models pretrained on proxy properties (using large datasets) are repurposed for target tasks with limited data. Integrating these attention-based models with traditional CFD methods (like finite element analysis) enables hybrid modeling for complex industrial systems such as turbines, pumps, and pipelines, while providing industrial-strength, extensible implementations."
|
| 97 |
+
}
|
| 98 |
+
]
|
requirements.txt
ADDED
|
@@ -0,0 +1,134 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
aiobotocore==3.6.0
|
| 2 |
+
aiohappyeyeballs==2.6.1
|
| 3 |
+
aiohttp==3.13.5
|
| 4 |
+
aioitertools==0.13.0
|
| 5 |
+
aiosignal==1.4.0
|
| 6 |
+
annotated-doc==0.0.4
|
| 7 |
+
annotated-types==0.7.0
|
| 8 |
+
anyio==4.13.0
|
| 9 |
+
attrs==26.1.0
|
| 10 |
+
babel==2.18.0
|
| 11 |
+
bert-score==0.3.13
|
| 12 |
+
botocore==1.43.0
|
| 13 |
+
brotli==1.2.0
|
| 14 |
+
certifi==2026.4.22
|
| 15 |
+
cffi==2.0.0
|
| 16 |
+
charset-normalizer==3.4.7
|
| 17 |
+
click==8.3.3
|
| 18 |
+
contourpy==1.3.3
|
| 19 |
+
courlan==1.3.2
|
| 20 |
+
cryptography==47.0.0
|
| 21 |
+
cycler==0.12.1
|
| 22 |
+
datasets==4.8.5
|
| 23 |
+
dateparser==1.4.0
|
| 24 |
+
dill==0.4.1
|
| 25 |
+
distro==1.9.0
|
| 26 |
+
evaluate==0.4.6
|
| 27 |
+
fastapi==0.136.1
|
| 28 |
+
fastembed==0.7.4
|
| 29 |
+
filelock==3.29.0
|
| 30 |
+
flatbuffers==25.12.19
|
| 31 |
+
fonttools==4.62.1
|
| 32 |
+
frozenlist==1.8.0
|
| 33 |
+
fsspec==2026.2.0
|
| 34 |
+
google-ai-generativelanguage==0.6.15
|
| 35 |
+
google-api-core==2.25.2
|
| 36 |
+
google-api-python-client==2.195.0
|
| 37 |
+
google-auth==2.50.0
|
| 38 |
+
google-auth-httplib2==0.3.1
|
| 39 |
+
google-genai==1.74.0
|
| 40 |
+
google-generativeai==0.8.6
|
| 41 |
+
googleapis-common-protos==1.74.0
|
| 42 |
+
groq==1.2.0
|
| 43 |
+
grpcio==1.80.0
|
| 44 |
+
grpcio-status==1.71.2
|
| 45 |
+
h11==0.16.0
|
| 46 |
+
h2==4.3.0
|
| 47 |
+
hf-xet==1.4.3
|
| 48 |
+
hpack==4.1.0
|
| 49 |
+
htmldate==1.9.4
|
| 50 |
+
httpcore==1.0.9
|
| 51 |
+
httplib2==0.31.2
|
| 52 |
+
httpx==0.28.1
|
| 53 |
+
huggingface_hub==1.13.0
|
| 54 |
+
hyperframe==6.1.0
|
| 55 |
+
idna==3.13
|
| 56 |
+
Jinja2==3.1.6
|
| 57 |
+
jmespath==1.1.0
|
| 58 |
+
joblib==1.5.3
|
| 59 |
+
jusText==3.0.2
|
| 60 |
+
kiwisolver==1.5.0
|
| 61 |
+
loguru==0.7.3
|
| 62 |
+
lxml==6.1.0
|
| 63 |
+
lxml_html_clean==0.4.4
|
| 64 |
+
markdown-it-py==4.0.0
|
| 65 |
+
MarkupSafe==3.0.3
|
| 66 |
+
matplotlib==3.10.9
|
| 67 |
+
mdurl==0.1.2
|
| 68 |
+
mmh3==5.2.1
|
| 69 |
+
mpmath==1.3.0
|
| 70 |
+
multidict==6.7.1
|
| 71 |
+
multiprocess==0.70.19
|
| 72 |
+
networkx==3.6.1
|
| 73 |
+
numpy==2.4.4
|
| 74 |
+
onnxruntime==1.25.1
|
| 75 |
+
openalex-official==0.3.3
|
| 76 |
+
packaging==26.2
|
| 77 |
+
pandas==3.0.2
|
| 78 |
+
pillow==11.3.0
|
| 79 |
+
portalocker==3.2.0
|
| 80 |
+
propcache==0.4.1
|
| 81 |
+
proto-plus==1.27.2
|
| 82 |
+
protobuf==5.29.6
|
| 83 |
+
py_rust_stemmers==0.1.5
|
| 84 |
+
pyarrow==24.0.0
|
| 85 |
+
pyasn1==0.6.3
|
| 86 |
+
pyasn1_modules==0.4.2
|
| 87 |
+
pycparser==3.0
|
| 88 |
+
pydantic==2.13.3
|
| 89 |
+
pydantic_core==2.46.3
|
| 90 |
+
Pygments==2.20.0
|
| 91 |
+
PyMuPDF==1.27.2.3
|
| 92 |
+
pyparsing==3.3.2
|
| 93 |
+
python-dateutil==2.9.0.post0
|
| 94 |
+
python-dotenv==1.2.2
|
| 95 |
+
pyTigerGraph==2.0.3
|
| 96 |
+
pytz==2026.2
|
| 97 |
+
PyYAML==6.0.3
|
| 98 |
+
qdrant-client==1.17.1
|
| 99 |
+
regex==2026.4.4
|
| 100 |
+
requests==2.33.1
|
| 101 |
+
rich==15.0.0
|
| 102 |
+
safetensors==0.7.0
|
| 103 |
+
scikit-learn==1.8.0
|
| 104 |
+
scipy==1.17.1
|
| 105 |
+
sentence-transformers==5.4.1
|
| 106 |
+
setuptools==81.0.0
|
| 107 |
+
shellingham==1.5.4
|
| 108 |
+
six==1.17.0
|
| 109 |
+
sniffio==1.3.1
|
| 110 |
+
starlette==1.0.0
|
| 111 |
+
sympy==1.14.0
|
| 112 |
+
tenacity==9.1.4
|
| 113 |
+
threadpoolctl==3.6.0
|
| 114 |
+
tiktoken==0.12.0
|
| 115 |
+
tld==0.13.2
|
| 116 |
+
tokenizers==0.22.2
|
| 117 |
+
toml==0.10.2
|
| 118 |
+
tomli_w==1.2.0
|
| 119 |
+
torch==2.11.0
|
| 120 |
+
tqdm==4.67.3
|
| 121 |
+
trafilatura==2.0.0
|
| 122 |
+
transformers==5.7.0
|
| 123 |
+
typer==0.25.1
|
| 124 |
+
typing-inspection==0.4.2
|
| 125 |
+
typing_extensions==4.15.0
|
| 126 |
+
tzlocal==5.3.1
|
| 127 |
+
uritemplate==4.2.0
|
| 128 |
+
urllib3==2.6.3
|
| 129 |
+
uvicorn==0.46.0
|
| 130 |
+
validators==0.35.0
|
| 131 |
+
websockets==16.0
|
| 132 |
+
wrapt==2.1.2
|
| 133 |
+
xxhash==3.7.0
|
| 134 |
+
yarl==1.23.0
|
run_benchmarks.py
ADDED
|
@@ -0,0 +1,274 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import json
|
| 2 |
+
import time
|
| 3 |
+
import logging
|
| 4 |
+
from pathlib import Path
|
| 5 |
+
# Imports are moved inside functions to avoid dependency issues on startup
|
| 6 |
+
|
| 7 |
+
# Configuration: Toggle which pipelines to run
|
| 8 |
+
RUN_BASELINE = False
|
| 9 |
+
RUN_HYBRID = True
|
| 10 |
+
RUN_GRAPH = False
|
| 11 |
+
START_QUESTION = 16 # 1-indexed. Set to resume from a specific question (e.g. 9 to skip Q1-8)
|
| 12 |
+
LIMIT_QUESTIONS = None # Set to None to run all questions from START_QUESTION onward
|
| 13 |
+
|
| 14 |
+
# Evaluation Config
|
| 15 |
+
GOLDEN_SET_PATH = "evaluation_set.json"
|
| 16 |
+
RESULTS_PATH = "benchmark_results.json"
|
| 17 |
+
|
| 18 |
+
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
|
| 19 |
+
logging.getLogger("httpx").setLevel(logging.WARNING)
|
| 20 |
+
logging.getLogger("httpcore").setLevel(logging.WARNING)
|
| 21 |
+
logging.getLogger("sentence_transformers").setLevel(logging.WARNING)
|
| 22 |
+
logging.getLogger("huggingface_hub").setLevel(logging.WARNING)
|
| 23 |
+
logger = logging.getLogger(__name__)
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
# ------------------------------------------------------------------
|
| 27 |
+
# Incremental file I/O β every question saves immediately
|
| 28 |
+
# ------------------------------------------------------------------
|
| 29 |
+
|
| 30 |
+
def _load_results() -> dict:
|
| 31 |
+
"""Load existing results from disk."""
|
| 32 |
+
if Path(RESULTS_PATH).exists():
|
| 33 |
+
try:
|
| 34 |
+
with open(RESULTS_PATH, "r") as f:
|
| 35 |
+
return json.load(f)
|
| 36 |
+
except (json.JSONDecodeError, KeyError):
|
| 37 |
+
pass
|
| 38 |
+
return {"detailed_results": {}}
|
| 39 |
+
|
| 40 |
+
|
| 41 |
+
def _save_one_result(pipeline_name: str, result: dict):
|
| 42 |
+
"""Append a single result to disk immediately (crash-safe).
|
| 43 |
+
|
| 44 |
+
Merges by question_number so re-runs overwrite stale entries.
|
| 45 |
+
"""
|
| 46 |
+
data = _load_results()
|
| 47 |
+
detailed = data.get("detailed_results", {})
|
| 48 |
+
|
| 49 |
+
# Get or create the list for this pipeline
|
| 50 |
+
entries = detailed.get(pipeline_name, [])
|
| 51 |
+
|
| 52 |
+
# Merge: replace if same question_number exists, else append
|
| 53 |
+
qn = result.get("question_number")
|
| 54 |
+
replaced = False
|
| 55 |
+
for i, existing in enumerate(entries):
|
| 56 |
+
if existing.get("question_number") == qn:
|
| 57 |
+
entries[i] = result
|
| 58 |
+
replaced = True
|
| 59 |
+
break
|
| 60 |
+
if not replaced:
|
| 61 |
+
entries.append(result)
|
| 62 |
+
|
| 63 |
+
# Sort by question number
|
| 64 |
+
entries.sort(key=lambda r: r.get("question_number", 0))
|
| 65 |
+
detailed[pipeline_name] = entries
|
| 66 |
+
|
| 67 |
+
# Write back (no summary yet β that comes at the end)
|
| 68 |
+
data["detailed_results"] = detailed
|
| 69 |
+
with open(RESULTS_PATH, "w") as f:
|
| 70 |
+
json.dump(data, f, indent=4)
|
| 71 |
+
|
| 72 |
+
|
| 73 |
+
def _compute_and_save_summary():
|
| 74 |
+
"""Read all results, compute batch-averaged metrics (Official Bulk BERTScore), and write summary."""
|
| 75 |
+
data = _load_results()
|
| 76 |
+
results_dict = data.get("detailed_results", {})
|
| 77 |
+
if not results_dict:
|
| 78 |
+
return
|
| 79 |
+
|
| 80 |
+
from services.metrics_service import MetricsService
|
| 81 |
+
metrics_svc = MetricsService()
|
| 82 |
+
|
| 83 |
+
summary = {}
|
| 84 |
+
for pipe_name, entries in results_dict.items():
|
| 85 |
+
if not entries:
|
| 86 |
+
continue
|
| 87 |
+
|
| 88 |
+
# Only score entries that have a judge result
|
| 89 |
+
valid_entries = [e for e in entries if e.get("judge") is not None]
|
| 90 |
+
if not valid_entries:
|
| 91 |
+
continue
|
| 92 |
+
|
| 93 |
+
passes = [e for e in valid_entries if e.get("judge") == "PASS"]
|
| 94 |
+
total_scored = len(valid_entries)
|
| 95 |
+
|
| 96 |
+
# Collective BERTScore (Batch mode as per Notion)
|
| 97 |
+
preds = [e["answer"] for e in valid_entries if e.get("answer")]
|
| 98 |
+
|
| 99 |
+
# Map questions to correct answers from evaluation_set for bulk score
|
| 100 |
+
with open(GOLDEN_SET_PATH, "r") as f:
|
| 101 |
+
eval_set = json.load(f)
|
| 102 |
+
ref_map = {item["question"]: (item.get("correct_answer") or item.get("exact_answer")) for item in eval_set}
|
| 103 |
+
|
| 104 |
+
refs = [ref_map.get(e["question"]) for e in valid_entries if e.get("answer")]
|
| 105 |
+
|
| 106 |
+
bulk_bert = 0.0
|
| 107 |
+
if preds and refs:
|
| 108 |
+
try:
|
| 109 |
+
# Direct batch calculation to satisfy Notion "Bulk" requirement
|
| 110 |
+
bert_results = metrics_svc.bertscore.compute(
|
| 111 |
+
predictions=preds,
|
| 112 |
+
references=refs,
|
| 113 |
+
lang="en",
|
| 114 |
+
rescale_with_baseline=True
|
| 115 |
+
)
|
| 116 |
+
bulk_bert = float(sum(bert_results["f1"]) / len(bert_results["f1"]))
|
| 117 |
+
except Exception as e:
|
| 118 |
+
logger.error(f"Error calculating bulk BERTScore for {pipe_name}: {e}")
|
| 119 |
+
|
| 120 |
+
avg_latency = sum(e.get("latency", 0) for e in valid_entries) / total_scored
|
| 121 |
+
|
| 122 |
+
summary[pipe_name] = {
|
| 123 |
+
"pass_rate": f"{(len(passes)/total_scored*100):.1f}%",
|
| 124 |
+
"avg_bert_score": f"{bulk_bert:.4f}",
|
| 125 |
+
"avg_latency": f"{avg_latency:.2f}s",
|
| 126 |
+
"total_questions": total_scored,
|
| 127 |
+
"scored_questions": total_scored
|
| 128 |
+
}
|
| 129 |
+
|
| 130 |
+
print("="*50)
|
| 131 |
+
for name, stats in summary.items():
|
| 132 |
+
print(f"\nPipeline: {name}")
|
| 133 |
+
print(f" - Pass Rate: {stats['pass_rate']}")
|
| 134 |
+
print(f" - Bulk BERTScore: {stats['avg_bert_score']} (Threshold for bonus: 0.55)")
|
| 135 |
+
print(f" - Avg Latency: {stats['avg_latency']}")
|
| 136 |
+
print("="*50)
|
| 137 |
+
print(f"Full results saved to {RESULTS_PATH}")
|
| 138 |
+
|
| 139 |
+
|
| 140 |
+
def _print_checkpoint(questions_done: int):
|
| 141 |
+
"""Print a quick running summary from whatever's on disk."""
|
| 142 |
+
data = _load_results()
|
| 143 |
+
detailed = data.get("detailed_results", {})
|
| 144 |
+
for name, entries in detailed.items():
|
| 145 |
+
valid = [r for r in entries
|
| 146 |
+
if r.get("judge") is not None and r.get("bert_score") is not None]
|
| 147 |
+
if not valid:
|
| 148 |
+
continue
|
| 149 |
+
passes = sum(1 for r in valid if r["judge"] == "PASS")
|
| 150 |
+
avg_bert = sum(float(r["bert_score"]) for r in valid) / len(valid)
|
| 151 |
+
pct = (passes / len(valid)) * 100
|
| 152 |
+
print(f"\nπ CHECKPOINT after {questions_done} questions "
|
| 153 |
+
f"({len(valid)} scored) β {name}")
|
| 154 |
+
print(f" Pass Rate: {pct:.1f}% | Avg BERT: {avg_bert:.4f}")
|
| 155 |
+
|
| 156 |
+
|
| 157 |
+
# ------------------------------------------------------------------
|
| 158 |
+
# Rate limit detection
|
| 159 |
+
# ------------------------------------------------------------------
|
| 160 |
+
|
| 161 |
+
def _is_rate_limit(text: str) -> bool:
|
| 162 |
+
text = str(text).lower()
|
| 163 |
+
return "429" in text and ("rate_limit" in text or "rate limit" in text)
|
| 164 |
+
|
| 165 |
+
|
| 166 |
+
# ------------------------------------------------------------------
|
| 167 |
+
# Main
|
| 168 |
+
# ------------------------------------------------------------------
|
| 169 |
+
|
| 170 |
+
def load_golden_set():
|
| 171 |
+
with open(GOLDEN_SET_PATH, "r") as f:
|
| 172 |
+
return json.load(f)
|
| 173 |
+
|
| 174 |
+
|
| 175 |
+
def run_benchmark():
|
| 176 |
+
golden_set = load_golden_set()
|
| 177 |
+
|
| 178 |
+
# Apply start offset (1-indexed β 0-indexed)
|
| 179 |
+
start_idx = max(0, START_QUESTION - 1)
|
| 180 |
+
golden_set = golden_set[start_idx:]
|
| 181 |
+
if LIMIT_QUESTIONS:
|
| 182 |
+
golden_set = golden_set[:LIMIT_QUESTIONS]
|
| 183 |
+
|
| 184 |
+
total = len(golden_set)
|
| 185 |
+
logger.info(f"Running benchmark on questions {START_QUESTION} to {START_QUESTION + total - 1} "
|
| 186 |
+
f"({total} questions)...")
|
| 187 |
+
|
| 188 |
+
# Initialize Pipelines
|
| 189 |
+
pipelines = {}
|
| 190 |
+
if RUN_BASELINE:
|
| 191 |
+
from core.pipeline_1.logic import PipelineLLMOnly
|
| 192 |
+
logger.info("Initializing Baseline Pipeline...")
|
| 193 |
+
pipelines["Baseline (LLM Only)"] = PipelineLLMOnly(top_n=5, max_full_text=5)
|
| 194 |
+
|
| 195 |
+
if RUN_HYBRID:
|
| 196 |
+
from core.pipeline_2.logic import PipelineRAG
|
| 197 |
+
logger.info("Initializing Hybrid RAG Pipeline...")
|
| 198 |
+
pipelines["Hybrid RAG (Pipeline 2)"] = PipelineRAG(retrieval_top_k=50, rerank_top_n=10)
|
| 199 |
+
|
| 200 |
+
if RUN_GRAPH:
|
| 201 |
+
from core.pipeline_3.logic import PipelineGraphRAG
|
| 202 |
+
logger.info("Initializing GraphRAG Pipeline...")
|
| 203 |
+
pipelines["GraphRAG (Pipeline 3)"] = PipelineGraphRAG(rerank_top_n=10)
|
| 204 |
+
|
| 205 |
+
logger.info("Warming up BERTScore model (avoids cold-start penalty on Q1)...")
|
| 206 |
+
from services.metrics_service import MetricsService as _MS
|
| 207 |
+
_ms_warmup = _MS()
|
| 208 |
+
_ms_warmup.bertscore.compute(predictions=["warmup"], references=["warmup"], lang="en", rescale_with_baseline=True)
|
| 209 |
+
logger.info("BERTScore ready.")
|
| 210 |
+
|
| 211 |
+
rate_limited = False
|
| 212 |
+
questions_done = 0
|
| 213 |
+
|
| 214 |
+
for i, item in enumerate(golden_set):
|
| 215 |
+
q_num = start_idx + i + 1
|
| 216 |
+
question = item["question"]
|
| 217 |
+
correct_answer = item.get("correct_answer") or item.get("exact_answer")
|
| 218 |
+
q_type = item.get("type", "factual")
|
| 219 |
+
|
| 220 |
+
logger.info(f"\n[Q{q_num}/{start_idx + total}] Testing: {question[:80]}...")
|
| 221 |
+
|
| 222 |
+
for name, pipe in pipelines.items():
|
| 223 |
+
logger.info(f"Running {name}...")
|
| 224 |
+
try:
|
| 225 |
+
output = pipe.run(question, ground_truth=correct_answer)
|
| 226 |
+
|
| 227 |
+
# Rate limit check
|
| 228 |
+
if _is_rate_limit(output.get("error", "")):
|
| 229 |
+
logger.warning(f"π Rate limit hit on Q{q_num}. Saving and stopping.")
|
| 230 |
+
rate_limited = True
|
| 231 |
+
break
|
| 232 |
+
|
| 233 |
+
# Extract metrics and save immediately
|
| 234 |
+
m = output.get("metrics", {})
|
| 235 |
+
res = {
|
| 236 |
+
"question_number": q_num,
|
| 237 |
+
"question": question,
|
| 238 |
+
"type": q_type,
|
| 239 |
+
"answer": output.get("answer"),
|
| 240 |
+
"judge": m.get("accuracy", {}).get("llm_judge"),
|
| 241 |
+
"bert_score": m.get("accuracy", {}).get("bert_score"),
|
| 242 |
+
"latency": m.get("latency_seconds"),
|
| 243 |
+
"cost": m.get("cost_usd"),
|
| 244 |
+
"tokens": m.get("tokens", {}).get("total"),
|
| 245 |
+
}
|
| 246 |
+
# π Save to disk RIGHT NOW β crash-safe
|
| 247 |
+
_save_one_result(name, res)
|
| 248 |
+
logger.info(f"β
Q{q_num} saved | {res['judge']} | BERT: {res['bert_score']}")
|
| 249 |
+
|
| 250 |
+
except Exception as e:
|
| 251 |
+
if _is_rate_limit(str(e)):
|
| 252 |
+
logger.warning(f"π Rate limit exception on Q{q_num}. Saving and stopping.")
|
| 253 |
+
rate_limited = True
|
| 254 |
+
break
|
| 255 |
+
logger.error(f"Error running {name} on Q{q_num}: {e}")
|
| 256 |
+
|
| 257 |
+
if rate_limited:
|
| 258 |
+
break
|
| 259 |
+
|
| 260 |
+
questions_done += 1
|
| 261 |
+
|
| 262 |
+
# ββ Checkpoint every 5 questions βοΏ½οΏ½οΏ½
|
| 263 |
+
if questions_done % 5 == 0:
|
| 264 |
+
_print_checkpoint(questions_done)
|
| 265 |
+
|
| 266 |
+
# Compute final averages
|
| 267 |
+
_compute_and_save_summary()
|
| 268 |
+
|
| 269 |
+
if rate_limited:
|
| 270 |
+
logger.info(f"\nπ‘ To resume, set START_QUESTION = {q_num} in run_benchmarks.py")
|
| 271 |
+
|
| 272 |
+
|
| 273 |
+
if __name__ == "__main__":
|
| 274 |
+
run_benchmark()
|
services/metrics_service.py
ADDED
|
@@ -0,0 +1,111 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import time
|
| 2 |
+
import logging
|
| 3 |
+
import os
|
| 4 |
+
import evaluate
|
| 5 |
+
from huggingface_hub import InferenceClient
|
| 6 |
+
|
| 7 |
+
logger = logging.getLogger(__name__)
|
| 8 |
+
|
| 9 |
+
# YOUR EXACT PROMPT
|
| 10 |
+
JUDGE_PROMPT = """Grade the system's answer.
|
| 11 |
+
Question: {q}
|
| 12 |
+
Correct answer: {correct}
|
| 13 |
+
System answer: {answer}
|
| 14 |
+
|
| 15 |
+
Reply with only PASS or FAIL.
|
| 16 |
+
PASS = the system answer correctly addresses the question with no major errors.
|
| 17 |
+
FAIL = the answer is wrong, missing, or contradicts the correct answer."""
|
| 18 |
+
|
| 19 |
+
class MetricsService:
|
| 20 |
+
def __init__(self, model_name="gemini-1.5-flash"):
|
| 21 |
+
self.model_name = model_name.lower()
|
| 22 |
+
self.hf_token = os.environ.get("HF_TOKEN")
|
| 23 |
+
|
| 24 |
+
# Setup as per your implementation
|
| 25 |
+
self.client = InferenceClient(
|
| 26 |
+
model="meta-llama/Llama-3.1-8B-Instruct",
|
| 27 |
+
token=self.hf_token
|
| 28 |
+
)
|
| 29 |
+
self.bertscore = evaluate.load("bertscore")
|
| 30 |
+
|
| 31 |
+
self.tiers = {
|
| 32 |
+
"flash": {"input": 0.075, "output": 0.30},
|
| 33 |
+
"pro": {"input": 3.50, "output": 10.50},
|
| 34 |
+
"groq_large": {"input": 0.59, "output": 0.79},
|
| 35 |
+
"groq_small": {"input": 0.05, "output": 0.08},
|
| 36 |
+
}
|
| 37 |
+
|
| 38 |
+
def calculate_cost(self, input_tokens, output_tokens):
|
| 39 |
+
if "70b" in self.model_name or "versatile" in self.model_name:
|
| 40 |
+
prices = self.tiers["groq_large"]
|
| 41 |
+
elif "8b" in self.model_name or "instant" in self.model_name:
|
| 42 |
+
prices = self.tiers["groq_small"]
|
| 43 |
+
elif "pro" in self.model_name:
|
| 44 |
+
prices = self.tiers["pro"]
|
| 45 |
+
else:
|
| 46 |
+
prices = self.tiers["flash"]
|
| 47 |
+
return ((input_tokens / 1_000_000) * prices["input"]) + ((output_tokens / 1_000_000) * prices["output"])
|
| 48 |
+
|
| 49 |
+
def process_metrics(self, client, query, answer, context, usage_metadata, start_time, abstracts_list=None, ground_truth=None):
|
| 50 |
+
"""EXACT implementation with NO truncation."""
|
| 51 |
+
# 1. Prepare "Correct Answer"
|
| 52 |
+
# If ground_truth is provided (from benchmark), use it. Otherwise fallback to abstracts/context.
|
| 53 |
+
final_ground_truth = ground_truth or ("\n\n".join(abstracts_list) if abstracts_list else context)
|
| 54 |
+
|
| 55 |
+
# 2. LLM-as-a-Judge β strip citation markers and disclaimer sentences before judging
|
| 56 |
+
import re as _re
|
| 57 |
+
clean_for_judge = _re.sub(r'\[Paper \d+\]', '', answer).strip()
|
| 58 |
+
# Strip sentences where the model hedges about missing context β these cause false judge FAILs
|
| 59 |
+
_DISCLAIMER = _re.compile(
|
| 60 |
+
r'[^.!?]*\b(not explicitly (mentioned|stated|provided)|'
|
| 61 |
+
r'not (mentioned|provided|specified) in (the )?(provided |given )?context|'
|
| 62 |
+
r'is not clear from|cannot be (determined|found) from|'
|
| 63 |
+
r'no (explicit |direct )?mention)\b[^.!?]*[.!?]?',
|
| 64 |
+
_re.IGNORECASE
|
| 65 |
+
)
|
| 66 |
+
clean_for_judge = _DISCLAIMER.sub('', clean_for_judge).strip()
|
| 67 |
+
prompt = JUDGE_PROMPT.format(
|
| 68 |
+
q=query,
|
| 69 |
+
correct=final_ground_truth,
|
| 70 |
+
answer=clean_for_judge
|
| 71 |
+
)
|
| 72 |
+
|
| 73 |
+
try:
|
| 74 |
+
verdict = self.client.chat_completion(
|
| 75 |
+
[{"role": "user", "content": prompt}],
|
| 76 |
+
max_tokens=20,
|
| 77 |
+
temperature=0.0
|
| 78 |
+
)
|
| 79 |
+
judge_pass = "PASS" in verdict.choices[0].message.content.upper()
|
| 80 |
+
except Exception as e:
|
| 81 |
+
logger.error(f"Judge error: {e}")
|
| 82 |
+
judge_pass = False
|
| 83 |
+
|
| 84 |
+
# 3. BERTScore β reuse the already-cleaned answer from above
|
| 85 |
+
clean_answer = clean_for_judge
|
| 86 |
+
try:
|
| 87 |
+
bert_results = self.bertscore.compute(
|
| 88 |
+
predictions=[clean_answer],
|
| 89 |
+
references=[final_ground_truth],
|
| 90 |
+
lang="en",
|
| 91 |
+
rescale_with_baseline=True
|
| 92 |
+
)
|
| 93 |
+
f1_score = float(bert_results["f1"][0])
|
| 94 |
+
except Exception as e:
|
| 95 |
+
logger.error(f"BERTScore error: {e}")
|
| 96 |
+
f1_score = 0.0
|
| 97 |
+
|
| 98 |
+
# 4. Latency and Cost
|
| 99 |
+
end_time = time.time()
|
| 100 |
+
input_tokens = usage_metadata.prompt_token_count if usage_metadata else 0
|
| 101 |
+
output_tokens = usage_metadata.candidates_token_count if usage_metadata else 0
|
| 102 |
+
|
| 103 |
+
return {
|
| 104 |
+
"tokens": {"input": input_tokens, "output": output_tokens, "total": input_tokens + output_tokens},
|
| 105 |
+
"latency_seconds": round(end_time - start_time, 2),
|
| 106 |
+
"cost_usd": f"{self.calculate_cost(input_tokens, output_tokens):.6f}",
|
| 107 |
+
"accuracy": {
|
| 108 |
+
"llm_judge": "PASS" if judge_pass else "FAIL",
|
| 109 |
+
"bert_score": f1_score
|
| 110 |
+
}
|
| 111 |
+
}
|
services/paper_fetcher.py
ADDED
|
@@ -0,0 +1,216 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import os
|
| 2 |
+
import requests
|
| 3 |
+
import fitz # PyMuPDF
|
| 4 |
+
import logging
|
| 5 |
+
import io
|
| 6 |
+
import trafilatura
|
| 7 |
+
from pathlib import Path
|
| 8 |
+
from dotenv import load_dotenv
|
| 9 |
+
|
| 10 |
+
# Load environment variables
|
| 11 |
+
load_dotenv()
|
| 12 |
+
|
| 13 |
+
logging.basicConfig(level=logging.INFO)
|
| 14 |
+
logger = logging.getLogger(__name__)
|
| 15 |
+
|
| 16 |
+
class PaperFetcher:
|
| 17 |
+
def __init__(self):
|
| 18 |
+
self.api_key = os.environ.get("OPENALEX_API_KEY")
|
| 19 |
+
self.headers = {
|
| 20 |
+
"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36",
|
| 21 |
+
"Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7",
|
| 22 |
+
"Accept-Language": "en-US,en;q=0.9",
|
| 23 |
+
"Accept-Encoding": "gzip, deflate, br",
|
| 24 |
+
"DNT": "1",
|
| 25 |
+
"Connection": "keep-alive",
|
| 26 |
+
"Upgrade-Insecure-Requests": "1",
|
| 27 |
+
"Sec-Fetch-Dest": "document",
|
| 28 |
+
"Sec-Fetch-Mode": "navigate",
|
| 29 |
+
"Sec-Fetch-Site": "none",
|
| 30 |
+
"Sec-Fetch-User": "?1"
|
| 31 |
+
}
|
| 32 |
+
self.session = requests.Session()
|
| 33 |
+
self.session.headers.update(self.headers)
|
| 34 |
+
|
| 35 |
+
def get_work_metadata(self, work_id):
|
| 36 |
+
"""Fetch metadata from OpenAlex."""
|
| 37 |
+
if not work_id.startswith("https://openalex.org/"):
|
| 38 |
+
if work_id.startswith("W"):
|
| 39 |
+
work_id = f"https://openalex.org/{work_id}"
|
| 40 |
+
|
| 41 |
+
url = f"https://api.openalex.org/works/{work_id.split('/')[-1]}"
|
| 42 |
+
params = {}
|
| 43 |
+
if self.api_key:
|
| 44 |
+
params["api_key"] = self.api_key
|
| 45 |
+
|
| 46 |
+
try:
|
| 47 |
+
response = requests.get(url, params=params, headers=self.headers)
|
| 48 |
+
response.raise_for_status()
|
| 49 |
+
return response.json()
|
| 50 |
+
except Exception as e:
|
| 51 |
+
logger.error(f"Failed to fetch metadata from OpenAlex: {e}")
|
| 52 |
+
return None
|
| 53 |
+
|
| 54 |
+
def get_semantic_scholar_pdf(self, doi):
|
| 55 |
+
"""Fallback: Fetch PDF URL from Semantic Scholar API."""
|
| 56 |
+
if not doi:
|
| 57 |
+
return None
|
| 58 |
+
|
| 59 |
+
clean_doi = doi.replace("https://doi.org/", "")
|
| 60 |
+
url = f"https://api.semanticscholar.org/graph/v1/paper/DOI:{clean_doi}?fields=openAccessPdf"
|
| 61 |
+
|
| 62 |
+
try:
|
| 63 |
+
logger.info(f"Checking Semantic Scholar for DOI: {clean_doi}...")
|
| 64 |
+
response = requests.get(url, headers=self.headers, timeout=15)
|
| 65 |
+
response.raise_for_status()
|
| 66 |
+
data = response.json()
|
| 67 |
+
oa_pdf = data.get("openAccessPdf")
|
| 68 |
+
if oa_pdf and oa_pdf.get("url"):
|
| 69 |
+
return oa_pdf["url"]
|
| 70 |
+
except Exception as e:
|
| 71 |
+
logger.error(f"Semantic Scholar fallback failed: {e}")
|
| 72 |
+
return None
|
| 73 |
+
|
| 74 |
+
def fetch_pdf_content(self, pdf_url):
|
| 75 |
+
"""Fetch PDF content with hardened redirect and session handling."""
|
| 76 |
+
try:
|
| 77 |
+
logger.info(f"Fetching PDF from {pdf_url}...")
|
| 78 |
+
# For PMC and similar, we need to be careful with redirects
|
| 79 |
+
response = self.session.get(pdf_url, timeout=30, allow_redirects=True)
|
| 80 |
+
|
| 81 |
+
# Handle the "Too many redirects" or meta-refresh redirects manually if needed
|
| 82 |
+
if response.status_code == 200 and 'application/pdf' in response.headers.get('Content-Type', ''):
|
| 83 |
+
return response.content
|
| 84 |
+
|
| 85 |
+
# If it's HTML, we might have been redirected to a challenge page
|
| 86 |
+
if 'text/html' in response.headers.get('Content-Type', ''):
|
| 87 |
+
logger.warning(f"PDF URL {pdf_url} returned HTML instead of PDF. Possibly a bot challenge.")
|
| 88 |
+
|
| 89 |
+
response.raise_for_status()
|
| 90 |
+
return response.content
|
| 91 |
+
except Exception as e:
|
| 92 |
+
logger.error(f"Failed to fetch PDF from {pdf_url}: {e}")
|
| 93 |
+
return None
|
| 94 |
+
|
| 95 |
+
def extract_text_from_bytes(self, pdf_bytes):
|
| 96 |
+
"""Extract all text content from PDF bytes."""
|
| 97 |
+
try:
|
| 98 |
+
doc = fitz.open(stream=pdf_bytes, filetype="pdf")
|
| 99 |
+
text = ""
|
| 100 |
+
for page in doc:
|
| 101 |
+
text += page.get_text()
|
| 102 |
+
doc.close()
|
| 103 |
+
return text.strip()
|
| 104 |
+
except Exception as e:
|
| 105 |
+
logger.error(f"Failed to extract text from PDF bytes: {e}")
|
| 106 |
+
return None
|
| 107 |
+
|
| 108 |
+
def extract_from_html(self, url):
|
| 109 |
+
"""Extract content from HTML landing page by downloading first to handle cookies/redirects."""
|
| 110 |
+
try:
|
| 111 |
+
logger.info(f"Attempting HTML extraction from {url}...")
|
| 112 |
+
# Use our session to get the HTML
|
| 113 |
+
response = self.session.get(url, timeout=20)
|
| 114 |
+
response.raise_for_status()
|
| 115 |
+
|
| 116 |
+
downloaded = response.text
|
| 117 |
+
if downloaded:
|
| 118 |
+
# Pass the HTML content directly to trafilatura
|
| 119 |
+
result = trafilatura.extract(downloaded, include_comments=False, include_tables=True)
|
| 120 |
+
if result:
|
| 121 |
+
logger.info(f"Successfully extracted {len(result)} chars from HTML.")
|
| 122 |
+
return result.strip()
|
| 123 |
+
except Exception as e:
|
| 124 |
+
logger.error(f"HTML extraction failed for {url}: {e}")
|
| 125 |
+
return None
|
| 126 |
+
|
| 127 |
+
def fetch_full_text(self, work_id):
|
| 128 |
+
"""Main method: arXiv -> OpenAlex -> Semantic Scholar -> HTML Scrape -> Abstract."""
|
| 129 |
+
metadata = self.get_work_metadata(work_id)
|
| 130 |
+
if not metadata:
|
| 131 |
+
return None
|
| 132 |
+
|
| 133 |
+
# 1. Gather all potential PDF candidates
|
| 134 |
+
pdf_candidates = []
|
| 135 |
+
|
| 136 |
+
# Check arXiv
|
| 137 |
+
ids = metadata.get("ids", {})
|
| 138 |
+
if "arxiv" in ids:
|
| 139 |
+
arxiv_id = ids["arxiv"].split("/")[-1].replace("abs/", "").replace("arxiv:", "")
|
| 140 |
+
pdf_candidates.append(f"https://arxiv.org/pdf/{arxiv_id}.pdf")
|
| 141 |
+
|
| 142 |
+
# Check OpenAlex locations
|
| 143 |
+
best_loc = metadata.get("best_oa_location")
|
| 144 |
+
if best_loc and best_loc.get("pdf_url"):
|
| 145 |
+
pdf_candidates.append(best_loc["pdf_url"])
|
| 146 |
+
|
| 147 |
+
for loc in metadata.get("locations", []):
|
| 148 |
+
if loc.get("pdf_url") and loc["pdf_url"] not in pdf_candidates:
|
| 149 |
+
pdf_candidates.append(loc["pdf_url"])
|
| 150 |
+
|
| 151 |
+
# 2. Try each PDF candidate
|
| 152 |
+
logger.info(f"Found {len(pdf_candidates)} PDF candidates for {work_id}")
|
| 153 |
+
for url in pdf_candidates:
|
| 154 |
+
pdf_bytes = self.fetch_pdf_content(url)
|
| 155 |
+
if pdf_bytes:
|
| 156 |
+
text = self.extract_text_from_bytes(pdf_bytes)
|
| 157 |
+
if text and len(text) > 200:
|
| 158 |
+
logger.info(f"Successfully extracted {len(text)} chars from {url}")
|
| 159 |
+
return text
|
| 160 |
+
else:
|
| 161 |
+
logger.warning(f"Extraction from {url} was too short or empty.")
|
| 162 |
+
|
| 163 |
+
# 3. If PDF fails, try Semantic Scholar for a new PDF link
|
| 164 |
+
ss_pdf_url = self.get_semantic_scholar_pdf(metadata.get("doi"))
|
| 165 |
+
if ss_pdf_url and ss_pdf_url not in pdf_candidates:
|
| 166 |
+
pdf_bytes = self.fetch_pdf_content(ss_pdf_url)
|
| 167 |
+
if pdf_bytes:
|
| 168 |
+
text = self.extract_text_from_bytes(pdf_bytes)
|
| 169 |
+
if text and len(text) > 200:
|
| 170 |
+
return text
|
| 171 |
+
|
| 172 |
+
# 4. If all PDFs fail, try HTML scraping from landing page
|
| 173 |
+
landing_page = metadata.get("landing_page_url")
|
| 174 |
+
if not landing_page and best_loc:
|
| 175 |
+
landing_page = best_loc.get("landing_page_url")
|
| 176 |
+
|
| 177 |
+
if landing_page:
|
| 178 |
+
text = self.extract_from_html(landing_page)
|
| 179 |
+
if text and len(text) > 500: # HTML extraction should be substantial
|
| 180 |
+
return text
|
| 181 |
+
|
| 182 |
+
# 5. Final Safety Net: Reconstruct Abstract
|
| 183 |
+
logger.warning(f"Could not get full text for {work_id}. Falling back to abstract.")
|
| 184 |
+
return self.get_abstract(metadata)
|
| 185 |
+
|
| 186 |
+
def get_abstract(self, metadata):
|
| 187 |
+
"""Reconstruct abstract from inverted index."""
|
| 188 |
+
inverted_index = metadata.get("abstract_inverted_index")
|
| 189 |
+
if not inverted_index:
|
| 190 |
+
return ""
|
| 191 |
+
|
| 192 |
+
max_index = 0
|
| 193 |
+
for indices in inverted_index.values():
|
| 194 |
+
if indices:
|
| 195 |
+
max_index = max(max_index, max(indices))
|
| 196 |
+
|
| 197 |
+
abstract_list = [""] * (max_index + 1)
|
| 198 |
+
for word, indices in inverted_index.items():
|
| 199 |
+
for index in indices:
|
| 200 |
+
abstract_list[index] = word
|
| 201 |
+
|
| 202 |
+
return " ".join(abstract_list).strip()
|
| 203 |
+
|
| 204 |
+
# Test Block
|
| 205 |
+
if __name__ == "__main__":
|
| 206 |
+
fetcher = PaperFetcher()
|
| 207 |
+
# Testing the PeerJ paper that has been failing
|
| 208 |
+
test_id = "W2741809807"
|
| 209 |
+
full_text = fetcher.fetch_full_text(test_id)
|
| 210 |
+
|
| 211 |
+
if full_text:
|
| 212 |
+
print("\n--- Success! First 200 characters ---")
|
| 213 |
+
print(full_text[:200])
|
| 214 |
+
print(f"\nTotal characters fetched: {len(full_text):,}")
|
| 215 |
+
else:
|
| 216 |
+
print("Failed to fetch any text.")
|
utils/download_aiml_data.py
ADDED
|
@@ -0,0 +1,80 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python3
|
| 2 |
+
import os
|
| 3 |
+
import subprocess
|
| 4 |
+
import sys
|
| 5 |
+
from dotenv import load_dotenv
|
| 6 |
+
|
| 7 |
+
def main():
|
| 8 |
+
# Load environment variables from .env file
|
| 9 |
+
load_dotenv()
|
| 10 |
+
|
| 11 |
+
# Configuration
|
| 12 |
+
output_dir = "data"
|
| 13 |
+
# Top AI/ML topics identified via API search
|
| 14 |
+
topic_ids = [
|
| 15 |
+
"T12072", # Machine Learning and Algorithms
|
| 16 |
+
"T11948", # Deep Learning and Neural Networks
|
| 17 |
+
]
|
| 18 |
+
topic_filter = "|".join(topic_ids)
|
| 19 |
+
|
| 20 |
+
pub_year = "2018-2024"
|
| 21 |
+
min_citations = "20"
|
| 22 |
+
|
| 23 |
+
# Check for API key
|
| 24 |
+
api_key = os.environ.get("OPENALEX_API_KEY")
|
| 25 |
+
if not api_key:
|
| 26 |
+
print("Error: OPENALEX_API_KEY environment variable is not set.")
|
| 27 |
+
print("Please get your API key from https://openalex.org/settings/api and set it:")
|
| 28 |
+
print("export OPENALEX_API_KEY='your-key-here'")
|
| 29 |
+
sys.exit(1)
|
| 30 |
+
|
| 31 |
+
# Ensure output directory exists
|
| 32 |
+
os.makedirs(output_dir, exist_ok=True)
|
| 33 |
+
|
| 34 |
+
filter_str = f"topics.id:{topic_filter},publication_year:{pub_year},cited_by_count:>{min_citations}"
|
| 35 |
+
|
| 36 |
+
print(f"π Starting OpenAlex download for AI/ML papers...")
|
| 37 |
+
print(f"π Output directory: {output_dir}")
|
| 38 |
+
print(f"π Filter: {filter_str}")
|
| 39 |
+
|
| 40 |
+
# Build the command
|
| 41 |
+
command = [
|
| 42 |
+
"openalex", "download",
|
| 43 |
+
"--api-key", api_key,
|
| 44 |
+
"--output", output_dir,
|
| 45 |
+
"--filter", filter_str,
|
| 46 |
+
"--resume",
|
| 47 |
+
"--workers", "10"
|
| 48 |
+
]
|
| 49 |
+
|
| 50 |
+
try:
|
| 51 |
+
# Run the command and pipe output to terminal
|
| 52 |
+
process = subprocess.Popen(
|
| 53 |
+
command,
|
| 54 |
+
stdout=subprocess.PIPE,
|
| 55 |
+
stderr=subprocess.STDOUT,
|
| 56 |
+
text=True,
|
| 57 |
+
bufsize=1
|
| 58 |
+
)
|
| 59 |
+
|
| 60 |
+
# Print progress in real-time
|
| 61 |
+
for line in process.stdout:
|
| 62 |
+
print(line, end="")
|
| 63 |
+
|
| 64 |
+
process.wait()
|
| 65 |
+
|
| 66 |
+
if process.returncode == 0:
|
| 67 |
+
print("\nβ
Download completed successfully.")
|
| 68 |
+
else:
|
| 69 |
+
print(f"\nβ Download failed with return code {process.returncode}.")
|
| 70 |
+
print("You can run this script again to resume from the last checkpoint.")
|
| 71 |
+
|
| 72 |
+
except KeyboardInterrupt:
|
| 73 |
+
print("\nπ Download interrupted by user. Run again to resume.")
|
| 74 |
+
sys.exit(1)
|
| 75 |
+
except Exception as e:
|
| 76 |
+
print(f"\nπ₯ An error occurred: {e}")
|
| 77 |
+
sys.exit(1)
|
| 78 |
+
|
| 79 |
+
if __name__ == "__main__":
|
| 80 |
+
main()
|
utils/token_count.py
ADDED
|
@@ -0,0 +1,183 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Calculate total token counts over data/ folder.
|
| 2 |
+
|
| 3 |
+
Reports three tiers:
|
| 4 |
+
1. Raw β full OpenAlex JSON as-is on disk
|
| 5 |
+
2. Optimised β only the fields the RAG pipeline actually reads
|
| 6 |
+
(id, title, abstract, authorships, topics, keywords,
|
| 7 |
+
referenced_works, cited_by_count, publication_year, doi,
|
| 8 |
+
best_oa_location.pdf_url)
|
| 9 |
+
3. LLM context β what actually reaches the embedding model and the LLM:
|
| 10 |
+
plain-text "Title: β¦\\n\\nAbstract: β¦" per paper
|
| 11 |
+
"""
|
| 12 |
+
|
| 13 |
+
import json
|
| 14 |
+
import sys
|
| 15 |
+
from pathlib import Path
|
| 16 |
+
|
| 17 |
+
# data/ lives at the project root, two levels up from utils/
|
| 18 |
+
DATA_DIR = Path(__file__).parents[1] / "data"
|
| 19 |
+
|
| 20 |
+
# Approximate token ratio (1 token β 4 chars for English text β GPT/Llama heuristic)
|
| 21 |
+
CHARS_PER_TOKEN = 4
|
| 22 |
+
|
| 23 |
+
# Fields the pipeline actually uses
|
| 24 |
+
_KEEP_KEYS = {
|
| 25 |
+
"id", "title", "display_name", "abstract_inverted_index",
|
| 26 |
+
"publication_year", "cited_by_count",
|
| 27 |
+
"authorships", "topics", "keywords", "referenced_works",
|
| 28 |
+
"primary_topic", "doi", "type",
|
| 29 |
+
}
|
| 30 |
+
|
| 31 |
+
|
| 32 |
+
def reconstruct_abstract(inv_index: dict) -> str:
|
| 33 |
+
"""Rebuild plain-text abstract from OpenAlex inverted index."""
|
| 34 |
+
if not inv_index:
|
| 35 |
+
return ""
|
| 36 |
+
pairs = []
|
| 37 |
+
for word, positions in inv_index.items():
|
| 38 |
+
for pos in positions:
|
| 39 |
+
pairs.append((pos, word))
|
| 40 |
+
pairs.sort()
|
| 41 |
+
return " ".join(w for _, w in pairs)
|
| 42 |
+
|
| 43 |
+
|
| 44 |
+
def slim_authorships(authorships: list) -> list:
|
| 45 |
+
"""Keep only author name + position (drop institutions, affiliations, etc.)."""
|
| 46 |
+
return [
|
| 47 |
+
{
|
| 48 |
+
"name": a.get("author", {}).get("display_name", ""),
|
| 49 |
+
"position": a.get("author_position", ""),
|
| 50 |
+
}
|
| 51 |
+
for a in (authorships or [])
|
| 52 |
+
]
|
| 53 |
+
|
| 54 |
+
|
| 55 |
+
def optimise_paper(raw: dict) -> dict:
|
| 56 |
+
"""Strip a raw OpenAlex JSON to only pipeline-relevant fields."""
|
| 57 |
+
optimised = {k: raw[k] for k in _KEEP_KEYS if k in raw}
|
| 58 |
+
|
| 59 |
+
# Replace inverted index with plain text abstract
|
| 60 |
+
if "abstract_inverted_index" in optimised:
|
| 61 |
+
optimised["abstract"] = reconstruct_abstract(optimised.pop("abstract_inverted_index"))
|
| 62 |
+
|
| 63 |
+
# Slim authorships
|
| 64 |
+
if "authorships" in optimised:
|
| 65 |
+
optimised["authorships"] = slim_authorships(optimised["authorships"])
|
| 66 |
+
|
| 67 |
+
# Slim topics to just display_name
|
| 68 |
+
if "topics" in optimised:
|
| 69 |
+
optimised["topics"] = [t.get("display_name", "") for t in (optimised["topics"] or [])]
|
| 70 |
+
|
| 71 |
+
# Slim keywords
|
| 72 |
+
if "keywords" in optimised:
|
| 73 |
+
optimised["keywords"] = [k.get("display_name", "") for k in (optimised["keywords"] or [])]
|
| 74 |
+
|
| 75 |
+
# Slim referenced_works to just IDs
|
| 76 |
+
if "referenced_works" in optimised:
|
| 77 |
+
optimised["referenced_works"] = [
|
| 78 |
+
r.rsplit("/", 1)[-1] if "/" in r else r
|
| 79 |
+
for r in (optimised["referenced_works"] or [])
|
| 80 |
+
]
|
| 81 |
+
|
| 82 |
+
return optimised
|
| 83 |
+
|
| 84 |
+
|
| 85 |
+
def count_tokens(text: str) -> dict:
|
| 86 |
+
chars = len(text)
|
| 87 |
+
words = len(text.split())
|
| 88 |
+
tokens_est = chars // CHARS_PER_TOKEN
|
| 89 |
+
return {"chars": chars, "words": words, "tokens_est": tokens_est}
|
| 90 |
+
|
| 91 |
+
|
| 92 |
+
def main():
|
| 93 |
+
if not DATA_DIR.exists():
|
| 94 |
+
print(f"ERROR: {DATA_DIR} not found. Expected: {DATA_DIR.resolve()}")
|
| 95 |
+
sys.exit(1)
|
| 96 |
+
|
| 97 |
+
files = sorted(DATA_DIR.glob("*.json"))
|
| 98 |
+
total = len(files)
|
| 99 |
+
print(f"Found {total} JSON files in {DATA_DIR.resolve()}\n")
|
| 100 |
+
|
| 101 |
+
raw_total = {"chars": 0, "words": 0, "tokens_est": 0, "bytes": 0}
|
| 102 |
+
opt_total = {"chars": 0, "words": 0, "tokens_est": 0, "bytes": 0}
|
| 103 |
+
llm_total = {"chars": 0, "words": 0, "tokens_est": 0, "bytes": 0}
|
| 104 |
+
|
| 105 |
+
for f in files:
|
| 106 |
+
raw_text = f.read_text(encoding="utf-8")
|
| 107 |
+
raw_total["bytes"] += len(raw_text.encode("utf-8"))
|
| 108 |
+
stats = count_tokens(raw_text)
|
| 109 |
+
raw_total["chars"] += stats["chars"]
|
| 110 |
+
raw_total["words"] += stats["words"]
|
| 111 |
+
raw_total["tokens_est"] += stats["tokens_est"]
|
| 112 |
+
|
| 113 |
+
try:
|
| 114 |
+
data = json.loads(raw_text)
|
| 115 |
+
|
| 116 |
+
# Tier 2: optimised JSON (all pipeline-read fields)
|
| 117 |
+
optimised = optimise_paper(data)
|
| 118 |
+
opt_text = json.dumps(optimised, ensure_ascii=False)
|
| 119 |
+
opt_total["bytes"] += len(opt_text.encode("utf-8"))
|
| 120 |
+
opt_stats = count_tokens(opt_text)
|
| 121 |
+
opt_total["chars"] += opt_stats["chars"]
|
| 122 |
+
opt_total["words"] += opt_stats["words"]
|
| 123 |
+
opt_total["tokens_est"] += opt_stats["tokens_est"]
|
| 124 |
+
|
| 125 |
+
# Tier 3: LLM context β the exact text sent to the embedding model and LLM
|
| 126 |
+
# Mirrors indexer.py _doc_text() and setup.py embed text construction
|
| 127 |
+
title = data.get("title") or ""
|
| 128 |
+
abstract = reconstruct_abstract(data.get("abstract_inverted_index") or {})
|
| 129 |
+
llm_text = f"Title: {title}\n\nAbstract: {abstract}"
|
| 130 |
+
llm_total["bytes"] += len(llm_text.encode("utf-8"))
|
| 131 |
+
llm_stats = count_tokens(llm_text)
|
| 132 |
+
llm_total["chars"] += llm_stats["chars"]
|
| 133 |
+
llm_total["words"] += llm_stats["words"]
|
| 134 |
+
llm_total["tokens_est"] += llm_stats["tokens_est"]
|
| 135 |
+
|
| 136 |
+
except json.JSONDecodeError:
|
| 137 |
+
pass
|
| 138 |
+
|
| 139 |
+
def fmt(n):
|
| 140 |
+
return f"{n:,}"
|
| 141 |
+
|
| 142 |
+
print("=" * 60)
|
| 143 |
+
print(" TIER 1 β RAW (full OpenAlex JSON as-is on disk)")
|
| 144 |
+
print("=" * 60)
|
| 145 |
+
print(f" Files: {fmt(total)}")
|
| 146 |
+
print(f" Total bytes: {fmt(raw_total['bytes'])} ({raw_total['bytes'] / 1e6:.1f} MB)")
|
| 147 |
+
print(f" Total chars: {fmt(raw_total['chars'])}")
|
| 148 |
+
print(f" Total words: {fmt(raw_total['words'])}")
|
| 149 |
+
print(f" Est. tokens: {fmt(raw_total['tokens_est'])} (~{raw_total['tokens_est'] / 1e6:.2f}M)")
|
| 150 |
+
print()
|
| 151 |
+
print("=" * 60)
|
| 152 |
+
print(" TIER 2 β OPTIMISED (pipeline-relevant fields only)")
|
| 153 |
+
print(" Fields: id, title, abstract, authorships, topics,")
|
| 154 |
+
print(" keywords, referenced_works, doi, year,")
|
| 155 |
+
print(" cited_by_count, best_oa_location.pdf_url")
|
| 156 |
+
print("=" * 60)
|
| 157 |
+
print(f" Files: {fmt(total)}")
|
| 158 |
+
print(f" Total bytes: {fmt(opt_total['bytes'])} ({opt_total['bytes'] / 1e6:.1f} MB)")
|
| 159 |
+
print(f" Total chars: {fmt(opt_total['chars'])}")
|
| 160 |
+
print(f" Total words: {fmt(opt_total['words'])}")
|
| 161 |
+
print(f" Est. tokens: {fmt(opt_total['tokens_est'])} (~{opt_total['tokens_est'] / 1e6:.2f}M)")
|
| 162 |
+
print()
|
| 163 |
+
print("=" * 60)
|
| 164 |
+
print(" TIER 3 β LLM CONTEXT (what reaches the embedding model + LLM)")
|
| 165 |
+
print(' Format: "Title: β¦\\n\\nAbstract: β¦" per paper')
|
| 166 |
+
print("=" * 60)
|
| 167 |
+
print(f" Files: {fmt(total)}")
|
| 168 |
+
print(f" Total bytes: {fmt(llm_total['bytes'])} ({llm_total['bytes'] / 1e6:.1f} MB)")
|
| 169 |
+
print(f" Total chars: {fmt(llm_total['chars'])}")
|
| 170 |
+
print(f" Total words: {fmt(llm_total['words'])}")
|
| 171 |
+
print(f" Est. tokens: {fmt(llm_total['tokens_est'])} (~{llm_total['tokens_est'] / 1e6:.2f}M)")
|
| 172 |
+
print()
|
| 173 |
+
|
| 174 |
+
r2_pct = (1 - opt_total["bytes"] / raw_total["bytes"]) * 100 if raw_total["bytes"] else 0
|
| 175 |
+
r3_pct = (1 - llm_total["bytes"] / raw_total["bytes"]) * 100 if raw_total["bytes"] else 0
|
| 176 |
+
print(f" Tier 1 β Tier 2 reduction: {r2_pct:.1f}% "
|
| 177 |
+
f"({fmt(raw_total['tokens_est'] - opt_total['tokens_est'])} tokens saved)")
|
| 178 |
+
print(f" Tier 1 β Tier 3 reduction: {r3_pct:.1f}% "
|
| 179 |
+
f"({fmt(raw_total['tokens_est'] - llm_total['tokens_est'])} tokens saved)")
|
| 180 |
+
|
| 181 |
+
|
| 182 |
+
if __name__ == "__main__":
|
| 183 |
+
main()
|