Embedding,到底嵌入了什么

516 阅读20分钟

Embedding,中文嵌入,听上去一脸懵,嵌入?把什么嵌入什么?嵌入到哪?Embedding到底是个啥,别急我们一步步来看

从推荐系统开始

现在你是一名旅游网站的开发人员,你接到了一个需求,需要根据旅客选择的酒店信息,推荐差不多的酒店。你表示产品一定是疯了,我又没去过这些酒店,我怎么知道谁和谁差不多。不过吐槽归吐槽,活还是得干。

于是你冥思苦想,从数据库里的酒店信息发现了一些端倪,如果说,能够把酒店信息中的价格,地址,服务项目都拆分出来,然后进行对比,这样是不是就可以找到差不多的酒店?

酒店信息在数据库中都有明确的字段,直接取就可以,但是如何进行对比呢?一个字段一个字段的匹配?肯定不行,那系统得慢成什么样子。于是,你想到可以先对酒店信息进行分词,然后统计每个词出现的词频,这样就可以通过计算两条数据的词频是否相近来判断两个酒店是不是类似。

余弦相似度

举个例子:

  • 信息A:我不爱吃蔬菜,我喜欢吃肉

  • 信息B:我不爱吃肉,我爱吃蔬菜 第一步:分词

  • 信息A:我/不爱/吃/蔬菜,我/喜欢/吃/肉

  • 信息B:我/不爱/吃/肉,我/爱/吃/蔬菜 第二步:列出所有的词

  • 我,不爱,爱,吃,肉,蔬菜,喜欢 第三步:统计词频

  • 信息A:我 2,不爱 1,爱 0,吃 2,肉 1,蔬菜 1,喜欢 1

  • 信息B:我 2,不爱 1,爱 2,吃 2,肉 1,蔬菜 1,喜欢 0

  • 然后我们就得到了两个数组,也可以说是两个向量

    • A:[2,1,0,2,1,1,1]
    • B:[2,1,2,2,1,1,0]

那么现在我们该怎么计算两个向量的相似度呢,总不能拿眼看吧。于是你想到了曾经学过的数学知识,余弦,我们可以通过计算两个向量的余弦夹角,来判断两个向量是不是相似的,具体公式如下:

图片.png

看上去有点头疼?没关系,我们可以通过numpy包来计算

import numpy as np

# 定义两个7维向量
vector_a = np.array([2,1,0,2,1,1,1])
vector_b = np.array([2,1,2,2,1,1,0])

# 计算点积
dot_product = np.dot(vector_a, vector_b)

# 计算向量的欧几里得范数
norm_a = np.linalg.norm(vector_a)
norm_b = np.linalg.norm(vector_b)

# 计算余弦相似度
cosine_similarity = dot_product / (norm_a * norm_b)

print("余弦相似度:", cosine_similarity)

最终结果:

0.819891591749923

根据余弦相似度的范围[-1,1]可以知道,这两个句子是相似度高达0.8。

但是,发现问题没有?这两个句子从语义上来说完全不是一回事啊,一个爱吃肉,一个爱吃蔬菜,哪里相似了?

N-Garm

通过思考你发现,分词统计词频的时候只是统计了每个词的数量,上下文并没有考虑进去,那如何把上下文考虑进去呢?既然单个词没有办法表达上下文,那么把分词的边界扩大不就可以了,比如:

  • 单个词划分:A/B/C/D/E/F
  • 两个词划分:AB,BC,CD,DE,EF
  • 三个词划分:ABC,BCD,CDE,DEF,
  • N个词划分。。。
import numpy as np
from collections import Counter

# 定义N-Gram分词函数
def n_gram_tokenize(text, n=3):
    return [text[i:i+n] for i in range(len(text) - n + 1)]

# 定义余弦相似度计算函数
def cosine_similarity(vec1, vec2):
    dot_product = np.dot(vec1, vec2)
    norm1 = np.linalg.norm(vec1)
    norm2 = np.linalg.norm(vec2)
    return dot_product / (norm1 * norm2)

# 定义两句话
sentence_a = "我不爱吃蔬菜,我喜欢吃肉"
sentence_b = "我不爱吃肉,我爱吃蔬菜"

