🌐

PLaMo翻訳のモデルで動画の字幕を翻訳する

に公開9

前日、Preferred Networks (PFN) が開発している LLM を使った PLaMo翻訳 がリリースされました。

PLaMo翻訳は LLM を使った翻訳なので、従来のニューラル機械翻訳に比べて、文脈から 空気を読んで いい感じに翻訳してくれるという特徴があります。X で見かけた以下のポストでは、構造を持ったテキストをうまく元の構造を維持したまま翻訳してくれているのが確認できます。

https://x.com/ibushi_maru/status/1927198535620894976

PLaMo 翻訳はデモサイトで簡単に試せるほか、モデル自体も PLaMo Community Licensehuggingface 上で公開されており、 8B モデルが動く GPU があればローカルで自分で動かすこともできるようになっています。

この記事では「huggingface 上で公開されているモデルを利用してローカルで実行する形で」「構造をもった大きいテキストを翻訳する」例として、動画の字幕 (.srt ファイル) の翻訳の実装を紹介します。今回 .srt を取り上げますが、format に関してはできるだけ LLM に丸投げする形で実装してみたので、他の format でも広く使えるような実装になっているのではないかと思います。

準備

必要な package の install

PLaMo 2 は vllm に実装がマージされているため vllm が入っていれば推論を実行できます。

pip install "vllm>=0.8.5"

vllm による推論は transformers による推論に比べて大幅に高速なため、追加で学習を行うような場合以外は vllm の利用をおすすめします。

.srt ファイルの用意

翻訳元の字幕が存在せず、動画の文字起こしから始める場合は OpenAI が公開している Whisper が便利です。

https://github.com/openai/whisper

これを使う場合は、ffmpeg なりを使って動画から音声を抽出し、whisper で書き起こして .srt ファイルを作成する、という流れになります。

# ffmpeg を使って音声を mp3 として抽出する
ffmpeg -i input_video.mp4 -vn -acodec libmp3lame -q:a 2 audio.mp3

# whisper を使って書き起こしをする
pip install -U openai-whisper
# この例では large-v3 という大きい強力なモデルを利用
whisper audio.mp3 --model large-v3 --language en

実装

適当な1文を翻訳させてみる

まずは vllm を import して、vllm.LLM の instance を作成します。

import vllm

llm = vllm.LLM(model="pfnet/plamo-2-translate", trust_remote_code=True, max_model_len=16 * 1024, max_num_seqs=16)

このとき、max_model_len は扱う最大 token 数、max_num_seqs は1バッチで並列に処理する request の数を制御します。両方とも、増やすと要求される GPU メモリが大きくなるパラメータで、今回は GeForce RTX 4090 (24GB VRAM) に載せられる最大の max_model_len を指定して、OOM しないように max_num_seqs を絞っています。

次に、公式のサンプルにある例を実行してみます

prompt = r'''<|plamo:op|>dataset
translation
<|plamo:op|>input lang=English
Write the text to be translated here.
<|plamo:op|>output lang=Japanese
'''

responses = llm.generate([prompt], sampling_params=vllm.SamplingParams(temperature=0, max_tokens=1024, stop=["<|plamo:op|>"]))
print(responses[0].outputs[0].text)  

prompt の <|plamo:op|>input lang=English\n の後の Write the text to be translated here. が翻訳対象の文字列になります。上記のコードを実行すると翻訳結果として ここに翻訳するテキストを入力してください。 と print されます。

ファイル全体を翻訳する

LLM による翻訳では、ある程度長い文章をまとめて翻訳させることで前後の文脈も加味した翻訳が可能になりますが、ファイルが大きくなると、GPU メモリやモデルの context-length の制約により、全文を一度に入力するのは難しくなります。

今回は、以下のようなアプローチで前後の文脈を加味しつつ入力を分割することにします。

  • 入力ファイルを行ごとに分割
  • lines_per_window = 64 行をまとめて入力
  • 入力のうち最初の数行 (lines_per_window - stride = 16 行) が前回の翻訳と被るようにして、前回の翻訳結果とうまくつながる翻訳を生成させるようにする

