Function generate_dense_embeddings

tests/test_hybrid.py:505–553

Generate dense embeddings for corpus and queries

(corpus: Dict, queries: Dict, dataset: str, embedding_model)

    return corpus, queries, qrels

def generate_dense_embeddings(corpus: Dict, queries: Dict, dataset: str, embedding_model) -> Tuple[List[str], np.ndarray, List[str], np.ndarray]:
    """Generate dense embeddings for corpus and queries"""
    cache_file = Path("datasets") / f"hybrid_{dataset}" / "dense_embeddings.npz"
    if cache_file.exists():
        print("Loading cached dense embeddings...")
        data = np.load(cache_file, allow_pickle=True)
        return (
            data["corpus_ids"].tolist(),
            data["corpus_embeddings"],
            data["query_ids"].tolist(),
            data["query_embeddings"],
        )

    print("Generating dense embeddings...")
    corpus_ids = list(corpus.keys())
    corpus_texts = [
        f"{corpus[doc_id].get('title', '')} {corpus[doc_id]['text']}"
        for doc_id in corpus_ids
    ]

    query_ids = list(queries.keys())
    query_texts = [queries[qid] for qid in query_ids]

    # Generate corpus embeddings
    corpus_embeddings = []
    for i in tqdm(
        range(0, len(corpus_texts), EMBEDDING_BATCH_SIZE), desc="Corpus embeddings"
    ):
        batch = corpus_texts[i : i + EMBEDDING_BATCH_SIZE]
        corpus_embeddings.extend(list(embedding_model.embed(batch)))
    corpus_embeddings = np.array(corpus_embeddings)

    # Generate query embeddings
    query_embeddings = []
    for i in tqdm(range(0, len(query_texts), EMBEDDING_BATCH_SIZE), desc="Query embeddings"):
        batch = query_texts[i : i + EMBEDDING_BATCH_SIZE]
        query_embeddings.extend(list(embedding_model.embed(batch)))
    query_embeddings = np.array(query_embeddings)

    # Cache the embeddings
    cache_file.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(
        cache_file,
        corpus_ids=corpus_ids,
        corpus_embeddings=corpus_embeddings,
        query_ids=query_ids,
        query_embeddings=query_embeddings,
    )
    return corpus_ids, corpus_embeddings, query_ids, query_embeddings


def brute_force_dense(