# 对两句话进行N-Gram分词
tokens_a = n_gram_tokenize(sentence_a, n=3)
tokens_b = n_gram_tokenize(sentence_b, n=3)

# 构建词表
vocab = list(set(tokens_a + tokens_b))

# 将句子转换为向量
def sentence_to_vector(tokens, vocab):
    counter = Counter(tokens)
    return [counter.get(word, 0) for word in vocab]

vector_a = sentence_to_vector(tokens_a, vocab)
vector_b = sentence_to_vector(tokens_b, vocab)

# 计算余弦相似度
similarity = cosine_similarity(vector_a, vector_b)
print(f"余弦相似度: {similarity:.4f}")

结果:

余弦相似度: 0.5000

可以发现好了很多,已经不是很像了。

这就是所谓的N-Gram也就是N元语法,N元语法有以下几个特征

  • 基于一个假设:第n个词出现与前n-1个词相关,与其他任何词不相关
  • N=1时为unigram N=2时为bigram N=3时为trigram
  • N-Gram是指给定一段文本,其中的N个item序列
  • 当一阶特征不够用的时候,不如处理文本特征的时候,一个关键词是一个特征,但是有的时候这样做不是很有用,采用N元语法,可以理解两个相邻关键词的特征组合

下面我们以酒店推荐系统为例,做一下实践

用到的数据集: github.com/susanli2016…

import pandas as pd
import numpy as np
from nltk.corpus import stopwords
from sklearn.metrics.pairwise import linear_kernel
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.decomposition import LatentDirichletAllocation
import re
import random
pd.options.display.max_columns = 30
import matplotlib.pyplot as plt
# 支持中文
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
df = pd.read_csv('Seattle_Hotels.csv', encoding="latin-1")
# 数据探索
print(df.head())
print('数据集中的酒店个数:', len(df))

图片.png 数据集中一共有152个酒店,我们选择第10个酒店

def print_description(index):
    example = df[df.index == index][['desc', 'name']].values[0]
    if len(example) > 0:
        print(example[0])
        print('Name:', example[1])
print('第10个酒店的描述:')
print_description(10)

图片.png

下一步我们对原始数据集进行清洗,考虑到英文中存在大量的重复性但区分度不高的词汇,所以我们设置停用词

# 创建英文停用词列表
ENGLISH_STOPWORDS = {
    'i', 'me', 'my', 'myself', 'we', 'our', 'ours', 'ourselves', 'you', "you're", "you've", "you'll", "you'd", 'your', 
    'yours', 'yourself', 'yourselves', 'he', 'him', 'his', 'himself', 'she', "she's", 'her', 'hers', 'herself', 'it', 
    "it's", 'its', 'itself', 'they', 'them', 'their', 'theirs', 'themselves', 'what', 'which', 'who', 'whom', 'this', 
    'that', "that'll", 'these', 'those', 'am', 'is', 'are', 'was', 'were', 'be', 'been', 'being', 'have', 'has', 'had', 
    'having', 'do', 'does', 'did', 'doing', 'a', 'an', 'the', 'and', 'but', 'if', 'or', 'because', 'as', 'until', 'while', 
    'of', 'at', 'by', 'for', 'with', 'about', 'against', 'between', 'into', 'through', 'during', 'before', 'after', 'above', 
    'below', 'to', 'from', 'up', 'down', 'in', 'out', 'on', 'off', 'over', 'under', 'again', 'further', 'then', 'once', 
    'here', 'there', 'when', 'where', 'why', 'how', 'all', 'any', 'both', 'each', 'few', 'more', 'most', 'other', 'some', 
    'such', 'no', 'nor', 'not', 'only', 'own', 'same', 'so', 'than', 'too', 'very', 's', 't', 'can', 'will', 'just', 'don', 
    "don't", 'should', "should've", 'now', 'd', 'll', 'm', 'o', 're', 've', 'y', 'ain', 'aren', "aren't", 'couldn', 
    "couldn't", 'didn', "didn't", 'doesn', "doesn't", 'hadn', "hadn't", 'hasn', "hasn't", 'haven', "haven't", 'isn', 
    "isn't", 'ma', 'mightn', "mightn't", 'mustn', "mustn't", 'needn', "needn't", 'shan', "shan't", 'shouldn', "shouldn't", 
    'wasn', "wasn't", 'weren', "weren't", 'won', "won't", 'wouldn', "wouldn't"
}
# 文本预处理
REPLACE_BY_SPACE_RE = re.compile('[/(){}[]|@,;]')
BAD_SYMBOLS_RE = re.compile('[^0-9a-z #+_]')
# 使用自定义的英文停用词列表替代nltk的stopwords
STOPWORDS = ENGLISH_STOPWORDS
# 对文本进行清洗
def clean_text(text):
    # 全部小写
    text = text.lower()
    # 用空格替代一些特殊符号,如标点
    text = REPLACE_BY_SPACE_RE.sub(' ', text)
    # 移除BAD_SYMBOLS_RE
    text = BAD_SYMBOLS_RE.sub('', text)
    # 从文本中去掉停用词
    text = ' '.join(word for word in text.split() if word not in STOPWORDS)
    return text
