C4AI Command-R Google Colab L4で動かしてみた

はじめに
手順
まとめ
ノートブック
参考サイト
1. 関連

はじめに

本記事では、Cohere For AIが開発した35億パラメータの大規模言語モデルであるC4AI Command-Rを使用して、Google Colab L4で実装する方法について解説します。Command-Rは、推論、要約、質問応答など様々なタスクに最適化されており、10言語でのマルチリンガル生成が可能です。

こちらの記事もおすすめ

日本語版：C4AI Command R+モデルカード

🚨 このモデルは、C4AI Command R+の非量子化バージョンです。bitsandbytesを使用したC4AI Command R+の量子化バージョンはこちらで見つけることができます。モデルの概要C4AI Command R+は、1,...

手順

環境設定

まず、必要なライブラリをインストールし、環境変数を設定します。

pip install 'git+https://github.com/huggingface/transformers.git' bitsandbytes accelerate

import os

os.environ['TORCH_SHOW_CPP_STACKTRACES'] = '1'
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'

モデルとトークナイザーの読み込み

BitsAndBytesConfigを使用して、4bitで量子化されたCommand-Rモデルを読み込みます。モデルIDを指定し、AutoTokenizerとAutoModelForCausalLMを使用してトークナイザーとモデルを読み込みます。

from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

bnb_config = BitsAndBytesConfig(llm_int8_enable_fp32_cpu_offload=True, load_in_4bit=True)
model_id = "CohereForAI/c4ai-command-r-v01-4bit"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    low_cpu_mem_usage=True,
    quantization_config=bnb_config,
    offload_folder="offload",
    offload_state_dict=True,
)

ユーザー入力のフォーマット

Command-R-plusチャットテンプレートを使用して、ユーザーのメッセージをフォーマットします。トークナイザーのapply_chat_templateメソッドを使用して、メッセージをトークン化し、生成プロンプトを追加します。

messages = [{"role": "user", "content": "こんにちは、お元気ですか？"}]
input_ids = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt")

応答の生成

model.generateメソッドを使用して、フォーマットされたユーザー入力から応答を生成します。max_new_tokens、do_sample、temperatureなどのパラメータを調整して、生成される応答の長さやランダム性を制御します。

gen_tokens = model.generate(
    input_ids,
    max_new_tokens=100,
    do_sample=True,
    temperature=0.3,
)

生成された応答の表示

トークナイザーのdecodeメソッドを使用して、生成されたトークンを人間が読める形式にデコードします。生成された応答をprint関数で表示します。

gen_text = tokenizer.decode(gen_tokens[0])
print(gen_text)

まとめ

C4AI Command-Rモデルを使用することで、高性能なチャットボットを簡単に実装することができます。量子化されたモデルを使用することで、メモリ使用量を削減しつつ、高品質な応答を生成できます。また、チャットテンプレートを使用することで、ユーザー入力を適切にフォーマットし、より自然なチャットボットとのやり取りが可能になります。

本記事で紹介した手順を参考に、C4AI Command-Rモデルを使用して独自のチャットボットを開発してみてください。