examples/embedding_search.py

import numpy as np


def cosine_similarity(query, matrix):
    query_norm = query / np.linalg.norm(query)
    matrix_norm = matrix / np.linalg.norm(matrix, axis=1, keepdims=True)
    return matrix_norm @ query_norm


def main():
    embedding_dim = 1024
    query = np.random.randn(embedding_dim).astype(np.float32)
    index = np.random.randn(5, embedding_dim).astype(np.float32)
    scores = cosine_similarity(query, index)
    ranked = np.argsort(scores)[::-1]

    print("Nearest neighbors:", ranked.tolist())
    print("Scores:", scores[ranked].tolist())


if __name__ == "__main__":
    main()