# 对desc字段进行清理,apply针对某列
df['desc_clean'] = df['desc'].apply(clean_text)
print(df['desc_clean'])

下一步我们开始对数据集进行建模

df.set_index('name', inplace = True)
# 使用TF-IDF提取文本特征,使用自定义停用词列表
tf = TfidfVectorizer(analyzer='word', ngram_range=(1, 3), min_df=0.01, stop_words=list(ENGLISH_STOPWORDS))
# 针对desc_clean提取tfidf
tfidf_matrix = tf.fit_transform(df['desc_clean'])
print('TFIDF feature names:')
#print(tf.get_feature_names_out())
print(len(tf.get_feature_names_out()))
#print('tfidf_matrix:')
#print(tfidf_matrix)
#print(tfidf_matrix.shape)
# 计算酒店之间的余弦相似度(线性核函数)
cosine_similarities = linear_kernel(tfidf_matrix, tfidf_matrix)
#print(cosine_similarities)
print(cosine_similarities.shape)
indices = pd.Series(df.index) #df.index是酒店名称

这里我们引入了另一个概念,TF-IDF,之前我们提到了N-Gram可以一定程度上抽取单词上下文的特征,但是仍然存在一个问题。 如果一个n-gram特征在一篇文档中频繁出现,我们可以说它是一个明显的特征,但是如果所有文档中都频繁出现,就说明这个词区分度不是很高,所以只用n-gram还是有些粗糙的,这里就要使用TF-IDF进一步计算

图片.png 也就是说TF-IDF值越大,那么这个特征就是区分度高,又很明显

# 建模
df.set_index('name', inplace = True)
# 使用TF-IDF提取文本特征,使用自定义停用词列表
tf = TfidfVectorizer(analyzer='word', ngram_range=(1, 3), min_df=0.01, stop_words=list(ENGLISH_STOPWORDS))
# 针对desc_clean提取tfidf
tfidf_matrix = tf.fit_transform(df['desc_clean'])
print('TFIDF feature names:')
#print(tf.get_feature_names_out())
print(len(tf.get_feature_names_out()))
#print('tfidf_matrix:')
#print(tfidf_matrix)
#print(tfidf_matrix.shape)

图片.png 我们可以看到一共抽取了3347个特征,然后我们来计算酒店信息的相似度矩阵

# 计算酒店之间的余弦相似度(线性核函数)
cosine_similarities = linear_kernel(tfidf_matrix, tfidf_matrix)
#print(cosine_similarities)
print(cosine_similarities.shape)
indices = pd.Series(df.index) #df.index是酒店名称

然后我们基于上一步的相似度矩阵,来尝试进行酒店推荐

def recommendations(name, cosine_similarities = cosine_similarities):
    recommended_hotels = []
    # 找到想要查询酒店名称的idx
    idx = indices[indices == name].index[0]
    print('idx=', idx)
    # 对于idx酒店的余弦相似度向量按照从大到小进行排序
    score_series = pd.Series(cosine_similarities[idx]).sort_values(ascending = False)
    # 取相似度最大的前10个(除了自己以外)
    top_10_indexes = list(score_series.iloc[1:11].index)
    # 放到推荐列表中
    for i in top_10_indexes:
        recommended_hotels.append(list(df.index)[i])
    return recommended_hotels
print(recommendations('Hilton Seattle Airport & Conference Center'))

图片.png 可以看到,结果还是ok的

