如何实现 Embedding 相似度查询

Embedding 是 LLM 的重要组件，如果通过 LlamaIndex 进行 Embedding 查询，使用 Vector Index 就可以方便的实现。如果我们想自己实现 Embedding 的相似度计算呢？Embedding 的流程是调用 Embedding 模型，模型返回向量，查询向量，自己实现的好处是不需要依赖 LlamaIndex，可以自己实现Vector存储、Vector 查询。这里我们可以借鉴 LlamaIndex 的实现方式，LlamaIndex 多种 Index，SummaryIndex 默认是返回所有的节点，同时也支持 Embedding 和 LLM 进行搜索。本文根据 SummaryIndex 的实现方式，实现自己的 Embedding 的相似度查询。

创建 Embedding

Embedding 是通过模型生成向量，LlamaIndex中，所有 Model 都是继承自 BaseEmbedding，通过get_query_embedding 获取向量，Ollama Embedding 创建 Vector。

from llm import get_ollama_embbedingembed_model = get_ollama_embbeding()
res = embed_model.get_query_embedding("北京")

LlamaIndex 调用 Ollama API 获取 Embedding，可以通过 Ollama 的 Python 库自行调用。

生成向量：

相似度

有了 Embedding 向量之后，接下来就是找相似度了，参考 LlamaIndex 中 SummaryIndex 中的 SummaryIndexEmbeddingRetriever，_retrieve 函数代码如下，首先获取所有的 Nodes，随后将生成 Vector，最后相似度计算：

def _retrieve(self,query_bundle: QueryBundle,) -> List[NodeWithScore]:"""Retrieve nodes."""node_ids = self._index.index_struct.nodes# top k nodesnodes = self._index.docstore.get_nodes(node_ids)query_embedding, node_embeddings = self._get_embeddings(query_bundle, nodes)top_similarities, top_idxs = get_top_k_embeddings(query_embedding,node_embeddings,similarity_top_k=self._similarity_top_k,embedding_ids=list(range(len(nodes))),)top_k_nodes = [nodes[i] for i in top_idxs]node_with_scores = []for node, similarity in zip(top_k_nodes, top_similarities):node_with_scores.append(NodeWithScore(node=node, score=similarity))logger.debug(f"> Top {len(top_idxs)} nodes:\n")nl = "\n"logger.debug(f"{nl.join([n.get_content() for n in top_k_nodes])}")return node_with_scores

下面的方法是是计算 Query 与 Node 之间的相似度，默认算法是余弦相似度算法，方法比较简单，相似度算法可以通过 similarity_fn 进行自定义。

def get_top_k_embeddings(query_embedding: List[float],embeddings: List[List[float]],similarity_fn: Optional[Callable[..., float]] = None,similarity_top_k: Optional[int] = None,embedding_ids: Optional[List] = None,similarity_cutoff: Optional[float] = None,
) -> Tuple[List[float], List]:"""Get top nodes by similarity to the query."""if embedding_ids is None:embedding_ids = list(range(len(embeddings)))similarity_fn = similarity_fn or default_similarity_fnembeddings_np = np.array(embeddings)query_embedding_np = np.array(query_embedding)similarity_heap: List[Tuple[float, Any]] = []for i, emb in enumerate(embeddings_np):similarity = similarity_fn(query_embedding_np, emb)if similarity_cutoff is None or similarity > similarity_cutoff:heapq.heappush(similarity_heap, (similarity, embedding_ids[i]))if similarity_top_k and len(similarity_heap) > similarity_top_k:heapq.heappop(similarity_heap)result_tups = sorted(similarity_heap, key=lambda x: x[0], reverse=True)result_similarities = [s for s, _ in result_tups]result_ids = [n for _, n in result_tups]return result_similarities, result_ids

similarity_fn 主要包括三种（欧氏距离、点积和余弦），LlamaIndex 中的实现如下，使用 NP 库实现。

def similarity(embedding1: Embedding,embedding2: Embedding,mode: SimilarityMode = SimilarityMode.DEFAULT,
) -> float:"""Get embedding similarity."""if mode == SimilarityMode.EUCLIDEAN:# Using -euclidean distance as similarity to achieve same ranking orderreturn -float(np.linalg.norm(np.array(embedding1) - np.array(embedding2)))elif mode == SimilarityMode.DOT_PRODUCT:return np.dot(embedding1, embedding2)else:product = np.dot(embedding1, embedding2)norm = np.linalg.norm(embedding1) * np.linalg.norm(embedding2)return product / norm

根据上面的逻辑，实现自己的相似度算法，这里只用余弦，这里使用英文做测试，nomic-embed-text 这个模型中文能力不行。

import numpy as np
def similarity(embedding1: list[float],embedding2: list[float]
) -> float:"""Get embedding similarity."""product = np.dot(embedding1, embedding2)norm = np.linalg.norm(embedding1) * np.linalg.norm(embedding2)return product / normfrom llm import get_ollama_embbedingembed_model = get_ollama_embbeding()
res1 = embed_model.get_query_embedding("The sun is shining brightly over the clear blue sky, and birds are chirping cheerfully.")
res2 = embed_model.get_query_embedding("The stock market faced a significant downturn today due to rising inflation concerns.")
res3 = embed_model.get_query_embedding("The cat is sitting on the windowsill, basking in the warm sunlight.")
res4 = embed_model.get_query_embedding("The feline rests on the windowsill, enjoying the warmth of the sun.")similarity(res1, res2), similarity(res3, res4)

总结

通过参考 LlamaIndex 的代码，本文实现了向量相似度算法，自己实现相似度算法比较简单，开发、测试也不用依赖过多的第三方库，如果使用其他语言，也容易实现。