Google Colab で ELYZA-japanese-Llama-2-7b を試す

# パッケージのインストール
!pip install transformers accelerate bitsandbytes

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# トークナイザーとモデルの準備
tokenizer = AutoTokenizer.from_pretrained(
    "elyza/ELYZA-japanese-Llama-2-7b-instruct"
)
model = AutoModelForCausalLM.from_pretrained(
    "elyza/ELYZA-japanese-Llama-2-7b-instruct", 
    torch_dtype=torch.float16,
    device_map="auto"
)

# プロンプトの準備
prompt = """<s>[INST] <<SYS>>
あなたは誠実で優秀な日本人のアシスタントです。
<</SYS>>

まどか☆マギカでは誰が一番かわいい？ [/INST]"""

# 推論の実行
with torch.no_grad():
    token_ids = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt")
    output_ids = model.generate(
        token_ids.to(model.device),
        max_new_tokens=256,
        pad_token_id=tokenizer.pad_token_id,
        eos_token_id=tokenizer.eos_token_id,
    )
output = tokenizer.decode(output_ids.tolist()[0][token_ids.size(1) :], skip_special_tokens=True)
print(output)

まどか☆マギカに登場するキャラクターの中で、誰が一番かわいいか、私の意見を述べます。

私は人間である冬木高校の暁美ほむらが一番かわいいと考えます。

ほむらは物語の冒頭、幼い頃に事故で両親を亡くし、兄の憂木凛と二人で暮らしています。彼女は明るく元気な性格で、周りの人からも愛されています。しかし、彼女は自分の運命を変えるため、魔法少女になるための試練に挑むことを決意します。

ほむらは魔法少女になるために、魔法少

from transformers import TextIteratorStreamer
from threading import Thread

# ストリーミング出力
with torch.no_grad():
    streamer = TextIteratorStreamer(tokenizer)
    token_ids = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt")
    generation_kwargs = dict(
        input_ids=token_ids.to(model.device),
        max_new_tokens=256,
        pad_token_id=tokenizer.pad_token_id,
        eos_token_id=tokenizer.eos_token_id,
        streamer=streamer, 
    )
    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()

# 出力
for new_text in streamer:
    print(new_text.replace(" ", ""), end="")