README.md

ResearchIT Phase 6 — LightGBM Reranker

Status: Pipeline tested on synthetic data. Awaiting real citation edges from Semantic Scholar API to produce the production model.
Parent project: siddhm11/ResearchIT
Replaces: Hand-tuned heuristic scorer in app/recommend/reranker.py

---

What This Is

A LightGBM lambdarank model that replaces the hand-tuned heuristic scorer in ResearchIT's recommendation pipeline. The heuristic uses 5 features with fixed weights. This model uses 37 features and learns optimal weights from data.

This repo contains:

• scripts/ — 3 Python scripts that build the full pipeline (data → train → evaluate)
• synthetic_model/ — A model trained on synthetic data (proof the pipeline works)
• test_results.json — Full benchmark results from the test suite

This repo does NOT yet contain:

• A production model trained on real citation data (that's the next step — run Script 1 with your S2 API key)

---

How It Works

The Problem

ResearchIT recommends arXiv papers using a multi-stage pipeline:

Qdrant ANN retrieval → Quota fusion → Reranking → MMR diversity → Feed

The reranking step currently uses a heuristic scorer:

score = 0.40 × cos(paper, long_term_profile)
      + 0.25 × cos(paper, short_term_profile)
      + 0.15 × recency_decay
      + 0.10 × retrieval_rank_confidence
      - 0.15 × cos(paper, negative_profile)

These weights are hand-tuned guesses. The model can't use citation count, co-citation, category match, or any feature interactions.

The Solution

Train a LightGBM lambdarank model on citation-graph pseudo-labels:

• Each arXiv paper acts as a "pseudo-user" — its reference list simulates what that researcher would "save"
• Direct citations → label 2 (strong positive)
• Co-citations → label 1 (weak positive — papers that share references)
• ANN-retrieved but not cited → label 0 (negative)

The model learns: given 37 features about a (user, paper) pair, which papers should rank higher?

---

The 3-Script Pipeline

Script 1: 01_fetch_citation_edges.py

What: Fetches the citation graph from Semantic Scholar API.
Input: arxiv_ids.txt (your 1.6M arXiv IDs, one per line)
Output: citations.parquet — (citing_arxiv_id, cited_arxiv_id, is_influential)

python scripts/01_fetch_citation_edges.py \
  --corpus-file arxiv_ids.txt \
  --output citations.parquet \
  --api-key YOUR_S2_API_KEY   # optional, faster with key

• Supports checkpoint/resume — safe to interrupt
• Supports both batch API (simple) and bulk download (faster for 1.6M papers)
• Filters to in-corpus edges only (both papers must be in your Qdrant collection)
• Expected output: ~20-40M citation edges

Script 2: 02_generate_training_triples.py

What: Generates labeled training data by running ANN searches against Qdrant.
Input: citations.parquet + Qdrant access + Turso access
Output: ltr_dataset/train.parquet + ltr_dataset/eval.parquet + feature_schema.json

python scripts/02_generate_training_triples.py \
  --citations citations.parquet \
  --corpus-file arxiv_ids.txt \
  --qdrant-url "https://YOUR_QDRANT_URL" \
  --qdrant-api-key "YOUR_KEY" \
  --qdrant-collection arxiv_bgem3_dense \
  --turso-url "https://YOUR_TURSO_URL" \
  --turso-token "YOUR_TOKEN" \
  --output-dir ./ltr_dataset \
  --num-queries 100000 \
  --candidates-per-query 50

• Enforces time-split: train on papers before 2023, eval on 2023+
• Asserts no temporal leakage: max(train.year) < min(eval.year)
• Uses scroll() + query_points() for Qdrant (correct API for qdrant-client 1.17+)
• Expected output: ~5M train rows, ~500K eval rows

Script 3: 03_train_lightgbm.py

What: Trains LightGBM and compares against the heuristic baseline.
Input: train.parquet + eval.parquet
Output: reranker_v1.txt (~100-300KB model file) + metrics

python scripts/03_train_lightgbm.py \
  --train-file ltr_dataset/train.parquet \
  --eval-file ltr_dataset/eval.parquet \
  --output-dir ./model_output \
  --num-boost-round 500 \
  --learning-rate 0.05

• Evaluates: nDCG@5, nDCG@10, nDCG@20, Recall@10, Recall@50, HR@10, MRR
• Compares LightGBM vs the exact heuristic from reranker.py
• Reports feature importance, latency benchmark, per-query win rates
• Outputs verdict: deploy / marginal / no improvement

---

Feature Schema (37 features)

The model uses 37 features organized in 4 groups:

Content/Retrieval Features (0-19) — Populated during pseudo-label training

#	Feature	Description	Source
0	qdrant_cosine_score	BGE-M3 cosine similarity from ANN search	Qdrant
1	candidate_position	Rank position in ANN results (0-indexed)	Qdrant
2	candidate_citation_count	Total citation count	Turso
3	candidate_log_citations	log(citation_count + 1)	Computed
4	candidate_influential_citations	Influential citation count (Semantic Scholar)	Turso
5	candidate_age_days	Days since publication	Turso
6	candidate_recency_score	exp(-0.002 × age_days) — matches heuristic formula	Computed
7	query_citation_count	Citation count of the query/user paper	Turso
8	query_age_days	Days since query paper was published	Turso
9	year_diff		query_year - candidate_year
10	same_primary_category	1 if same primary arXiv category, else 0	Turso
11	co_citation_count	Papers that cite BOTH query and candidate	Citation graph
12	shared_author_count	Number of shared authors (case-insensitive)	Turso
13	candidate_is_newer	1 if candidate published after query	Computed
14	query_log_citations	log(query_citation_count + 1)	Computed
15	citation_count_ratio	candidate_citations / (query_citations + 1)	Computed
16	age_ratio	candidate_age / (query_age + 1)	Computed
17	candidate_citations_per_year	citation_count / max(age_years, 0.5)	Computed
18	query_num_references	How many papers the query paper cites (in-corpus)	Citation graph
19	candidate_num_cited_by	How many corpus papers cite the candidate	Citation graph

User Behavior Features (20-30) — Zero-filled for pseudo-labels, active for real users

#	Feature	Description	Source in ResearchIT
20	ewma_longterm_similarity	cos(candidate, user long-term EWMA profile)	profiles.py α=0.03
21	ewma_shortterm_similarity	cos(candidate, user short-term EWMA profile)	profiles.py α=0.40
22	ewma_negative_similarity	cos(candidate, user negative EWMA profile)	profiles.py α=0.15
23	cluster_importance	Importance weight of the serving cluster	clustering.py
24	cluster_distance_to_medoid	cos(candidate, cluster medoid embedding)	clustering.py
25	is_suppressed_category	1 if candidate's category is suppressed (≥3 dismissals in 14 days)	db.py
26	onboarding_category_match	1 if candidate matches user's onboarding selections	db.py
27	user_total_saves	Total papers user has saved	interactions table
28	user_total_dismissals	Total papers user has dismissed	interactions table
29	user_days_since_last_save	Days since user's most recent save	interactions table
30	user_session_save_count	Saves in current session	In-memory state

Cross Features (31-36) — Feature interactions

#	Feature	Formula
31	cosine_x_recency	qdrant_cosine_score × candidate_recency_score
32	cosine_x_citations	qdrant_cosine_score × candidate_log_citations
33	category_x_recency	same_primary_category × candidate_recency_score
34	cosine_x_cocitation	qdrant_cosine_score × log(co_citation_count + 1)
35	position_inverse	1 / (candidate_position + 1)
36	citations_x_recency	candidate_log_citations × candidate_recency_score

Why 37 Features Instead of 5?

The heuristic uses features 0, 1, 5 (via recency), 20, 21 — that's it. LightGBM additionally uses:

• Citation count (#2, #3, #17) — papers with more citations are generally higher quality
• Co-citation count (#11) — the strongest signal in our tests. Papers cited by the same community as your interests are highly relevant
• Category match (#10) — same arXiv category = topically aligned
• Shared authors (#12) — self-citation / collaborator network signal
• Feature interactions (#31-36) — e.g., a paper that is both semantically similar AND frequently co-cited is much more relevant than either signal alone

---

Benchmark Results (Synthetic Data)

These results are from realistic synthetic data that mimics citation graph patterns. Real data results will differ.

LightGBM vs Heuristic Baseline

Metric	Heuristic	LightGBM	Improvement
nDCG@3	0.9305	0.9979	+7.2%
nDCG@5	0.9138	0.9980	+9.2%
nDCG@10	0.9111	0.9985	+9.6%
nDCG@20	0.9472	0.9988	+5.5%

Per-Query Win Rate (500 eval queries)

• LightGBM wins: 91.4%
• Heuristic wins: 0.4%
• Ties: 8.2%

Production Metrics

Metric	Value	Target	Status
Latency (100 candidates)	0.088ms	<1ms	✅ 11× under budget
Model size	286 KB	<500 KB	✅
Model reload	Identical predictions	—	✅
Handles NaN input	Graceful	—	✅
Handles extreme values	No crash	—	✅
Train-eval gap	0.0008	<0.05	✅ No overfitting

Top 10 Features by Importance

Rank	Feature	Importance
1	cosine_x_cocitation	95,386
2	qdrant_cosine_score	37,300
3	co_citation_count	14,124
4	candidate_citation_count	3,399
5	candidate_position	2,615
6	cosine_x_citations	1,623
7	year_diff	1,229
8	candidate_log_citations	894
9	age_ratio	564
10	position_inverse	427

The #1 feature is cosine_x_cocitation — the interaction between embedding similarity and co-citation count. This is a signal the heuristic cannot access.

---

How to Produce the Production Model

Prerequisites

pip install httpx pyarrow tqdm numpy qdrant-client lightgbm

Step 1: Export corpus IDs

Export your 1.6M arXiv IDs from Turso to a text file:

SELECT arxiv_id FROM papers;

Save as arxiv_ids.txt, one ID per line.

Step 2: Run the pipeline

# Fetch citation edges (~30-90 min)
python scripts/01_fetch_citation_edges.py \
  --corpus-file arxiv_ids.txt \
  --output citations.parquet

# Generate training data (~1-2 hours, needs Qdrant + Turso access)
python scripts/02_generate_training_triples.py \
  --citations citations.parquet \
  --corpus-file arxiv_ids.txt \
  --qdrant-url "$QDRANT_URL" \
  --qdrant-api-key "$QDRANT_API_KEY" \
  --turso-url "$TURSO_URL" \
  --turso-token "$TURSO_DB_TOKEN" \
  --output-dir ./ltr_dataset

# Train (~5 min)
python scripts/03_train_lightgbm.py \
  --train-file ltr_dataset/train.parquet \
  --eval-file ltr_dataset/eval.parquet \
  --output-dir ./model_output

Step 3: Deploy

Copy model_output/reranker_v1.txt to your ResearchIT Space.

---

Integration Into ResearchIT

The model integrates into app/recommend/reranker.py with minimal changes:

# Load once at startup
import lightgbm as lgb

_lgb_model = None
try:
    _lgb_model = lgb.Booster(model_file="reranker_v1.txt")
    print("[reranker] LightGBM model loaded")
except Exception:
    print("[reranker] Using heuristic fallback")

# In rerank_candidates():
features = compute_features_v2(...)  # expanded to 37 features
if _lgb_model is not None:
    scores = _lgb_model.predict(features)
else:
    scores = heuristic_score(features)  # existing fallback — always works

The heuristic scorer remains as a fallback. If the model file is missing or fails to load, the system silently uses the heuristic. No user-facing impact.

This is LightGBM-1 in the Doc 07 multi-stage architecture:

Qdrant ANN → LightGBM-1 (this model) → [TinyBERT → LightGBM-2] (Phase 8b, future)

---

Time-Split Evaluation

Training data uses a strict temporal split to prevent leakage:

• Train: Query papers published before 2023-01-01
• Eval: Query papers published on or after 2023-01-01
• Assertion: max(train.year) < min(eval.year) — verified in Script 2

This means the model never sees future papers citing past papers during training — matching how it would perform in production on newly published papers.

---

Known Limitations

Citation ≠ User Interest

Citation pseudo-labels ("this author cited that paper in their bibliography") are different from real user signals ("this user saved that paper in their feed"). A foundational paper like "Attention Is All You Need" gets label=2 in citation data but might be dismissed by users who've already read it.

Mitigation: The candidate_log_citations and candidate_citations_per_year features help LightGBM learn a popularity penalty. When 500+ real user interactions accumulate, retrain on actual save/dismiss data — the 11 user behavior features (20-30) activate and the model learns real preferences.

Synthetic vs Real

The benchmark numbers above are from synthetic data. Real citation data will have:

• Noisier labels (not all citations are "relevant" — some are obligatory references)
• Lower nDCG improvement over heuristic (expect +3-5% on real data vs +9.6% on synthetic)
• Different feature importance rankings

---

Dependencies

lightgbm>=4.0
httpx>=0.24
pyarrow>=12.0
numpy>=1.24
qdrant-client>=1.17
tqdm>=4.65

---

References

• ResearchIT Doc 06 §3.1 — LightGBM lambdarank architecture decision
• ResearchIT Doc 07 §A6 — Time-split evaluation protocol
• ResearchIT Doc 07 §B.4 — Multi-stage reranker architecture
• PinnerSage (Pal et al., KDD 2020) — Ward clustering + importance-weighted retrieval
• Taobao ULIM (Meng et al., RecSys 2025) — Quota allocation, +5.54% CTR
• YouTube (Xia et al., 2023) — 3× gain from negative signals in reranking
• Bruch et al. (SIGIR 2022) — RRF optimizes Recall not nDCG