Skip to content

RRF: batch-encode query variants instead of one embed call per query #95

Description

@VarunGitGood

Why

_rrf() embeds query variants one at a time in a Python loop:

# repi/retrieval/rrf.py:108-113
query_embeddings = []
for q in queries:
    q_emb = self.embedding_func([q])[0]   # N serial model calls
    if hasattr(q_emb, 'tolist'):
        q_emb = q_emb.tolist()
    query_embeddings.append(q_emb)

embedding_func is backed by SentenceTransformers, which batch-encodes far more efficiently than N separate single-item calls (one forward pass, better tensor utilisation). With query expansion producing ~5 variants, this is 5 inference calls where 1 would do — 4× the avoidable model overhead per search, and search_logs runs RRF on most ReAct turns.

Scope (in)

  • Encode all variants in a single batched call:
embeddings = self.embedding_func(queries)
query_embeddings = [e.tolist() if hasattr(e, "tolist") else e for e in embeddings]
  • Preserve the per-row .tolist() normalisation (pgvector store expects plain lists).
  • Preserve ordering (zip(queries, query_embeddings) downstream depends on it).

Scope (out)

  • Changing the embedding model or the embedding_func signature.

Acceptance

  • _rrf produces identical rankings to before (same embeddings, just batched).
  • A unit test confirms embedding_func is invoked once with the full query list rather than once per query.

Files

  • repi/retrieval/rrf.py_rrf, lines ~108-113.

Metadata

Metadata

Assignees

No one assigned

    Labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions