elec_rag_ui/resources/asr-worker.cjs

/**
 * ASR worker process
 *
 * 独立进程运行 sherpa-onnx-node，避免与 onnxruntime-node 发生 DLL 冲突。
 * 通过 process.send / process.on('message') 与主进程通信。
 */

const { existsSync } = require('fs')
const { join } = require('path')

process.on('uncaughtException', (err) => {
  console.error('[asr-worker] 未捕获异常：', err)
  if (process.send) {
    process.send({ type: 'error', message: `ASR worker 未捕获异常：${err.message || err}` })
  }
  process.exit(1)
})

const { OnlineRecognizer, OfflinePunctuation } = require('sherpa-onnx-node')

let recognizer = null
let activeStream = null
let loadedModelDir = ''
let punctuation = null
let loadedPunctuationModelDir = ''
let chunkCount = 0
let resultCount = 0

function resolveModelFiles(modelDir) {
  if (!modelDir || !existsSync(modelDir)) return null

  const encoderCandidates = ['encoder.int8.onnx', 'encoder.onnx']
  let encoder = ''
  for (const name of encoderCandidates) {
    const p = join(modelDir, name)
    if (existsSync(p)) {
      encoder = p
      break
    }
  }

  const decoder = join(modelDir, 'decoder.onnx')
  const tokens = join(modelDir, 'tokens.txt')

  if (!encoder || !existsSync(decoder) || !existsSync(tokens)) return null
  return { encoder, decoder, tokens }
}

function resolvePunctuationModel(modelDir) {
  if (!modelDir || !existsSync(modelDir)) return null
  const candidates = ['model.int8.onnx', 'model.onnx']
  for (const name of candidates) {
    const p = join(modelDir, name)
    if (existsSync(p)) return p
  }
  return null
}

function send(type, payload = {}) {
  if (process.send) {
    process.send({ type, ...payload })
  }
}

function addPunctuation(text) {
  if (!text || !punctuation) return text
  try {
    return punctuation.addPunct(text)
  } catch (e) {
    console.error('[asr-worker] 标点恢复失败：', e)
    return text
  }
}

const CJK_PUNCTUATION_RE = /[，。！？、；：""''（）【】《》]$/

