webassembly009 transformers.js 网页端侧推理 NLLB翻译模型

应用页面

• 运行代码：npm i , npm run dev
  

App.jsx

import { useEffect, useRef, useState } from 'react';
import LanguageSelector from './components/LanguageSelector'; // 语言选择器组件
import Progress from './components/Progress'; // 进度条组件

import './App.css'; // 应用样式

function App() {

  // 状态管理
  const [ready, setReady] = useState(null); // 模型准备状态
  const [disabled, setDisabled] = useState(false); // 翻译按钮是否禁用
  const [progressItems, setProgressItems] = useState([]); // 进度条
  // 输入输出状态管理
  const [input, setInput] = useState('I love walking my dog.'); // 输入文本
  const [sourceLanguage, setSourceLanguage] = useState('eng_Latn'); // 源语言
  const [targetLanguage, setTargetLanguage] = useState('fra_Latn'); // 目标语言
  const [output, setOutput] = useState(''); // 输出文本


  // 创建一个对worker对象的引用
  const worker = useRef(null);
  // 使用`useEffect`钩子在组件挂载时设置worker
  useEffect(() => {
    if (!worker.current) {
      // 如果worker不存在，则创建一个新的Worker实例
      worker.current = new Worker(new URL('./worker.js', import.meta.url), {
        type: 'module' // 指定worker脚本类型为ES模块
      });
    }

    // 给worker添加消息事件监听器
    worker.current.addEventListener('message', onMessageReceived);
    
    // 定义从worker接收消息后的回调处理函数
    const onMessageReceived = (e) => {
      switch (e.data.status) {
        case 'initiate':     // 开始加载模型文件时，更新状态并添加进度项
          setReady(false);
          setProgressItems(prev => [...prev, e.data]);
          break;
        case 'progress':    // 更新模型文件加载进度
          setProgressItems(prev => prev.map(item => item.file === e.data.file ? { ...item, progress: e.data.progress } : item));
          break;
        case 'done':        // 模型文件加载完成，从进度列表中移除该项目
          setProgressItems(prev => prev.filter(item => item.file !== e.data.file));
          break;
        case 'ready':       // 工作线程准备好接受消息
          setReady(true);
          break;
        case 'update':      // 更新输出文本
          setOutput(e.data.output);
          break;
        case 'complete':    // 翻译完成，启用“翻译”按钮
          setDisabled(false);
          break;
      }
    };

    // 返回一个清理函数，在组件卸载时移除事件监听器
    return () => worker.current.removeEventListener('message', onMessageReceived);
  });

  // 翻译函数
  const translate = () => {
    setDisabled(true); // 禁用翻译按钮
    worker.current.postMessage({ // 发送消息给worker开始翻译
      text: input,
      src_lang: sourceLanguage,
      tgt_lang: targetLanguage,
    });
  }


  // 渲染UI
  return (
    <>
      <h1>Transformers.js</h1>
      <h2>ML-powered multilingual translation in React!</h2>

      {/* UI布局，关键在于调用translate函数 */}
      <div className='container'>
        <div className='language-container'>
          <LanguageSelector type={"Source"} defaultLanguage={"eng_Latn"} onChange={x => setSourceLanguage(x.target.value)} />
          <LanguageSelector type={"Target"} defaultLanguage={"fra_Latn"} onChange={x => setTargetLanguage(x.target.value)} />
        </div>
        <div className='textbox-container'>
          <textarea value={input} rows={3} onChange={e => setInput(e.target.value)}></textarea>
          <textarea value={output} rows={3} readOnly></textarea>
        </div>
      </div>

      <button disabled={disabled} onClick={translate}>Translate</button>

      <div className='progress-bars-container'>
        {ready === false && <label>Loading models... (only run once)</label>}
        {progressItems.map(data => (
          <div key={data.file}>
            <Progress text={data.file} percentage={data.progress} />
          </div>
        ))}
      </div>
    </>
  );
}

export default App; // 导出默认组件

worker.js

• 代码定义了一个使用单例模式的MyTranslationPipeline，并通过Web Worker监听主线程的消息来执行翻译任务（Worker线程消息传递）：

import { pipeline } from '@xenova/transformers'; // 导入pipeline函数用于加载模型


 // 使用单例模式确保仅加载一次翻译管道实例。由于加载管道是一个昂贵的操作，我们不希望每次翻译句子时都重新加载。
class MyTranslationPipeline {
    static task = 'translation'; // 任务类型
    static model = 'Xenova/nllb-200-distilled-600M'; // 模型名称
    static instance = null; // 静态变量：存储单例实例

    /**
     * 获取或创建一个翻译管道的实例。
     * @param progress_callback - 进度回调函数，用于追踪模型加载进度。@returns {Promise<*>} - 返回已加载的翻译管道实例。
     */
    static async getInstance(progress_callback = null) {
        if (this.instance === null) { // 如果尚未创建实例，则加载新的翻译管道并保存
            // 1. pipeline函数transformers.js模型加载：默认情况下，模型将从 Hugging Face Hub 下载并存储在 浏览器缓存 中，有关更多信息，请参见https://hugging-face.cn/docs/transformers.js/custom_usage。
            this.instance = await pipeline(this.task, this.model, { progress_callback });
        }
        return this.instance; // 返回实例
    }
}

// 监听来自主线程的消息事件
self.addEventListener('message', async (event) => {
    // 获取翻译管道实例，首次调用时将加载管道并保存以备将来使用
    let translator = await MyTranslationPipeline.getInstance(x => {
        // 添加进度回调函数，以便跟踪模型加载进度
        self.postMessage(x);
    });

    // 2. transformers.js模型执行实际的翻译操作
    let output = await translator(event.data.text, {
        tgt_lang: event.data.tgt_lang, // 目标语言
        src_lang: event.data.src_lang, // 源语言

        // 允许部分输出，并通过回调函数发送更新
        callback_function: x => {
            self.postMessage({
                status: 'update', // 状态更新
                output: translator.tokenizer.decode(x[0].output_token_ids, { skip_special_tokens: true }) // 解码输出文本
            });
        }
    });

    // 将翻译结果发送回主线程
    self.postMessage({
        status: 'complete', // 完成状态
        output: output, // 最终的翻译输出
    });
});