CTranslate2で推論をStream出力する on Databricks

今回もDatabricks上で検証しています。

# 導入

以下の記事でCTranslate2変換モデルをlangchainで使う方法を書きました。
今回はCTranslate2でストリーミングする方法です。
ちなみに今回はlangchainを使いません。langchainを使う中でのストリームは次回を予定。

https://qiita.com/isanakamishiro2/items/96956be54444dd3e8cba

# ストリーミングって？

全てのトークン生成が終わってから結果を返すのではなく、
ChatGPTのようにトークン単位で推論結果が次々出力されていくアレ。
そうすることで体感的な待ち時間を減らせる。

ちなみにHuggingface transformersでのStream実装については、[こちら](https://huggingface.co/docs/transformers/main/en/internal/generation_utils#transformers.TextStreamer)などに情報があるようです。（まだ開発中？）

# CTranslate2でストリームを実装

以下の内容が参考になります。

https://opennmt.net/CTranslate2/generation.html?highlight=stream#token-streaming

二つほどオプションがあります。

## Option1. Callbackを使う

ctranslate2.Generatorのメソッド```generate_batch```は引数としてコールバック関数を渡すことが出来ます。
トークンの生成ごとにコールバックが走るため、ストリーミング処理を簡単に実装することができます。

サンプルコードは以下のようになります。

```python
# generatorとtokenizerは既にインスタンスができてる前提

from ctranslate2 import GenerationStepResult

# プロンプトの準備
prompt = "ユーザー: まどか☆マギカでは誰が一番かわいい？\nシステム: "

# コールバック関数
# 引数にGenerationStepResultをとり、bool値を返す関数となっている必要がある
def _callback(step_result:GenerationStepResult) -> bool:
    # step_result引数に生成したトークンのIDが含まれるので、それを使ってデコード
    word = tokenizer.decode([step_result.token_id])

    # step_resultのtokenが▁から始まる場合、スペースを付与する
    if step_result.token.startswith("▁"):
        word = " " + word

    # 表示
    print(word, end="", flush=True)    
    # 次のトークン生成に移る場合、Falseを返す。Trueを返すとそこで推論処理が終了する
    return False

# いつもの推論実行
tokens = tokenizer.convert_ids_to_tokens(
    tokenizer.encode(prompt, add_special_tokens=False)
)
results = generator.generate_batch(
    [tokens],
    asynchronous=False, # Trueにすると非同期実行できる
    max_length=256,
    sampling_topk=10,
    sampling_temperature=0.1,
    repetition_penalty=1.1,
    include_prompt_in_result=False,
    return_scores=True,
    callback=_callback, # コールバック関数を指定
)

# 最後に全出力
print(tokenizer.decode(results[0].sequences_ids[0]), end="", flush=True)

```

```_callback```の中を変更すれば、生成トークン単位に様々な処理を実行できます。
また、この場合```include_prompt_in_result=False```は必須の指定となります。

## Option2. generate_tokensメソッドを使う

ctranslate2.Generatorには```generate_tokens```というメソッドがあり、トークン生成単位でイテレーションを回すことができます。


サンプルコードは以下のようになります。

```python
# generatorとtokenizerは既にインスタンスができてる前提
# プロンプトの準備
prompt = "ユーザー: まどか☆マギカでは誰が一番かわいい？\nシステム: "

# 推論の実行
tokens = tokenizer.convert_ids_to_tokens(
    tokenizer.encode(prompt, add_special_tokens=False)
)

# step_resultsは非同期のイテレーションオブジェクト
step_results = generator.generate_tokens(
    tokens,
    max_length=256,
    sampling_topk=10,
    sampling_temperature=0.1,
    repetition_penalty=1.1,
    return_log_prob=True,
)

# 全ての生成トークンの記録用
output_ids = []

# イテレーションを回して処理
for step_result in step_results:
    output_ids.append(step_result.token_id)

    word = tokenizer.decode([step_result.token_id])
    # step_resultのtokenが▁から始まる場合、スペースを付与する
    if step_result.token.startswith("▁"):
        word = " " + word

    print(word, end="", flush=True)

print(tokenizer.decode(output_ids), end="", flush=True)

```

どちらの方法でも同様の結果を得ることができます。
(Option1のコールバックもgenerate_tokensメソッドを内部的に使っているようです）

## 注意点：

上のサンプル実装だと、ストリーム表示中に文字化けした文字が表示されるケースがあります。
日本語特化モデルは大丈夫だと思いますが、Vicunaなどの多言語モデルを使うとだいたい起きます。
これはトークナイザにSentencepieceを使う場合などで、１トークンが日本語の一文字ではなく、Unicodeの1バイト分を生成するようなケースのときに、トークン単位では正常な文字を生成できないためです。

この場合、有効な文字を生成できるまでトークンをバッファして、まとめてデコードするなどの工夫が必要になります。
この対応は次回やってみます。

# まとめ

CTranslate2だと、ストリーミング処理も比較的簡単にできるのではないかと思います。
特に生成速度が遅いLLMを対話型チャットボットに使う場合など、こういった工夫は大事だと思います。

次回はlangchainを使うケースでのストリーミングをやってみます。