導入

以前、AutoAWQを使ってAWQフォーマットのモデルを読み込んだりしてみました。
langchainで使えた方がいろいろ便利なのですが、AutoAWQは見たところlangchainのLLM/ChatModel対応がされていないようなので、カスタムChat Modelクラスを作って見ます。

なお、以前CTranslate2対応のカスタムChat Modelクラス作成をした記事はこちら。
基本的にはこれを踏襲します。

実行・検証はDatabricks上で行っています。

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

AutoAWQやlangchainなど、必要なモジュールをインストール。

%pip install -U autoawq transformers accelerate langchain

dbutils.library.restartPython()

カスタムLLMモデルを作る

最低限の実装はざっくりと以下のようなコードになります。
エラー処理等、かなり省略しているので注意。一応、ストリーミング出力対応です。

from awq import AutoAWQForCausalLM
from awq.models.base import BaseAWQForCausalLM
from transformers import (
    AutoTokenizer,
    PreTrainedTokenizer,
    PreTrainedTokenizerFast,
    TextIteratorStreamer,
)
from threading import Thread
import asyncio

from typing import (
    Any,
    List,
    Union,
    Mapping,
    Tuple,
    Optional,
)
from langchain.callbacks.manager import (
    AsyncCallbackManagerForLLMRun,
    CallbackManagerForLLMRun,
)
from langchain.chat_models.base import BaseChatModel
from langchain.pydantic_v1 import Field, root_validator
from langchain.schema.messages import (
    AIMessage,
    AIMessageChunk,
    BaseMessage,
    ChatMessage,
    HumanMessage,
    SystemMessage,
)
from langchain.schema import (
    ChatGeneration,
    ChatResult,
)


def load_pretrained_model(
    model_path: str,
    max_new_tokens: int = None,
    fuse_layers: bool = True,
    safetensors: bool = True,
    use_fast: bool = False,
    trust_remote_code: bool = False,
) -> Tuple[BaseAWQForCausalLM, Union[PreTrainedTokenizer, PreTrainedTokenizerFast]]:
    """事前学習済みモデルとTokenizerを返す"""
    # Load model
    model = AutoAWQForCausalLM.from_quantized(
        model_path,
        max_new_tokens=max_new_tokens,
        fuse_layers=fuse_layers,
        trust_remote_code=trust_remote_code,
        safetensors=safetensors,
    )
    tokenizer = AutoTokenizer.from_pretrained(
        model_path,
        use_faset=use_fast,
        trust_remote_code=trust_remote_code,
    )

    return model, tokenizer


class ChatAutoAWQ(BaseChatModel):
    """AWQ(AutoAWQ) models.

    To use, you should have the ``autoawq`` python package installed.

    Example:
        .. code-block:: python

            generator, tokenizer = load_pretrained_model(model_path)
            hf = ChatAutoAWQ(generator=generator, tokenizer=tokenizer)
    """

    generator: BaseAWQForCausalLM
    """ AutoAWQ Model """
    tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast]
    """ Tokenizer """

    # メッセージテンプレート
    human_message_template: str = "ユーザ: {}"
    ai_message_template: str = "システム: {}"
    system_message_template: str = "{}"

    do_sample: bool = True
    max_new_tokens: int = 64
    repetition_penalty: float = 1.1
    temperature: float = 1
    top_k: int = 1
    top_p: float = 0.95

    prompt_line_separator: str = "\n"

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

    def _format_message_as_text(self, message: BaseMessage) -> str:
        if isinstance(message, ChatMessage):
            message_text = f"{self.prompt_line_separator}{message.role.capitalize()}: {message.content}"
        elif isinstance(message, HumanMessage):
            message_text = self.human_message_template.format(message.content)
        elif isinstance(message, AIMessage):
            message_text = self.ai_message_template.format(message.content)
        elif isinstance(message, SystemMessage):
            message_text = self.system_message_template.format(message.content)
        else:
            raise ValueError(f"Got unknown type {message}")
        return message_text

    def _format_messages_as_text(self, messages: List[BaseMessage]) -> str:
        return self.prompt_line_separator.join(
            [self._format_message_as_text(message) for message in messages]
        )

    def _generate_stream(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> TextIteratorStreamer:

        prompt = self._format_messages_as_text(messages)

        tokens = self.tokenizer(
            prompt,
            # add_special_tokens=True,
            return_tensors="pt",
        ).input_ids.cuda()

        streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=True, skip_special_tokens=True)

        # Generate output
        generation_kwargs = dict(
            inputs=tokens,
            do_sample=self.do_sample,
            temperature=self.temperature,
            top_p=self.top_p,
            top_k=self.top_k,
            max_new_tokens=self.max_new_tokens,
            streamer=streamer,
            num_return_sequences=1,
        )
        thread = Thread(target=self.generator.generate, kwargs=generation_kwargs)
        thread.start()

        return streamer

    def _generate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> ChatResult:

        streamer = self._generate_stream(messages, stop, run_manager, **kwargs)

        text = ""
        count = 0
        
        for new_text in streamer:
            if not new_text:
                continue

            text += new_text
            count += 1

            if run_manager:
                run_manager.on_llm_new_token(
                    new_text,
                    verbose=self.verbose,
                )

        chat_generation = ChatGeneration(message=AIMessage(content=text))
        return ChatResult(
            generations=[chat_generation],
            llm_output={"completion_tokens": count},
        )

    async def _agenerate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> ChatResult:
        streamer = self._generate_stream(messages, stop, run_manager, **kwargs)

        text = ""
        count = 0

        for new_text in streamer:

            if not new_text:
                continue

            text += new_text
            count += 1

            # await asyncio.sleep(0)
            if run_manager:
                await run_manager.on_llm_new_token(
                    new_text,
                    verbose=self.verbose,
                )

        chat_generation = ChatGeneration(message=AIMessage(content=text))
        return ChatResult(
            generations=[chat_generation],
            llm_output={"completion_tokens": count},
        )

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

