DatabricksのAI提案コメントを日本語化してみる（ネタ）

# 導入

以下の記事にあるように、DatabricksのUnity Catalogで、テーブルに対するAI提案コメントがサポートされました。

https://qiita.com/taka_yayoi/items/d001e0b1b48a336fa65f

待ち望んでいた機能なので喜び勇んで試しているのですが、記事内にもあるように現時点では英語で作成されます。

> 正直なところ、言語設定に合わせてコメントの言語も切り替えてほしいですよね…。是非、フィードバックお願いします！

というわけで私もフィードバックは送ったのですが、実装されるまで待ってられない！
というわけで、きっと誰もが考えたであろう「生成されたAIコメントを日本語に翻訳して登録し直す」をやってみました。

かなりお遊び＆実験成分多めです。

:::note info
今回、翻訳処理にローカルLLMを使っていますが、はっきり言って過剰手段です。
Google翻訳のAPIを叩いたりする方がいいと思います。
実務で万が一マネしようとする場合、気を付けましょう。
:::

# Step1. テーブルを準備

まず、適当にコメントを入れるためのテーブルを作ります。

今回は、標準で提供されているサンプルテーブルである```samples.nyctaxi.trips```を、一旦適当なカタログ/スキーマである```training.examples```にコピーして使います。

```sql
%sql

CREATE TABLE IF NOT EXISTS training.examples.trips DEEP CLONE samples.nyctaxi.trips;
```

その後、AI提案コメント機能を使って、このテーブルのコメントおよび列コメントを生成します。

以下のようになりました。全部英語ですね。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/898617ac-b845-dcc1-6fd5-0d7edb2ccf1c.png)

# Step2. 翻訳用のローカルLLMを準備

翻訳するためにLLMを準備します。
上にも書きましたが、過剰手段だと理解しつつ、せっかくなのでLLMを使って翻訳させます。
あと、OpenAI等のプロプライエタリなAPIを使ったりするとお金もかかるので、ローカルLLMを使います。Databricksのクラスタ利用料はお布施ですので払いましょう。
All you need is LLM.

というわけで必要なモジュールをインストール。


```python
%pip install -U autoawq transformers accelerate langchain

dbutils.library.restartPython()
```

[以前の記事](https://qiita.com/isanakamishiro2/items/acbad2dc0f3bc8176bbd)で作成したAutoAWQをlangchainのチャットモデルとして利用する関数やクラスを再利用してモデルを読み込みます。

<details><summary>全コード(autoawq_chat.py)</summary>

```python
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,
        }
```
</details>



また、今回は以下のモデルを利用しました。

https://huggingface.co/TheBloke/Xwin-LM-13B-v0.2-AWQ

```python
from autoawq_chat import load_pretrained_model

# モデルは事前にUnity Catalogのボリュームにダウンロード済みとする
model_path = "/Volumes/モデルの保管場所/models--TheBloke--Xwin-LM-13B-v0.2-AWQ"

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

langchainを使ってプロンプトテンプレートやLLMChainを作成します。

```python
from autoawq_chat import ChatAutoAWQ
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
from langchain.prompts.chat import (
    ChatPromptTemplate,
    SystemMessagePromptTemplate,
    AIMessagePromptTemplate,
    HumanMessagePromptTemplate,
    MessagesPlaceholder,
)

# チャット履歴(history)を含むようにプロンプトテンプレートを作成
chat_prompt = ChatPromptTemplate.from_messages(
    [
        SystemMessagePromptTemplate.from_template("ユーザからの英文を日本語に翻訳してください。翻訳結果だけ出力してください。"),
        HumanMessagePromptTemplate.from_template("{comment}"),
        AIMessagePromptTemplate.from_template(""),
    ]
)

llm = ChatAutoAWQ(
    generator=model,
    tokenizer=tokenizer,
    system_message_template="{}",
    human_message_template="USER: {}",
    ai_message_template="ASSISTANT: {}",
    max_new_tokens=1024,
)

# Chain作成
llm_chain = LLMChain(
    llm=llm,
    prompt=chat_prompt,
    verbose=False,
)
```

これで準備完了です。

# Step3. テーブル情報の取得

AIコメントが付与されたテーブルの情報を取得します。

```python
# 件数少ないので、全てpandasのデータフレーム化する

table = "training.examples.trips"

# テーブル自体の情報取得
sql = f"DESCRIBE DETAIL {table}"
table_detail_pdf = spark.sql(sql).toPandas()

# テーブルの各カラムの情報取得
sql = f"DESCRIBE {table}"
col_detail_pdf = spark.sql(sql).toPandas()
col_detail_pdf["table"] = table
```

# Step4. 翻訳する

取得したテーブル情報を使って翻訳します。

まずは翻訳用関数を定義。

```python
# 翻訳用関数
def translate(comment:str) -> str:
    return llm_chain.predict(comment=comment)
```

テーブルコメントの翻訳

```python
table_detail_pdf["description_jp"] = table_detail_pdf["description"].apply(translate)

display(table_detail_pdf[["description","description_jp"]])
```

結果：
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/adb5e0b0-e7b9-ef06-1dd0-7fb6eaca3217.png)

それっぽく翻訳されていますね。

次に各カラムのコメントを翻訳します。

```python
col_detail_pdf["comment_jp"] = col_detail_pdf["comment"].apply(translate)

# 。は後処理で除くことにする
col_detail_pdf["comment_jp"] = col_detail_pdf["comment_jp"].apply(lambda x: x.replace("。", ""))

display(col_detail_pdf[["comment","comment_jp"]])
```

結果：
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/5dfe2cbb-8308-076a-e668-a232ceeac9dc.png)

表現の統一感がマチマチですが、こちらも翻訳されました。
「距離（マイル）」など、LLMならではの翻訳表現は結構好きです。

# Step5. テーブル情報の更新

翻訳結果を生成できたので、テーブル情報に書き戻します。

まずはテーブルコメント。

```python

def update_table_comment(comment:str):
    table = "training.examples.trips"
    sql = f"COMMENT ON TABLE {table} IS {{comment}}"
    spark.sql(sql, comment=comment)

table_detail_pdf["description_jp"].apply(update_table_comment)
```

次にカラムコメント。

```python
def update_table_column_comment(col:str, comment:str):
    table = "training.examples.trips"
    sql = f"ALTER TABLE {table} ALTER COLUMN {col} COMMENT {{comment}}"

    spark.sql(sql, comment=comment)

col_detail_pdf.apply(lambda x: update_table_column_comment(x["col_name"], x["comment_jp"]), axis=1)
```

問題なければこれで更新されます。

更新結果をカタログエクスプローラ上で確認してみましょう。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/f7203271-fd7e-ae60-d65d-69ad26b03975.png)

日本語コメントに置き換わっていますね！

# まとめ

英語で生成されたAI提案コメントを日本語化して書き戻してみました。

とはいえ、AI提案コメントの多言語対応はすぐにDatabricks側に実装されそうなので、正直完全にネタだと思います。
私のように待ちきれない人は試してみてください。

しかし、なんでもかんでもLLMは良くないと思いつつ、ついつい便利なのでやっちゃうなあ。。。いい時代になりました。