DatabricksでrinnaのファインチューニングからGUI構築までやってみる

ある意味、GWの宿題的な取り組みとして。

やりたかったのは以下のフローです。

1. rinnaのファインチューニング
1. MLflowによるメトリクスやモデルのトラッキング
1. モデルサービングによるRESTエンドポイントの構築
1. Streamlitからモデルを呼び出し

結論として、上のエンドツーエンドのフローをDatabricksで完結させることができました。嬉しい。

自分の参考資料はこちらです。途中で触れていますが、他の方の記事も参考にさせていただいています。

- [DatabricksのモデルサービングでLLMを用いたチャットボットを動かす](https://qiita.com/taka_yayoi/items/5f1db6447055a79284b8)
- [MLflow 2\.3のHugging Faceトランスフォーマーのサポートを試す](https://qiita.com/taka_yayoi/items/ad370a7f57c4eae58800)
- [\[翻訳\] Hugging Face transformersのクイックスタート](https://qiita.com/taka_yayoi/items/f3c74c8a242368412615)
- [Databricksで日本語GPT\-2モデルをファインチューニングして文章生成をやってみる](https://qiita.com/taka_yayoi/items/17be2b5f7136c4af86e9)
- [Databricksでrinnaの日本語GPT\-2モデルのファインチューニングを試す](https://qiita.com/taka_yayoi/items/89d7b078e80fed931243)
- [MLflow 2\.3のご紹介：ネイティブLLMのサポートと新機能による強化](https://qiita.com/taka_yayoi/items/431fa69430c5c6a5e741)


# ライブラリのインストール

transformersやMLflowは最新バージョンをインストールします。Databricksランタイムも最新の13.0MLです。

```
%pip install transformers==4.25.1
%pip install sentencepiece
%pip install mlflow==2.3.1
```

DBRを13.0にするとライブラリインストール時に`Note: you may need to restart the kernel using dbutils.library.restartPython() to use updated packages.`といったメッセージが出るので念の為kernelを再起動します。

```py:Python
dbutils.library.restartPython()
```

# Tokenizerの設定

(多分)ベースモデルに含まれていないDatabricksの用語などを追加します。ここはまだ試行錯誤の段階です。

```py:Python
from transformers import T5Tokenizer,AutoModelForCausalLM

tokenizer = T5Tokenizer.from_pretrained("rinna/japanese-gpt2-medium")
DATABRICKS_TOKENS = ["MLflow", "Databricks", "Delta Lake", "Spark"]
num_added_toks = tokenizer.add_tokens(DATABRICKS_TOKENS)

model = AutoModelForCausalLM.from_pretrained("rinna/japanese-gpt2-medium")

print('We have added', num_added_toks, 'tokens')
model.resize_token_embeddings(len(tokenizer)) # Notice: resize_token_embeddings expect to receive the full size of the new vocabulary, i.e. the length of the tokenizer.
```

# rinnaのファインチューニング

引き続き、こちらの記事を参考にさせていただいています。

https://zenn.dev/robes/articles/24ee45dd81636a

```py:Python
from transformers import TextDataset, DataCollatorForLanguageModeling
from transformers import Trainer, TrainingArguments, AutoModelWithLMHead

train_data_path = "/dbfs/FileStore/shared_uploads/takaaki.yayoi@databricks.com/dolly/taka_qiita_cleansed.csv"

# データセットの設定
train_dataset = TextDataset(
    tokenizer = tokenizer,
    file_path = train_data_path,
    block_size = 128 # 文章の長さを揃える必要がある
)

# データの入力に関する設定
data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm= False
)

# 訓練に関する設定
training_args = TrainingArguments(
    output_dir="/dbfs/tmp/takaaki.yayoi@databricks.com/rinna/output_20230502/",  # 関連ファイルを保存するパス
    overwrite_output_dir=True,  # ファイルを上書きするかどうか
    num_train_epochs=3,  # エポック数
    per_device_train_batch_size=8,  # バッチサイズ
    logging_steps=100,  # 途中経過を表示する間隔
    save_steps=800  # モデルを保存する間隔
)

# トレーナーの設定
trainer = Trainer(
    model =model,
    args=training_args,
    data_collator = data_collator,
    train_dataset = train_dataset
)
```

以下が今回の肝です。MLflowにメトリクスとモデルをトラッキングしてもらいます。MLflowでtransformersがサポートされたので実装は簡単です。これで、今後のチューニングが捗ります。GPUクラスターの場合はパイプラインで`device=0`の設定が必要です。

https://mlflow.org/docs/latest/python_api/mlflow.transformers.html

```py:Python
import mlflow
from transformers import pipeline

with mlflow.start_run():
  trainer.train()

  task = "text-generation"
  architecture = "rinna/japanese-gpt2-small"

  # fine tune後のモデルとトークナイザー、GPUを指定
  sentence_pipeline = pipeline(
    task=task, tokenizer=tokenizer, model=model, device=0
  )

  # Validate that the overrides function
  prompts = ["生成型モデルは", "今日は天気がいいので"]

  # validation of config prior to save or log
  inference_config = {
    "top_k": 2,
    "num_beams": 5,
    "max_length": 30,
    "temperature": 0.62,
    "top_p": 0.85,
    "repetition_penalty": 1.15,
  }

  # Verify that no exceptions are thrown
  sentence_pipeline(prompts, **inference_config)

  mlflow.transformers.log_model(
    transformers_model=sentence_pipeline,
    artifact_path="my_sentence_generator",
    task=task,
    inference_config=inference_config,
    registered_model_name="taka-hugging-face" # モデルレジストリに登録
  )
```

エポック`3`にしているのでしばし待ちます。なお、DBR13.0ですとメトリクスの画面がアップデートされています。
![Screenshot 2023-05-07 at 10.07.59.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/51cf082d-fcd4-e754-4605-70a3c148366d.png)

トラッキングされているモデルの画面でlossも確認できます。
![Screenshot 2023-05-07 at 11.18.44.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/70e7ad0a-2ecb-e80e-3d96-4e6b5a8dd15f.png)

最後にモデルレジストリに登録しているので、すぐにREST APIでモデルを呼び出せるようになります。前のバージョンがあるので今回はバージョン3になっています。LLMのチューニングは従来の機械学習より奥深いものですので、これまで以上に試行錯誤が発生すると思います。その際に、MLflowのトラッキングやモデルレジストリを活用することで、パラメーター、メトリクス、モデルのバージョンなどを簡単に管理できるようになります。
![Screenshot 2023-05-07 at 11.30.46.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/24eee46e-f5c1-7f2f-d160-3e3e7380655e.png)

# バッチ推論

デバッグの目的で、トラッキングされたモデルをロードして推論してみます。

```py:Python
loaded_model = mlflow.transformers.load_model('runs:/4dec269b93074f59a75fb7dcb66b0768/my_sentence_generator')
loaded_model("Databricksとは")
```

```
Out[6]: [{'generated_text': 'Databricksとは?parkはビッグデータ構造化データや、大規模データを取り扱う際に直面する最大の課題であるの課題に正面から取り組んでいきます。データ処理とデータサイエンスのゴールは同一です: 機械'}]
```

まだ怪しいですが、ファイチューニングの成果が現れています。ファインチューニングで使用しているデータは、私のQiitaの記事です。

# GUI経由モデルの活用

ここまででは、ノートブックからモデルを呼び出せるだけで、広くユーザーの方に活用いただくというところまで到達していません。そこで、サービングエンドポイントの出番です。これを活用することで、モデルをREST API経由で呼び出せるようになります。詳細手順はこちらに。

https://qiita.com/taka_yayoi/items/5f1db6447055a79284b8

エンドポイントに上記バージョン3のモデルをデプロイします。これで`URL`経由でモデルにアクセスできるようになります。
![Screenshot 2023-05-07 at 11.38.17.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/e283e21b-9334-85db-8479-44a926db647c.png)

GUI、前処理、後処理をstreamlitを使って実装します。これまでに作ったものの使い回しですが。

```py:chatbot_rinna.py
import streamlit as st 
import numpy as np 
from PIL import Image
import base64
import io

import os
import requests
import numpy as np
import pandas as pd

import json

st.header('Databricks Q&A Chatbot on Databrikcs')
st.write('''
- [MLflow 2\.3のご紹介：ネイティブLLMのサポートと新機能による強化 \- Qiita](https://qiita.com/taka_yayoi/items/431fa69430c5c6a5e741)
- [MLflow 2\.3のHugging Faceトランスフォーマーのサポートを試す \- Qiita](https://qiita.com/taka_yayoi/items/ad370a7f57c4eae58800)
''')

def score_model(prompt):
  # 1. パーソナルアクセストークンを設定してください
  # 今回はデモのため平文で記載していますが、実際に使用する際には環境変数経由で取得する様にしてください。
  token = "<パーソナルアクセストークン>"

  # 2. モデルエンドポイントのURLを設定してください
  url = '<URL>'
  headers = {'Authorization': f'Bearer {token}'}
  #st.write(token)
  
  data_json_str = f"""
  {{
  "inputs": [
    [
      "{prompt}"
    ]
  ]
}}
  """

  data_json = json.loads(data_json_str)
   
  response = requests.request(method='POST', headers=headers, url=url, json=data_json)
  if response.status_code != 200:
    raise Exception(f'Request failed with status {response.status_code}, {response.text}')
  return response.json()

prompt = st.text_input("プロンプト")

if prompt != "":
  response = score_model(prompt)
  st.write(response['predictions'])
```

```
streamlit run chatbot_rinna.py
```

動きました！
![Screenshot 2023-05-07 at 11.43.41.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/144244a3-5f5a-b7bd-0657-83e7ba69a5af.png)

これで一連のフローを実装できたので、モデルのファインチューニングの沼に潜ります。

### Databricksクイックスタートガイド

[Databricksクイックスタートガイド](https://www.amazon.co.jp/dp/B09V1YXFVQ/)


### Databricks無料トライアル

[Databricks無料トライアル](https://databricks.com/jp/try-databricks)