| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Example of using ColBERT late interaction models for reranking and scoring. |
| 4 | |
| 5 | ColBERT (Contextualized Late Interaction over BERT) uses per-token embeddings |
| 6 | and MaxSim scoring for document reranking, providing better accuracy than |
| 7 | single-vector models while being more efficient than cross-encoders. |
| 8 | |
| 9 | vLLM supports ColBERT with multiple encoder backbones. Start the server |
| 10 | with one of the following: |
| 11 | |
| 12 | # BERT backbone (works out of the box) |
| 13 | vllm serve answerdotai/answerai-colbert-small-v1 |
| 14 | |
| 15 | # ModernBERT backbone |
| 16 | vllm serve lightonai/GTE-ModernColBERT-v1 \ |
| 17 | --hf-overrides '{"architectures": ["ColBERTModernBertModel"]}' |
| 18 | |
| 19 | # Jina XLM-RoBERTa backbone |
| 20 | vllm serve jinaai/jina-colbert-v2 \ |
| 21 | --hf-overrides '{"architectures": ["ColBERTJinaRobertaModel"]}' \ |
| 22 | --trust-remote-code |
| 23 | |
| 24 | Then run this script: |
| 25 | python colbert_rerank_online.py |
| 26 | """ |
| 27 | |
| 28 | import json |
| 29 | |
| 30 | import requests |
| 31 | |
| 32 | # Change this to match the model you started the server with |
| 33 | MODEL = "answerdotai/answerai-colbert-small-v1" |
| 34 | BASE_URL = "http://127.0.0.1:8000" |
| 35 | |
| 36 | headers = {"accept": "application/json", "Content-Type": "application/json"} |
| 37 | |
| 38 | documents = [ |
| 39 | "Machine learning is a subset of artificial intelligence.", |
| 40 | "Python is a programming language.", |
| 41 | "Deep learning uses neural networks for complex tasks.", |
| 42 | "The weather today is sunny.", |
| 43 | ] |
| 44 | |
| 45 | |
| 46 | def rerank_example(): |
| 47 | """Use the /rerank endpoint to rank documents by query relevance.""" |
| 48 | print("=== Rerank Example ===") |
| 49 | |
| 50 | data = { |
| 51 | "model": MODEL, |
| 52 | "query": "What is machine learning?", |
| 53 | "documents": documents, |
| 54 | } |
| 55 | |
| 56 | response = requests.post(f"{BASE_URL}/rerank", headers=headers, json=data) |
| 57 | result = response.json() |
| 58 | print(json.dumps(result, indent=2)) |
| 59 | |
| 60 | print("\nRanked documents (most relevant first):") |
| 61 | for item in result["results"]: |
| 62 | doc_idx = item["index"] |
| 63 | score = item["relevance_score"] |
| 64 | print(f" Score {score:.4f}: {documents[doc_idx]}") |
| 65 | |
| 66 | |
| 67 | def score_example(): |
| 68 | """Use the /score endpoint for pairwise query-document scoring.""" |
| 69 | print("\n=== Score Example ===") |
| 70 | |
| 71 | data = { |
| 72 | "model": MODEL, |
| 73 | "text_1": "What is machine learning?", |
| 74 | "text_2": [ |
| 75 | "Machine learning is a subset of AI.", |
| 76 | "The weather is sunny.", |
| 77 | ], |
| 78 | } |
| 79 | |
| 80 | response = requests.post(f"{BASE_URL}/score", headers=headers, json=data) |
| 81 | result = response.json() |
| 82 | print(json.dumps(result, indent=2)) |
| 83 | |
| 84 | |
| 85 | def main(): |
| 86 | rerank_example() |
| 87 | score_example() |
| 88 | |
| 89 | |
| 90 | if __name__ == "__main__": |
| 91 | main() |