test_examples.py

"""Test all examples from the README to verify they work as documented."""

from rerank import MLXReranker

print("="*80)
print("TEST 1: Basic Example (from README)")
print("="*80)

# Initialize the reranker
reranker = MLXReranker()

# Your query and documents
query = "What are the health benefits of green tea?"
documents = [
    "Green tea contains antioxidants called catechins that may help reduce inflammation and protect cells from damage.",
    "El precio del café ha aumentado un 20% este año debido a problemas en la cadena de suministro.",
    "Studies show that drinking green tea regularly can improve brain function and boost metabolism.",
    "Basketball is one of the most popular sports in the United States.",
    "绿茶富含儿茶素等抗氧化剂，可以降低心脏病风险，还有助于控制体重。",
    "Le thé vert est riche en antioxydants et peut améliorer la fonction cérébrale.",
]

# Rerank documents
results = reranker.rerank(query, documents)

# Results are sorted by relevance score (highest first)
for result in results:
    print(f"Score: {result['relevance_score']:.4f}")
    print(f"Document: {result['document'][:100]}...")
    print()

print("="*80)
print("TEST 2: Top-N results (from README)")
print("="*80)

# Get only top 3 results
top_results = reranker.rerank(query, documents, top_n=3)
print(f"Requested top 3, got {len(top_results)} results")
for i, result in enumerate(top_results, 1):
    print(f"{i}. Score: {result['relevance_score']:.4f}, Index: {result['index']}")

print("\n" + "="*80)
print("TEST 3: With embeddings (from README)")
print("="*80)

# Get embeddings for further processing
results_with_embeddings = reranker.rerank(query, documents, return_embeddings=True)
for i, result in enumerate(results_with_embeddings[:2], 1):  # Just show first 2
    embedding = result['embedding']  # numpy array of shape (512,)
    print(f"Doc {i}: embedding shape = {embedding.shape}, first 5 values = {embedding[:5]}")

print("\n" + "="*80)
print("ALL TESTS PASSED ✓")
print("="*80)