Function main

tests/test_hybrid.py:1096–1254

Main function to run the hybrid raw search benchmark

(
    dataset: str = "arguana",
    top_k: int = 10,
)

    return filename

def main(
    dataset: str = "arguana",
    top_k: int = 10,
):
    """Main function to run the hybrid raw search benchmark"""
    print("=" * 60)
    print("Raw Hybrid Search Benchmark (dense + server-side sparse)")
    print("=" * 60)

    try:
        # Setup
        embedding_model = setup_embedding_model()
        create_session()

        # Load dataset
        corpus, queries, qrels = get_beir_dataset(dataset)
        print(f"Loaded {dataset}: {len(corpus)} docs, {len(queries)} queries, {len(qrels)} qrels")

        # Calculate average query length (for raw text)
        avg_query_len = sum(len(q.split()) for q in queries.values()) / len(queries)
        print(f"Average query length: {avg_query_len:.1f} tokens")

        # Generate dense embeddings
        corpus_ids, corpus_embeddings, query_ids, query_embeddings = generate_dense_embeddings(
            corpus, queries, dataset, embedding_model
        )

        # Build brute-force baselines for evaluation
        print("\nBuilding brute-force baselines...")
        bf_dense = brute_force_dense(
            query_embeddings, corpus_embeddings, corpus_ids, query_ids, top_k, dataset
        )

        # Build sparse vectors for brute-force comparison
        sparse_vectors = build_sparse_vectors_for_bf(corpus, dataset)

        # Get corpus stats for sparse queries
        total_docs = len(sparse_vectors)
        term_doc_freq = defaultdict(int)
        total_len = sum(len(v["text"].split()) for v in sparse_vectors)
        avg_len = total_len / total_docs
        for v in sparse_vectors:
            for idx in v["indices"]:
                term_doc_freq[idx] += 1

        sparse_queries = build_sparse_queries_for_bf(
            queries, dataset, (total_docs, term_doc_freq, avg_len)
        )
        bf_sparse = brute_force_sparse(sparse_queries, sparse_vectors, top_k, dataset)

        # Create collection and indexes
        collection_name = f"{COLLECTION_PREFIX}_{dataset}"
        collection = ensure_collection(collection_name)

        # Upsert hybrid vectors (dense + raw text)
        start_time = time.time()
        upsert_hybrid_vectors(collection, corpus_ids, corpus_embeddings, corpus)
        indexing_time = time.time() - start_time