基于Embedding模型微调的中文意图识别模型(18种意图)

轻量级意图识别模型，基于BGE-M3 Embedding模型搭建

训练数据 SetFit/amazon_massive_scenario_zh-CN

核心模型结构

import torch from torch import nn from sentence_transformers import SentenceTransformer class EmbeddingBasedIntentModel(torch.nn.Module): def __init__(self, embedding_model, device) -> None: super().__init__() self.n_classes = 18 self.embedding = SentenceTransformer(embedding_model, trust_remote_code=True).to(device) self.fc = nn.Sequential( nn.Linear(1024, 128), nn.ReLU(), nn.Dropout(0.3), nn.Linear(128, self.n_classes), ).to(device) def forward(self, input_ids, attention_mask): x = { "input_ids": input_ids, "attention_mask": attention_mask } x = self.embedding(x)["sentence_embedding"] x = self.fc(x)

调用方式

from inference import EmbeddingBasedIntentModelWrapper device = "cpu" embedding_path = 'YOUR_PATH_TO_BGE_EMBEDDING' model_checkpoint = "YOUR_PATH_TO_THE_MODEL" model = EmbeddingBasedIntentModelWrapper(embedding_path, model_checkpoint, device) while True: input_text = input("Enter input: ") result = model.classify(input_text) print(result) ''' 以下为输出效果 Enter input: 帮我开个灯 iot Enter input: 青花瓷 play Enter input: 外面冷不冷 weather Enter input: 点个汉堡王 takeaway Enter input: 买张去东京的机票 transport Enter input: 英国伦敦现在几点 datetime Enter input: 给谢老板发个邮件 email Enter input: 提醒我下周六和小王出去玩 calendar Enter input: 定个明天早上9点的闹钟 alarm Enter input: 音量调到最小 audio '''

训练脚本,模型 checkpoint请见Github

训练分数