mirror of
https://github.com/hwchase17/langchain.git
synced 2025-07-16 17:53:37 +00:00
add get_top_k_cosine_similarity method to get max top k score and index (#5059)
# Row-wise cosine similarity between two equal-width matrices and return the max top_k score and index, the score all greater than threshold_score. Co-authored-by: Dev 2049 <dev.dev2049@gmail.com>
This commit is contained in:
parent
039f8f1abb
commit
e57ebf3922
@ -1,5 +1,5 @@
|
||||
"""Math utils."""
|
||||
from typing import List, Union
|
||||
from typing import List, Optional, Tuple, Union
|
||||
|
||||
import numpy as np
|
||||
|
||||
@ -23,3 +23,34 @@ def cosine_similarity(X: Matrix, Y: Matrix) -> np.ndarray:
|
||||
similarity = np.dot(X, Y.T) / np.outer(X_norm, Y_norm)
|
||||
similarity[np.isnan(similarity) | np.isinf(similarity)] = 0.0
|
||||
return similarity
|
||||
|
||||
|
||||
def cosine_similarity_top_k(
|
||||
X: Matrix,
|
||||
Y: Matrix,
|
||||
top_k: Optional[int] = 5,
|
||||
score_threshold: Optional[float] = None,
|
||||
) -> Tuple[List[Tuple[int, int]], List[float]]:
|
||||
"""Row-wise cosine similarity with optional top-k and score threshold filtering.
|
||||
|
||||
Args:
|
||||
X: Matrix.
|
||||
Y: Matrix, same width as X.
|
||||
top_k: Max number of results to return.
|
||||
score_threshold: Minimum cosine similarity of results.
|
||||
|
||||
Returns:
|
||||
Tuple of two lists. First contains two-tuples of indices (X_idx, Y_idx),
|
||||
second contains corresponding cosine similarities.
|
||||
"""
|
||||
if len(X) == 0 or len(Y) == 0:
|
||||
return [], []
|
||||
score_array = cosine_similarity(X, Y)
|
||||
sorted_idxs = score_array.flatten().argsort()[::-1]
|
||||
top_k = top_k or len(sorted_idxs)
|
||||
top_idxs = sorted_idxs[:top_k]
|
||||
score_threshold = score_threshold or -1.0
|
||||
top_idxs = top_idxs[score_array.flatten()[top_idxs] > score_threshold]
|
||||
ret_idxs = [(x // score_array.shape[1], x % score_array.shape[1]) for x in top_idxs]
|
||||
scores = score_array.flatten()[top_idxs].tolist()
|
||||
return ret_idxs, scores
|
||||
|
@ -2,8 +2,19 @@
|
||||
from typing import List
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from langchain.math_utils import cosine_similarity
|
||||
from langchain.math_utils import cosine_similarity, cosine_similarity_top_k
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def X() -> List[List[float]]:
|
||||
return [[1.0, 2.0, 3.0], [0.0, 1.0, 0.0], [1.0, 2.0, 0.0]]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def Y() -> List[List[float]]:
|
||||
return [[0.5, 1.0, 1.5], [1.0, 0.0, 0.0], [2.0, 5.0, 2.0], [0.0, 0.0, 0.0]]
|
||||
|
||||
|
||||
def test_cosine_similarity_zero() -> None:
|
||||
@ -27,13 +38,41 @@ def test_cosine_similarity_empty() -> None:
|
||||
assert len(cosine_similarity(empty_list, np.random.random((3, 3)))) == 0
|
||||
|
||||
|
||||
def test_cosine_similarity() -> None:
|
||||
X = [[1.0, 2.0, 3.0], [0.0, 1.0, 0.0], [1.0, 2.0, 0.0]]
|
||||
Y = [[0.5, 1.0, 1.5], [1.0, 0.0, 0.0], [2.0, 5.0, 2.0]]
|
||||
def test_cosine_similarity(X: List[List[float]], Y: List[List[float]]) -> None:
|
||||
expected = [
|
||||
[1.0, 0.26726124, 0.83743579],
|
||||
[0.53452248, 0.0, 0.87038828],
|
||||
[0.5976143, 0.4472136, 0.93419873],
|
||||
[1.0, 0.26726124, 0.83743579, 0.0],
|
||||
[0.53452248, 0.0, 0.87038828, 0.0],
|
||||
[0.5976143, 0.4472136, 0.93419873, 0.0],
|
||||
]
|
||||
actual = cosine_similarity(X, Y)
|
||||
assert np.allclose(expected, actual)
|
||||
|
||||
|
||||
def test_cosine_similarity_top_k(X: List[List[float]], Y: List[List[float]]) -> None:
|
||||
expected_idxs = [(0, 0), (2, 2), (1, 2), (0, 2), (2, 0)]
|
||||
expected_scores = [1.0, 0.93419873, 0.87038828, 0.83743579, 0.5976143]
|
||||
actual_idxs, actual_scores = cosine_similarity_top_k(X, Y)
|
||||
assert actual_idxs == expected_idxs
|
||||
assert np.allclose(expected_scores, actual_scores)
|
||||
|
||||
|
||||
def test_cosine_similarity_score_threshold(
|
||||
X: List[List[float]], Y: List[List[float]]
|
||||
) -> None:
|
||||
expected_idxs = [(0, 0), (2, 2)]
|
||||
expected_scores = [1.0, 0.93419873]
|
||||
actual_idxs, actual_scores = cosine_similarity_top_k(
|
||||
X, Y, top_k=None, score_threshold=0.9
|
||||
)
|
||||
assert actual_idxs == expected_idxs
|
||||
assert np.allclose(expected_scores, actual_scores)
|
||||
|
||||
|
||||
def test_cosine_similarity_top_k_and_score_threshold(
|
||||
X: List[List[float]], Y: List[List[float]]
|
||||
) -> None:
|
||||
expected_idxs = [(0, 0), (2, 2), (1, 2), (0, 2)]
|
||||
expected_scores = [1.0, 0.93419873, 0.87038828, 0.83743579]
|
||||
actual_idxs, actual_scores = cosine_similarity_top_k(X, Y, score_threshold=0.8)
|
||||
assert actual_idxs == expected_idxs
|
||||
assert np.allclose(expected_scores, actual_scores)
|
||||
|
Loading…
Reference in New Issue
Block a user