下面我们介绍如何使用BGE-M3模型实现短文本的分类,比如用户输入一段文字,然后对其进行分类,从而得到其对应的主题的需求。
这里训练的数据使用的是今日头条的新闻,其数据格式类似如下:
6551700932705387022_!_101_!_news_culture_!_京城最值得你来场文化之旅的博物馆_!_保利集团,马未都,中国科学技术馆,博物馆,新中国 6552368441838272771_!_101_!_news_culture_!_发酵床的垫料种类有哪些？哪种更好？_!_ 6552407965343678723_!_101_!_news_culture_!_上联：黄山黄河黄皮肤黄土高原。怎么对下联？_!_ 6552332417753940238_!_101_!_news_culture_!_林徽因什么理由拒绝了徐志摩而选择梁思成为终身伴侣？_!_ 6552475601595269390_!_101_!_news_culture_!_黄杨木是什么树？_!_ ...
现在要做的是就是根据后面的标题预测前面的分类。 首先需要进行预处理,对应的代码如下:
import zhconv
import jieba.analyse
from collections import defaultdict
filename = "toutiao_cat_data.txt"
train_file = "toutiao.txt"
data = defaultdict(list)
stop_words = []
with open("stop_word.txt",encoding="utf-8") as f:
for line in f:
line = line.strip()
stop_words.append(line)
def parse_clue_line(line):
parts = line.split('!_')
if len(parts) >= 5:
title = parts
keywords = parts
# 构造适合 BGE-M3 的描述文本
return parts[1], f"{title} {keywords}"
return parts[1], ""
def extract_keywords(text, top_k=5):
"""提取单条文本的关键词"""
# allowPOS 限制为名词、动词等实词
kws = jieba.analyse.extract_tags(text, topK=top_k, withWeight=False, allowPOS=('n', 'ns', 'nr', 'nt', 'v'))
return kws
with open(filename, encoding="utf-8") as f:
for line in f:
label, text = parse_clue_line(line)
text = zhconv.convert(text, "zh-hans")
data[label].append(text)
# 4. 核心逻辑：提取关键词 -> 去重 -> 拼接 -> 向量化
label_embeddings = {}
MAX_KEYWORDS_PER_LABEL = 200 # 限制每个标签最终拼接的关键词总数，防止超出模型窗口
with open(train_file, "w", encoding="utf-8") as f:
for label, texts in data.items():
all_keywords = []
seen_keywords = set()
# 遍历该标签下的所有文本，提取关键词并去重
for text in texts:
kws = extract_keywords(text, top_k=5)
for kw in kws:
if kw not in seen_keywords and kw not in stop_words and len(kw) > 1:
all_keywords.append(kw)
seen_keywords.add(kw)
# 达到上限则停止，保证向量质量不被稀释
if len(all_keywords) >= MAX_KEYWORDS_PER_LABEL:
break
if len(all_keywords) >= MAX_KEYWORDS_PER_LABEL:
break
# 拼接成主题描述字符串
merged_desc = " ".join(all_keywords)
f.write("__label__{},{}\n".format(label,merged_desc))
其输出结果类似如下:
__label__101_,博物馆 来场 中国科学技术馆 保利 京城 垫料 发酵 种类 黄皮肤 上联 下联 黄山 黄土高原 林徽因 终身伴侣 梁思成 徐志摩 拒绝 黄杨木 草根 星光 登上 写实 绘画 松涛 婉转 骑牛 老子 读书 醉人 ...
之后加载BGE-M3模型,其相关的代码如下:
import numpy as np
from FlagEmbedding import BGEM3FlagModel
def get_top3_categories(query: str, top_k=3):
"""
返回用户问题最匹配的 Top-K 分类及分数
"""
# 生成查询向量
query_embedding = model.encode([query], batch_size=128)['dense_vecs']
# 计算余弦相似度 (BGE-M3 dense vectors 已归一化，点积即余弦相似度)
similarities = np.dot(topic_embeddings, query_embedding.T).flatten()
# 获取排序后的索引 (从高到低)
sorted_indices = np.argsort(similarities)[::-1]
results = []
for i in range(min(top_k, len(topics))):
idx = sorted_indices[i]
results.append({
"rank": i + 1,
"category_id": topics[idx]['id'],
"category_name": topics[idx]['name'],
"score": float(similarities[idx]) # 转换为 Python float
})
return results
if __name__ == "__main__":
# 1. 初始化模型
model = BGEM3FlagModel('BAAI/bge-m3', use_fp16=False)
# 2. 预设主题及其描述（已提取关键词拼接）
topics = []
classify_dict = {
"100": "民生","101": "文化","102": "娱乐","103": "体育", "106": "房产","104": "财经", "107": "汽车","108": "教育","109": "科技","110": "军事","112": "旅游","113": "国际", "114": "证券股票","115": "农业","116": "电竞游戏"
}
with open("toutiao.txt",encoding="utf-8") as f:
for line in f:
if line:
label,text = line.split(",")
_id = label.replace("__label__","").replace("_","").strip()
topics.append({"desc": text, "id":_id, "name": classify_dict[_id]})
# 3. 预生成主题向量 (离线完成)
topic_descs = [t['desc'] for t in topics]
topic_embeddings = model.encode(topic_descs, batch_size=128)['dense_vecs']
# --- 测试 ---
while 1:
try:
query = input("> ")
#query = "我想了解 Milvus 和 MySQL 在存储向量时的区别"
top3 = get_top3_categories(query)
for item in top3:
print(f"排名: {item['rank']} | 分类: {item['category_name']} | 相似度: {item['score']:.4f}")
except KeyboardInterrupt:
break
之后是其测试结果:
$ 八一建军节前夕这次集体学习，传递了哪些重要信息？ 排名: 1 | 分类: 军事 | 相似度: 0.4640 排名: 2 | 分类: 国际 | 相似度: 0.4147 排名: 3 | 分类: 教育 | 相似度: 0.4112 $ 新势力7月洗牌：零跑独破10万，“蔚小理米”困在3万档 排名: 1 | 分类: 电竞游戏 | 相似度: 0.4557 排名: 2 | 分类: 汽车 | 相似度: 0.4516 排名: 3 | 分类: 财经 | 相似度: 0.3863 $ 日媒透露：高市早苗已有两月未与任何外国领导人通电话；安倍晋三当年每月与特朗普通一次话，即使没有大事，也会问候特朗普妻子 排名: 1 | 分类: 国际 | 相似度: 0.4725 排名: 2 | 分类: 军事 | 相似度: 0.4083 排名: 3 | 分类: 科技 | 相似度: 0.3925 $ 美财长称按照特朗普指示，正在全球范围内追查伊朗资产：如今伊朗甚至无力支付军饷 排名: 1 | 分类: 国际 | 相似度: 0.4821 排名: 2 | 分类: 科技 | 相似度: 0.4519 排名: 3 | 分类: 军事 | 相似度: 0.4241 $ AI大模型龙头，大跌超20%，资金逆势抢筹超20亿港元 排名: 1 | 分类: 证券股票 | 相似度: 0.5167 排名: 2 | 分类: 科技 | 相似度: 0.4370 排名: 3 | 分类: 财经 | 相似度: 0.4087 $ 日媒试玩国产《锦衣卫》：高速激战很上头！ 排名: 1 | 分类: 汽车 | 相似度: 0.4776 排名: 2 | 分类: 电竞游戏 | 相似度: 0.4366 排名: 3 | 分类: 军事 | 相似度: 0.4262 $ 通关仅需1h！性感八尺大人"温暖"新游可亲密大姐姐 排名: 1 | 分类: 电竞游戏 | 相似度: 0.4293 排名: 2 | 分类: 文化 | 相似度: 0.4217 排名: 3 | 分类: 汽车 | 相似度: 0.4204 $ 超强台风 “白海豚”或成今年“风王”！官方：关键在6日到8日前后，不要只盯登陆点 排名: 1 | 分类: 汽车 | 相似度: 0.4433 排名: 2 | 分类: 电竞游戏 | 相似度: 0.4344 排名: 3 | 分类: 财经 | 相似度: 0.4316
可以看到对于新闻的分类,效果还是可以的。
如果喜欢这篇文章或对您有帮助，可以：[☕] 请我喝杯咖啡 | [💓] 小额赞助
