【LLM】Google ColabでStableLMを実装する

このシリーズではChatGPTを中心とした最新の大規模言語モデル（LLM）に関する情報をまとめています。最新AI技術を使用した「ChatGPT」をはじめとした、自然言語処理技術の概要や活用方法について紹介しておりますので、ぜひ最後までご覧ください。

※大規模言語モデルは技術の進歩がとても速い分野となりますので、最新の情報もこちらからご確認ください。

今回の記事ではLLMの1つであるStableLMの実装を紹介します。
Google Colabを使用して簡単に実装できますので、ぜひ最後までご覧ください。

1. StableLMの概要
2. 導入
3. 英語の場合
4. まとめ

StableLMの概要

「StableLM」とは、Stability AIが開発したオープンソースの言語モデルで、30億パラメータと70億パラメータのモデルがあり、将来はさらに大きなモデルも用意される予定です。このモデルは、商用や研究目的で自由に使えます。学習に使われるデータセットは「The Pile」の3倍の大きさで、1.5兆トークンのコンテンツがあります。そのため、小さいパラメータ数でも高い性能が発揮できます。

また、研究用にファインチューニングされたモデルも公開されており、5つのオープンソースデータセットを使って会話型エージェントが開発される予定です。これらのモデルは研究目的でのみ使用でき、非営利のライセンスで公開されています。

Stability AIは、オープンで透明なAI技術を目指し、StableLMなどのモデルを開発しています。このモデルは、透明性が高く、アクセスしやすく、協力的なAI技術を提供することを目指しています。これにより、研究者や開発者がより効率的で実用的なAIツールを開発し、創造性や生産性を向上させ、新しい経済機会を開拓できるようになります。

StableLMはGitHubリポジトリで利用でき、今後も技術報告書の公開や開発者とのコラボレーションが予定されています。また、クラウドソーシング型プログラムやAIアシスタント向けのオープンソースデータセットの作成など、コミュニティ活動も展開していく予定です。

詳細は以下のリンクからご確認ください。

詳細:https://ja.stability.ai/blog/stability-ai-launches-the-first-of-its-stablelm-suite-of-language-models
モデル:https://huggingface.co/stabilityai/stablelm-base-alpha-7b

導入

ここからはGoogle colabを使用して実装していきます。
（Google colabの使用方法はこちら⇨使い方）

※今回紹介するコードは以下のリンクからもご覧いただけます。

まずはGPUを使用できるように設定をします。

「ランタイムのタイプを変更」→「ハードウェアアクセラレータ」をGPUに変更・「GPUクラス」をプレミアムに変更・「ランタイムの仕様」をハイメモリに変更

必要なライブラリをインストールします。

!pip install accelerate transformers

from transformers import AutoModelForCausalLM, AutoTokenizer

# トークナイザーとモデルの準備
tokenizer = AutoTokenizer.from_pretrained("StabilityAI/stablelm-tuned-alpha-7b")
model = AutoModelForCausalLM.from_pretrained("StabilityAI/stablelm-tuned-alpha-7b")
model.half().cuda()

以上で準備が完了しました。

英語の場合

「What are the famous sightseeing spots in Japan?」に続く文章を生成します。

from transformers import StoppingCriteria
import torch

# トークン停止クラスの定義
class StopOnTokens(StoppingCriteria):
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
        stop_ids = [50278, 50279, 50277, 1, 0]
        for stop_id in stop_ids:
            if input_ids[0][-1] == stop_id:
                return True
        return False

# システムプロンプトの定義
system_prompt = """<|SYSTEM|># StableLM Tuned (Alpha version)
- StableLM is a helpful and harmless open-source AI language model developed by StabilityAI.
- StableLM is excited to be able to help the user, but will refuse to do anything that could be considered harmful to the user.
- StableLM is more than just an information source, StableLM is also able to write poetry, short stories, and make jokes.
- StableLM will refuse to participate in anything that could harm a human.
"""

from transformers import StoppingCriteriaList

# プロンプトの準備
prompt = "What are some famous tourist destinations in Japan?"
prompt = f"{system_prompt}<|USER|>{prompt}<|ASSISTANT|>"

# 推論の実行
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
tokens = model.generate(
  **inputs,
  max_new_tokens=64,
  temperature=0.7,
  do_sample=True,
  stopping_criteria=StoppingCriteriaList([StopOnTokens()])
)
print("response :", tokenizer.decode(tokens[0], skip_special_tokens=False))

出力結果:

<|USER|>What are some famous tourist destinations in Japan?<|ASSISTANT|>
1. Tokyo
2. Osaka
3. Kyoto
4. Hiroshima
5. Nara
6. Sapporo
7. Hiroshima Castle
8. Mount Fuji
9. Hiroshima Peace Memorial Park
10. Tsukiji Fish Market.
<|USER|>

# 日本語訳
1. 東京
2. 大阪
3. 京都
4. 広島
5. 奈良
6. 札幌
7. 広島城
8. 富士山
9. 広島平和記念公園
10. 築地魚河岸