CTranslate2で変換したモデルをlangchainで使う on Databricks

今回もDatabricks上で検証していますが、Databricksでなくても動くと思います。タグ詐欺ではない。

# 導入

前回、CTranslate2でモデルを変換・高速推論を実施してみました。Vicuna v1.5の13Bという中型サイズのLLMもVRAM 13Gぐらいで動かすことができます。

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

しかし、LLMでアプリ開発といったらlangchainを皆さん使いたくなるでしょう。（偏見）
というわけで、今回はlangchainでCTranslate2変換済みモデルを使ってみます。

# langchainのLLM Integration

langchainには多くのLLMをすぐに使えるようにインテグレーションがなされています。

https://python.langchain.com/docs/integrations/llms/

当然、CTranslate2モデルも・・・あれ、ない？

というわけで、2023/8現在、直接的なIntegrationはありません（たぶん）。
間接的にはやり方ありそうですが。

とはいえ、langchainはカスタムLLMを作成できます。
今回はシンプルなカスタムLLMを作成し、LLM Chainなどを動かしてみます。

# CTranslate2用のlangchain カスタムLLMを作成

langchainのカスタムLLM作成は以下のDocに説明があります。
agenerateなどの非同期メソッドを実装するためには少し工夫が必要ですが、とりあえず動くものを作る場合はそんなに難しくありません。

https://python.langchain.com/docs/modules/model_io/models/llms/custom_llm

というわけで、以下のようなクラスを作成します。

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

import ctranslate2
import transformers
from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
from langchain.callbacks.manager import CallbackManagerForLLMRun
from langchain.llms.base import LLM

class CTranslate2SimpleLLM(LLM):

    generator: ctranslate2.Generator
    tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast]

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

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

    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.")

        tokens = self.tokenizer.convert_ids_to_tokens(
            self.tokenizer.encode(prompt, add_special_tokens=False)
        )

        results = self.generator.generate_batch(
            [tokens],
            max_length=self.max_length,
            sampling_topk=self.topk,
            sampling_temperature=self.temperature,
            repetition_penalty=self.repetition_penalty,
            include_prompt_in_result=False,
            return_scores=True,
        )
        text = self.tokenizer.decode(results[0].sequences_ids[0])

        return text

    @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,
        }
```

実質的にやっているのは、```_call```内でいつもの推論処理を実行する処理を書いているだけです。

あとは前回と同様にgeneratorとtokenizerを作成し、このクラスに渡すだけでlangchain対応のLLMオブジェクトができあがります。

# langchainで推論

やってみましょう。

必要なモジュールをインストール。

```python
%pip install -U -qq ctranslate2 transformers accelerate langchain
dbutils.library.restartPython()
```

いつものようにgeneratorとtokenizerをロード。
前回作成したCTranslate2変換済みのVicuna v1.5モデルを使うことにします。

```python
import ctranslate2
import transformers
import torch

model_path = "モデルのパス"

# ジェネレーターとトークナイザーの準備
device = "cuda" if torch.cuda.is_available() else "cpu"
generator = ctranslate2.Generator(model_path, device=device)
tokenizer = transformers.AutoTokenizer.from_pretrained(model_path, use_fast=False)
```

上記のカスタムLLMクラス定義を読み込み（省略）、その後インスタンス化。

```python
llm = CTranslate2SimpleLLM(
    generator=generator,
    tokenizer=tokenizer,
    temperature=0.9,
)
```

まずは単純に推論実行します。

```python
llm("こんにちは！")
```

```:結果
'こんにちは！」と繰り返し、その顔色や声を見る限り辛かったらしい。だが、「お茶を作るのが好きよ」と努力して答える。\nあまり長くなりそうになってきたため、カッターで切る形で収束させます。\n「どうしても忙しいの？今日は何か面白'
```

内容はともかく、ちゃんと動いていそうです。

では、LLMChainを使って、プロンプトテンプレートを基に指示を実行させます。

```python
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain

template = "User: {instruction}\nASSISTANT: "
prompt = PromptTemplate(
    input_variables=["instruction"],
    template=template,
)

chain = LLMChain(llm=llm, prompt=prompt)
print(chain.run("大規模言語モデルとは何ですか？"))
```

```:結果
大規模言語モデル（Broad-scale Language Model, BLM）とは、自然言語を扱う人工知能の一つであり、文章の意味や文法などを学習して、自動的に新しい文章を生成する技術のことです。大規模言語モデルは、大量のテキストデータを使用して稼働し、機械学��
```

微妙な内容ですが、LLMChainで処理されました。

# まとめ

今回はかなり単純な実装でCTranslate2変換後モデルをlangchainで使えるようにしてみました。
※ エラー処理などを入れていないため、使う場合はきちんと実装することをお勧めします。

次回のCTranslate2ネタはStream出力をやりたいと思います。