64 lines
2.1 KiB
Python
64 lines
2.1 KiB
Python
"""
|
|
Test script for the Jina Reranker integration.
|
|
This script tests the reranker functionality by comparing results with and without reranking.
|
|
"""
|
|
|
|
import json
|
|
import time
|
|
import os
|
|
from pathlib import Path
|
|
from typing import Dict, List, Any, Optional
|
|
|
|
# Import just what we need for the simple test
|
|
from ranking.jina_reranker import JinaReranker, get_jina_reranker
|
|
|
|
def test_simple_reranker():
|
|
"""Test the Jina Reranker with a simple query and documents"""
|
|
# Initialize the reranker directly without parameters (it will read from config)
|
|
try:
|
|
reranker = get_jina_reranker()
|
|
print("Successfully initialized Jina Reranker")
|
|
except Exception as e:
|
|
print(f"Error initializing Jina Reranker: {str(e)}")
|
|
return
|
|
|
|
# Simple query and documents
|
|
query = "What is quantum computing?"
|
|
documents = [
|
|
"Quantum computing is a type of computation that harnesses quantum mechanics.",
|
|
"Classical computers use bits, while quantum computers use qubits.",
|
|
"Machine learning is a subset of artificial intelligence.",
|
|
"Quantum computers can solve certain problems faster than classical computers."
|
|
]
|
|
|
|
print(f"Testing reranker with query: {query}")
|
|
print(f"Documents: {documents}")
|
|
|
|
# Rerank the documents
|
|
try:
|
|
reranked = reranker.rerank(query, documents)
|
|
print(f"Reranked results: {json.dumps(reranked, indent=2)}")
|
|
|
|
# Save the results to a file for analysis
|
|
results_dir = Path("results")
|
|
results_dir.mkdir(exist_ok=True)
|
|
results_file = results_dir / f"reranked_results_{int(time.time())}.json"
|
|
|
|
with open(results_file, "w") as f:
|
|
json.dump(reranked, f, indent=2)
|
|
|
|
print(f"Results saved to {results_file}")
|
|
return True
|
|
except Exception as e:
|
|
print(f"Error reranking: {str(e)}")
|
|
return False
|
|
|
|
if __name__ == "__main__":
|
|
# Just run the simple test
|
|
success = test_simple_reranker()
|
|
|
|
if success:
|
|
print("Jina Reranker test completed successfully!")
|
|
else:
|
|
print("Jina Reranker test failed.")
|