但是,但是仍然不太好,为什么?我们可以看到前面的特征抽取抽出了3347个特征,这还是原始数据集的内容不是很长,如果塞一篇小说进去。。。。会导致维度爆炸

这个时候就要提到我们今天要说的Embedding,也就是嵌入

Embedding

刚才我们发现,抽取出来的特征矩阵维度太多,计算量可能过大,那咋办?既然维度多,那就降维。 我们可以将维度过多的特征嵌入到一个维度固定的矩阵里,转换成维度相同的向量,这就是嵌入,把一堆东西塞到一个固定大小的盒子里。

维度降下来了,向量之间是可以通过余弦夹角计算相似度的

Word2Vec

Word2Vec 是Embedding的一种方式,市面上还有很多Embedding模型,这里只是用来举个例子,方便大家理解原理

图片.png

这张图是Word2Vec简化版的神经网络,分为Input(输入),Hidden(隐藏),Output(输出)三层,假设我们现在有一本《三国演义》,我把它交给Word2Vec,Word2Vec会将整本书进行压缩并向量化,并且抽取某些词上下文的特征,举个例子,张飞属于蜀国,刘备属于蜀国,张飞,刘备和蜀国的在语义上就会更接近,但是曹操和蜀国在语义上就会更远,基于这个特点,隐藏层会形成一个矩阵,这个矩阵的大小是n*m,n是有多少个词,m是每个词的维度数量。

然后我们就得到了一个查找表

图片.png 给到对应的词,然后在隐藏层矩阵中找到对应的向量,之后由输出层输出,这样我们就可以拿着包含了某个词上下文特征的向量,去做接下来的任务

Word2Vec的两种模式
  1. Skip-Gram,跟定输入的词预测上下文

图片.png 2. CBOW给定上下文,预测相对应的词

图片.png

原理上大致就是这样,下面我们来看下实际的操作,我们使用Gensim工具来完成 安装方式

pip install gensim

gensim是一个开源的工具包,可 以 从非结构化文本中,无监督地学习到隐层的主题向 量表达,每一个向量变换的操作都对应着一个主题模型,支持TF-IDF,LDA, LSA, word2vec 等 多种主题模型算 法

  • 关键参数
    • window,句子中当前单词和被预测单词的最大距离
    • min_count,需要训练词语的最小出现次数,默认为5
    • size,向量维度,默认为100
    • worker,训练使用的线程数,默认为1即不使用多线程 我们以西游记作为基本语料,计算小说中人物的相似度

首先我们使用jieba库对小说进行分词

# -*-coding: utf-8 -*-
# 对txt文件进行中文分词
import jieba
import os
from utils import files_processing

# 源文件所在目录
source_folder = './journey_to_the_west/source'
segment_folder = './journey_to_the_west/segment'

# 字词分割,对整个文件内容进行字词分割
def segment_lines(file_list,segment_out_dir,stopwords=[]):
    for i,file in enumerate(file_list):
        segment_out_name=os.path.join(segment_out_dir,'segment_{}.txt'.format(i))
        with open(file, 'rb') as f:
            document = f.read()
            document_cut = jieba.cut(document)
            sentence_segment=[]
            for word in document_cut:
                if word not in stopwords:
                    sentence_segment.append(word)
            result = ' '.join(sentence_segment)
            result = result.encode('utf-8')
            with open(segment_out_name, 'wb') as f2:
                f2.write(result)

# 对source中的txt文件进行分词,输出到segment目录中
file_list=files_processing.get_files_list(source_folder, postfix='*.txt')
segment_lines(file_list, segment_folder)

图片.png

然后我们将处理好的文件转化为一个sentence迭代器

from gensim.models import word2vec
import multiprocessing

# 如果目录中有多个文件,可以使用PathLineSentences
segment_folder = './journey_to_the_west/segment'
sentences = word2vec.PathLineSentences(segment_folder)

最后,我们使用word2vec进行训练

# 设置模型参数,进行训练
model = word2vec.Word2Vec(sentences, vector_size=100, window=3, min_count=1)
print(model.wv.similarity('孙悟空', '猪八戒')) # 孙悟空和猪八戒的相似度
print(model.wv.similarity('孙悟空', '孙行者')) # 孙悟空和孙行者的相似度
print(model.wv.most_similar(positive=['孙悟空', '唐僧'], negative=['孙行者'])) # 孙悟空+唐僧-孙行者=?

