Function compute_brute_force_results

tests/test-sparse-vector.py:192–217

Compute dot product similarity using GPU acceleration with CuPy, fallback to CPU

(vectors, query_vectors, dimension, top_k=10)

def compute_brute_force_results(vectors, query_vectors, dimension, top_k=10):
    """Compute dot product similarity using GPU acceleration with CuPy, fallback to CPU"""
    if os.path.exists(BRUTE_FORCE_RESULTS_FILE):
        print(f"Loading existing brute force results from {BRUTE_FORCE_RESULTS_FILE}")
        with open(BRUTE_FORCE_RESULTS_FILE, "rb") as f:
            return pickle.load(f)

    print(
        f"Computing brute force dot product similarity for {len(query_vectors)} queries..."
    )
    results = []

    try:
        # Try GPU computation first
        print("Attempting GPU computation...")
        results = _compute_brute_force_gpu(vectors, query_vectors, dimension, top_k)
    except Exception as e:
        print(f"GPU computation failed: {e}")
        print("Falling back to CPU computation...")
        results = _compute_brute_force_cpu(vectors, query_vectors, dimension, top_k)

    # Save to disk
    with open(BRUTE_FORCE_RESULTS_FILE, "wb") as f:
        pickle.dump(results, f)
    print(f"Brute force results computed and saved to {BRUTE_FORCE_RESULTS_FILE}")
    return results


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