AudacityTorch项目实战：语音转文本标签模型集成指南

前言

在音频处理领域，将语音转换为文本标签是一个常见且重要的任务。AudacityTorch项目为开发者提供了将深度学习模型集成到Audacity音频编辑软件中的能力。本文将详细介绍如何利用Facebook的Wav2Vec2模型，创建一个能够将语音波形转换为文本标签的Audacity插件。

环境准备

依赖安装

首先需要安装必要的Python包：

!pip install torchaudio==0.9.0
!pip install transformers
!pip install audacitorch

这些包分别提供：

• torchaudio：音频处理基础功能
• transformers：预训练模型加载
• audacitorch：Audacity模型集成接口

基础导入

import torch
from torch import nn
from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC
import torchaudio
import json

# 禁用梯度计算以提高效率
torch.set_grad_enabled(False)

模型封装原理

AudacityTorch要求模型必须遵循特定的接口规范，核心是WaveformToLabelsBase基类。这个基类确保了模型输入输出的标准化，使Audacity能够正确解析模型结果。

模型封装结构

我们需要构建两个主要组件：

1. SubModels：包含实际的语音识别模型和处理器
2. ModelWrapper：继承自WaveformToLabelsBase，实现模型接口

from audacitorch.core import WaveformToLabelsBase

class SubModels(nn.Module):
    def __init__(self):
        super().__init__()
        # 加载预训练模型和处理器
        self._model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-base-960h", torchscript=True)
        self._processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h", torchscript=True)
        # 创建token到索引的映射
        self.token_to_idx = {val:key for key, val in self._processor.tokenizer.decoder.items()}

关键辅助函数

@torch.jit.script_if_tracing
def get_timestamps(num_preds: int, total_time: int):
    """处理时间戳生成，避免除零错误"""
    if num_preds == 0:
        return torch.empty(1, 0)
    else:
        equal_size_timestamp = total_time / num_preds
        timestamps = torch.zeros(num_preds, 2)
        return timestamps

@torch.jit.script_if_tracing
def check_empty_output(preds, timestamps):
    """处理空输出情况"""
    if preds.shape[0] == 0:
        return torch.tensor([0]), torch.tensor([[0., 0.01]])
    else:
        return preds, timestamps

核心封装类

class ModelWrapper(WaveformToLabelsBase):
    def do_forward_pass(self, _input):
        # 音频预处理
        input_values = self.model._processor(_input, return_tensors="pt", padding="longest").input_values[0]
        
        # 模型推理
        logits = self.model._model(input_values)[0]
        predicted_ids = torch.argmax(logits, dim=-1)
        transcription = self.model._processor.decode(predicted_ids[0])
        num_preds = len(transcription)

        # 将预测转换为one-hot编码
        preds_onehot = torch.FloatTensor(num_preds, len(self.model.token_to_idx))
        preds_onehot.zero_()
        for i, token in enumerate(transcription):
            if token == ' ':
                token = '<s>'
            token_idx = self.model.token_to_idx[token]
            preds_onehot[i][token_idx] = 0.99
        
        # 计算时间戳
        total_time = _input.shape[1] / 16000
        timestamps = get_timestamps(num_preds, total_time)
        
        # 返回标准化输出
        preds = torch.argmax(preds_onehot, dim=-1, keepdim=False) if preds_onehot.numel() else preds_onehot
        preds, timestamps = check_empty_output(preds, timestamps)
        return (preds, timestamps)

模型元数据配置

为了让Audacity正确识别和使用模型，需要提供详细的元数据：

vocab = [str(letter) for letter in sub_models._processor.tokenizer.decoder.values()]

metadata = {
    'sample_rate': 16000,  # 采样率
    'domain_tags': ['speech'],  # 适用领域
    'short_description': '语音转文本标签模型',
    'long_description': '基于Facebook的Wav2Vec2模型实现的语音转文本功能...',
    'tags': ['speech-to-text'],  # 功能标签
    'effect_type': 'waveform-to-labels',  # 效果类型
    'multichannel': False,  # 是否支持多声道
    'labels': vocab,  # 输出标签集
}

模型导出与保存

模型追踪与保存

from audacitorch.utils import save_model, get_example_inputs

# 获取示例输入并追踪模型
example_inputs = get_example_inputs()
traced_model = torch.jit.trace(torchscript_model, example_inputs[0])

# 保存模型和元数据
save_model(traced_model, metadata, Path('audacity-Wav2Vec2-Base'))

导出文件结构

保存后将生成两个文件：

1. audacity-Wav2Vec2-Base.pt：模型权重文件
2. audacity-Wav2Vec2-Base.json：模型元数据文件

技术要点解析

1. 模型兼容性：当前Huggingface的transformers模块对torchscript的支持有限，Wav2Vec2系列模型兼容性较好

2. 输入输出规范：

   • 输入：单声道音频波形，采样率16kHz
   • 输出：包含预测标签和对应时间戳的元组

3. 异常处理：

   • 空输入处理
   • 除零保护
   • 无效输出处理

结语

通过本文的步骤，开发者可以将先进的语音识别模型集成到Audacity中，为用户提供高质量的语音转文本功能。这种集成方式不仅限于Wav2Vec2模型，也可以扩展到其他音频处理模型，为Audacity生态带来更多可能性。