以下为你介绍几种常用的中文开源 Embedding 模型,以及它们的输入输出调用方式和代码示例。
1. Sentence Transformers(包含中文预训练模型)
- 简介:Sentence Transformers 是一个用于生成句子嵌入的 Python 库,提供了多种预训练模型,其中有适用于中文的模型,能够将句子转换为固定长度的向量,方便进行语义相似度计算等任务。
- 输入输出:
- 输入:字符串形式的单个句子或者句子列表。
- 输出:与输入句子对应的向量表示,输出为 Numpy 数组或者 PyTorch 张量。
- 代码示例:
from sentence_transformers import SentenceTransformer
import numpy as np
# 加载中文预训练模型
model = SentenceTransformer('shibing624/text2vec-base-chinese')
# 输入句子
sentences = ["这是一个测试句子", "另一个测试句子"]
# 生成句子嵌入
embeddings = model.encode(sentences)
# 打印结果
for sentence, embedding in zip(sentences, embeddings):
print(f"句子: {sentence}")
print(f"嵌入向量: {embedding[:5]}...(长度: {len(embedding)})")
print()
2. ERNIE(百度)
- 简介:ERNIE(Enhanced Representation through Knowledge Integration)是百度研发的预训练模型,它在大规模语料上进行预训练,融合了大量的知识信息,在中文自然语言处理任务中表现出色。可以用于生成文本的 Embedding 表示。
- 输入输出:
- 输入:中文文本,使用分词器处理后的词元列表。
- 输出:每个词元对应的向量表示,或者经过池化后的句子向量表示。
- 代码示例:
from transformers import AutoTokenizer, AutoModel
import torch
# 加载 ERNIE 分词器和模型
tokenizer = AutoTokenizer.from_pretrained("ernie-1.0")
model = AutoModel.from_pretrained("ernie-1.0")
# 输入中文文本
text = "这是一段中文文本。"
# 分词
inputs = tokenizer(text, return_tensors='pt')
# 模型推理
with torch.no_grad():
outputs = model(**inputs)
# 获取词元的嵌入向量
last_hidden_states = outputs.last_hidden_state
print("词元嵌入向量形状:", last_hidden_states.shape)
# 简单的池化操作获取句子向量(这里使用 [CLS] 标记的向量)
sentence_embedding = last_hidden_states[:, 0, :]
print("句子嵌入向量形状:", sentence_embedding.shape)
3. BERT - 中文预训练模型
- 简介:BERT 是 Google 提出的预训练语言模型,Hugging Face 提供了中文预训练的 BERT 模型,能够捕捉中文文本的语义信息,生成高质量的词元嵌入和句子嵌入。
- 输入输出:
- 输入:中文文本,经过分词器处理后添加特殊标记(如
[CLS]
和[SEP]
)。 - 输出:每个词元对应的向量表示,以及经过池化后的句子向量表示。
- 输入:中文文本,经过分词器处理后添加特殊标记(如
- 代码示例:
from transformers import BertTokenizer, BertModel
import torch
# 加载中文 BERT 分词器和模型
tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')
model = BertModel.from_pretrained('bert-base-chinese')
# 输入中文文本
text = "这是一个中文测试。"
# 分词并添加特殊标记
inputs = tokenizer(text, return_tensors='pt')
# 模型推理
with torch.no_grad():
outputs = model(**inputs)
# 获取词元的嵌入向量
last_hidden_states = outputs.last_hidden_state
print("词元嵌入向量形状:", last_hidden_states.shape)
# 简单的池化操作获取句子向量(这里使用 [CLS] 标记的向量)
sentence_embedding = last_hidden_states[:, 0, :]
print("句子嵌入向量形状:", sentence_embedding.shape)
4. XLM - RoBERTa(支持中文)
- 简介:XLM - RoBERTa 是一种多语言预训练模型,支持包括中文在内的多种语言,能够生成高质量的文本嵌入,适用于跨语言和单语言的自然语言处理任务。
- 输入输出:
- 输入:中文文本,经过分词器处理后添加特殊标记。
- 输出:每个词元对应的向量表示,以及经过池化后的句子向量表示。
- 代码示例:
from transformers import XLMRobertaTokenizer, XLMRobertaModel
import torch
# 加载 XLM - RoBERTa 分词器和模型
tokenizer = XLMRobertaTokenizer.from_pretrained('xlm - roberta - base')
model = XLMRobertaModel.from_pretrained('xlm - roberta - base')
# 输入中文文本
text = "中文文本处理示例。"
# 分词并添加特殊标记
inputs = tokenizer(text, return_tensors='pt')
# 模型推理
with torch.no_grad():
outputs = model(**inputs)
# 获取词元的嵌入向量
last_hidden_states = outputs.last_hidden_state
print("词元嵌入向量形状:", last_hidden_states.shape)
# 简单的池化操作获取句子向量(这里使用 [CLS] 标记的向量)
sentence_embedding = last_hidden_states[:, 0, :]
print("句子嵌入向量形状:", sentence_embedding.shape)
在运行上述代码之前,需要确保已经安装了相应的库,如 transformers
、sentence - transformers
等,可以使用 pip install
命令进行安装。同时,部分模型需要从 Hugging Face 的模型库中下载,可能需要一定的网络环境支持。