让我们看一下输出结果

图片.png

效果好像不是很好,让我们调整参数再来一次

# 设置模型参数,进行训练
model2 = word2vec.Word2Vec(sentences, vector_size=128, window=5, min_count=5, workers=multiprocessing.cpu_count())
# 保存模型
model2.save('./models/word2Vec.model')
print(model2.wv.similarity('孙悟空', '猪八戒'))
print(model2.wv.similarity('孙悟空', '孙行者'))
print(model2.wv.most_similar(positive=['孙悟空', '唐僧'], negative=['孙行者']))

图片.png 可以看到好了一些些,但还不是特别理想,因为word2vec还是一个相对粗糙的方式,

但是基于这个原理,可以窥探推荐系统的一角,将商品,视频,小说等特征代换文本特征,我们可以通过抽取商品,视频,小说的特征,当一个用户购买某个商品,刷过某个视频之后,可以去匹配特征相似的对象,然后输出相似度最大的top-n个。。。。

Embedding模型选择

huggingface 上我们可以看到各种embedding模型的排名

图片.png

也可以在 魔搭 找到需要的embeding模型

图片.png

下面是一些常见embedding模型的介绍:

  1. 通用文本嵌入模型
    • BGE-M3(智源研究院)
      • 特点:支持100+语言,输入长度达8192 tokens,融合密集、稀疏、 多向量混合检索,适合跨语言长文档检索。
      • 适用场景:跨语言长文档检索、高精度RAG应用。
    • text-embedding-3-large(OpenAI)
      • 特点:向量维度3072,长文本语义捕捉能力强,英文表现优秀。
      • 适用场景:英文内容优先的全球化应用。
    • Jina-embeddings-v2(Jina AI)
      • 特点:参数量仅35M,支持实时推理(RT<50ms),适合轻量化 部署。
      • 适用场景:轻量级文本处理、实时推理任务。
  2. 中文嵌入模型
    • xiaobu-embedding-v2
      • 特点:针对中文语义优化,语义理解能力强。
      • 适用场景:中文文本分类、语义检索。
    • M3E-Turbo
      • 特点:针对中文优化的轻量模型,适合本地私有化部署。
      • 适用场景:中文法律、医疗领域检索任务。
    • stella-mrl-large-zh-v3.5-1792
      • 特点:处理大规模中文数据能力强,捕捉细微语义关系。
      • 适用场景:中文文本高级语义分析、自然语言处理任务。
  3. 指令驱动与复杂任务模型
    • gte-Qwen2-7B-instruct(阿里巴巴)
      • 特点:基于Qwen大模型微调,支持代码与文本跨模态检索。
      • 适用场景:复杂指令驱动任务、智能问答系统。
    • E5-mistral-7B(Microsoft)
      • 特点:基于Mistral架构,Zero-shot任务表现优异。
      • 适用场景:动态调整语义密度的复杂系统。
  4. 企业级与复杂系统
    • BGE-M3(智源研究院)
      • 特点:适合企业级部署,支持混合检索。
      • 适用场景:企业级语义检索、复杂RAG应用。
    • E5-mistral-7B(Microsoft)
      • 特点:适合企业级部署,支持指令微调。
      • 适用场景:需要动态调整语义密度的复杂系统。

下面我们简单看一下BGE-M3和gte-qwen2的使用

BGE-M3

#模型下载
from modelscope import snapshot_download
model_dir = snapshot_download('BAAI/bge-m3', cache_dir='/root/autodl-tmp/models')


from FlagEmbedding import BGEM3FlagModel

model = BGEM3FlagModel('./models/BAAI/bge-m3',  
                       use_fp16=True) # Setting use_fp16 to True speeds up computation with a slight performance degradation

sentences_1 = ["What is BGE M3?", "Defination of BM25"]
sentences_2 = ["BGE M3 is an embedding model supporting dense retrieval, lexical matching and multi-vector interaction.", 
               "BM25 is a bag-of-words retrieval function that ranks a set of documents based on the query terms appearing in each document"]

embeddings_1 = model.encode(sentences_1, 
                            batch_size=12, 
                            max_length=8192, # If you don't need such a long length, you can set a smaller value to speed up the encoding process.
                            )['dense_vecs']
