watashiha/watashiha-gpt-6bをGoogleColobで動かす

AI

準備

以下をインストールします

!pip install torch
!pip install transformers
!pip install sentencepiece

推論

モデルのロード

サンプルコード通りbf16でロードします

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "watashiha/watashiha-gpt-6b"
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16)

if torch.cuda.is_available():
    model = model.to("cuda")

サンプルプロンプト

推論

text = "お題:ホラー映画の「○○○から逃げろ！」<SEP>回答:"
token_ids = tokenizer.encode(text, add_special_tokens=False, return_tensors="pt").to(model.device)

output_ids = model.generate(
    token_ids,
    do_sample=True,
    max_new_tokens=32,
    top_p=0.9,
    top_k=50,
    pad_token_id=tokenizer.pad_token_id,
    eos_token_id=tokenizer.eos_token_id,
)
output = tokenizer.decode(output_ids.tolist()[0], skip_special_tokens=True)
print(output)

結果

お題:ホラー映画の「○○○から逃げろ!」 回答:ゾンビ映画

面白いセリフ

推論

text = "お題:アニメのセリフの続き「月に代わって◯◯◯」<SEP>回答:"
token_ids = tokenizer.encode(text, add_special_tokens=False, return_tensors="pt").to(model.device)

output_ids = model.generate(
    token_ids,
    do_sample=True,
    max_new_tokens=32,
    top_p=0.9,
    top_k=50,
    pad_token_id=tokenizer.pad_token_id,
    eos_token_id=tokenizer.eos_token_id,
)
output = tokenizer.decode(output_ids.tolist()[0], skip_special_tokens=True)
print(output)

結果

お題:アニメのセリフの続き「月に代わって◯◯◯」 回答:、今週のジャンプは休載します

まどマギプロンプト改

推論





結果



text = "お題:まどマギに絡めた面白いこと<SEP>回答:"
token_ids = tokenizer.encode(text, add_special_tokens=False, return_tensors="pt").to(model.device)

output_ids = model.generate(
    token_ids,
    do_sample=True,
    max_new_tokens=32,
    top_p=0.9,
    top_k=50,
    pad_token_id=tokenizer.pad_token_id,
    eos_token_id=tokenizer.eos_token_id,
)
output = tokenizer.decode(output_ids.tolist()[0], skip_special_tokens=True)
print(output)





結果



お題:まどマギに絡めた面白いこと 回答:どら焼きの食べ放題





使用リソース