From 464002a86329357140064c30d6339f75ffb53c45 Mon Sep 17 00:00:00 2001 From: Windpicker-owo <3431391539@qq.com> Date: Sat, 13 Dec 2025 16:59:47 +0800 Subject: [PATCH] =?UTF-8?q?feat(similarity):=20=E9=87=8D=E6=9E=84=E7=9B=B8?= =?UTF-8?q?=E4=BC=BC=E5=BA=A6=E8=AE=A1=E7=AE=97=E5=87=BD=E6=95=B0=EF=BC=8C?= =?UTF-8?q?=E4=BC=98=E5=8C=96=E6=80=A7=E8=83=BD=E5=B9=B6=E5=A2=9E=E5=8A=A0?= =?UTF-8?q?=E6=96=87=E6=A1=A3=E6=B3=A8=E9=87=8A?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/memory_graph/perceptual_manager.py | 41 +-------- src/memory_graph/utils/similarity.py | 119 ++++++++++++++----------- 2 files changed, 70 insertions(+), 90 deletions(-) diff --git a/src/memory_graph/perceptual_manager.py b/src/memory_graph/perceptual_manager.py index 6d58f8c26..2cfd1305e 100644 --- a/src/memory_graph/perceptual_manager.py +++ b/src/memory_graph/perceptual_manager.py @@ -21,6 +21,7 @@ import numpy as np from src.common.logger import get_logger from src.memory_graph.models import MemoryBlock, PerceptualMemory from src.memory_graph.utils.embeddings import EmbeddingGenerator +from src.memory_graph.utils.similarity import _compute_similarities_sync logger = get_logger(__name__) @@ -406,47 +407,9 @@ class PerceptualMemoryManager: ) -> np.ndarray: """在后台线程中向量化计算相似度，避免阻塞事件循环。""" return await asyncio.to_thread( - self._compute_similarities_sync, query_embedding, block_embeddings, block_norms + _compute_similarities_sync, query_embedding, block_embeddings, block_norms ) - @staticmethod - def _compute_similarities_sync( - query_embedding: np.ndarray, - block_embeddings: list[np.ndarray], - block_norms: list[float] | None = None, - ) -> np.ndarray: - import numpy as np - - if not block_embeddings: - return np.zeros(0, dtype=np.float32) - - query = np.asarray(query_embedding, dtype=np.float32) - blocks = np.asarray(block_embeddings, dtype=np.float32) - - if blocks.ndim == 1: - blocks = blocks.reshape(1, -1) - - query_norm = np.linalg.norm(query) - if query_norm == 0.0: - return np.zeros(blocks.shape[0], dtype=np.float32) - - if block_norms is None: - block_norms_array = np.linalg.norm(blocks, axis=1) - else: - block_norms_array = np.asarray(block_norms, dtype=np.float32) - if block_norms_array.shape[0] != blocks.shape[0]: - block_norms_array = np.linalg.norm(blocks, axis=1) - - valid_mask = block_norms_array > 0 - similarities = np.zeros(blocks.shape[0], dtype=np.float32) - - if valid_mask.any(): - normalized_blocks = blocks[valid_mask] / block_norms_array[valid_mask][:, None] - normalized_query = query / query_norm - similarities[valid_mask] = normalized_blocks @ normalized_query - - return np.clip(similarities, 0.0, 1.0) - async def recall_blocks( self, query_text: str, diff --git a/src/memory_graph/utils/similarity.py b/src/memory_graph/utils/similarity.py index b1d8c0d69..f1d62dc03 100644 --- a/src/memory_graph/utils/similarity.py +++ b/src/memory_graph/utils/similarity.py @@ -5,12 +5,69 @@ """ import asyncio -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any if TYPE_CHECKING: import numpy as np +def _compute_similarities_sync( + query_embedding: "np.ndarray", + block_embeddings: "np.ndarray | list[np.ndarray] | list[Any]", + block_norms: "np.ndarray | list[float] | None" = None, +) -> "np.ndarray": + """ + 计算 query 向量与一组向量的余弦相似度（同步/向量化实现）。 + + - 返回 float32 ndarray + - 输出范围裁剪到 [0.0, 1.0] + - 支持可选的 block_norms 以减少重复 norm 计算 + """ + import numpy as np + + if block_embeddings is None: + return np.zeros(0, dtype=np.float32) + + query = np.asarray(query_embedding, dtype=np.float32) + + if isinstance(block_embeddings, (list, tuple)) and len(block_embeddings) == 0: + return np.zeros(0, dtype=np.float32) + + blocks = np.asarray(block_embeddings, dtype=np.float32) + if blocks.dtype == object: + blocks = np.stack( + [np.asarray(vec, dtype=np.float32) for vec in block_embeddings], + axis=0, + ) + + if blocks.size == 0: + return np.zeros(0, dtype=np.float32) + + if blocks.ndim == 1: + blocks = blocks.reshape(1, -1) + + query_norm = float(np.linalg.norm(query)) + if query_norm == 0.0: + return np.zeros(blocks.shape[0], dtype=np.float32) + + if block_norms is None: + block_norms_array = np.linalg.norm(blocks, axis=1).astype(np.float32, copy=False) + else: + block_norms_array = np.asarray(block_norms, dtype=np.float32) + if block_norms_array.shape[0] != blocks.shape[0]: + block_norms_array = np.linalg.norm(blocks, axis=1).astype(np.float32, copy=False) + + dot_products = blocks @ query + denom = block_norms_array * np.float32(query_norm) + + similarities = np.zeros(blocks.shape[0], dtype=np.float32) + valid_mask = denom > 0 + if valid_mask.any(): + np.divide(dot_products, denom, out=similarities, where=valid_mask) + + return np.clip(similarities, 0.0, 1.0) + + def cosine_similarity(vec1: "np.ndarray", vec2: "np.ndarray") -> float: """ 计算两个向量的余弦相似度 @@ -25,23 +82,16 @@ def cosine_similarity(vec1: "np.ndarray", vec2: "np.ndarray") -> float: try: import numpy as np - # 确保是numpy数组 - if not isinstance(vec1, np.ndarray): - vec1 = np.array(vec1) - if not isinstance(vec2, np.ndarray): - vec2 = np.array(vec2) + vec1 = np.asarray(vec1, dtype=np.float32) + vec2 = np.asarray(vec2, dtype=np.float32) - # 归一化 - vec1_norm = np.linalg.norm(vec1) - vec2_norm = np.linalg.norm(vec2) + vec1_norm = float(np.linalg.norm(vec1)) + vec2_norm = float(np.linalg.norm(vec2)) - if vec1_norm == 0 or vec2_norm == 0: + if vec1_norm == 0.0 or vec2_norm == 0.0: return 0.0 - # 余弦相似度 - similarity = np.dot(vec1, vec2) / (vec1_norm * vec2_norm) - - # 确保在 [0, 1] 范围内（处理浮点误差） + similarity = float(np.dot(vec1, vec2) / (vec1_norm * vec2_norm)) return float(np.clip(similarity, 0.0, 1.0)) except Exception: @@ -74,43 +124,10 @@ def batch_cosine_similarity(vec1: "np.ndarray", vec_list: list["np.ndarray"]) -> 相似度列表 """ try: - import numpy as np + if not vec_list: + return [] - # 确保是numpy数组 - if not isinstance(vec1, np.ndarray): - vec1 = np.array(vec1) - - # 批量转换为numpy数组 - vec_list = [np.array(vec) for vec in vec_list] - - # 计算归一化 - vec1_norm = np.linalg.norm(vec1) - if vec1_norm == 0: - return [0.0] * len(vec_list) - - # 计算所有向量的归一化 - vec_norms = np.array([np.linalg.norm(vec) for vec in vec_list]) - - # 避免除以0 - valid_mask = vec_norms != 0 - similarities = np.zeros(len(vec_list)) - - if np.any(valid_mask): - # 批量计算点积 - valid_vecs = np.array(vec_list)[valid_mask] - dot_products = np.dot(valid_vecs, vec1) - - # 计算相似度 - valid_norms = vec_norms[valid_mask] - valid_similarities = dot_products / (vec1_norm * valid_norms) - - # 确保在 [0, 1] 范围内 - valid_similarities = np.clip(valid_similarities, 0.0, 1.0) - - # 填充结果 - similarities[valid_mask] = valid_similarities - - return similarities.tolist() + return _compute_similarities_sync(vec1, vec_list).tolist() except Exception: return [0.0] * len(vec_list) @@ -134,5 +151,5 @@ __all__ = [ "batch_cosine_similarity", "batch_cosine_similarity_async", "cosine_similarity", - "cosine_similarity_async" + "cosine_similarity_async", ]