similarity.py

import numpy as np

def cosine_similarity(query_vector: np.ndarray, corpus_vectors: np.ndarray) -> np.ndarray:
    query_norm = query_vector / np.linalg.norm(query_vector)
    corpus_norm = corpus_vectors / np.linalg.norm(corpus_vectors, axis=1, keepdims=True)
    return np.dot(corpus_norm, query_norm)