使用Universal Sentence Encoder计算AG-NEWS数据集类别相似度咨询

核心实现步骤 #

• 步骤1：数据集分组采样
  先把AG-NEWS数据集按标签拆分为World、Sports、Business、Sci/Tech四个独立子集，每个子集均匀采样1000~5000条文本即可，样本量足够的情况下结果稳定性高，也无需全量计算浪费资源。
• 步骤2：生成类别语义表征
  有两种成熟方案可选：
  • 方案1（简单易用）：用Universal Sentence Encoder对每个类别下的所有采样文本生成句子嵌入，再对该类别所有嵌入向量取算术平均，得到对应类别的唯一中心表征向量
  • 方案2（贴合类别语义）：额外生成四个类别名称的文本嵌入，和上述样本平均向量做加权融合（比如样本向量权重取0.8，类别名称向量权重取0.2），可以减少样本噪音带来的表征偏差
• 步骤3：计算4×4相似度矩阵
  对4个类别的表征向量两两计算余弦相似度，最终输出的4×4矩阵就是你需要的类别相似度结果，矩阵对角线值恒为1（类别与自身的相似度最高）

简化代码示例 #

import tensorflow_hub as hub
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
from datasets import load_dataset

# 加载USE模型和AG-NEWS数据集
use_model = hub.load("https://tfhub.dev/google/universal-sentence-encoder/4")
ag_news = load_dataset("ag_news", split="train")

# 按标签分组存储对应类别的文本
class_texts = [[], [], [], []]
for item in ag_news:
    class_texts[item["label"]].append(item["text"])

# 每个类别采样2000条文本生成类别中心嵌入
class_embeddings = []
for texts in class_texts:
    sampled_texts = np.random.choice(texts, 2000, replace=False)
    embeddings = use_model(sampled_texts)
    class_center = np.mean(embeddings, axis=0)
    class_embeddings.append(class_center)

# 生成4*4相似度矩阵
similarity_matrix = cosine_similarity(class_embeddings)
print(similarity_matrix)

可选优化方向 #

• 如果生成的相似度矩阵和预期的类别语义偏差较大，可以替换余弦相似度为PLDA等适配分类任务的度量方法，更贴合AG-NEWS的实际分类边界
• 可以对样本做预处理，过滤掉长度过短、噪音过大的文本，提升类别表征的准确性

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