# 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](https://huggingface.co/spaces/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:
```python
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)

```bash
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`

```bash
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

```bash
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| | Computed |
| 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
```bash
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:
```sql
SELECT arxiv_id FROM papers;
```
Save as `arxiv_ids.txt`, one ID per line.

### Step 2: Run the pipeline
```bash
# 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:

```python
# 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