mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-09-08 22:21:15 +00:00
39 lines
1.5 KiB
Python
39 lines
1.5 KiB
Python
"""Cosine similarity for single vectors and batched matrices."""
|
|
|
|
import numpy as np
|
|
|
|
|
|
def cosine_similarity(vec1: list[float], vec2: list[float]) -> float:
|
|
"""Cosine similarity of two equal-length vectors; returns 0.0 if either has zero norm."""
|
|
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:
|
|
"""Pairwise cosine similarity matrix between two batches; output shape (N1, N2)."""
|
|
if nd_array1.ndim != 2 or nd_array2.ndim != 2:
|
|
raise ValueError(
|
|
f"Expected 2D arrays, got shapes {nd_array1.shape} and {nd_array2.shape}",
|
|
)
|
|
if nd_array1.shape[1] != nd_array2.shape[1]:
|
|
raise ValueError(
|
|
f"Embedding dimensions must match: {nd_array1.shape[1]} != {nd_array2.shape[1]}",
|
|
)
|
|
|
|
dot_products = np.dot(nd_array1, nd_array2.T)
|
|
norms1 = np.linalg.norm(nd_array1, axis=1)
|
|
norms2 = np.linalg.norm(nd_array2, axis=1)
|
|
norm_products = np.outer(norms1, norms2)
|
|
# Guard against zero-norm rows to keep division finite.
|
|
norm_products = np.where(norm_products == 0, 1e-10, norm_products)
|
|
|
|
return dot_products / norm_products
|