Function _compute_brute_force_cpu

tests/test-sparse-vector.py:220–262

CPU-based brute force computation using scipy.sparse

(vectors, query_vectors, dimension, top_k=10)

def _compute_brute_force_cpu(vectors, query_vectors, dimension, top_k=10):
    """CPU-based brute force computation using scipy.sparse"""
    from scipy.sparse import csr_matrix as cpu_csr_matrix

    results = []
    n_vectors = len(vectors)
    print("Building dataset sparse matrix on CPU...")

    # Build dataset matrix
    data_list = []
    indices_list = []
    indptr = [0]

    for vec in tqdm(vectors):
        data_list.extend(vec["values"])
        indices_list.extend(vec["indices"])
        indptr.append(len(data_list))

    A = cpu_csr_matrix((data_list, indices_list, indptr), shape=(n_vectors, dimension))

    print("Computing similarities...")
    for i, query in enumerate(tqdm(query_vectors)):
        # Build query vector
        q_data = query["values"]
        q_indices = query["indices"]
        q_indptr = [0, len(q_data)]

        Q = cpu_csr_matrix((q_data, q_indices, q_indptr), shape=(1, dimension))

        # Compute dot products
        scores = A.dot(Q.T).toarray().flatten()

        # Get top-k
        top_indices = np.argpartition(scores, -top_k)[-top_k:]
        top_indices = top_indices[np.argsort(scores[top_indices])[::-1]]

        top_results = []
        for idx in top_indices:
            top_results.append({"id": vectors[idx]["id"], "score": float(scores[idx])})

        results.append({"query_id": query["id"], "top_results": top_results})

    return results


def _compute_brute_force_gpu(vectors, query_vectors, dimension, top_k=10):