MPT-7B-StoryWriter-65k+を走らせてみる

mosaicml/mpt-7b-storywriter · Hugging Face

今回は商用利用可能なLLMとして公開された、MPT-7Bのうちトークン数が大きいフィクションストーリーを読み書きするために設計されたモデルMPT-7B-StoryWriter-65k＋をGoogle Colabで動かしてみました。

今回試したモデルMPT-7B-StoryWriterは、2048のシーケンス長のMPT-7を65k トークンのコンテキスト長で 微調整することによって構築されたおり、推論時にALiBiを利用すると65k トークンを超えても推論を実行できるそうです。

npakaさんの記事を参考にさせてもらいました。ありがとうございます。

import torch
import transformers
from transformers import AutoTokenizer, AutoModelForCausalLM

# トークナイザーとモデルの準備
tokenizer = AutoTokenizer.from_pretrained(
    "mosaicml/mpt-7b-storywriter"
)
model = AutoModelForCausalLM.from_pretrained(
    "mosaicml/mpt-7b-storywriter", 
    torch_dtype=torch.float16,
    trust_remote_code=True
).to("cuda:0")

太宰治 走れメロス

今回は、走れメロスのエピローグを生成させてみました。 日本語だとうまく動かなかったため、原文をdeeplで英訳したものを入力しました。

with open("./Run, Meros!.txt") as f:
    text = f.read()
prompt = text + "\n\n epilogue:"
inputs = tokenizer(prompt, return_tensors='pt').to(model.device)
input_length = inputs.input_ids.shape[1]
input_length

5794

# 推論の実行
inputs = tokenizer(prompt, return_tensors='pt').to(model.device)
input_length = inputs.input_ids.shape[1]
outputs = model.generate(
    **inputs, 
    max_new_tokens=256, 
    do_sample=True, 
    temperature=0.7, 
    top_p=0.7, 
    top_k=50, 
    return_dict_in_generate=True
)
token = outputs.sequences[0, input_length:]
output_str = tokenizer.decode(token)

# 確認
print("output :", output_str)

A100 を用いてGPU RAM は40GBギリギリ、実行時間は8分でした。

群衆の端に立っていた王の娘は、抱き合っている二人を見つめた。彼女は驚いた。

"あなたたちは何をしているのですか？なぜ、抱き合うのですか？"

「友達だから抱き合ってるんだよ」とメロスは言った。

「友達なの？友達なんですか？…


The king's daughter, who was standing at the edge of the crowd, looked at the two men who were embracing each other. She was surprised.

"What are you doing? Why are you embracing each other?"
"We are embracing each other because we are friends," said Melos.

"You are friends? You are friends? You are friends? You are friends? You are friends? You are friends? You are friends?（以下リピート）

出力をdeeplで日本語訳

うーん。例によって満足いく出力をだすにはパラメータチューニングが難しそうです。今回は大きなトークンを入力して、なんらかの、ぽい出力が得られたのでひとまず満足です。ああ、A100 80GBがほしいぞよ。

おしまい