jina-embeddings-v3是一种多语言多任务文本嵌入模型,专为各种NLP应用程序而设计。该模型基于 Jina-XLM-RoBERTa 架构,支持旋转位置嵌入,可以处理8192个标记的长输入序列。此外,它还具有5个LoRA适配器,可高效生成特定于任务的嵌入。
特点
-
扩展序列程度:使用RoPE支持高达8192令牌
-
特定任务的嵌入:通过 task 参数自定义嵌入
-
retrieval.query :用于非对称检索任务中的查询嵌入
-
retrieval.passage :用于非对称检索任务中的段落嵌入
-
separation :用于聚类和重新排序应用程序中的嵌入
-
classification :用于分类任务中的嵌入
-
text-matching :用于量化两个文本之间相似性的任务中的嵌入,例如STS或对称检索任务
-
-
Matryoshka 嵌入向量:支持灵活的嵌入向量大小( 32, 64, 128, 256, 512, 768, 1024 ),允许截断嵌入向量以适应应用程序
用法
在集成模型时应用均值汇聚
为什么用 Mean Pooling
均值汇聚从模型的输出中获取所有标记嵌入,并在句子或段落级别对他们进行平均。这种方法已被证明可以产生高质量的句子嵌入。
现有 encode 函数,自动处理这个问题。如果直接在 encode 函数之外使用模型,需要手动应用均值池化。
from transformers import AutoModel
# Initialize the model
model = AutoModel.from_pretrained("jinaai/jina-embeddings-v3", trust_remote_code=True)
texts = [
"Follow the white rabbit.", # English
"Sigue al conejo blanco.", # Spanish
"Suis le lapin blanc.", # French
"跟着白兔走。", # Chinese
"اتبع الأرنب الأبيض.", # Arabic
"Folge dem weißen Kaninchen.", # German
]
# When calling the `encode` function, you can choose a `task` based on the use case:
# 'retrieval.query', 'retrieval.passage', 'separation', 'classification', 'text-matching'
# Alternatively, you can choose not to pass a `task`, and no specific LoRA adapter will be used.
embeddings = model.encode(texts, task="text-matching")
# Compute similarities
print(embeddings[0] @ embeddings[1].T)
上述代码使用了 transformers 库和 jinaai/jina-embeddings-v3 模型生成文本句子的嵌入向量(embedding)常用于语义搜索、文本相似度比较或下游NLP任务。
整体流程梳理:
-
加载模型和分词器
-
对句子进行编码
-
通过模型向前传播生成输出
-
使用平均池化(mean pooling)+ 掩码(mask)聚合token embeddings
-
L2 归一化 得到最终的句向量
代码详细分析:
-
mean_pooling(model_output, attention_mask) 对 Transformer 输出的 token 级嵌入进行平均池化(考虑 attention_mask 避免 padding token 干扰)
token_embeddings = model_output[0] # shape: (batch_size, seq_len, hidden_dim)
input_mask_expanded = (
attention_mask.unsqueeze(-1) # shape: (batch_size, seq_len, 1)
.expand(token_embeddings.size()) # shape: (batch_size, seq_len, hidden_dim)
.float()
)
# 加权求和,再除以有效 token 的个数 => 得到句向量
return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(
input_mask_expanded.sum(1), min=1e-9
)
-
输入句子和编码
sentences = ["How is the weather today?", "What is the current weather like today?"]
encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors="pt")
将两个句子编码为 token IDs,自动添加 padding 和 truncation
return_tensors="pt" 返回 PyTorch 张量格式
-
加载 Jina Embedding 模型
tokenizer = AutoTokenizer.from_pretrained("jinaai/jina-embeddings-v3")
model = AutoModel.from_pretrained("jinaai/jina-embeddings-v3", trust_remote_code=True)
这是一个多任务嵌入模型,可以根据不同的任务自动切换 adapter
trust_remote_code=True 允许加载模型自定以逻辑
-
指定任务 retrieval.query
task = 'retrieval.query'
task_id = model._adaptation_map[task]
adapter_mask = torch.full((len(sentences),), task_id, dtype=torch.int32)
adapter_mask 是这个模型自定以的任务切换机制
每个句子打上该任务的ID标签,使模型使用对应的 adapter
-
模型向前传播
with torch.no_grad():
model_output = model(**encoded_input, adapter_mask=adapter_mask)
不计算梯度,节省显存
模型返回的第一个输出通常是 last_hidden_state
-
Mean Pooling + L2 Normalization
embeddings = mean_pooling(model_output, encoded_input["attention_mask"])
embeddings = F.normalize(embeddings, p=2, dim=1)
mean_pooling :从 token 级转成句级向量
F.normalize :对每个向量做 L2 归一化,使得后续计算余弦相似度变得合理
jina-embeddings-v3 是一个 multi-adapter embedding model,支持多种下游任务(如检索、分类、对话等)每种任务使用不同的 adapter。
这相当于一次加载多个任务专家模块,在保持基础语言理解能力的同时微调适配不同目标任务,无需每个任务单独 finetune 整个模型。
其他情况
默认情况下,该模型支持的最大序列长度为8192个令牌。但是,如果要将输入文本截断为更短的长度,则可以将 max_length 参数传递给 encode 函数:
embeddings = model.encode(["Very long ... document"], max_length=2048)
如果你想使用 Matryoshka 嵌入并切换到不同的维度,你可以通过将 truncate_dim 参数传递给 encode 函数来调整:
embeddings = model.encode(['Sample text'], truncate_dim=256)
最新版本的(3.1.0)的 SentenceTransformers 还支持 jina-embeddings-v3 :
pip install sentence-transformers
from sentence_transformers import SentenceTransformer
model = SentenceTransformer("jinaai/jina-embeddings-v3", trust_remote_code=True)
task = "retrieval.query"
embeddings = model.encode(
["What is the weather like in Berlin today?"],
task=task,
prompt_name=task,
)
可以使用 SentenceTransformers 微调 jina-embeddings-v3 。要针对特定任务进行微调,应该在将模型传递给 ST Trainer 之前 (在初始化期间)设置任务:
model = SentenceTransformer("jinaai/jina-embeddings-v3", trust_remote_code=True, model_kwargs={'default_task': 'classification'})
或之后:
model = SentenceTransformer("jinaai/jina-embeddings-v3", trust_remote_code=True)
model[0].default_task = 'classification'
这样,可以针对所选任务微调 LoRA 适配器。
但是,如果要微调整个模型,要确保在加载模型时将主要参数设置为可训练:
model = SentenceTransformer("jinaai/jina-embeddings-v3", trust_remote_code=True, model_kwargs={'lora_main_params_trainable': True})
是允许微调整个模型,而不仅仅是 LoRA 适配器
ONNX inference. ONNX 推理
可以使用ONNX通过 jina-embeddings-v3 进行高效推理
import onnxruntime
import numpy as np
from transformers import AutoTokenizer, PretrainedConfig
# Mean pool function
def mean_pooling(model_output: np.ndarray, attention_mask: np.ndarray):
token_embeddings = model_output
input_mask_expanded = np.expand_dims(attention_mask, axis=-1)
input_mask_expanded = np.broadcast_to(input_mask_expanded, token_embeddings.shape)
sum_embeddings = np.sum(token_embeddings * input_mask_expanded, axis=1)
sum_mask = np.clip(np.sum(input_mask_expanded, axis=1), a_min=1e-9, a_max=None)
return sum_embeddings / sum_mask
# Load tokenizer and model config
tokenizer = AutoTokenizer.from_pretrained('jinaai/jina-embeddings-v3')
config = PretrainedConfig.from_pretrained('jinaai/jina-embeddings-v3')
# Tokenize input
input_text = tokenizer('sample text', return_tensors='np')
# ONNX session
model_path = 'jina-embeddings-v3/onnx/model.onnx'
session = onnxruntime.InferenceSession(model_path)
# Prepare inputs for ONNX model
task_type = 'text-matching'
task_id = np.array(config.lora_adaptations.index(task_type), dtype=np.int64)
inputs = {
'input_ids': input_text['input_ids'],
'attention_mask': input_text['attention_mask'],
'task_id': task_id
}
# Run model
outputs = session.run(None, inputs)[0]
# Apply mean pooling and normalization to the model outputs
embeddings = mean_pooling(outputs, input_text["attention_mask"])
embeddings = embeddings / np.linalg.norm(embeddings, ord=2, axis=1, keepdims=True)