試してみる

今回はせっかくなので、HuggingfaceのZephyr 7B AlphaのAWQモデルを使ってみます。
Zephyrについては以下の記事を参照してください。

Mistral 7Bベースでファインチューニングしたモデルです。

こちらから事前にモデルをダウンロードしておき、今回作ったクラスで読み込みます。

UC_VOLUME = "/Volumes/モデルのスナップショット保管場所"
model_dir = "/models--TheBloke--zephyr-7B-alpha-AWQ"

model_path = f"{UC_VOLUME}{model_dir}"

from chat_models.autoawq_chat import load_pretrained_model

model, tokenizer = load_pretrained_model(model_path, fuse_layers=False)

from chat_models.autoawq_chat import ChatAutoAWQ
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler


llm = ChatAutoAWQ(
    generator=model,
    tokenizer=tokenizer,
    system_message_template="<|system|>\n{}</s>",
    human_message_template="<|user|>\n{}</s>",
    ai_message_template="<|assistant|>{}",
    max_new_tokens=256,
)

result = llm.predict("東京の明日の天気は？", callbacks=[StreamingStdOutCallbackHandler()])

出力


<|assistant|>
今日は東京の天気は、晴れです。明日の天気は、晴れと少しの雨が予測されています。

動きました。実際にはストリームで出力されます。

もう少し、いろいろ試してみましょう。
Chat Modelなので、それらしく指定します。

from langchain.schema.messages import (
    AIMessage,
    HumanMessage,
) 

output = llm(
    [
        HumanMessage(content="架空の生き物をデザインし、その特徴や能力について日本語で説明してください。"),
        AIMessage(content="\n"),
    ]
)
print(output.content)

出力

「エレクトリック・フロー」という名前の生き物は、電気の力を活用している、幻的な生き物です。

この生き物は、電気の流れを感知し、その流れを利用して動きます。電気の強さに応じて、速度や方向を変えることができます。

エレクトリック・フローは、電気の強さに応じて、縁を張り、電気の強さに応じて、縁を張り、電気の強さに応じて、縁を張り、電気の強さに応じて、縁を張り、電気の強さに応じて、縁を張り、電気の強さに応じて、縁を張り、電気の強さに応じて、縁を張り、電気の強さに応じて、縁を張

なんか繰り返してますね。ペナルティの設定が必要かな。

コードも生成してみます。

output = llm(
    [
        HumanMessage(content="ランダムな10個の要素からなるリストを作成してソートするコードをPythonで書いてください。"),
        AIMessage(content="\n"),
    ]
)
print(output.content)

出力

Here's a Python code to generate a list of 10 random numbers and sort them:

```python
import random

numbers = [random.randint(1, 100) for _ in range(10)]
numbers.sort()

print(numbers)
``

Explanation:

1. We first import the `random` module to generate random numbers.

2. We create an empty list `numbers` and use a list comprehension to generate 10 random numbers between 1 and 100 (inclusive) using the `randint()` function from the `random` module.

3. We then sort the list using the `sort()` method.

4. Finally, we print the sorted list.

できました。生成されるコードの動作確認もしてみましたが、きちんと動きます。

まとめ

AWQモデルをlangchainで使えるようにするためにカスタムChat Modelを作って見ました。

しかし、Mistral 7Bのファインチューニングモデルがワンサカ出てきていて、比較的少ないパラメータで高い性能を発揮しているのが面白いです。
全体的に日本語がイマイチなところがあり、日本語が得意なファインチューニングモデルが出てくることを願ってます。