embeddings_2 = model.encode(sentences_2)['dense_vecs']

similarity = embeddings_1 @ embeddings_2.T
print('embeddings_1:',embeddings_1)
print('embeddings_2:',embeddings_2)
print('embeddings_2.T:',embeddings_2.T)
print('similarity:',similarity)

similarity = embeddings_1 @ embeddings_2.T 用于计算两组嵌入向量(embeddings)之间的相似度矩阵

embedding_1 包含的是第一组句子sentences_1 的嵌入向量,形状是[sentences_1的长度(2),嵌入维度(1024)] embeddubg_2 包含的是第二组句子sentences_2 的嵌入向量,形状是[sentences_2的长度(2),嵌入维度(1024)]

embeddings_2.T 是对 embeddings_2 进行转置操作,形状变为 [嵌入维度(1024), sentences_2的数量(2)]

@符号用于矩阵乘法运算

总的来说,这段代码计算了两组句子的预想相似度矩阵,结果的形状是[[sentences_1的数量(2), sentences_2的数量(2)]

我们来看一下输出:

===============================embeddings_1===============================
[[-0.03411702 -0.04707835 -0.00089448 ...  0.04828532  0.00755428
  -0.02961648]
 [-0.0104174  -0.04479254 -0.02429204 ... -0.008193    0.01503994
   0.01113799]]
===============================embeddings_2===============================
[[-0.00011139 -0.06657311 -0.00018451 ...  0.04317546 -0.02131795
   0.01383034]
 [ 0.00191871 -0.05004179 -0.00031009 ... -0.01365599  0.00120368
   0.0082929 ]]
===============================embeddings_2.T===============================
[[-0.00011139  0.00191871]
 [-0.06657311 -0.05004179]
 [-0.00018451 -0.00031009]
 ...
 [ 0.04317546 -0.01365599]
 [-0.02131795  0.00120368]
 [ 0.01383034  0.0082929 ]]
===============================similarity===============================
[[0.6259036  0.34749582]
 [0.34986773 0.6782464 ]]

可以看到: "What is BGE M3?" 与 "BGE M3 is an embedding model..." 的相似度为0.6265(较高) "What is BGE M3?" 与 "BM25 is a bag-of-words retrieval function..." 的相似度为 0.3477(较低) "Defination of BM25" 与 "BGE M3 is an embedding model..." 的 相似度为0.3499(较低) "Defination of BM25" 与 "BM25 is a bag-of-words retrieval function..." 的相似度为 0.678(较高)

GTE-QW2

示例一

from modelscope import snapshot_download
model_dir = snapshot_download('iic/gte_Qwen2-1.5B-instruct', cache_dir='./models')

from sentence_transformers import SentenceTransformer

model_dir = "./qw_models/iic/gte_Qwen2-1___5B-instruct"
model = SentenceTransformer(model_dir, trust_remote_code=True)
# In case you want to reduce the maximum length:
model.max_seq_length = 8192

queries = [
    "how much protein should a female eat",
    "summit define",
]
documents = [
    "As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.",
    "Definition of summit for English Language Learners. : 1  the highest point of a mountain : the top of a mountain. : 2  the highest level. : 3  a meeting or series of meetings between the leaders of two or more governments.",
]

query_embeddings = model.encode(queries, prompt_name="query")
document_embeddings = model.encode(documents)

scores = (query_embeddings @ document_embeddings.T) * 100
print(scores.tolist())

输出:

[[78.49691772460938, 17.04286003112793], [14.924489974975586, 75.37960815429688]]

示例二

import torch  # PyTorch深度学习库
import torch.nn.functional as F  # PyTorch函数式接口,包含各种神经网络函数

from torch import Tensor  # 导入Tensor类型,用于类型提示
from modelscope import AutoTokenizer, AutoModel  # 从modelscope导入自动分词器和模型加载器


# 定义最后一个token池化函数
# 该函数从最后的隐藏状态中提取每个序列的最后一个有效token的表示
def last_token_pool(last_hidden_states: Tensor,
                 attention_mask: Tensor) -> Tensor:
    # 检查是否为左侧填充(即所有序列最后一个位置都有效)
    left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
    if left_padding:
        # 如果是左侧填充,直接返回最后一个位置的隐藏状态
        return last_hidden_states[:, -1]
    else:
        # 如果是右侧填充,计算每个序列的实际长度(减1是因为索引从0开始)
        sequence_lengths = attention_mask.sum(dim=1) - 1
        batch_size = last_hidden_states.shape[0]
        # 返回每个序列最后一个有效token的隐藏状态
        return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]


