Add 02_generate_training_triples.py

scripts/02_generate_training_triples.py

@@ -0,0 +1,748 @@
+"""
+Step 2: Generate LightGBM training triples from citation edges.
+
+Produces: train.parquet + eval.parquet
+  Each row = (query_arxiv_id, candidate_arxiv_id, label, feature_1, ..., feature_N)
+
+Labels:
+  2 = directly cited by query paper (strong positive)
+  1 = co-cited with query paper (weak positive)
+  0 = retrieved but not cited (negative)
+
+Time-split:
+  train: query papers published before 2023-01-01
+  eval:  query papers published on or after 2023-01-01
+
+Usage:
+  python 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
+
+Prerequisites:
+  - citations.parquet from Step 1
+  - Qdrant Cloud access (ANN search + embedding retrieval)
+  - Turso access (paper metadata)
+  - pip install httpx pyarrow qdrant-client tqdm numpy
+
+Feature Schema (37 features):
+  See FEATURE_SCHEMA below for the full list.
+  Features 1-20 are populated from citation graph + metadata.
+  Features 21-27 are zero-filled (EWMA/cluster/suppression — need real users).
+  All 37 feature columns are present so the model schema is stable.
+
+Author: ResearchIT ML Pipeline — Phase 6, Step 2
+"""
+from __future__ import annotations
+
+import argparse
+import asyncio
+import json
+import os
+import random
+import time
+from collections import defaultdict
+from datetime import datetime, timezone
+from pathlib import Path
+
+import httpx
+import numpy as np
+import pyarrow as pa
+import pyarrow.parquet as pq
+from tqdm import tqdm
+
+try:
+    from qdrant_client import QdrantClient
+    from qdrant_client.models import Filter, FieldCondition, MatchValue
+except ImportError:
+    print("ERROR: pip install qdrant-client")
+    raise
+
+
+# ── Feature Schema ───────────────────────────────────────────────────────────
+# This defines ALL 37 features. Features 21-27 are zero-filled for pseudo-label
+# training but will be populated when real user data is available.
+#
+# The schema is designed so that the LightGBM model trained on pseudo-labels
+# can be retrained on real data without changing the feature layout.
+
+FEATURE_SCHEMA = [
+    # === Content/Retrieval features (populated during pseudo-label training) ===
+    "qdrant_cosine_score",           # 0: ANN cosine similarity
+    "candidate_position",            # 1: rank position in ANN results (0-indexed)
+    "candidate_citation_count",      # 2: citation count of candidate paper
+    "candidate_log_citations",       # 3: log(citation_count + 1)
+    "candidate_influential_citations",  # 4: influential citation count
+    "candidate_age_days",            # 5: days since candidate was published
+    "candidate_recency_score",       # 6: exp(-0.002 * age_days) — matches heuristic
+    "query_citation_count",          # 7: citation count of query/user paper
+    "query_age_days",                # 8: days since query paper was published
+    "year_diff",                     # 9: |query_year - candidate_year|
+    "same_primary_category",         # 10: 1 if same primary arXiv category, else 0
+    "co_citation_count",             # 11: papers that cite BOTH query and candidate
+    "shared_author_count",           # 12: number of shared authors
+    "candidate_is_newer",            # 13: 1 if candidate published after query, else 0
+    "query_log_citations",           # 14: log(query_citation_count + 1)
+    "citation_count_ratio",          # 15: candidate_citations / (query_citations + 1)
+    "age_ratio",                     # 16: candidate_age / (query_age + 1)
+    "candidate_citations_per_year",  # 17: citation_count / max(age_years, 0.5)
+    "query_num_references",          # 18: how many papers the query paper cites (in-corpus)
+    "candidate_num_cited_by",        # 19: how many corpus papers cite the candidate
+
+    # === User behavior features (zero-filled for pseudo-labels, active for real users) ===
+    "ewma_longterm_similarity",      # 20: cos(candidate, user long-term EWMA profile)
+    "ewma_shortterm_similarity",     # 21: cos(candidate, user short-term EWMA profile)
+    "ewma_negative_similarity",      # 22: cos(candidate, user negative EWMA profile)
+    "cluster_importance",            # 23: importance weight of serving cluster
+    "cluster_distance_to_medoid",    # 24: cos(candidate, cluster medoid)
+    "is_suppressed_category",        # 25: 1 if candidate's category is suppressed
+    "onboarding_category_match",     # 26: 1 if candidate matches user's onboarding categories
+
+    # === Interaction features (zero-filled for pseudo-labels) ===
+    "user_total_saves",              # 27: total papers user has saved
+    "user_total_dismissals",         # 28: total papers user has dismissed
+    "user_days_since_last_save",     # 29: days since user's last save
+    "user_session_save_count",       # 30: saves in current session
+
+    # === Cross features (computed from combinations) ===
+    "cosine_x_recency",             # 31: qdrant_cosine_score × candidate_recency_score
+    "cosine_x_citations",           # 32: qdrant_cosine_score × candidate_log_citations
+    "category_x_recency",           # 33: same_primary_category × candidate_recency_score
+    "cosine_x_cocitation",          # 34: qdrant_cosine_score × log(co_citation_count + 1)
+    "position_inverse",             # 35: 1 / (candidate_position + 1)
+    "citations_x_recency",          # 36: candidate_log_citations × candidate_recency_score
+]
+
+NUM_FEATURES = len(FEATURE_SCHEMA)  # 37
+assert NUM_FEATURES == 37, f"Expected 37 features, got {NUM_FEATURES}"
+
+# Time split cutoff
+EVAL_CUTOFF = "2023-01-01"
+EVAL_CUTOFF_DATE = datetime(2023, 1, 1, tzinfo=timezone.utc)
+
+
+# ── Citation Graph Loading ───────────────────────────────────────────────────
+
+def load_citation_graph(citations_path: str) -> tuple[dict, dict, dict]:
+    """
+    Load citation edges and build lookup structures.
+
+    Returns:
+        references: {citing_id: set(cited_ids)} — outgoing references
+        cited_by:   {cited_id: set(citing_ids)} — incoming citations
+        co_citation_counts: precomputed co-citation matrix (lazily computed per query)
+    """
+    table = pq.read_table(citations_path)
+    citing_col = table.column("citing_arxiv_id").to_pylist()
+    cited_col = table.column("cited_arxiv_id").to_pylist()
+
+    references: dict[str, set[str]] = defaultdict(set)
+    cited_by: dict[str, set[str]] = defaultdict(set)
+
+    for citing, cited in zip(citing_col, cited_col):
+        references[citing].add(cited)
+        cited_by[cited].add(citing)
+
+    print(f"Loaded citation graph:")
+    print(f"  {len(references)} papers with outgoing references")
+    print(f"  {len(cited_by)} papers with incoming citations")
+    print(f"  {sum(len(v) for v in references.values())} total edges")
+
+    return dict(references), dict(cited_by), {}
+
+
+def compute_co_citation_count(
+    query_id: str,
+    candidate_id: str,
+    cited_by: dict[str, set[str]],
+) -> int:
+    """Count papers that cite BOTH query and candidate."""
+    citing_query = cited_by.get(query_id, set())
+    citing_candidate = cited_by.get(candidate_id, set())
+    return len(citing_query & citing_candidate)
+
+
+# ── Turso Metadata Fetching ─────────────────────────────────────────────────
+
+async def fetch_turso_metadata_batch(
+    arxiv_ids: list[str],
+    turso_url: str,
+    turso_token: str,
+) -> dict[str, dict]:
+    """Fetch paper metadata from Turso DB."""
+    if not arxiv_ids:
+        return {}
+
+    pipeline_url = turso_url.rstrip("/")
+    if pipeline_url.startswith("libsql://"):
+        pipeline_url = "https://" + pipeline_url[len("libsql://"):]
+    elif not pipeline_url.startswith("https://"):
+        pipeline_url = "https://" + pipeline_url
+
+    placeholders = ", ".join(["?" for _ in arxiv_ids])
+    sql = f"""SELECT arxiv_id, title, authors, primary_topic, update_date,
+                     citation_count, influential_citations
+              FROM papers WHERE arxiv_id IN ({placeholders})"""
+
+    args = [{"type": "text", "value": aid} for aid in arxiv_ids]
+
+    payload = {
+        "requests": [
+            {"type": "execute", "stmt": {"sql": sql, "args": args}},
+            {"type": "close"},
+        ]
+    }
+
+    headers = {
+        "Authorization": f"Bearer {turso_token}",
+        "Content-Type": "application/json",
+    }
+
+    async with httpx.AsyncClient(timeout=15) as client:
+        resp = await client.post(f"{pipeline_url}/v2/pipeline", json=payload, headers=headers)
+        resp.raise_for_status()
+
+    data = resp.json()
+    results = data.get("results", [])
+    if not results:
+        return {}
+
+    execute_result = results[0]
+    if execute_result.get("type") == "error":
+        print(f"[turso] Query error: {execute_result.get('error')}")
+        return {}
+
+    response = execute_result.get("response", {})
+    result_data = response.get("result", {})
+    cols = [c["name"] for c in result_data.get("cols", [])]
+    rows = result_data.get("rows", [])
+
+    output = {}
+    for row in rows:
+        values = {}
+        for i, col in enumerate(cols):
+            cell = row[i]
+            values[col] = None if cell.get("type") == "null" else cell.get("value", "")
+
+        arxiv_id = values.get("arxiv_id")
+        if not arxiv_id:
+            continue
+
+        # Parse citation counts
+        try:
+            citation_count = int(values.get("citation_count") or 0)
+        except (ValueError, TypeError):
+            citation_count = 0
+        try:
+            influential = int(values.get("influential_citations") or 0)
+        except (ValueError, TypeError):
+            influential = 0
+
+        # Parse authors
+        authors_raw = values.get("authors") or ""
+        if authors_raw.startswith("["):
+            try:
+                author_list = json.loads(authors_raw)
+            except json.JSONDecodeError:
+                author_list = [a.strip() for a in authors_raw.split(",") if a.strip()]
+        else:
+            author_list = [a.strip() for a in authors_raw.split(",") if a.strip()]
+
+        output[arxiv_id] = {
+            "arxiv_id": arxiv_id,
+            "primary_topic": values.get("primary_topic") or "",
+            "update_date": values.get("update_date") or "",
+            "citation_count": citation_count,
+            "influential_citations": influential,
+            "authors": author_list,
+        }
+
+    return output
+
+
+# ── Feature Computation ──────────────────────────────────────────────────────
+
+def compute_paper_age_days(published_str: str) -> int:
+    """Compute age in days from a YYYY-MM-DD date string."""
+    now = datetime.now(timezone.utc)
+    try:
+        pub_date = datetime.strptime(published_str[:10], "%Y-%m-%d").replace(tzinfo=timezone.utc)
+        return max(0, (now - pub_date).days)
+    except (ValueError, TypeError):
+        return 365  # default 1 year
+
+
+def parse_year(published_str: str) -> int:
+    """Extract year from YYYY-MM-DD string."""
+    try:
+        return int(published_str[:4])
+    except (ValueError, TypeError, IndexError):
+        return 2020  # default
+
+
+def compute_shared_authors(authors_a: list[str], authors_b: list[str]) -> int:
+    """Count shared authors between two papers (case-insensitive)."""
+    set_a = {a.lower().strip() for a in authors_a if a.strip()}
+    set_b = {b.lower().strip() for b in authors_b if b.strip()}
+    return len(set_a & set_b)
+
+
+def compute_features_for_pair(
+    query_meta: dict,
+    candidate_meta: dict,
+    qdrant_score: float,
+    candidate_position: int,
+    co_citation_count: int,
+    query_num_references: int,
+    candidate_num_cited_by: int,
+) -> np.ndarray:
+    """
+    Compute the full 37-feature vector for a (query, candidate) pair.
+
+    Features 20-30 (user behavior) are zero-filled for pseudo-label training.
+    """
+    features = np.zeros(NUM_FEATURES, dtype=np.float32)
+
+    # --- Content/Retrieval features (0-19) ---
+
+    # 0: qdrant_cosine_score
+    features[0] = qdrant_score
+
+    # 1: candidate_position
+    features[1] = float(candidate_position)
+
+    # 2: candidate_citation_count
+    cand_citations = candidate_meta.get("citation_count", 0)
+    features[2] = float(cand_citations)
+
+    # 3: candidate_log_citations
+    features[3] = np.log(cand_citations + 1)
+
+    # 4: candidate_influential_citations
+    features[4] = float(candidate_meta.get("influential_citations", 0))
+
+    # 5: candidate_age_days
+    cand_age = compute_paper_age_days(candidate_meta.get("update_date", ""))
+    features[5] = float(cand_age)
+
+    # 6: candidate_recency_score (matches heuristic in reranker.py)
+    features[6] = np.exp(-0.002 * cand_age)
+
+    # 7: query_citation_count
+    query_citations = query_meta.get("citation_count", 0)
+    features[7] = float(query_citations)
+
+    # 8: query_age_days
+    query_age = compute_paper_age_days(query_meta.get("update_date", ""))
+    features[8] = float(query_age)
+
+    # 9: year_diff
+    query_year = parse_year(query_meta.get("update_date", ""))
+    cand_year = parse_year(candidate_meta.get("update_date", ""))
+    features[9] = abs(query_year - cand_year)
+
+    # 10: same_primary_category
+    query_cat = query_meta.get("primary_topic", "")
+    cand_cat = candidate_meta.get("primary_topic", "")
+    features[10] = 1.0 if (query_cat and cand_cat and query_cat == cand_cat) else 0.0
+
+    # 11: co_citation_count
+    features[11] = float(co_citation_count)
+
+    # 12: shared_author_count
+    features[12] = float(compute_shared_authors(
+        query_meta.get("authors", []),
+        candidate_meta.get("authors", []),
+    ))
+
+    # 13: candidate_is_newer
+    features[13] = 1.0 if cand_year > query_year else 0.0
+
+    # 14: query_log_citations
+    features[14] = np.log(query_citations + 1)
+
+    # 15: citation_count_ratio
+    features[15] = cand_citations / (query_citations + 1)
+
+    # 16: age_ratio
+    features[16] = cand_age / (query_age + 1)
+
+    # 17: candidate_citations_per_year
+    cand_age_years = max(cand_age / 365.0, 0.5)
+    features[17] = cand_citations / cand_age_years
+
+    # 18: query_num_references
+    features[18] = float(query_num_references)
+
+    # 19: candidate_num_cited_by
+    features[19] = float(candidate_num_cited_by)
+
+    # --- User behavior features (20-30): zero-filled for pseudo-labels ---
+    # features[20] = ewma_longterm_similarity   → 0.0
+    # features[21] = ewma_shortterm_similarity   → 0.0
+    # features[22] = ewma_negative_similarity    → 0.0
+    # features[23] = cluster_importance           → 0.0
+    # features[24] = cluster_distance_to_medoid   → 0.0
+    # features[25] = is_suppressed_category       → 0.0
+    # features[26] = onboarding_category_match    → 0.0
+    # features[27] = user_total_saves             → 0.0
+    # features[28] = user_total_dismissals        → 0.0
+    # features[29] = user_days_since_last_save    → 0.0
+    # features[30] = user_session_save_count      → 0.0
+
+    # --- Cross features (31-36) ---
+
+    # 31: cosine_x_recency
+    features[31] = features[0] * features[6]
+
+    # 32: cosine_x_citations
+    features[32] = features[0] * features[3]
+
+    # 33: category_x_recency
+    features[33] = features[10] * features[6]
+
+    # 34: cosine_x_cocitation
+    features[34] = features[0] * np.log(co_citation_count + 1)
+
+    # 35: position_inverse
+    features[35] = 1.0 / (candidate_position + 1)
+
+    # 36: citations_x_recency
+    features[36] = features[3] * features[6]
+
+    return features
+
+
+# ── Main Pipeline ────────────────────────────────────────────────────────────
+
+async def generate_triples(
+    citations_path: str,
+    corpus_ids: list[str],
+    qdrant_url: str,
+    qdrant_api_key: str,
+    qdrant_collection: str,
+    turso_url: str,
+    turso_token: str,
+    output_dir: str,
+    num_queries: int,
+    candidates_per_query: int,
+    seed: int = 42,
+):
+    """Main pipeline: load graph → sample queries → ANN search → compute features."""
+
+    output_path = Path(output_dir)
+    output_path.mkdir(parents=True, exist_ok=True)
+
+    # ── Step 1: Load citation graph ──────────────────────────────────────
+    print("=" * 60)
+    print("STEP 1: Loading citation graph...")
+    references, cited_by, _ = load_citation_graph(citations_path)
+
+    corpus_set = set(corpus_ids)
+    print(f"Corpus size: {len(corpus_set)}")
+
+    # Pre-compute per-paper stats
+    num_references = {pid: len(refs) for pid, refs in references.items()}
+    num_cited_by = {pid: len(citers) for pid, citers in cited_by.items()}
+
+    # ── Step 2: Connect to Qdrant ────────────────────────────────────────
+    print("\nSTEP 2: Connecting to Qdrant...")
+    qdrant = QdrantClient(url=qdrant_url, api_key=qdrant_api_key, timeout=30)
+    collection_info = qdrant.get_collection(qdrant_collection)
+    print(f"  Collection: {qdrant_collection}")
+    print(f"  Points: {collection_info.points_count}")
+
+    # ── Step 3: Sample query papers ──────────────────────────────────────
+    print("\nSTEP 3: Sampling query papers...")
+
+    # Only sample papers that have references (otherwise no positive labels)
+    papers_with_refs = [pid for pid in corpus_ids if pid in references and len(references[pid]) >= 3]
+    print(f"  Papers with ≥3 in-corpus references: {len(papers_with_refs)}")
+
+    rng = random.Random(seed)
+    if len(papers_with_refs) > num_queries:
+        sampled_queries = rng.sample(papers_with_refs, num_queries)
+    else:
+        sampled_queries = papers_with_refs
+        print(f"  Warning: only {len(sampled_queries)} papers have enough references")
+
+    print(f"  Sampled {len(sampled_queries)} query papers")
+
+    # ── Step 4: Fetch metadata for all relevant papers ───────────────────
+    print("\nSTEP 4: Fetching metadata from Turso...")
+
+    # Collect all paper IDs we'll need metadata for
+    all_needed_ids = set(sampled_queries)
+    for qid in sampled_queries:
+        all_needed_ids.update(references.get(qid, set()))
+    # We'll also need metadata for ANN candidates, but we fetch those per-batch
+
+    # Fetch in batches of 500 (Turso limit)
+    metadata_cache: dict[str, dict] = {}
+    needed_list = list(all_needed_ids & corpus_set)
+    batch_size = 500
+    for i in tqdm(range(0, len(needed_list), batch_size), desc="Fetching metadata"):
+        batch = needed_list[i:i + batch_size]
+        try:
+            meta = await fetch_turso_metadata_batch(batch, turso_url, turso_token)
+            metadata_cache.update(meta)
+        except Exception as e:
+            print(f"  Warning: metadata batch failed: {e}")
+    print(f"  Cached metadata for {len(metadata_cache)} papers")
+
+    # ── Step 5: Time-split the queries ───────────────────────────────────
+    print(f"\nSTEP 5: Applying time-split (eval cutoff: {EVAL_CUTOFF})...")
+
+    train_queries = []
+    eval_queries = []
+    skipped = 0
+
+    for qid in sampled_queries:
+        meta = metadata_cache.get(qid)
+        if not meta:
+            skipped += 1
+            continue
+        pub_date = meta.get("update_date", "")
+        year = parse_year(pub_date)
+        if year < 2023:
+            train_queries.append(qid)
+        else:
+            eval_queries.append(qid)
+
+    print(f"  Train queries (pre-2023): {len(train_queries)}")
+    print(f"  Eval queries (2023+):     {len(eval_queries)}")
+    print(f"  Skipped (no metadata):    {skipped}")
+
+    # Verify no temporal leakage
+    if train_queries and eval_queries:
+        max_train_year = max(parse_year(metadata_cache[q].get("update_date", "")) for q in train_queries if q in metadata_cache)
+        min_eval_year = min(parse_year(metadata_cache[q].get("update_date", "")) for q in eval_queries if q in metadata_cache)
+        print(f"  Max train year: {max_train_year}")
+        print(f"  Min eval year:  {min_eval_year}")
+        assert max_train_year < min_eval_year, "TEMPORAL LEAKAGE DETECTED!"
+        print(f"  ✅ No temporal leakage")
+
+    # ── Step 6: Generate triples ─────────────────────────────────────────
+    print("\nSTEP 6: Generating training triples...")
+
+    for split_name, query_ids in [("train", train_queries), ("eval", eval_queries)]:
+        if not query_ids:
+            print(f"  Skipping {split_name} — no queries")
+            continue
+
+        print(f"\n  Processing {split_name} split ({len(query_ids)} queries)...")
+
+        all_query_ids = []
+        all_candidate_ids = []
+        all_labels = []
+        all_features = []
+
+        for qi, qid in enumerate(tqdm(query_ids, desc=f"  {split_name}")):
+            query_meta = metadata_cache.get(qid, {})
+            query_refs = references.get(qid, set())
+
+            # Build co-cited set: papers that share references with query
+            co_cited = set()
+            for ref_id in query_refs:
+                co_cited.update(references.get(ref_id, set()))
+            co_cited -= query_refs  # exclude direct citations
+            co_cited.discard(qid)    # exclude self
+
+            # ANN search from Qdrant
+            try:
+                # Look up query paper by arxiv_id payload field
+                # retrieve() takes point IDs (integers), not payload values.
+                # Use scroll() with a FieldCondition filter to find by arxiv_id.
+                scroll_results, _ = qdrant.scroll(
+                    collection_name=qdrant_collection,
+                    scroll_filter=Filter(
+                        must=[FieldCondition(key="arxiv_id", match=MatchValue(value=qid))]
+                    ),
+                    limit=1,
+                    with_vectors=True,
+                    with_payload=True,
+                )
+                if not scroll_results:
+                    continue
+
+                query_vector = scroll_results[0].vector
+                if query_vector is None:
+                    continue
+
+                # ANN search using the query paper's embedding
+                results = qdrant.query_points(
+                    collection_name=qdrant_collection,
+                    query=query_vector,
+                    limit=candidates_per_query,
+                    with_payload=True,
+                )
+
+                candidates = []
+                for hit in results.points:
+                    cand_id = hit.payload.get("arxiv_id") if hit.payload else None
+                    if cand_id and cand_id != qid and cand_id in corpus_set:
+                        candidates.append((cand_id, hit.score))
+
+            except Exception as e:
+                if qi < 3:  # Only print first few errors
+                    print(f"    Warning: Qdrant query failed for {qid}: {e}")
+                continue
+
+            if not candidates:
+                continue
+
+            # Fetch metadata for candidates not yet cached
+            uncached = [cid for cid, _ in candidates if cid not in metadata_cache]
+            if uncached:
+                try:
+                    meta_batch = await fetch_turso_metadata_batch(
+                        uncached[:500], turso_url, turso_token
+                    )
+                    metadata_cache.update(meta_batch)
+                except Exception:
+                    pass
+
+            # Compute features and labels for each candidate
+            for pos, (cand_id, qdrant_score) in enumerate(candidates):
+                cand_meta = metadata_cache.get(cand_id, {})
+
+                # Label assignment
+                if cand_id in query_refs:
+                    label = 2  # direct citation
+                elif cand_id in co_cited:
+                    label = 1  # co-cited
+                else:
+                    label = 0  # not cited
+
+                # Co-citation count
+                cocite_count = compute_co_citation_count(qid, cand_id, cited_by)
+
+                # Feature vector
+                feat = compute_features_for_pair(
+                    query_meta=query_meta,
+                    candidate_meta=cand_meta,
+                    qdrant_score=qdrant_score,
+                    candidate_position=pos,
+                    co_citation_count=cocite_count,
+                    query_num_references=num_references.get(qid, 0),
+                    candidate_num_cited_by=num_cited_by.get(cand_id, 0),
+                )
+
+                all_query_ids.append(qid)
+                all_candidate_ids.append(cand_id)
+                all_labels.append(label)
+                all_features.append(feat)
+
+        # ── Save to parquet ──────────────────────────────────────────────
+        if not all_features:
+            print(f"  No data for {split_name} split!")
+            continue
+
+        feature_matrix = np.array(all_features, dtype=np.float32)
+
+        # Build parquet table
+        columns = {
+            "query_arxiv_id": pa.array(all_query_ids, type=pa.string()),
+            "candidate_arxiv_id": pa.array(all_candidate_ids, type=pa.string()),
+            "label": pa.array(all_labels, type=pa.int32()),
+        }
+
+        # Add each feature as a named column
+        for fi, fname in enumerate(FEATURE_SCHEMA):
+            columns[fname] = pa.array(feature_matrix[:, fi].tolist(), type=pa.float32())
+
+        # Add group_size info (candidates per query, needed for LightGBM)
+        # We track this separately
+        table = pa.table(columns)
+
+        out_file = output_path / f"{split_name}.parquet"
+        pq.write_table(table, str(out_file), compression="snappy")
+
+        # Print stats
+        label_counts = {0: 0, 1: 0, 2: 0}
+        for l in all_labels:
+            label_counts[l] = label_counts.get(l, 0) + 1
+
+        num_queries_actual = len(set(all_query_ids))
+        print(f"\n  {split_name} split saved to {out_file}")
+        print(f"    Rows: {len(all_labels)}")
+        print(f"    Queries: {num_queries_actual}")
+        print(f"    Avg candidates/query: {len(all_labels) / max(num_queries_actual, 1):.1f}")
+        print(f"    Labels: 0={label_counts[0]}, 1={label_counts[1]}, 2={label_counts[2]}")
+        print(f"    Label 2 rate: {100*label_counts[2]/max(len(all_labels),1):.2f}%")
+        print(f"    Label 1 rate: {100*label_counts[1]/max(len(all_labels),1):.2f}%")
+        print(f"    Features: {NUM_FEATURES}")
+
+    # ── Save feature schema ──────────────────────────────────────────────
+    schema_file = output_path / "feature_schema.json"
+    with open(schema_file, "w") as f:
+        json.dump({
+            "features": FEATURE_SCHEMA,
+            "num_features": NUM_FEATURES,
+            "pseudo_label_features": list(range(0, 20)) + list(range(31, 37)),
+            "user_features_zero_filled": list(range(20, 31)),
+            "eval_cutoff": EVAL_CUTOFF,
+            "description": "37-feature schema for ResearchIT LightGBM reranker. "
+                           "Features 20-30 are zero-filled during pseudo-label training "
+                           "and will be populated when real user data is available.",
+        }, f, indent=2)
+    print(f"\nFeature schema saved to {schema_file}")
+
+
+# ── CLI ──────────────────────────────────────────────────────────────────────
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Generate LightGBM training triples from citation graph"
+    )
+    parser.add_argument("--citations", required=True, help="citations.parquet from Step 1")
+    parser.add_argument("--corpus-file", required=True, help="Text file with arXiv IDs")
+    parser.add_argument("--qdrant-url", required=True)
+    parser.add_argument("--qdrant-api-key", required=True)
+    parser.add_argument("--qdrant-collection", default="arxiv_bgem3_dense")
+    parser.add_argument("--turso-url", required=True)
+    parser.add_argument("--turso-token", required=True)
+    parser.add_argument("--output-dir", default="./ltr_dataset")
+    parser.add_argument("--num-queries", type=int, default=100000)
+    parser.add_argument("--candidates-per-query", type=int, default=50)
+    parser.add_argument("--seed", type=int, default=42)
+
+    args = parser.parse_args()
+
+    # Load corpus IDs
+    corpus_ids = []
+    with open(args.corpus_file) as f:
+        for line in f:
+            line = line.strip()
+            if line and not line.startswith("#"):
+                if line.startswith("arXiv:"):
+                    line = line[6:]
+                corpus_ids.append(line)
+    print(f"Loaded {len(corpus_ids)} corpus IDs")
+
+    asyncio.run(generate_triples(
+        citations_path=args.citations,
+        corpus_ids=corpus_ids,
+        qdrant_url=args.qdrant_url,
+        qdrant_api_key=args.qdrant_api_key,
+        qdrant_collection=args.qdrant_collection,
+        turso_url=args.turso_url,
+        turso_token=args.turso_token,
+        output_dir=args.output_dir,
+        num_queries=args.num_queries,
+        candidates_per_query=args.candidates_per_query,
+        seed=args.seed,
+    ))
+
+    print("\n✅ Done! Training triples generated.")
+
+
+if __name__ == "__main__":
+    main()