MistralをGoogle Colabで実行

はじめに

LLMを実行するためのサンプルコードです。
jupyter notebookは下記にあります。
https://github.com/toshi-4886/LLM/blob/main/notebook/1_mistral.ipynb

概要

  • 小型で性能の高いMistralを動かします。

  • 認証などについても説明します。

事前準備

  1. Hugging Faceのアカウント作成

  2. アクセストークンの作成
    Hugging Faceのアカウントにログインし、「自分のアイコン→Settings→Access Tokens→+ Create new token」からトークンを作成します。
    「Read access to contents of all public gated repos you can access」の権限が必要となります。

3 . モデルの利用条件の同意
Hugging Faceで使用したいモデルのページにアクセスして、利用条件に同意します。

実装

1. ライブラリのインポート

import sys
import os

!pip install -U bitsandbytes
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
import torch

2. 認証

Hugging Faceからモデルをダウンロードするために、アカウント認証を行います。
login()を実行するとトークンを入力するUIが出てくるので、そこにトークンを入力します。

また、colabのシークレットにトークンを保存しておけば、それを利用することもできます。

from huggingface_hub import login
# login()

# シークレットにトークンを保存しておけば下記で認証可能
from google.colab import userdata
login(token=userdata.get('HF_TOKEN'))

3. モデルの準備

トークナイザーとモデルを準備します。
モデルはv0.3を4bit量子化したものを使用します。

tokenizer = AutoTokenizer.from_pretrained(
    "mistralai/Mistral-7B-Instruct-v0.3",
)

quantization_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16)

model = AutoModelForCausalLM.from_pretrained(
    "mistralai/Mistral-7B-Instruct-v0.3",
    torch_dtype=torch.float16,
    device_map="auto",
    trust_remote_code=False,
    quantization_config=quantization_config,
)

4. 推論

プロンプトをモデルに入力して、得られた出力を表示します。

# 推論時間を計測
%%time

# プロンプトの準備
conversation = [{"role": "user", "content": "What's the weather like in Paris?"}]

# 推論の実行
with torch.no_grad():
    token_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True,
            return_dict=True, return_tensors="pt").to(model.device)
    output_ids = model.generate(
        **token_ids,
        pad_token_id=tokenizer.eos_token_id,
        temperature=0.1,
        do_sample=True,
        top_p=0.95,
        top_k=40,
        max_new_tokens=256,
    )
output = tokenizer.decode(output_ids[0], skip_special_tokens=True)
print(output)
What's the weather like in Paris? I don't have real-time capabilities, so I can't provide the current weather in Paris. However, Paris, like much of France, has a temperate oceanic climate, characterized by mild winters and relatively warm summers. Rainfall is relatively evenly distributed throughout the year. For a more accurate and up-to-date forecast, I recommend checking a reliable weather service.
CPU times: user 5.51 s, sys: 4.68 ms, total: 5.51 s
Wall time: 5.51 s

おわりに

今回の結果

Mistralでプロンプトの回答を生成することができました。
天気に関する質問に答えていないので、ハルシネーションを防ぐための調整もされていることが確認できました。

次にやること

プロンプトの設定方法について様々な方法を試してみたいと思います。

参考資料

いいなと思ったら応援しよう!