Meshyboi commited on
Commit
90645a4
Β·
verified Β·
1 Parent(s): 2c06739

Upload 27 files

Browse files
.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()