Implementing ZeroEntropy Reranking with turbopuffer Retrieval

!pip -q install turbopuffer zeroentropy sentence-transformers python-dotenv

import os
from dotenv import load_dotenv
load_dotenv()

# ↓ Optionally set directly for quick testing in Colab (replace the placeholders)
# os.environ['TURBOPUFFER_API_KEY'] = 'tpuf_...'
# os.environ['ZEROENTROPY_API_KEY'] = 'ze_...'

assert os.getenv('TURBOPUFFER_API_KEY') is not None, 'Please set TURBOPUFFER_API_KEY'
assert os.getenv('ZEROENTROPY_API_KEY') is not None, 'Please set ZEROENTROPY_API_KEY'
print('Keys detected ✅')

from turbopuffer import Turbopuffer
from sentence_transformers import SentenceTransformer
import uuid, time

# Init SDK
tpuf = Turbopuffer(api_key=os.environ['TURBOPUFFER_API_KEY'], region="gcp-us-central1",)
ns_name = f'ze_tpuf_demo_{int(time.time())}'
ns = tpuf.namespace(ns_name)

# Tiny corpus
docs = [
  {
    "id": str(uuid.uuid4()),
    "title": "Henry Cavendish",
    "text": "In 1798, Cavendish used a torsion balance to measure the Earth's density and effectively determine the gravitational constant G. His experiment measured the tiny gravitational attraction between lead spheres."
  },
  {
    "id": str(uuid.uuid4()),
    "title": "Isaac Newton",
    "text": "Newton formulated the law of universal gravitation, F = G m1 m2 / r^2, explaining that gravitational force obeys an inverse-square dependence on distance."
  },
  {
    "id": str(uuid.uuid4()),
    "title": "Charles-Augustin de Coulomb",
    "text": "Coulomb used a torsion balance to study electrostatic forces, establishing Coulomb's law for charges, which also follows an inverse-square relationship."
  },
  {
    "id": str(uuid.uuid4()),
    "title": "Gravitational Potential Energy",
    "text": "The potential energy between two masses is U = -G m1 m2 / r; near Earth's surface U ≈ m g h."
  }
]


# Build vectors (local, no external API)
model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
embs = model.encode([d['text'] for d in docs], normalize_embeddings=True)

# Upsert into Turbopuffer: text (BM25) + vector
rows = []
for d, v in zip(docs, embs):
    rows.append({
        'id': d['id'],
        'vector': v.tolist(),
        'title': d['title'],
        'text': d['text'],
    })

res = ns.write(
    upsert_rows=rows,
    distance_metric='cosine_distance',
    # Explicit schema to enable BM25 on `text`
    schema={
        'text': {'type': 'string', 'full_text_search': True},
        'title': {'type': 'string', 'full_text_search': True}
    }
)
print('Upserted:', res.rows_upserted, 'rows to', ns_name)

from typing import List, Dict

query = "Who first measured the gravitational constant — not the person who formulated the law of universal gravitation?"

top_k = 5

# ANN vector search
vec = model.encode([query], normalize_embeddings=True)[0].tolist()
ann = ns.query(
    rank_by=("vector", "ANN", vec),
    top_k=top_k,
    include_attributes=["title","text"]
).rows

# BM25 search
bm25 = ns.query(
    rank_by=("text", "BM25", query),
    top_k=top_k,
    include_attributes=["title","text"]
).rows

def rrf_with_scores(list_of_rankings: List[List[Dict]], k: int = 60):
    """Return (fused_list, rrf_scores_dict)."""
    scores = {}
    for ranking in list_of_rankings:
        for rank, item in enumerate(ranking, start=1):
            scores[item.id] = scores.get(item.id, 0.0) + 1.0 / (k + rank)
    id2item = {item.id: item for ranking in list_of_rankings for item in ranking}
    ordered_ids = sorted(scores, key=scores.get, reverse=True)
    fused = [id2item[i] for i in ordered_ids]
    return fused, scores

# Fuse and print only RRF results
fused_results, rrf_scores = rrf_with_scores([ann, bm25])
fused_results = fused_results[:top_k]

for i, r in enumerate(fused_results, 1):
    print(f"Rank {i}. {r.title} — rrf_score={rrf_scores[r.id]:.6f}")

import os
from zeroentropy import ZeroEntropy

def zeroentropy_rerank_or_unranked(results, query, model="zerank-1-small", k=None, return_scores=False):
    """Return results reordered by ZeroEntropy; optionally include scores."""
    if not os.getenv("ZEROENTROPY_API_KEY"):
        print("Warning: ZEROENTROPY_API_KEY not set, returning unranked results")
        return results

    ze = ZeroEntropy(api_key=os.environ["ZEROENTROPY_API_KEY"])
    documents = [getattr(r, "text", str(r)) for r in results]  # list[str]
    out = ze.models.rerank(
        query=query,
        documents=documents,
        model=model,                 # "zerank-1" or "zerank-1-small"
        top_n=k or len(documents),
    )

    # out.results is sorted by relevance (desc)
    if return_scores:
        return [(results[r.index], r.relevance_score) for r in out.results]
    else:
        return [results[r.index] for r in out.results]

# Use it
reranked_with_scores = zeroentropy_rerank_or_unranked(
    fused_results, query, model="zerank-1-small", return_scores=True
)

for rank, (doc, score) in enumerate(reranked_with_scores, 1):
    print(f"Rank {rank}. {doc.title} — {score:.4f}")