Custom Embeddings & Retrieval Optimization

import json
from typing import List, Dict

def build_retrieval_benchmark(
    queries: List[str],
    relevant_docs: Dict[str, List[str]],
    corpus: List[str]
) -> Dict:
    """Build an offline benchmark for retrieval evaluation.

    Args:
        queries: list of real user queries
        relevant_docs: mapping query -> list of relevant doc IDs
        corpus: list of all documents
    """
    benchmark = {
        "queries": [],
        "corpus_size": len(corpus),
    }

    for query in queries:
        entry = {
            "query": query,
            "relevant": relevant_docs.get(query, []),
            "num_relevant": len(relevant_docs.get(query, [])),
        }
        benchmark["queries"].append(entry)

    # Save for reproducible evaluation
    with open("retrieval_benchmark.json", "w") as f:
        json.dump(benchmark, f, indent=2)

    print(f"Benchmark: {len(queries)} queries, "
          f"{len(corpus)} docs, "
          f"avg {sum(len(v) for v in relevant_docs.values())/len(queries):.1f} "
          f"relevant per query")
    return benchmark

from openai import OpenAI

client = OpenAI()

def generate_training_queries(document_chunk: str, n: int = 3):
    """Generate synthetic queries for a document chunk.

    This 'doc2query' approach creates training pairs
    when labeled data is scarce.
    """
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{
            "role": "user",
            "content": f"""Generate {n} diverse search queries that
this document chunk would be the ideal answer for.
Make queries realistic (how a user would actually search).
Return one query per line, no numbering.

Document chunk:
{document_chunk}"""
        }],
        temperature=0.8
    )

    queries = response.choices[0].message.content.strip().split("\n")
    # Each (query, chunk) pair becomes a positive training example
    return [{"query": q.strip(), "positive": document_chunk}
            for q in queries if q.strip()]

import numpy as np
from sentence_transformers import SentenceTransformer

def mine_hard_negatives(
    queries, positives, corpus,
    model_name="BAAI/bge-base-en-v1.5",
    top_k=10, n_negatives=3
):
    """Mine hard negatives using current embedding model.

    For each query, find the top-k nearest corpus items
    that are NOT in the positive set.
    """
    model = SentenceTransformer(model_name)

    # Encode everything
    q_emb = model.encode(queries, normalize_embeddings=True)
    c_emb = model.encode(corpus, normalize_embeddings=True)

    # Compute cosine similarities
    sims = q_emb @ c_emb.T  # [n_queries, n_corpus]

    triplets = []
    for i, query in enumerate(queries):
        # Get indices sorted by similarity (descending)
        ranked = np.argsort(-sims[i])
        pos_set = set(positives[i])
        # Hard negatives: highest-similarity non-positives
        hard_negs = [
            corpus[idx] for idx in ranked
            if idx not in pos_set
        ][:n_negatives]

        for neg in hard_negs:
            triplets.append({
                "query": query,
                "positive": corpus[positives[i][0]],
                "negative": neg,
            })
    return triplets

from sentence_transformers import (
    SentenceTransformer, InputExample, losses
)
from torch.utils.data import DataLoader

# Load base model to fine-tune
model = SentenceTransformer("BAAI/bge-base-en-v1.5")

# Prepare training data as (query, positive_doc) pairs
# MNRL uses in-batch negatives automatically
train_examples = [
    InputExample(texts=[
        "metformin side effects in elderly",
        "Common adverse effects of metformin in patients over 65..."
    ]),
    InputExample(texts=[
        "dosing guidelines for lisinopril",
        "Recommended starting dose of lisinopril is 10mg daily..."
    ]),
    # ... more (query, positive) pairs
]

# Use MNRL: in-batch negatives, no mining needed
train_dataloader = DataLoader(train_examples, batch_size=32,
                              shuffle=True)
train_loss = losses.MultipleNegativesRankingLoss(model)

# Fine-tune for 1 epoch
model.fit(
    train_objectives=[(train_dataloader, train_loss)],
    epochs=1,
    warmup_steps=100,
    output_path="./domain-adapted-embedding",
)

from typing import List
import tiktoken

def chunk_document(
    text: str,
    max_tokens: int = 400,
    overlap_tokens: int = 50,
    doc_title: str = "",
) -> List[str]:
    """Split a document into overlapping chunks with metadata prefix."""
    enc = tiktoken.get_encoding("cl100k_base")
    tokens = enc.encode(text)

    chunks = []
    start = 0
    while start < len(tokens):
        end = start + max_tokens
        chunk_tokens = tokens[start:end]
        chunk_text = enc.decode(chunk_tokens)

        # Prefix with document metadata for richer embedding
        if doc_title:
            chunk_text = f"[{doc_title}] {chunk_text}"

        chunks.append(chunk_text)

        # Advance by (max_tokens - overlap) to create overlap
        start += max_tokens - overlap_tokens

    print(f"Document: {len(tokens)} tokens -> {len(chunks)} chunks")
    return chunks

from sentence_transformers import SentenceTransformer
import numpy as np

model = SentenceTransformer("intfloat/multilingual-e5-large")

# Test cross-lingual alignment:
# Query in English, documents in multiple languages
query = "query: What are the side effects of aspirin?"

