RustからWhisperを使ってみる

Rustで音声認識

はじめに

今回はWhisperを用いて音声認識を行います。 オリジナルのWhisperはPythonを使って実装されていますが、世の中にはWhisperをCで実装したライブラリwhisper.cppがあります。 さらに、それをRustで使うためにラッピングしたライブラリwhisper-rsを用いることになります。

Whisperとは

WhisperはOpenAIが開発した汎用的な音声認識モデルです。 多様な音声の大規模データセットで学習され、音声翻訳や言語識別だけでなく、多言語音声認識を行うことができるマルチタスクモデルでもあります。

したがって、今回の用途である音声認識以外にも認識言語の検出なども行える、非常に多機能なモデルとなっています。

ソースコードとモデルの重みは公開されており、簡単に使用できるようになっています。

Whisper.cppを使う上での注意点

Whisper.cppは以下の制限があります。

• モノラル、32bit float PCMしか対応していないため、例えば、一般的な16bit 整数、ステレオ音声を入力すると意図しない結果が得られる。
• モデルのデータはPyTorch標準のpt形式ではなく、独自のGGML形式で保存されている。

1点目はWhisper-rsに備え付けの関数を用いることで容易に対応することができます。

2点目のGGML形式のモデルはwhisper.cppのmodelsフォルダにあるdownload-ggml-modelsスクリプトを実行することでダウンロードできます。 Windowsであれば、download-ggml-models.cmd, Linuxであれば、download-ggml-models.shを実行することでダウンロードできます。

Whisper-rsを使う

Whisper-rsのexamplesにはファイルを読み込んで使う例はありません。 今回はrodioを使ってmp3ファイルを読み出して、Whisperによって音声認識するところまでをお試ししてみます。

use whisper_rs::{FullParams, SamplingStrategy, WhisperContext};

// 認識したスクリプトを保存する構造体
pub struct ScriptSegment {
  pub script: String,
  pub start: i64,
  pub end: i64,
}

// 認識機能のinterface
pub trait Scripter {
  fn parse(&mut self, audio_data: &[f32]) -> Vec<ScriptSegment>;
}

// Whisper実装
pub struct WhisperScripter {
  ctx: WhisperContext,
}

impl WhisperScripter {
  fn new(model_path: &str) -> Self {
    WhisperScripter {
      ctx: WhisperContext::new(model_path).expect("Failed to load model"),
    }
  }
}

impl Scripter for WhisperScripter {
  fn parse(&mut self, audio_data: &[f32]) -> Vec<ScriptSegment> {
    let mut params = FullParams::new(SamplingStrategy::Greedy { n_past: 0 });
    params.set_language("ja");

    self
      .ctx
      .full(params, audio_data)
      .expect("Failed to recognize audio.");
    let num_segments = self.ctx.full_n_segments();

    let mut scripts = Vec::new();
    for i in 0..num_segments {
      scripts.push(ScriptSegment {
        script: self
          .ctx
          .full_get_segment_text(i)
          .expect("failed to get segment"),
        start: self.ctx.full_get_segment_t0(i),
        end: self.ctx.full_get_segment_t1(i),
      });
    }
    scripts
  }
}

#[cfg(test)]
mod tests {
  use rodio;
  use std::io::BufReader;

  use crate::scripter::*;

  #[test]
  fn test_whisper_scripter() {
    // mp3を読み出し・デコード
    let file = std::fs::File::open("assets/example-audio/common_voice_ja_31833274.mp3").unwrap();
    let decoder = rodio::Decoder::new(BufReader::new(file)).unwrap();
    let data = decoder.into_iter().collect::<Vec<i16>>();

    // Whisperで使えるように変換(i16->f32)
    let data =
      whisper_rs::convert_stereo_to_mono_audio(&whisper_rs::convert_integer_to_float_audio(&data));

    assert!(!data.is_empty());

    // 音声認識
    let scripts = WhisperScripter::new("assets/model/ggml-tiny.bin").parse(&data);
    assert!(!scripts.is_empty());

    // let ans = "人類学者のいう所によれば、原始社会の生産作用も広義において法律的に支配せられているのである。"
    let mut all_script = "".to_owned();
    for script in scripts {
      all_script.push_str(&script.script);
    }
    assert!(
      all_script.contains("人類学者")
        || all_script.contains("原始社会")
        || all_script.contains("支配")
    );
  }
}