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

今回もDatabricks上で検証していますが、そうでなくても動くと思います。


# 導入

以下の記事では、CTranslate2変換モデルでストリーミングする方法を書きました。

https://qiita.com/isanakamishiro2/items/a659aaf6971f85c05b55

今回はさらにlangchain上でもストリーミングする方法です。

# langchainでストリーミング

langchainでは、CTranslate2のgenerate_batch同様、トークン生成後のコールバックを受け取る仕組みがあります。

https://python.langchain.com/docs/modules/model_io/models/chat/streaming

単純に言えば、このコールバックの仕組に乗っかることで、langchainのストリーミングを利用できます。

# langchainのストリーミング対応カスタムLLM作成

ちょっと長いですが、以下サンプルコード。

```python
from typing import Any, List, Union, Mapping, Optional, Iterable

import ctranslate2
from ctranslate2 import GenerationStepResult
import transformers
from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast

from langchain.llms.base import LLM
from langchain.callbacks.manager import CallbackManagerForLLMRun


class CTranslate2StreamLLM(LLM):

    generator: Generator
    tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast]

    max_length: int = 128
    repetition_penalty: float = 1.1
    temperature: float = 0.1
    topk: int = 1

    @property
    def _llm_type(self) -> str:
        return "CTranslate2"

    def _generate_tokens(
        self,
        prompt: str,
    ) -> Iterable[GenerationStepResult]:
        
        # 推論の実行
        tokens = self.tokenizer.convert_ids_to_tokens(
            self.tokenizer.encode(prompt, add_special_tokens=False)
        )

        step_results = self.generator.generate_tokens(
            tokens,
            max_length=self.max_length,
            sampling_topk=self.topk,
            sampling_temperature=self.temperature,
            repetition_penalty=self.repetition_penalty,
            return_log_prob=True,
        )

        return step_results

    def _decode_with_buffer(
        self, step_result: GenerationStepResult, token_buffer: list
    ) -> Union[str, None]:
        token_buffer.append(step_result.token_id)
        word = self.tokenizer.decode(token_buffer)

        # 全て変換不能文字の場合、終了
        if all(c == "�" for c in word):
            return None

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

        # 正常な文字が生成できた場合、バッファをクリア
        token_buffer.clear()

        return word

    def _call(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
    ) -> str:
        if stop is not None:
            raise ValueError("stop kwargs are not permitted.")

        step_results = self._generate_tokens(prompt)

        output_ids = []
        token_buffer = []

        for step_result in step_results:

            output_ids.append(step_result.token_id)

            if run_manager:
                if word := self._decode_with_buffer(step_result, token_buffer):

                    run_manager.on_llm_new_token(
                        word,
                        verbose=self.verbose,
                        logprobs=step_result.log_prob if step_result.log_prob else None,
                    )

        if output_ids:
            text = self.tokenizer.decode(output_ids)
            return text

        return ""

    async def _acall(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
    ) -> str:
        if stop is not None:
            raise ValueError("stop kwargs are not permitted.")

        step_results = self._generate_tokens(prompt)

        output_ids = []
        token_buffer = []

        for step_result in step_results:

            output_ids.append(step_result.token_id)

            if run_manager:
                if word := self._decode_with_buffer(step_result, token_buffer):

                    await run_manager.on_llm_new_token(
                        word,
                        verbose=self.verbose,
                        logprobs=step_result.log_prob if step_result.log_prob else None,
                    )

        if output_ids:
            text = self.tokenizer.decode(output_ids)
            return text

        return ""

    @property
    def _identifying_params(self) -> Mapping[str, Any]:
        """Get the identifying parameters."""
        return {
            "generator": self.generator,
            "tokenizer": self.tokenizer,
            "max_length": self.max_length,
            "repetition_penalty": self.repetition_penalty,
            "temperature": self.temperature,
            "topk": self.topk,
        }
```

CTranslate2を使ったシンプルなカスタムモデルについては、前の記事を参照ください。

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

ポイントは、以下です。

- トークン生成後に```_call```内でrun_managerのon_llm_new_tokenを呼び出すことで、langchain標準のコールバックを実行しています。
- ```_call```だけではなく、```_acall```も実装しました。
    - ```_acall```はagenerateなどの非同期メソッドで呼ばれる処理です。
    - ```_acall```を実装しない場合、代わりに```_call```が呼ばれますが、今回の処理では都合が悪いため```_acall```も実装しています。
    - ```_call```と```_acall```の違いは、run_manager.on_llm_new_tokenの呼び出しにawaitを付けているかどうかだけです。もうちょっと綺麗に書けると思いますが、今回は割愛。
- トークン単位のデコード処理は```_decode_with_buffer```にまとめています。
    - 前回記事での注意点として挙げた、「複数トークンが揃わないと文字化けする」問題に対しても雑ですが対応しています。


# langchainで推論

やってみます。
関連モジュールのインストールやgenerator/tokenizerの作成は以前の記事と同様です。

今回は[rinna/japanese-gpt-neox-3.6b-instruction-ppo](https://huggingface.co/rinna/japanese-gpt-neox-3.6b-instruction-ppo)をCTranslate2で変換したモデルを使ってみました。

```python
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler

llm = CTranslate2StreamLLM(
    generator=generator,
    tokenizer=tokenizer,
    callbacks=[StreamingStdOutCallbackHandler()],
)

llm("ユーザー: 大規模言語モデルとは？<NL>システム: ")
```
```:結果
大規模言語モデルとは、大規模なコーパスと語彙の統計情報を使用して学習や推論を実行する高度なアルゴリズムです。大規模言語モデルは、様々な種類の言語学習や教育に使用されています。

```

実際には、ストリームで標準出力に出力されます。


```_acall```も実装したので、非同期生成(agenerate)も動きます。

```python
await llm.agenerate(["ユーザー: まどか☆マギカでは誰が一番かわいい？<NL>システム: "])
```
```:結果
私はキュゥべえが一番かわいいと思います。彼はとてもかわいく、彼の言葉や行動はとても魅力的です。また、彼が見せる優しさや共感にとても惹かれます。そして、彼が作る世界もとても良いです。彼と関わって、一緒に過ごすことで、とても楽しい時間を過ごすことができます。

```

こちらも実際にはストリームで出力されます。

# (参考) MLflowへの保管はできるのか？

Databricksと言えばMLflow。
MLflowはlangchainのトラッキングもサポートしているので、上のクラスの保管もできそうに見えるのですが、残念ながらカスタムクラスは対象外です。
※ 最新版の2.6.0では試していないので、シレっとできるようになっているかもしれませんが。。。
CTranslate2モデルのMLflow管理については、また別の記事でやりたいと思います。

# まとめ

というわけでlangchainのストリーミング実装でした。
自己流のため、おかしなところが正直ある気がしていますので、ツッコミお待ちしております。

さて、実際にLLMを本番環境として使う際はREST対応などのAPIサーバを立てることが多いと思いますが、次回はストリーミング対応のREST APIサーバの実装をしたいと思います。

※　夏休みの間に書けるかしら。。。