mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-09-23 00:43:18 +00:00
88 lines
3 KiB
Python
88 lines
3 KiB
Python
"""Vector similarity computation utilities.
|
|
|
|
Provides functions for calculating cosine similarity between vectors,
|
|
with support for both single vectors and batch operations using NumPy.
|
|
"""
|
|
|
|
import numpy as np
|
|
|
|
|
|
def cosine_similarity(vec1: list[float], vec2: list[float]) -> float:
|
|
"""Calculate the cosine similarity between two numeric vectors.
|
|
|
|
Cosine similarity measures the cosine of the angle between two vectors,
|
|
returning a value between -1 (opposite) and 1 (identical direction).
|
|
|
|
Args:
|
|
vec1: First vector as a list of floats.
|
|
vec2: Second vector as a list of floats.
|
|
|
|
Returns:
|
|
Cosine similarity value in range [-1.0, 1.0].
|
|
Returns 0.0 if either vector has zero magnitude.
|
|
|
|
Raises:
|
|
ValueError: If vectors have different lengths.
|
|
|
|
Examples:
|
|
>>> cosine_similarity([1.0, 0.0], [1.0, 0.0])
|
|
1.0
|
|
>>> cosine_similarity([1.0, 0.0], [0.0, 1.0])
|
|
0.0
|
|
>>> cosine_similarity([1.0, 1.0], [-1.0, -1.0])
|
|
-1.0
|
|
"""
|
|
if len(vec1) != len(vec2):
|
|
raise ValueError(f"Vectors must have same length: {len(vec1)} != {len(vec2)}")
|
|
|
|
dot_product = sum(a * b for a, b in zip(vec1, vec2))
|
|
magnitude1 = sum(a * a for a in vec1) ** 0.5
|
|
magnitude2 = sum(b * b for b in vec2) ** 0.5
|
|
|
|
if magnitude1 == 0 or magnitude2 == 0:
|
|
return 0.0
|
|
|
|
return dot_product / (magnitude1 * magnitude2)
|
|
|
|
|
|
def batch_cosine_similarity(nd_array1: np.ndarray, nd_array2: np.ndarray) -> np.ndarray:
|
|
"""Calculate cosine similarity matrix between two batches of vectors.
|
|
|
|
Efficiently computes pairwise cosine similarities using matrix operations.
|
|
|
|
Args:
|
|
nd_array1: Matrix of shape (batch_size1, emb_size) representing
|
|
the first batch of embedding vectors.
|
|
nd_array2: Matrix of shape (batch_size2, emb_size) representing
|
|
the second batch of embedding vectors.
|
|
|
|
Returns:
|
|
Similarity matrix of shape (batch_size1, batch_size2) where
|
|
result[i, j] is the cosine similarity between nd_array1[i] and
|
|
nd_array2[j]. Values are in range [-1.0, 1.0].
|
|
|
|
Raises:
|
|
ValueError: If embedding dimensions don't match between arrays.
|
|
"""
|
|
if nd_array1.shape[1] != nd_array2.shape[1]:
|
|
raise ValueError(
|
|
f"Embedding dimensions must match: {nd_array1.shape[1]} != {nd_array2.shape[1]}",
|
|
)
|
|
|
|
# Compute dot products: (batch_size1, emb_size) @ (emb_size, batch_size2)
|
|
# Result shape: (batch_size1, batch_size2)
|
|
dot_products = np.dot(nd_array1, nd_array2.T)
|
|
|
|
# Compute L2 norms for each vector
|
|
norms1 = np.linalg.norm(nd_array1, axis=1) # Shape: (batch_size1,)
|
|
norms2 = np.linalg.norm(nd_array2, axis=1) # Shape: (batch_size2,)
|
|
|
|
# Compute outer product of norms: (batch_size1, 1) @ (1, batch_size2)
|
|
# Result shape: (batch_size1, batch_size2)
|
|
norm_products = np.outer(norms1, norms2)
|
|
|
|
# Avoid division by zero
|
|
norm_products = np.where(norm_products == 0, 1e-10, norm_products)
|
|
|
|
# Compute cosine similarities
|
|
return dot_products / norm_products
|