docs = [
    "passage: Aspirin side effects include stomach bleeding.",  # EN
    "passage: Nebenwirkungen von Aspirin sind Magenblutungen.",  # DE
    "passage: Les effets secondaires de l'aspirine incluent...",  # FR
    "passage: The history of bicycle manufacturing.",  # EN irrelevant
]

# Encode with instruction prefix (required for E5)
q_emb = model.encode([query], normalize_embeddings=True)
d_emb = model.encode(docs, normalize_embeddings=True)

# Cosine similarities
sims = (q_emb @ d_emb.T)[0]
for doc, sim in sorted(zip(docs, sims), key=lambda x: -x[1]):
    print(f"  {sim:.3f}  {doc[:60]}...")
# Expected: all three aspirin docs rank above the bicycle doc

import faiss
import numpy as np

dim = 768
n_vectors = 1_000_000

# Generate sample embeddings (replace with real embeddings)
vectors = np.random.randn(n_vectors, dim).astype("float32")
faiss.normalize_L2(vectors)  # normalize for cosine similarity

# Option 1: Flat index (exact, but uses full memory)
index_flat = faiss.IndexFlatIP(dim)
index_flat.add(vectors)
print(f"Flat: {index_flat.ntotal * dim * 4 / 1e9:.2f} GB")

# Option 2: IVF + PQ (compressed, approximate)
n_lists = 1024   # number of Voronoi cells
m = 48            # number of subquantizers
n_bits = 8        # bits per subquantizer
quantizer = faiss.IndexFlatIP(dim)
index_pq = faiss.IndexIVFPQ(quantizer, dim, n_lists, m, n_bits)
index_pq.train(vectors)
index_pq.add(vectors)
index_pq.nprobe = 32  # search 32 of 1024 cells

# Compare recall: search both indexes for same queries
queries = np.random.randn(100, dim).astype("float32")
faiss.normalize_L2(queries)
_, exact = index_flat.search(queries, 10)
_, approx = index_pq.search(queries, 10)

# Recall@10: fraction of true top-10 found by PQ
recall = np.mean([
    len(np.intersect1d(e, a)) / 10
    for e, a in zip(exact, approx)
])
print(f"PQ Recall@10: {recall:.1%}")

import numpy as np
from sklearn.metrics import precision_recall_curve

def find_optimal_threshold(
    similarities: np.ndarray,
    labels: np.ndarray,
    beta: float = 1.0
) -> dict:
    """Find the threshold that maximizes F-beta score.

    Args:
        similarities: cosine similarities for each (query, doc) pair
        labels: 1 = relevant, 0 = not relevant
        beta: F-beta weight (beta > 1 favors recall,
              beta < 1 favors precision)
    """
    precisions, recalls, thresholds = precision_recall_curve(
        labels, similarities
    )

    # Compute F-beta at each threshold
    f_scores = (
        (1 + beta**2) * precisions * recalls /
        (beta**2 * precisions + recalls + 1e-8)
    )

    best_idx = np.argmax(f_scores)
    return {
        "threshold": float(thresholds[best_idx]),
        "precision": float(precisions[best_idx]),
        "recall": float(recalls[best_idx]),
        "f_score": float(f_scores[best_idx]),
    }

import json, time
from typing import List, Dict

def run_canary_check(
    search_fn,
    canary_set: List[Dict],
    k: int = 10,
) -> Dict:
    """Check retrieval quality against known-good pairs.

    Args:
        search_fn: function(query) -> list of doc_ids
        canary_set: [{"query": str, "expected_doc_id": str}, ...]
        k: number of results to check
    """
    hits = 0
    scores = []
    failures = []

    for canary in canary_set:
        results = search_fn(canary["query"], top_k=k)
        result_ids = [r["id"] for r in results]

        if canary["expected_doc_id"] in result_ids:
            hits += 1
            rank = result_ids.index(canary["expected_doc_id"]) + 1
            scores.append(1.0 / rank)  # reciprocal rank
        else:
            failures.append(canary["query"])
            scores.append(0.0)

    result = {
        "timestamp": time.time(),
        "recall_at_k": hits / len(canary_set),
        "mrr": sum(scores) / len(scores),
        "failures": failures,
    }

    # Alert if recall drops below threshold
    if result["recall_at_k"] < 0.9:
        print(f"ALERT: Canary recall dropped to "
              f"{result['recall_at_k']:.1%}")
    return result

from typing import List, Dict, Callable
import json

def compare_indexes(
    query: str,
    old_search: Callable,
    new_search: Callable,
    k: int = 10,
) -> Dict:
    """Compare retrieval results from old and new embedding indexes.

    Returns overlap metrics and both result sets for analysis.
    """
    old_results = old_search(query, top_k=k)
    new_results = new_search(query, top_k=k)

    old_ids = [r["id"] for r in old_results]
    new_ids = [r["id"] for r in new_results]

    # Measure overlap between result sets
    overlap = len(set(old_ids) & set(new_ids))

    # Rank correlation (Jaccard similarity of top-k)
    jaccard = overlap / len(set(old_ids) | set(new_ids))

    # Top-1 agreement
    top1_match = old_ids[0] == new_ids[0] if old_ids and new_ids else False

    return {
        "query": query,
        "overlap_at_k": overlap / k,
        "jaccard": jaccard,
        "top1_match": top1_match,
        "old_top3": old_ids[:3],
        "new_top3": new_ids[:3],
    }

# Run over benchmark queries and log for analysis
# Look for patterns in disagreements to understand
# where the new model differs from the old one