File size: 600 Bytes
14e9a9f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 | 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()
|