import numpy as np def cosine_similarity( query_vector: np.ndarray, corpus_vectors: np.ndarray ) -> np.ndarray: """ Calculate cosine similarity between prompt vectors. Args: query_vector: Vectorized prompt query of shape (1, D). corpus_vectors: Vectorized prompt corpus of shape (N, D). Returns: The vector of shape (N,) with values in range [-1, 1] where 1 is max similarity i.e., two vectors are the same. """ query_norm = np.linalg.norm(query_vector, axis=1)[0] corpus_norms = np.linalg.norm(corpus_vectors, axis=1) dot_products = np.dot(corpus_vectors, query_vector.T).flatten() similarities = dot_products / (query_norm * corpus_norms) return similarities