78 lines
2.5 KiB
TypeScript
78 lines
2.5 KiB
TypeScript
import { pipeline, env } from "@xenova/transformers";
|
|
|
|
// 配置 Transformers.js
|
|
// 在 Vercel 或容器环境中，通常建议使用远程模型以避免文件系统问题
|
|
// env.allowLocalModels = false;
|
|
// env.useBrowserCache = false;
|
|
|
|
// 定义意图类别
|
|
// 英文标签通常在这些预训练模型上表现更好
|
|
const candidateLabels = [
|
|
"generate video", // 视频生成
|
|
"generate image", // 图片生成
|
|
"edit image", // 图片编辑
|
|
"chat", // 普通聊天
|
|
];
|
|
|
|
// 单例模式保存 pipeline
|
|
let classifierPromise: Promise<any> | null = null;
|
|
|
|
const getClassifier = () => {
|
|
if (!classifierPromise) {
|
|
// 使用量化版本的 Zero-shot Classification 模型
|
|
// xenova/nli-deberta-v3-xsmall 是一个平衡了体积(~80MB)和精度的模型
|
|
// 注意：首次运行时会从 Hugging Face Hub 下载模型，可能需要几秒钟
|
|
console.log("[Classifier] Initializing pipeline with Xenova/nli-deberta-v3-xsmall...");
|
|
classifierPromise = pipeline(
|
|
"zero-shot-classification",
|
|
"Xenova/nli-deberta-v3-xsmall"
|
|
);
|
|
}
|
|
return classifierPromise;
|
|
};
|
|
|
|
/**
|
|
* 识别用户输入的意图
|
|
* @param text 用户输入的文本
|
|
* @returns 意图标签：'video', 'image_gen', 'image_edit', 'chat'
|
|
*/
|
|
export async function classifyIntent(text: string): Promise<"video" | "image_gen" | "image_edit" | "chat"> {
|
|
try {
|
|
const classifier = await getClassifier();
|
|
|
|
// 限制输入长度，防止内存溢出或处理过慢
|
|
const truncatedText = text.slice(0, 500);
|
|
|
|
const result = await classifier(truncatedText, candidateLabels);
|
|
|
|
// result 结构: { sequence: string, labels: string[], scores: number[] }
|
|
// labels 和 scores 是按 score 降序排列的
|
|
const topLabel = result.labels[0];
|
|
const topScore = result.scores[0];
|
|
|
|
console.log(`[Classifier] Input: "${truncatedText.slice(0, 50)}...", Top Label: ${topLabel}, Score: ${topScore}`);
|
|
|
|
// 如果置信度太低，回退到 chat
|
|
if (topScore < 0.2) {
|
|
return "chat";
|
|
}
|
|
|
|
// 映射回应用内部使用的意图代码
|
|
switch (topLabel) {
|
|
case "generate video":
|
|
return "video";
|
|
case "generate image":
|
|
return "image_gen";
|
|
case "edit image":
|
|
return "image_edit";
|
|
default:
|
|
return "chat";
|
|
}
|
|
} catch (error) {
|
|
console.error("[Classifier] Classification failed:", error);
|
|
// 发生错误时回退到 chat (或者调用者可以回退到正则匹配)
|
|
return "chat";
|
|
}
|
|
}
|
|
|