基于Hugging Face Hub的Dolly模型计算Embeddings余弦相似度遇问题求助

问题根源 #

你用feature-extraction pipeline获取的是每个token的embedding，而非整个文本的统一embedding：

• 输入的单条文本返回形状为(1, 7, 2560)：1是样本数，7是句子的token数量，2560是模型隐藏层维度
• 数据集的3条文本返回形状为(3, N, 2560)，每条文本的token数N各不相同
  这种三维数组直接做归一化或点积运算，必然会出现维度不匹配的错误。

修正方案 #

核心是先对每个文本的token embeddings做池化处理，将其转换为单向量的文本级embedding，再进行后续计算。这里用最常用的均值池化：

# Import Pipeline
from transformers import pipeline
import torch
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.preprocessing import normalize

# Create Feature Extraction Object
feature_extraction = pipeline('feature-extraction',
                              model='databricks/dolly-v2-3b', 
                              torch_dtype=torch.bfloat16,
                              trust_remote_code=True, 
                              device_map="auto")

# Define Inputs
text = ["hey how are you doing today love"]
dataset = ["hey how are you doing today love", "I am doing great", "What about you?"]

# Create Embeddings
text_embeddings = feature_extraction(text)
dataset_embeddings = feature_extraction(dataset)

# 1. 均值池化：将每个文本的所有token embedding取平均，得到文本级embedding
# 处理查询文本
text_embedding = np.mean(text_embeddings[0], axis=0).reshape(1, -1)  # 形状变为 (1, 2560)
# 处理数据集文本
dataset_embeddings_pooled = []
for emb in dataset_embeddings:
    pooled_emb = np.mean(emb, axis=0)
    dataset_embeddings_pooled.append(pooled_emb)
dataset_embedding = np.array(dataset_embeddings_pooled)  # 形状变为 (3, 2560)

# 2. L2归一化
text_embedding_norm = normalize(text_embedding, norm='l2')
dataset_embedding_norm = normalize(dataset_embedding, norm='l2')

# 3. 计算余弦相似度与角度距离
cos_sim = cosine_similarity(text_embedding_norm, dataset_embedding_norm)
angular_distance = np.arccos(cos_sim) / np.pi

print("余弦相似度：", cos_sim)
print("角度距离：", angular_distance)

关键说明 #

• 均值池化是处理GPT类模型（无[CLS] token）文本embedding的通用方案，能将可变长度的token序列转换为固定维度的文本向量
• 若偏好其他池化方式，也可以尝试最大值池化或加权池化，效果差异不大
• 使用sklearn的cosine_similarity可以自动处理维度匹配，比手动点积更简洁可靠

内容的提问来源于stack exchange，提问作者john_mon