setup_reranker.py

"""
setup_reranker.py
- Downloads Qwen3-Reranker-0.6B
- Exports it to ONNX and saves locally
- Loads ONNX model with GPU and runs a test
"""

import os
import numpy as np
from pathlib import Path
import time
import torch
from sentence_transformers import CrossEncoder

# -------------------------------------------------------------------
# 1. Configuration
# -------------------------------------------------------------------
MODEL_ID = "Qwen/Qwen3-Reranker-0.6B"
ONNX_DIR = "./qwen3-reranker-onnx"          # where to save converted model
QUERY = "Which planet is known as the Red Planet?"
PASSAGES = [
    "Venus is often called Earth's twin because of its similar size and proximity.",
    "Mars, known for its reddish appearance, is often referred to as the Red Planet.",
    "Jupiter, the largest planet in our solar system, has a prominent red spot."
]

t0 = time.time()
device = os.getenv("LEXGRAPH_RERANKER_DEVICE") or ("cuda" if torch.cuda.is_available() else "cpu")
try:
    model = CrossEncoder(
        "Qwen/Qwen3-Reranker-0.6B",
        device=device,
    )
except Exception as exc:
    if device != "cuda":
        raise

    error_text = str(exc).lower()
    if not any(marker in error_text for marker in ("cuda", "nvml", "cudacachingallocator", "gpu", "oom", "out of memory")):
        raise

    print("⚠️  Reranker CUDA load failed; falling back to CPU:", exc)
    model = CrossEncoder("Qwen/Qwen3-Reranker-0.6B", device="cpu")

print(f"Model loaded in {time.time() - t0:.1f}s")

print(device := model.device)  # Should print "cuda" or "cuda:0"
t1 = time.time()
scores = model.predict([(QUERY, passage) for passage in PASSAGES], batch_size=1)
print(f"Prediction took {time.time() - t1:.3f}s")
print(scores)
# print("Relevance Scores:")
# for passage, score in zip(PASSAGES, scores):
#     print(f"Score: {score:.4f} | Passage: {passage[:60]}...")