jina-embeddings-v3模型本地GPU使用教程

如何在本地GPU环境中使用jinaai/jina-embeddings-v3模型生成文本嵌入向量。jina-embeddings-v3是一款拥有5.7亿参数的先进文本嵌入模型，支持89种语言，最大输入长度为8192个token，特别适合多语言和长文本检索任务。

环境准备

在使用jina-embeddings-v3之前，需要确保您的GPU环境满足要求：

硬件需求

支持CUDA的GPU
对于最佳性能，推荐使用支持FlashAttention-2的GPU（如A100、RTX 3090、RTX 4090、H100等Ampere、Ada或Hopper架构的GPU）(huggingface.co/jinaai/jina...)

安装必要依赖

yaml 复制代码

python
# 安装基本依赖
!pip install transformers torch einops
!pip install 'numpy<2'

# 如果GPU支持FlashAttention-2，可以安装以提升性能
!pip install flash-attn --no-build-isolation

使用Transformers库加载模型

基本加载方法

ini 复制代码

python
from transformers import AutoModel, AutoTokenizer
import torch

# 确保将模型加载到GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用设备: {device}")

# 加载模型和分词器
model = AutoModel.from_pretrained("jinaai/jina-embeddings-v3", trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("jinaai/jina-embeddings-v3")

# 将模型移至GPU
model = model.to(device)

生成文本嵌入向量

方法一：使用encode方法（推荐）

ini 复制代码

python
# 准备文本
texts = [
    "人工智能正在改变我们的生活方式", # 中文
    "Artificial intelligence is changing the way we live", # 英文
    "La inteligencia artificial está cambiando nuestra forma de vida" # 西班牙文
]

# 使用encode方法生成嵌入向量
# task参数可选："retrieval.query", "retrieval.passage", "separation", "classification", "text-matching"
embeddings = model.encode(texts, task="text-matching")

# 查看嵌入向量形状
print(f"嵌入向量形状: {embeddings.shape}")

# 计算相似度矩阵
similarity_matrix = embeddings @ embeddings.T
print("相似度矩阵:")
print(similarity_matrix)

方法二：使用tokenizer和模型的forward方法

ini 复制代码

python
# 准备文本
text = "人工智能正在改变我们的生活方式"

# 使用tokenizer处理文本
inputs = tokenizer(text, return_tensors="pt").to(device)

# 获取模型输出
with torch.no_grad():
    outputs = model(**inputs)

# 提取句子嵌入（使用平均池化）
# 这里使用最后一层隐藏状态的平均值作为句子嵌入
embeddings = outputs.last_hidden_state.mean(dim=1)
print(f"嵌入向量形状: {embeddings.shape}")

任务特定的嵌入生成

jina-embeddings-v3的一个主要特点是内置了特定任务的LoRA适配器，可以根据不同任务生成优化的嵌入向量

python 复制代码

python
# 针对不同任务生成嵌入向量
texts = ["这是一个测试文档，用于展示jina-embeddings-v3的性能"]

# 1. 检索任务 - 文档段落
passage_embeddings = model.encode(texts, task="retrieval.passage")

# 2. 检索任务 - 查询
query = ["如何展示jina-embeddings-v3的性能"]
query_embeddings = model.encode(query, task="retrieval.query")

# 3. 聚类任务
cluster_embeddings = model.encode(texts, task="separation")

# 4. 分类任务
classification_embeddings = model.encode(texts, task="classification")

# 5. 文本匹配任务
matching_embeddings = model.encode(texts, task="text-matching")

# 比较不同任务嵌入向量的形状
print(f"文档嵌入向量形状: {passage_embeddings.shape}")
print(f"查询嵌入向量形状: {query_embeddings.shape}")
print(f"聚类嵌入向量形状: {cluster_embeddings.shape}")
print(f"分类嵌入向量形状: {classification_embeddings.shape}")
print(f"文本匹配嵌入向量形状: {matching_embeddings.shape}")

调整嵌入向量维度

jina-embeddings-v3支持灵活调整输出嵌入向量的维度，这得益于Matryoshka表示学习（MRL）8。

ini 复制代码

python
# 使用不同维度生成嵌入向量
text = ["这是一个测试样例"]

# 默认维度为1024
default_dim_embeddings = model.encode(text, task="text-matching")
print(f"默认维度嵌入向量形状: {default_dim_embeddings.shape}")

# 自定义维度（如果使用HuggingFace Transformers API）
# 注意：如果直接使用HuggingFace的API，可能需要在encode方法中添加dimensions参数
# 这里假设encode方法支持dimensions参数
custom_dim_embeddings = model.encode(text, task="text-matching", dimensions=128)
print(f"自定义维度嵌入向量形状: {custom_dim_embeddings.shape}")

处理长文本

jina-embeddings-v3支持高达8192个token的输入长度，并引入了迟分（Late Chunking）技术

ini 复制代码

python
# 处理长文本
long_text = ["这是一个非常长的文档..." * 100]  # 创建一个长文本

# 使用迟分技术处理长文本
# 注意：如果使用API，需要设置late_chunking=True
# 这里假设encode方法支持late_chunking参数
long_text_embeddings = model.encode(long_text, task="retrieval.passage")
print(f"长文本嵌入向量形状: {long_text_embeddings.shape}")

完整示例：文本相似度计算

ini 复制代码

python
import torch
from transformers import AutoModel, AutoTokenizer

# 加载模型和分词器
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = AutoModel.from_pretrained("jinaai/jina-embeddings-v3", trust_remote_code=True).to(device)
tokenizer = AutoTokenizer.from_pretrained("jinaai/jina-embeddings-v3")

# 准备文本
texts = [
    "人工智能正在改变我们的生活方式",
    "机器学习是人工智能的一个子领域",
    "深度学习是机器学习的一种方法",
    "自然语言处理是人工智能的重要应用",
    "计算机视觉也是人工智能的重要分支"
]

# 生成嵌入向量
embeddings = model.encode(texts, task="text-matching")
print(f"生成的嵌入向量形状: {embeddings.shape}")

# 计算相似度矩阵
similarity_matrix = embeddings @ embeddings.T
print("相似度矩阵:")
print(similarity_matrix)

# 查找与第一个文本最相似的其他文本
similarities = similarity_matrix[0]
most_similar_idx = torch.argsort(similarities, descending=True)[1]  # 排除自身
print(f"与'{texts[0]}'最相似的文本是: {texts[most_similar_idx]}")

结论

jina-embeddings-v3是一个功能强大的多语言嵌入模型，支持多种语言和多种任务类型1 2。本文介绍了如何在本地GPU环境中使用该模型生成嵌入向量，以及如何根据不同任务选择适当的LoRA适配器。

对于特定任务，建议选择相应的任务类型：

检索任务：使用retrieval.query和retrieval.passage
聚类任务：使用separation
分类任务：使用classification
语义相似性任务：使用text-matching7