また、生成結果が必ず元の format に従うとは限らないため、チェックして、うまく言っていない場合は再生成を行ってあげるようにします。具体的なチェック方法は翻訳したい対象によって様々な実装が考えられますが、今回はできるだけ汎用な方法で行いたいため、生成結果の行数のチェックだけ行うことにします。

具体的な実装は以下のとおりです。

https://github.com/zaburo-ch/plamo-2-translate-srt/blob/main/main.py#L2-L84

.srt ファイルは 4行ごとにまとまりのあるデータになっているので、被せる行数を 4 の倍数にすることで、変な分割にならないようにしています。逆に、それ以外の .srt 固有の工夫はしておらず.srt format 固有の、番号や timestamp などの情報を翻訳後も保持するというところは LLM にまる投げしています。そのため、.srt 以外のフォーマットでも利用可能なアプローチになっているのではないかと思います。

結果のサンプル

https://www.youtube.com/watch?v=_rfUsZ4QjKE

音声は、Gemini に「翻訳モデルのデモとして、英語のプレゼンの書き起こしを翻訳する、ということを行いたいと思っています。そのために適当な英語のプレゼンの原稿が必要です。適当な原稿を考えていただけないでしょうか?」という prompt で作成してもらった原稿を、GCP の Text-to-Speech AI で音声に変換したものです。これを上記の手順で、whisper で英語 .srt に変換、plamo-2-translate で日本語 .srt に変換して、字幕付きの動画を作成しました。

英語 .srt
https://github.com/zaburo-ch/plamo-2-translate-srt/blob/main/sample_src.srt

翻訳後の日本語 .srt
https://github.com/zaburo-ch/plamo-2-translate-srt/blob/main/sample_translated.srt

一部、英語と日本語の語順の関係で、字幕としては不自然な翻訳になっていたりする部分もありますが、全体的に機械翻訳としては十分満足できるレベルのものができているのではないかと思います。

Discussion

abeqabeq

タイムリーで有益な記事の公開をありがとうございます。自分の手持ちの専用メモリ6GBのノートブック用3060ではmax_model_lenやmax_num_seqsを小さくしても動作することができませんでしたが正しいインストール方法がわかってよかったです。

ShilikShilik

mlx_lmバージョンのmain.pyで試してみてください。

import fire
from mlx_lm import load, generate
from tqdm import tqdm
import random

def main(
    src_file_path: str,
    dst_file_path: str,
    lines_per_window: int = 8,
    stride: int = 6,
    max_tokens: int = 4096,  
) -> None:
    with open(src_file_path) as fp:
        src_lines = fp.read().splitlines()

    # load MLX model
    model, tokenizer = load(
        "mlx-community/plamo-2-translate",
        tokenizer_config={"trust_remote_code": True}
    )
    tokenizer.add_eos_token("<|plamo:op|>")

    template_en_to_ja = r"""<|plamo:op|>dataset
translation
<|plamo:op|>input lang=English
{english}
<|plamo:op|>output lang=Japanese
"""

    translated_lines = []

    for i in tqdm(range(0, len(src_lines), stride), desc="[Translation Progress]"):
        # 空行を翻訳しなくていいように足してあげる
        while len(translated_lines) < len(src_lines) and src_lines[len(translated_lines)] == "":
            translated_lines.append("")

        if len(translated_lines) >= min(i + lines_per_window, len(src_lines)):
            continue

        src_text = "\n".join(src_lines[i : i + lines_per_window])
        prompt = template_en_to_ja.format(english=src_text)

        # 今注目している範囲で、すでに翻訳されているものがあれば context として加えてそこから先を翻訳する
        if len(translated_lines[i : i + lines_per_window]) > 0:
            context = "\n".join(translated_lines[i : i + lines_per_window]) + "\n"
        else:
            context = ""
        prompt += context

        # 行数など簡単にチェックできる範囲でおかしい結果のときは retry する。最大5回やってだめなら raise する
        for trial in range(5):
            # MLX-LM では複数候補を一度に生成できないため、必要に応じて複数回呼び出す
            candidates = []
            num_candidates = 1 if trial == 0 else 8

            for _ in range(num_candidates):
                # 温度設定とシード設定
                #temp = 0.0 if trial == 0 else 0.7

                # MLX-LM での生成
                response = generate(
                    model,
                    tokenizer,
                    prompt=prompt,
                    max_tokens=max_tokens // 2,
                    verbose=False,
                )

                # レスポンスから翻訳結果を抽出(<|plamo:op|> で停止)
                result = response.split("<|plamo:op|>")[0].rstrip()
                candidates.append(result)

            # outputs の中から最も良いものを選ぶ
            # (簡単のため、ここでは末尾以外の行数が揃っていれば採用とする。タスクに合わせていい感じの評価を実装する)
            best_result: str | None = None
            for result in candidates:
                if (context + result).count("\n") == src_text.rstrip().count("\n"):
                    best_result = result
                    break

            if best_result is not None:
                translated_lines.extend(best_result.split("\n"))

                # 空行を翻訳しなくていいように足してあげる
                while len(translated_lines) < len(src_lines) and src_lines[len(translated_lines)] == "":
                    translated_lines.append("")

                break
        else:
            raise RuntimeError("Translation failed after 5 attempts")

    with open(dst_file_path, "w") as fp:
        fp.write("\n".join(translated_lines))


