LLMのダウンロード

# キャッシュの場所を指定したいとき
import os
os.environ['HF_HOME'] = './cache'

# Huggingfaceのトークン認証が必要なとき(Gemmaなど)
import getpass
os.environ['HF_TOKEN'] = getpass.getpass('HF_TOKEN:')

from transformers import AutoTokenizer, AutoModelForCausalLM

# 公開されているモデルの名前
model_name = 'google/gemma-1.1-2b-it'

# トークナイザとモデルの読み込み
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map='cuda',
)

LLMのローカル保存

この方法でローカルに保存しておけば、モデルの公開が終了した場合や更新されてしまった場合などでも固定のモデルを使用し続けることができます。

# model_nameと同じ名前(ディレクトリ名)で保存
tokenizer.save_pretrained(model_name)
model.save_pretrained(model_name)

ローカル保存したLLMを実行

# 今回は公開されているモデルの名前と同じ名前のディレクトリへ格納したため
# model_nameをgoogle/gemma-1.1-2b-itとします
# ちなみにHuggingfaceに公開されているものとローカルにあるもので同じ名前がある場合は
# ローカルの方が優先されてロードされるみたいです
model_name = 'google/gemma-1.1-2b-it'
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map='cuda',
)

# 今回はchat_templateがあるchatモデルであるため、
# tokenizer.apply_chat_templateを使用
chat = [
    {'role': 'user', 'content': '徳川幕府の歴代将軍を表形式で教えてください。'},
]
input_ids = tokenizer.apply_chat_template(chat, tokenize=True, return_tensors='pt').to(model.device)
outputs = model.generate(input_ids, max_new_tokens=1000)
print(tokenizer.decode(outputs[0]))

出力

<bos><start_of_turn>user
徳川幕府の歴代将軍を表形式で教えてください。<end_of_turn>
* 1333年 - 1335年：足利義昭
* 1335年 - 1346年：足利義元
* 1346年 - 1351年：足利義昭
* 1351年 - 1368年：足利義元
* 1368年 - 1374年：足利義昭
* 1374年 - 1380年：足利義元
* 1380年 - 1382年：足利義昭
* 1382年 - 1385年：足利義元
* 1385年 - 1394年：足利義昭<eos>