TensorFlow.js聊天机器人开发：如何为文本分词添加特殊令牌？

自定义令牌插入+基础分词器扩展 #

既然你已经在用Universal Sentence Encoder，可以在分词前手动插入自定义特殊令牌，再结合专用分词工具处理令牌序列。

首先定义专属特殊令牌，用于区分上下文、历史消息和当前提问，然后在文本预处理阶段把这些令牌插入对应位置，再用tfjs-tokenizers库完成分词：

// 定义特殊令牌
const SPECIAL_TOKENS = {
  context: '<context>',
  pastMsg: '<past_msg>',
  prompt: '<prompt>'
};

// 预处理聊天文本，插入特殊令牌
function buildInputText(context, pastMessages, currentPrompt) {
  let input = `${SPECIAL_TOKENS.context} ${context} `;
  pastMessages.forEach(msg => {
    input += `${SPECIAL_TOKENS.pastMsg} ${msg} `;
  });
  input += `${SPECIAL_TOKENS.prompt} ${currentPrompt}`;
  return input;
}

// 结合tfjs-tokenizers完成分词
import { WordPieceTokenizer } from '@tensorflow-models/tfjs-tokenizers';

async function tokenizeChatInput(context, pastMsgs, prompt) {
  // 加载WordPiece分词器（适配BERT类预训练模型）
  const tokenizer = await WordPieceTokenizer.load();
  // 预处理文本
  const processedText = buildInputText(context, pastMsgs, prompt);
  // 编码得到令牌序列
  const encoded = tokenizer.encode(processedText);
  return encoded.tokens;
}

扩展现有分词器词汇表 #

如果需要让分词器原生识别特殊令牌，避免令牌被拆分，可以直接向tfjs-tokenizers的词汇表中添加自定义令牌：

import { WordPieceTokenizer } from '@tensorflow-models/tfjs-tokenizers';

async function createCustomTokenizer() {
  // 加载预训练分词器
  const baseTokenizer = await WordPieceTokenizer.load();
  // 获取现有词汇表
  const vocab = baseTokenizer.getVocab();
  // 添加自定义特殊令牌
  const newTokens = ['<context>', '<past_msg>', '<prompt>'];
  newTokens.forEach(token => {
    if (!vocab.has(token)) {
      vocab.set(token, vocab.size);
    }
  });
  // 用更新后的词汇表初始化新的分词器
  return new WordPieceTokenizer(vocab);
}

// 使用示例
const customTokenizer = await createCustomTokenizer();
const testText = `<context> 周末计划 <past_msg> 想去爬山 <prompt> 推荐一座近郊的山？`;
const tokenized = customTokenizer.encode(testText);
console.log(tokenized.tokens); // 特殊令牌会被完整保留

关键注意点 #

• 若基于预训练模型开发聊天机器人，添加新令牌后需要对模型进行微调，让模型学习到这些令牌的语义作用
• Universal Sentence Encoder主打句子嵌入，若需要令牌级别的细粒度处理，优先搭配tfjs-tokenizers使用
• 特殊令牌尽量用独特格式（如尖括号包裹），避免和普通词汇冲突

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