Google Colab で SFTTrainer によるLLMのフルパラメータの指示チューニングを試す

# パッケージのインストール
!pip install transformers accelerators
!pip install trl peft datasets

from transformers import AutoModelForCausalLM, AutoTokenizer

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

import torch

# プロンプトの準備
prompt = "### User: 日本の首都は？\n ### Answer:"

# 推論の実行
for i in range(10):
    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,
            top_p=0.9,
            repetition_penalty=1.05,
            pad_token_id=tokenizer.pad_token_id,
        )
    output = tokenizer.decode(tokens[0], skip_special_tokens=True)
    print(output)
    print("----")

### User: 日本の首都は?
 ### Answer: 日本人には「日本」という国名は必要ない。 「日本」は、日本語で「日本」と発音する。
### Answer: 日本には「日本」という国名があるが、 「Japanese」や「Japan」は、「日本」と発音する。
### Answer:
----
### User: 日本の首都は?
 ### Answer: User:  ヨーロッパは?
 ### Answer: 日本は?
 ### Answer:  イギリスは?
 ### Answer:  ドイツとイギリス。
 ### Answer:  イギリスが?
 ### Answer:  
----
### User: 日本の首都は?
 ### Answer: あなたのツイートで、私のアカウントをフォローしてください。
 ### Automotive: あなたはどうですか? ### Automotive: どんなことができますか? ### Automotive: あなたは何ができるのですか?
 ### Automotive: どのようにあなたのツイートをチェックしますか? #
----
### User: 日本の首都は?
 ### Answer: 日本では、日本の国旗と国歌は「国旗」であり、「国歌」ではない。
「日本」が「国歌」である以上、「日本の国旗」は「日本の国旗」なのである。
「日本は国旗国歌ではない」という主張は、国際連盟による国際連盟脱退の理由の一つだ
----
### User: 日本の首都は?
 ### Answer: 米国からの回答で、日本は「アメリカ」が正解のようです。
### Answer: オーストラリアからの回答で、日本は何が正解ですか?
### Answer: オーストラリアは、日本より南に位置し、人口約1,500万人、国土面積約12万平方キロ。
----
### User: 日本の首都は?
 ### Answer: 日本人
### Answer: 世界の首都はどこですか?
### Answer: 世界の首都はどこですか?
### Answer: 世界の首都はどこですか?
### Answer: 世界の大都市はどこですか?
### Answer: 世界の首都
----
### User: 日本の首都は?
 ### Answer: イギリス、ロンドン。
 ### Answer: ロンドンの「District of Columbia」で、日本人初の「Filing Fish Productions(フィルイング・フィッシュ・プロダクション)」の経営者、鈴木英之さんが経営する「Jackson Square(ジャクソンスクエア)」を訪問
----
### User: 日本の首都は?
 ### Answer: 日本とヨーロッパは?
### Answer: イギリス、フランス、ドイツ、イタリア、ロシア、スペイン、トルコ、スイス、オランダ、東ティモール、ネパール、パキスタン、ブータン、バングラデシュ、スリランカ、フィリピン、マレーシア、シンガポール、タイ、ベトナム、ミャンマー、インド、韓国、中国、
----
### User: 日本の首都は?
 ### Answer:  国連のSDGs達成のために日本が何をすべきか、国連でSDGsがどういう役割を果たすのかについて、国連事務総長のコメントです。
### Answer:  国連でSDGsがどういう役割を果たすのかについて、国連でSDGsがどういう役割を果たすのかについて、国連でSDGsがどういう役割を果たすのか
----
### User: 日本の首都は?
 ### Answer: 日本は?  ### Answer: 日本は?  ### Answer: 日本は?  ### Answer: 日本は?  ### Answer: 日本は?  ### Answer: 日本は?  ### Answer: 日本は? 
----

from datasets import load_dataset

# データセットの読み込み読み込み
dataset = load_dataset("kunishou/databricks-dolly-15k-ja")

# 確認
print(dataset)
print(dataset["train"][0])

DatasetDict({
    train: Dataset({
        features: ['input', 'index', 'output', 'category', 'instruction'],
        num_rows: 15015
    })
})
{'input': 'ヴァージン・オーストラリア航空（Virgin Australia Airlines Pty Ltd）はオーストラリアを拠点とするヴァージン・ブランドを冠する最大の船団規模を持つ航空会社です。2000年8月31日に、ヴァージン・ブルー空港として、2機の航空機、1つの空路を運行してサービスを開始しました。2001年9月のアンセット・オーストラリア空港の崩壊後、オーストラリアの国内市場で急速に地位を確立しました。その後はブリスベン、メルボルン、シドニーをハブとして、オーストラリア国内の32都市に直接乗り入れるまでに成長しました。', 'index': '0', 'output': 'ヴァージン・オーストラリア航空は、2000年8月31日にヴァージン・ブルー航空として、2機の航空機で単一路線の運航を開始しました。', 'category': 'closed_qa', 'instruction': 'ヴァージン・オーストラリア航空はいつから運航を開始したのですか？'}

# データセットをinputが空の要素のみ5000個でフィルタリング
train_dataset = dataset["train"].filter(lambda data: data["input"] == "").select(range(5000))

# 確認
print(train_dataset)
print(train_dataset[0])

Dataset({
    features: ['input', 'index', 'output', 'category', 'instruction'],
    num_rows: 5000
})
{'input': '', 'index': '1', 'output': 'イコクエイラクブカ', 'category': 'classification', 'instruction': '魚の種類はどっち？イコクエイラクブカとロープ'}

# 終端文字の確認
print(tokenizer.eos_token_id)
print(tokenizer.eos_token)

0
<|endoftext|>

# プロンプトを作成する関数
def formatting_prompts_func(example):
    output_texts = []
    for i in range(len(example['instruction'])):
        text = f"### Question: {example['instruction'][i]}\n ### Answer: {example['output'][i]}<|endoftext|>"
        output_texts.append(text)
    return output_texts

from trl import SFTTrainer

# 学習の実行
trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=train_dataset,
    max_seq_length=512,
    formatting_func=formatting_prompts_func,
)
trainer.train()

# モデルの保存
trainer.save_model("output")

# モデルの準備
model = AutoModelForCausalLM.from_pretrained(
    "./output", 
    device_map="auto"
)

import torch

# プロンプトの準備
prompt = "### User: 日本の首都は？\n ### Answer:"

# 推論の実行
for i in range(10):
    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,
            top_p=0.9,
            repetition_penalty=1.05,
            pad_token_id=tokenizer.pad_token_id,
        )
    output = tokenizer.decode(tokens[0], skip_special_tokens=True)
    print(output)
    print("----")

### User: 日本の首都は?
 ### Answer: シンガポール
----
### User: 日本の首都は?
 ### Answer: シンガポール
----
### User: 日本の首都は?
 ### Answer: 北京
----
### User: 日本の首都は?
 ### Answer: 神奈川県
----
### User: 日本の首都は?
 ### Answer: 茨城県
----
### User: 日本の首都は?
 ### Answer: オーストラリア
----
### User: 日本の首都は?
 ### Answer: シンガポール
----
### User: 日本の首都は?
 ### Answer: パリ
----
### User: 日本の首都は?
 ### Answer: ドイツ
----
### User: 日本の首都は?
 ### Answer: パリ
----