| """Minimal load + score example. Run from this repo dir: python load_example.py""" | |
| import torch | |
| from literank.config import ModelConfig | |
| from literank.model import Ranker | |
| from literank.checkpoint import load_checkpoint | |
| ckpt = torch.load("model.pt", map_location="cpu", weights_only=False) | |
| ranker = Ranker(ModelConfig(**ckpt["config"])) | |
| load_checkpoint("model.pt", ranker) | |
| ranker.eval() | |
| query = "what is late interaction in retrieval?" | |
| docs = [ | |
| "LITE is a learnable late-interaction re-ranker for document retrieval.", | |
| "Bananas are a good source of potassium.", | |
| ] | |
| with torch.no_grad(): | |
| scores = ranker.score([query] * len(docs), docs) | |
| for s, d in sorted(zip(scores.tolist(), docs), reverse=True): | |
| print(f"{s:8.3f} {d}") | |