"""Retrieve with pplx-embed, rerank with zerank: the recommended pair from the findings.

  pip install "onnxruntime>=1.31" tokenizers numpy huggingface_hub
  python pipeline_example.py

Downloads the two recommended variants (about 1.9 GB) on first use.
"""
import numpy as np

from inference import load

docs = [
    "To reset the router, hold the reset button on the back for ten seconds until the light blinks.",
    "The router has four LAN ports and one WAN port.",
    "Firmware updates are installed from the admin page under System > Update.",
    "Para reiniciar el router, mantén pulsado el botón de reinicio durante diez segundos.",
    "The warranty covers hardware defects for two years.",
]
query = "How do I factory-reset my router?"

embedder = load("nynefold/pplx-embed-v1-0.6b-ONNX")
reranker = load("nynefold/zerank-1-small-ONNX")

# 1. embed and take the top k (the findings rerank the top 20; 15 costs ~0.4 nDCG points)
d = embedder.embed_documents(docs)
q = embedder.embed_queries([query])
top = np.argsort(-(q @ d.T)[0])[:20]

# 2. rerank the candidates
scores = reranker.rerank(query, [docs[i] for i in top])
for s, i in sorted(zip(scores, top), reverse=True):
    print(f"{s:.3f}  {docs[i]}")