首次提交：上传本地文件夹

utils/embedding.py

@@ -0,0 +1,95 @@
+import numpy as np
+from typing import List, Dict, Any
+from .llm import embedding as embedding_model
+
+class TextEmbedder:
+    def __init__(self):
+        """初始化文本嵌入器，使用自定义的embedding模型"""
+        pass
+    
+    def get_embedding(self, text: str) -> np.ndarray:
+        """
+        获取文本的嵌入向量
+        
+        参数:
+            text: 输入文本
+        
+        返回:
+            嵌入向量
+        """
+        return embedding_model.embed(text)
+    
+    def get_embeddings(self, texts: List[str]) -> List[np.ndarray]:
+        """
+        批量获取文本的嵌入向量
+        
+        参数:
+            texts: 输入文本列表
+        
+        返回:
+            嵌入向量列表
+        """
+        return [embedding_model.embed(text) for text in texts]
+    
+    def compute_similarity(self, text1: str, text2: str) -> float:
+        """
+        计算两段文本的相似度
+        
+        参数:
+            text1: 第一段文本
+            text2: 第二段文本
+        
+        返回:
+            相似度分数 (0-1)
+        """
+        emb1 = self.get_embedding(text1)
+        emb2 = self.get_embedding(text2)
+        return self._cosine_similarity(emb1, emb2)
+    
+    def find_most_similar(self, query: str, candidates: List[str], top_k: int = 5) -> List[Dict[str, Any]]:
+        """
+        找出与查询最相似的候选文本
+        
+        参数:
+            query: 查询文本
+            candidates: 候选文本列表
+            top_k: 返回的最相似文本数量
+        
+        返回:
+            最相似文本及其相似度的列表
+        """
+        query_emb = self.get_embedding(query)
+        candidate_embs = self.get_embeddings(candidates)
+        
+        similarities = []
+        for i, emb in enumerate(candidate_embs):
+            similarity = self._cosine_similarity(query_emb, emb)
+            similarities.append((i, similarity))
+        
+        # 按相似度降序排序
+        similarities.sort(key=lambda x: x[1], reverse=True)
+        
+        results = []
+        for i, sim in similarities[:top_k]:
+            results.append({
+                "text": candidates[i],
+                "similarity": float(sim)
+            })
+        
+        return results
+    
+    def _cosine_similarity(self, vec1: np.ndarray, vec2: np.ndarray) -> float:
+        """
+        计算两个向量的余弦相似度
+        
+        参数:
+            vec1: 第一个向量
+            vec2: 第二个向量
+        
+        返回:
+            余弦相似度 (0-1)
+        """
+        dot_product = np.dot(vec1, vec2)
+        norm1 = np.linalg.norm(vec1)
+        norm2 = np.linalg.norm(vec2)
+        return dot_product / (norm1 * norm2)