OpenCALM-7Bをローカルでいい感じに使うpythonスクリプト

概要

サイバーエージェントが公開した日本語LLMのOpenCALM-7BをWSLでいい感じに使うスクリプトをChatGPTに書いてもらったので残しておきます。
自分はPythonのことほんとに何もわからないでやっていますのでご了承ください。
参考にしたページ
https://note.com/npaka/n/n2185b422a2f2

使い方

1. transformerのインストール
   1. pip install transformers accelerate
2. モデルのあるリポジトリのclone?（すべてのファイルは必要ないかも）
3. 下記スクリプトをCloneしたフォルダにおいて実行
4. checkpointのロードが終わったら質問を入力してエンター
5. やめる場合はquit

素のtransformerの挙動なので？文章を補完させる形で聞くのがいいらしい。

コード

run.py

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# モデルとトークナイザーの準備
model = AutoModelForCausalLM.from_pretrained(
    "cyberagent/open-calm-7b", 
    device_map="auto", 
    torch_dtype=torch.float16
)
tokenizer = AutoTokenizer.from_pretrained("cyberagent/open-calm-7b")

while True:
    # ユーザーからの質問を受け取る
    question = input("Enter your question (or 'quit' to stop): ")
    
    # ループを抜けるための条件
    if question.lower() == 'quit':
        break
    
    # プロンプトの作成
    prompt = f"{question}\nA:"
    
    # 推論の実行
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        tokens = model.generate(
            **inputs,
            max_new_tokens=64,
            do_sample=True,
            temperature=0.7,
            pad_token_id=tokenizer.pad_token_id,
        )
        
    output = tokenizer.decode(tokens[0], skip_special_tokens=True)
    print(output)