if __name__ == "__main__":
    fire.Fire(main)
abeqabeq

ご回答、ご教示をありがとうございます。
mlx-lmはApple Silicon Macに対してのことですね、こちらはwsl ubuntuでやっていました。
失礼しました。
<del>
pip install mlx-lm
としたら
libmlx.so: cannot open shared object file
と表示されたので
pip install "mlx[cpu]"
としたらメッセージは消えて、処理が始まりましたが
corrupted double-linked list
Aborted (core dumped)
と表示され処理が中断されたので、max_tokens: int = 4096を2048にしてみましたが
Segmentation fault (core dumped)
と表示され処理が中断されました。
何か思い当たりがありましたらまたアドバイスをいただけるととてもうれしいです。
</del>

ShilikShilik

長文の場合は、以下のようなコマンドで-c を大きくしてみたらいかがでしょうか。

build/bin/llama-cli -m 'plamo-2-translate.gguf'  -n 128 -c 2048 -ngl 32 -b 32 -p  '<your_prompt>' -no-cnv

追記:
長文を渡すと誤動作するというエラーはubuntu環境のみ見られました。MAC Mini M4では問題なく訳文を生成したことを確認しました。

abeqabeq

Shilik様 ご教示をありがとうございます。
パラメータ -no-cnv -n 128 -b 32 -c 2048 で試したところ正しい結果は得られませんでした。

dataset
translation
input lang=English
This PR supports PLaMo2 model in llama.cpp, which was also requested on a related discussion thread: #13874. This model uses a custom-implemented tokenizer, so this PR includes both the model itself (which uses an architecture combining Mamba and Attention, similar to Jamba) as well as implementing the new custom tokenizer.
output lang=Japanese
BBBoooooooooooooooooooooooooooooooooooooooooooooo(以下oが続く)

パラメタ-ngl 32を追加するとHello, how are you?でも誤動作しました。

AMD Ryzen9 5900HX with Radeon Graphics+Windows11にVulkanをインストールしたPC上でllama-b5899-bin-win-vulkan-x64.zipを解凍したプログラムを使っています。
以上、報告します。報告の内容に不足があればご指摘いただけますようお願いします。

abeqabeq

Shilik様 ご教示をありがとうございます。
自分の使えるPCでRTX3060 Laptop GPU搭載のPCがあったので、llama-b6018-bin-win-cuda-12.4-x64.zipで試したところ、前述程度の英文は問題なく翻訳できました。なのでvulkan版に起因する問題だったかもしれません。次はそのPC上でCPU版を試してみます

abeqabeq

GPUを搭載しないPCでllama-b6027-bin-win-cpu-x64.zipを試したところ長文の英語でも問題なく、かなり速く結果を出しました。llama-b6026-bin-win-vulkan-x64.zipを試してみたら同様症状だったので今までの問題はvulkan版の問題だと思いました。何かのLLMを試したときvulkan版の方がcpu版よりは速いみたいなことがあったのでvulkanを使っていましたが、plamoの場合はcpu版でも速いのでcpu版を使おうと思います。