Jina-embeddings 使用教程

Jina-embeddings 是一个强大的工具，可以用来生成文本的嵌入向量（embeddings），这些向量可用于相似度搜索、分类、重排序（reranking）等任务。在这个教程中，我将展示如何使用 jina-embeddings 来生成嵌入向量，并进行重排序操作。

1. 安装依赖

首先，确保你已经安装了 jina 和 jina-embeddings：

pip install jina jina-embeddings

2. 加载 Jina-embeddings 模型

Jina-embeddings 提供了一些预训练的嵌入模型。我们可以加载一个特定的模型来生成文本嵌入。

例如，我们可以加载 jina-embeddings-v2-base 模型，以下是如何加载模型并生成嵌入向量的示例：

from jina import Document, DocumentArray
from jina.embeddings import TransformerEncoder

# 初始化 Jina-embeddings 模型
model = TransformerEncoder()

# 创建文档数组，包含需要生成嵌入的文本
texts = ['How is the weather today?', '今天天气怎么样?']
docs = DocumentArray([Document(text=text) for text in texts])

# 使用 Jina-embeddings 模型生成嵌入向量
model.encode(docs)

# 输出嵌入结果
for doc in docs:
    print(f"Text: {doc.text}")
    print(f"Embedding: {doc.embedding}")

在这段代码中：

  • 我们首先加载了 jina.embeddings.TransformerEncoder 模型，这个模型可以生成文本的嵌入向量。
  • 接着，我们将输入文本转换为 DocumentArray，然后通过调用 model.encode() 生成每个文本的嵌入。
  • 嵌入可以通过 doc.embedding 来获取。

3. 计算余弦相似度

生成嵌入后，我们可以使用余弦相似度来比较不同文本之间的相似性。

import numpy as np

# 计算余弦相似度
def cos_sim(a, b):
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

# 比较两条文本的相似度
similarity = cos_sim(docs[0].embedding, docs[1].embedding)
print(f"Cosine Similarity: {similarity}")

4. 文本重排序（Reranking）

文本重排序（Reranking）是在初步检索后，对返回结果基于其嵌入向量进行重新排序。通常，我们根据嵌入向量和查询向量之间的相似度来排序。

假设我们已经有一个查询文本和多个候选文本，我们可以使用生成的嵌入来对这些候选文本进行重排序。

# 示例查询
query_text = "What's the weather like today?"

# 候选文本
candidate_texts = ['How is the weather today?', 'It is sunny today.', 'What is the capital of France?', '今天天气怎么样?']

# 创建文档数组
query_doc = Document(text=query_text)
candidate_docs = DocumentArray([Document(text=text) for text in candidate_texts])

# 生成查询和候选文本的嵌入
model.encode(DocumentArray([query_doc]))
model.encode(candidate_docs)

# 计算查询和候选文本之间的相似度
similarities = [cos_sim(query_doc.embedding, doc.embedding) for doc in candidate_docs]

# 根据相似度进行重排序
sorted_candidates = sorted(zip(candidate_texts, similarities), key=lambda x: x[1], reverse=True)

# 输出重排序后的结果
for text, score in sorted_candidates:
    print(f"Candidate: {text}, Similarity: {score}")

在这个示例中：

  • 我们将查询文本与多个候选文本一起生成嵌入向量。
  • 然后计算查询文本与每个候选文本的相似度。
  • 最后，根据相似度对候选文本进行重排序，并输出排序后的结果。

5. 集成到搜索系统

你可以将 Jina-embeddings 嵌入和重排序的流程集成到你的搜索系统中，以增强搜索体验。具体步骤如下：

  1. 初步检索：从数据库或搜索引擎中返回初步的候选文本。
  2. 生成嵌入向量：使用 Jina-embeddings 为查询和候选文本生成嵌入向量。
  3. 相似度计算：基于嵌入向量，计算查询和候选文本之间的相似度。
  4. 重排序：根据相似度对候选结果进行重排序，并将最相关的结果返回给用户。

6. 使用 GPU 加速（可选）

如果你的系统支持 GPU，可以通过在加载模型时指定设备为 cuda 来加速嵌入的生成。

model = TransformerEncoder(device='cuda')  # 使用 GPU 加速

这样，模型将使用 GPU 来加速推理过程，特别是在处理大量文本时，GPU 可以显著提高性能。

总结

  • 生成文本嵌入：通过 Jina-embeddings 模型，可以轻松地为任意文本生成高维向量嵌入。
  • 相似度计算：使用余弦相似度比较文本嵌入，进行相似性判断。
  • 重排序：根据嵌入向量和查询之间的相似度，对候选文本进行重排序，提高检索结果的相关性。
  • GPU 支持：在加载模型时，可以选择使用 GPU 加速推理。

通过这些步骤，你可以利用 Jina-embeddings 生成高效的文本嵌入，并基于这些嵌入来优化文本检索和重排序的效果。