function endsWithPunctuation(text) {
  return CJK_PUNCTUATION_RE.test(text) || /[,.!?;:\"'\)\]}$]$/.test(text)
}

function fallbackSegmentPunctuation(text) {
  if (!text) return text
  if (endsWithPunctuation(text)) return text
  // 段末加逗号，表示停顿但未结束
  return text + '，'
}

function fallbackFinalPunctuation(text) {
  if (!text) return text
  if (endsWithPunctuation(text)) return text
  // 整句结束加句号
  return text + '。'
}

function initModel(modelDir, numThreads = 2, punctuationModelDir = '') {
  // 加载/重载标点模型（如果配置变化）
  if (punctuationModelDir && punctuationModelDir !== loadedPunctuationModelDir) {
    const punctModel = resolvePunctuationModel(punctuationModelDir)
    if (!punctModel) {
      console.warn(`[asr-worker] 标点模型目录不完整，未找到 model.int8.onnx 或 model.onnx：${punctuationModelDir}`)
      punctuation = null
      loadedPunctuationModelDir = ''
    } else {
      try {
        punctuation = new OfflinePunctuation({
          model: {
            ctTransformer: punctModel,
            numThreads,
            provider: 'cpu',
            debug: 0
          }
        })
        loadedPunctuationModelDir = punctuationModelDir
        console.log(`[asr-worker] 标点模型加载成功：${punctModel}`)
      } catch (e) {
        console.error('[asr-worker] 加载标点模型失败：', e)
        punctuation = null
        loadedPunctuationModelDir = ''
      }
    }
  } else if (!punctuationModelDir) {
    punctuation = null
    loadedPunctuationModelDir = ''
  }

  if (loadedModelDir === modelDir && recognizer) {
    send('ready')
    return
  }

  const files = resolveModelFiles(modelDir)
  if (!files) {
    send('error', { message: `ASR 模型目录不完整，需包含 encoder(.int8).onnx、decoder.onnx 与 tokens.txt：${modelDir}` })
    return
  }

  if (recognizer) {
    try {
      if (activeStream) {
        activeStream.inputFinished?.()
        activeStream = null
      }
      recognizer = null
    } catch (e) {
      // ignore
    }
  }

  const config = {
    featConfig: {
      sampleRate: 16000,
      featureDim: 80
    },
    modelConfig: {
      paraformer: {
        encoder: files.encoder,
        decoder: files.decoder
      },
      tokens: files.tokens,
      numThreads,
      provider: 'cpu',
      debug: 0,
      modelType: 'paraformer'
    },
    decodingMethod: 'greedy_search',
    maxActivePaths: 4,
    enableEndpoint: true,
    rule1MinTrailingSilence: 800,
    rule2MinTrailingSilence: 1200,
    rule3MinUtteranceLength: 20000
  }

  try {
    recognizer = new OnlineRecognizer(config)
    loadedModelDir = modelDir
    send('ready')
  } catch (e) {
    send('error', { message: `加载 ASR 模型失败：${e.message || e}` })
  }
}

function startStream() {
  if (!recognizer) {
    send('error', { message: 'ASR 模型未加载' })
    return
  }
  if (activeStream) {
    try {
      activeStream.inputFinished?.()
    } catch {
      // ignore
    }
  }
  activeStream = recognizer.createStream()
  send('streamStarted')
}

function feedChunk(samples) {
  if (!recognizer || !activeStream) {
    send('error', { message: 'ASR 未初始化或未开始识别流' })
    return
  }

  chunkCount++
  const floatSamples = new Float32Array(samples)
  const rms = Math.sqrt(floatSamples.reduce((s, v) => s + v * v, 0) / floatSamples.length)

  // 每 50 个 chunk 打印一次统计
  if (chunkCount % 50 === 1) {
    console.log(
      `[asr-worker] 收到 chunk #${chunkCount}, 样本数=${floatSamples.length}, rms=${rms.toFixed(5)}`
    )
  }

  try {
    activeStream.acceptWaveform({ sampleRate: 16000, samples: floatSamples })

    while (recognizer.isReady(activeStream)) {
      recognizer.decode(activeStream)
    }

    const isEndpoint = recognizer.isEndpoint(activeStream)
    if (isEndpoint) {
      recognizer.reset(activeStream)
    }

    const result = recognizer.getResult(activeStream)
    let text = result.text?.trim() ?? ''
    // 端点处对当前段落后处理：优先使用标点模型，否则用简单规则断句
    if (text && isEndpoint) {
      if (punctuation) {
        text = addPunctuation(text)
      } else {
        text = fallbackSegmentPunctuation(text)
      }
    }
    if (text) {
      resultCount++
      console.log(`[asr-worker] 识别结果 #${resultCount}: "${text}" isEndpoint=${isEndpoint}`)
    }
    send('result', { text, isEndpoint })
  } catch (e) {
    console.error('[asr-worker] 识别失败：', e)
    send('error', { message: `识别失败：${e.message || e}` })
  }
}

function stopStream() {
  if (!recognizer || !activeStream) {
    send('stopped', { text: '' })
    return
  }

  try {
    const result = recognizer.getResult(activeStream)
    let text = result.text?.trim() ?? ''
    if (text) {
      if (punctuation) {
        text = addPunctuation(text)
      } else {
        text = fallbackFinalPunctuation(text)
      }
    }
    try {
      activeStream.inputFinished?.()
    } catch {
      // ignore
    }
    activeStream = null
    send('stopped', { text })
  } catch (e) {
    activeStream = null
    send('stopped', { text: '' })
    send('error', { message: `停止识别失败：${e.message || e}` })
  }
}

function unloadModel() {
  try {
    activeStream = null
    recognizer = null
    punctuation = null
    loadedModelDir = ''
    loadedPunctuationModelDir = ''
  } catch {
    // ignore
  }
}

process.on('message', (msg) => {
  if (!msg || typeof msg !== 'object') return

  switch (msg.type) {
    case 'init':
      initModel(msg.modelDir, msg.numThreads, msg.punctuationModelDir)
      break
    case 'startStream':
      startStream()
      break
    case 'feedChunk':
      feedChunk(msg.samples)
      break
    case 'stopStream':
      stopStream()
      break
    case 'unloadModel':
      unloadModel()
      break
    default:
      break
  }
})

// 保持进程运行，直到主进程发送 unloadModel 或断开连接
process.on('disconnect', () => {
  unloadModel()
  process.exit(0)
})