GTE中文通用文本向量表示模型

GTE中文通用文本表示模型

文本表示是自然语言处理(NLP)领域的核心问题, 其在很多NLP、信息检索的下游任务中发挥着非常重要的作用。近几年, 随着深度学习的发展，尤其是预训练语言模型的出现极大的推动了文本表示技术的效果, 基于预训练语言模型的文本表示模型在学术研究数据、工业实际应用中都明显优于传统的基于统计模型(词袋法、TF-IDF) 或者浅层神经网络的文本表示模型。这里, 我们主要关注基于预训练语言模型的文本表示。GTE项目地址https://modelscope.cn/models/damo/nlp_gte_sentence-embedding_chinese-small/summary

文本表示示例, 输入一个句子, 输入一个固定维度的连续向量:

• 输入: 吃完海鲜可以喝牛奶吗?
• 输出: [0.27162,-0.66159,0.33031,0.24121,0.46122,...]

文本的向量表示通常可以用于文本聚类、文本相似度计算等下游任务中。

二、文本表示模型

基于监督数据训练的文本表示模型通常采用Dual Encoder框架, 如下图所示。在Dual Encoder框架中, Query和Document文本通过预训练语言模型编码后, 通常采用预训练语言模型[CLS]位置的向量作为最终的文本向量表示。基于标注数据的标签, 通过计算query-document之间的cosine距离度量两者之间的相关性。

GTE-zh模型使用retromae初始化训练模型，之后利用两阶段训练方法训练模型：第一阶段利用大规模弱弱监督文本对数据训练模型，第二阶段利用高质量精标文本对数据以及挖掘的难负样本数据训练模型。具体训练方法请参考论文Towards General Text Embeddings with Multi-stage Contrastive Learning。

二、使用方式和范围

使用方式:

• 直接推理, 对给定文本计算其对应的文本向量表示，向量维度512

使用范围:

• 本模型可以使用在通用领域的文本向量表示及其下游应用场景, 包括 两文档间文本相似度计算、query&多doc候选的相似度排序

2.1 如何使用

在ModelScope框架上，提供输入文本(默认最长文本长度为128)，即可以通过简单的Pipeline调用来使用GTE文本向量表示模型。ModelScope封装了统一的接口对外提供单文档向量表示、双文档文本相似度、多候选相似度计算等功能

2.2 安装

pip3 install torch
pip3 install transformers
pip3 install modelscope

三、代码示例

为方便实验，选择体积较小的模型文件damo/nlp_gte_sentence-embedding_chinese-smallfrom modelscope.models import Model
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks


model_id = "damo/nlp_gte_sentence-embedding_chinese-small" #57M
#model_id = "damo/nlp_gte_sentence-embedding_chinese-large #621M
pipeline_se = pipeline(Tasks.sentence_embedding,
                       model=model_id)

# 当输入包含“soure_sentence”与“sentences_to_compare”时，会输出source_sentence中首个句子与sentences_to_compare中每个句子的向量表示，以及source_sentence中首个句子与sentences_to_compare中每个句子的相似度。
inputs = {
        "source_sentence": ["吃完海鲜可以喝牛奶吗?"],
        "sentences_to_compare": [
            "不可以，早晨喝牛奶不科学",
            "吃了海鲜后是不能再喝牛奶的，因为牛奶中含得有维生素C，如果海鲜喝牛奶一起服用会对人体造成一定的伤害",
            "吃海鲜是不能同时喝牛奶吃水果，这个至少间隔6小时以上才可以。",
            "吃海鲜是不可以吃柠檬的因为其中的维生素C会和海鲜中的矿物质形成砷"
        ]
    }

result = pipeline_se(input=inputs)
print (result)

Run{'text_embedding': array([[-0.03317244, -0.0419106 , -0.03626636, ..., -0.0132677 ,
         -0.02028614, -0.01542077],
        [-0.04563809, -0.06220782, -0.03775004, ...,  0.01267119,
         -0.01111769, -0.03390383],
        [-0.02073098, -0.04639562, -0.04818704, ..., -0.00754705,
         -0.00731624, -0.02740852],
        [-0.00037597, -0.05922904, -0.0459275 , ..., -0.00697823,
         -0.02154762, -0.02951157],
        [-0.00491675, -0.02552056, -0.03427778, ..., -0.00760836,
         -0.00404084, -0.0509829 ]], dtype=float32),
 'scores': [0.8542333245277405,
  0.9613471031188965,
  0.947378396987915,
  0.8620702028274536]}

默认向量维度512,  两个向量做内积距离计算得到score

# 当输入仅含有soure_sentence时，会输出source_sentence中每个句子的向量表示。
inputs2 = {
        "source_sentence": [
            "不可以，早晨喝牛奶不科学",
            "吃了海鲜后是不能再喝牛奶的，因为牛奶中含得有维生素C，如果海鲜喝牛奶一起服用会对人体造成一定的伤害",
            "吃海鲜是不能同时喝牛奶吃水果，这个至少间隔6小时以上才可以。",
            "吃海鲜是不可以吃柠檬的因为其中的维生素C会和海鲜中的矿物质形成砷"
        ]
}
result = pipeline_se(input=inputs2)
print (result)

Run{'text_embedding': array([[-0.04563809, -0.06220782, -0.03775004, ...,  0.01267119,
        -0.01111769, -0.03390383],
       [-0.02073098, -0.04639562, -0.04818704, ..., -0.00754705,
        -0.00731624, -0.02740852],
       [-0.00037597, -0.05922904, -0.0459275 , ..., -0.00697823,
        -0.02154762, -0.02951157],
       [-0.00491675, -0.02552056, -0.03427778, ..., -0.00760836,
        -0.00404084, -0.0509829 ]], dtype=float32), 'scores': []}