# 定义获取详细指令的函数
# 将任务描述和查询组合成特定格式的指令
def get_detailed_instruct(task_description: str, query: str) -> str:
    return f'Instruct: {task_description}\nQuery: {query}'


# 每个查询都必须附带一个描述任务的简短指令
# 定义任务描述:给定网络搜索查询,检索相关的回答段落
task = 'Given a web search query, retrieve relevant passages that answer the query'
# 创建查询列表,每个查询都通过get_detailed_instruct函数添加了任务描述
queries = [
    get_detailed_instruct(task, 'how much protein should a female eat'),  # 女性应该摄入多少蛋白质
    get_detailed_instruct(task, 'summit define')  # summit(顶峰)的定义
]
# 检索文档不需要添加指令
documents = [
    "As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.",  # 关于女性蛋白质摄入量的文档
    "Definition of summit for English Language Learners. : 1  the highest point of a mountain : the top of a mountain. : 2  the highest level. : 3  a meeting or series of meetings between the leaders of two or more governments."  # 关于summit定义的文档
]
# 将查询和文档合并为一个输入文本列表
input_texts = queries + documents

# 设置模型路径
model_dir = "./models/iic/gte_Qwen2-1___5B-instruct"
# 加载分词器,trust_remote_code=True允许使用远程代码
tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
# 加载模型
model = AutoModel.from_pretrained(model_dir, trust_remote_code=True)

# 设置最大序列长度
max_length = 8192

# 对输入文本进行分词处理
# padding=True:对较短的序列进行填充,使批次中所有序列长度一致
# truncation=True:截断超过max_length的序列
# return_tensors='pt':返回PyTorch张量
batch_dict = tokenizer(input_texts, max_length=max_length, padding=True, truncation=True, return_tensors='pt')
# 将分词后的输入传入模型,获取输出
outputs = model(**batch_dict)
# 使用last_token_pool函数从最后的隐藏状态中提取每个序列的表示
embeddings = last_token_pool(outputs.last_hidden_state, batch_dict['attention_mask'])

# 对嵌入向量进行L2归一化,使其长度为1
# p=2表示L2范数,dim=1表示在第1维(特征维度)上进行归一化
embeddings = F.normalize(embeddings, p=2, dim=1)
# 计算查询和文档之间的相似度分数
# embeddings[:2]:查询的嵌入向量(前两个)
# embeddings[2:]:文档的嵌入向量(后两个)
# .T:转置操作
# * 100:将相似度分数缩放到0-100的范围
scores = (embeddings[:2] @ embeddings[2:].T) * 100
# 打印相似度分数
print(scores.tolist())
# 结果解释:
# [[70.00666809082031, 8.184867858886719], [14.62420654296875, 77.71405792236328]]
# 70.00: 第一个查询(蛋白质摄入)与第一个文档(蛋白质指南)的相似度(高相关)
# 8.18: 第一个查询(蛋白质摄入)与第二个文档(summit定义)的相似度(低相关)
# 14.62: 第二个查询(summit定义)与第一个文档(蛋白质指南)的相似度(低相关)
# 77.71: 第二个查询(summit定义)与第二个文档(summit定义)的相似度(高相关)

gte-Qwen2-7B-instruct 是基于 Qwen2的指令优化型嵌入模型

  • 指令优化:经过大量指令-响应对的训练,特别擅长理解和生 成高质量的文本。
  • 性能表现:在文本生成、问答系统、文本分类、情感分析、 命名实体识别和语义匹配等任务中表现优异。
  • 适合场景:适合复杂问答系统,处理复杂的多步推理问题, 能够生成准确且自然的答案。

优势:

  • 指令理解和执行能力强,适合复杂的指令驱动任务。
  • 多语言支持,能够处理多种语言的文本。
  • 在文本生成和语义理解任务中表现优异。

局限:

  • 计算资源需求较高,适合资源充足的环境。