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