【PyTorch】GPTをPrompt-tuningしてみた

exp
├── run.py
├── data/
│   ├── aio_02_train.jsonl
│   └── aio_02_dev_v1.0.jsonl
└── outputs/
    ├── soft_prompt.pt （これから生成されるファイル）
    └── out.txt （これから生成されるファイル）

import json
import logging
import os
import random
import sys
from dataclasses import dataclass

import numpy as np
import torch
import torch.nn as nn
from torch.optim import AdamW
from tqdm import tqdm
from transformers import AutoConfig, AutoTokenizer, AutoModelForCausalLM

seed = 42
train_file = 'data/aio_02_train.jsonl'
eval_file = 'data/aio_02_dev_v1.0.jsonl'
eval_out_file = 'outputs/out.txt'
output_dir = 'outputs'
num_train_epochs = 1
n_prompt_tokens = 100
train_batch_size = 16
learning_rate = 3e-4
max_new_tokens = 10
model_name = "rinna/japanese-gpt2-medium"

class PromptTuningLM(nn.Module):
    def __init__(
        self,
        model_name: str,
        n_prompt_tokens: int,
        config: AutoConfig,
        soft_prompt_path: str = None,
    ):
        super(PromptTuningLM, self).__init__()
        self.n_prompt_tokens = n_prompt_tokens 
        # 事前学習済みのGPTの呼び出し
        self.lm = AutoModelForCausalLM.from_pretrained(model_name, config=config)
        # Promptに対する埋め込みベクトルの作成
        self.soft_prompt = nn.Embedding(n_prompt_tokens, config.hidden_size)
        torch.nn.init.xavier_uniform_(self.soft_prompt.weight)

        # GPTの重みを固定
        for param in self.lm.parameters():
            param.requires_grad = False

        # [推論時] Promptに対する学習済みの埋め込みベクトルをロード
        if soft_prompt_path is not None: 
            logger.info(f"Set soft prompt. ({n_prompt_tokens} tokens)")
            self.soft_prompt = torch.load(soft_prompt_path)

    def _extend_inputs(self, input_ids) -> torch.Tensor:
        """
        Promptに対する埋め込みベクトルを付与する
        """
        # input_idsをベクトルに変換する（事前学習モデルが異なる場合は変更する必要あり）
        inputs_embeds = self.lm.transformer.wte(input_ids)
        if len(list(inputs_embeds.shape)) == 2:
            inputs_embeds = inputs_embeds.unsqueeze(0)
        # Promptに対する埋め込みベクトルとinputs_embedsを連結する
        batch_size = inputs_embeds.size(0)
        learned_embeds = self.soft_prompt.weight.repeat(batch_size, 1, 1)
        extended_embeds = torch.cat([learned_embeds, inputs_embeds], dim=1)
        return extended_embeds

    def _extend_labels(self, labels, ignore_index=-100) -> torch.Tensor:
        """
        inputに合わせて正解ラベルにPromptに対するラベルを付与する
        """
        if len(list(labels.shape)) == 1:
            labels = labels.unsqueeze(0)
        n_batches = labels.shape[0]
        # Promptに対してignore_indexを付与（-100に設定していれば損失が計算されない）
        prompt_labels = torch.full((n_batches, self.n_prompt_tokens), 
                                    ignore_index).to(labels.device)
        # Promptに対するラベルと元の正解ラベルを連結する
        extended_labels = torch.cat([prompt_labels, labels], dim=1)
        return extended_labels

    def save_soft_prompt(self, path: str, filename: str):
        """
        Promptに対する埋め込みベクトルの保存
        """
        torch.save(self.soft_prompt, os.path.join(path, filename))
        logger.info(f"Saved soft prompt: {os.path.join(path, filename)}")

    def forward(self, input_ids, labels=None, return_dict=None):
        # Promptを付与したベクトル
        inputs_embeds = self._extend_inputs(input_ids)
        if labels is not None:
            labels = self._extend_labels(labels)

        return self.lm(
            inputs_embeds=inputs_embeds,
            labels=labels,
            return_dict=return_dict,
        )

    def generate(self, input_text, tokenizer, max_new_tokens, eos_token_id, device):
        """
        [推論時]自己回帰で回答を生成する
        """
        input_ids = tokenizer.encode(input_text, add_special_tokens=False)
        cur_ids = torch.tensor(input_ids).unsqueeze(0).to(device)
        # 最大でmax_new_tokensだけ単語を生成する
        for _ in range(max_new_tokens):
            outputs = self.forward(cur_ids)
            softmax_logits = torch.softmax(outputs.logits[0,-1], dim=0)
            # 最大確率の単語を次の単語として一意に決定
            next_token_id = int(softmax_logits.argmax().to('cpu'))
            # もし選択された単語がeos_tokenなら生成を終了する
            if next_token_id == eos_token_id:
                break
            # 選択された単語をcur_idsに追加して次の処理を行う
            next_token_id = torch.tensor([[next_token_id]]).to(device)
            cur_ids = torch.cat([cur_ids, next_token_id], dim=1)

        # 生成した単語ID列をテキストに変換する
        output_list = list(cur_ids.squeeze().to('cpu').numpy())
        output_text = tokenizer.decode(output_list)
        return output_text

logging.basicConfig(level=logging.INFO, stream=sys.stdout)
logger = logging.getLogger(__name__)

os.makedirs(output_dir, exist_ok=True)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
n_gpu = torch.cuda.device_count()

random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
if n_gpu > 0:
    torch.cuda.manual_seed_all(seed)

tokenizer = AutoTokenizer.from_pretrained(model_name)
config = AutoConfig.from_pretrained(model_name)
model = PromptTuningLM(
    model_name,
    n_prompt_tokens=n_prompt_tokens,
    config=config,
)

model.to(device)

# LayerNorm.{weight, bias}に対してweight_decay=0.01を設定
param_optimizer = list(model.named_parameters())
optimizer_grouped_parameters = [
    {'params': [p for n, p in param_optimizer if not ('ln' in n)],
     'weight_decay': 0.01},
    {'params': [p for n, p in param_optimizer if 'ln' in n],
     'weight_decay': 0.0}
]
optimizer = AdamW(optimizer_grouped_parameters, lr=learning_rate)

@dataclass
class InputExample():
    question: str
    answer: str

def create_examples(filename):
    examples = []
    with open(filename, 'r') as f:
        for line in f:
            example = json.loads(line)
            examples.append(InputExample(
                question = example['question'],
                answer = example['answers'][0]))
    return examples

class CustomDataset(torch.utils.data.IterableDataset):
    def __init__(self, tokenizer, generator):
        super().__init__()
        self._tokenizer = tokenizer
        self._generator = generator

    @classmethod
    def from_texts(cls, tokenizer, texts):
        return cls(tokenizer=tokenizer, generator=lambda: texts)

    def __iter__(self):
        for text in self._generator():
            ids = self._tokenizer.encode(text)
            yield {"input_ids": ids, "labels": ids}

def collate_fn(samples):
    batch = {'input_ids': [], 'labels': []}
    for sample in samples:
        batch['input_ids'].append(torch.tensor(sample['input_ids']))
        batch['labels'].append(torch.tensor(sample['labels']))
    batch['input_ids'] = torch.nn.utils.rnn.pad_sequence(
                    batch['input_ids'], batch_first=True, padding_value=3)
    batch['labels'] = torch.nn.utils.rnn.pad_sequence(
                    batch['labels'], batch_first=True, padding_value=3)
    return batch

logger.info("***** Running training *****")
train_examples = create_examples(train_file)
train_texts = [example.question + tokenizer.sep_token + example.answer
                for example in train_examples]
train_data = CustomDataset.from_texts(tokenizer, texts=train_texts)
train_dataloader = torch.utils.data.DataLoader(dataset=train_data,
                    batch_size=train_batch_size, collate_fn=collate_fn)

model.train()
for epoch in range(int(num_train_epochs)):
    logger.info(f'Epoch: {epoch+1}')
    for batch in tqdm(train_dataloader, desc="Iteration"):
        input_ids = batch['input_ids'].to(device)
        labels = batch['labels'].to(device)
        outputs = model(input_ids, labels=labels)
        loss = outputs.loss
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

# Promptに対する埋め込みベクトルのみ保存する
model_to_save = model.module if hasattr(model, 'module') else model
model_to_save.save_soft_prompt(output_dir, 'soft_prompt.pt')

logger.info("***** Running evaluation *****")
tokenizer = AutoTokenizer.from_pretrained(model_name)
config = AutoConfig.from_pretrained(model_name)
model = PromptTuningLM(
    model_name,
    n_prompt_tokens=n_prompt_tokens,
    soft_prompt_path=os.path.join(output_dir, 'soft_prompt.pt'),
    config=config,
)
model.to(device)
model.eval()
output_texts = []

with open(eval_file, 'r') as f:
    for line in f:
        data = json.loads(line)
        input_text = data['question'] + tokenizer.sep_token
        output_text = model.generate(input_text, tokenizer,
                        max_new_tokens, tokenizer.eos_token_id, device)
        output_texts.append(output_text+'\n')
with open(eval_out_file, 'w') as f:
    f.writelines(output_texts)

# 事例１
入力：映画『ウエスト・サイド物語』に登場する2つの少年グループといえば、シャーク団と何団?
[Prompt-tuning] GPT-2の出力  ：ウエスト・サイド物語の登場キャラクターのうち、2
[Prompt-tuning] GPT(1b)の出力：スクルージ・マクダック
[自作Prompt] GPT-2の出力      ：「シャーク団」
[自作Prompt] GPT(1b)の出力    ：、何?
答え：ジェット団

# 事例2
入力：氷った海に穴を開けて漁をすることから、漢字で「氷の下の魚」と書くタラ科の魚は何?
[Prompt-tuning] GPT-2の出力  ：タラ
[Prompt-tuning] GPT(1b)の出力：カレイ
[自作Prompt] GPT-2の出力      ：タラ
[自作Prompt] GPT(1b)の出力    ：、このページ、このページ、このページ、
答え：コマイ

# 事例３
入力：アメリカとロシア間で結ばれた「戦略兵器削減条約」のことをアルファベット5文字の略称で何という?
[Prompt-tuning] GPT-2の出力  ：戦略兵器削減条約
[Prompt-tuning] GPT(1b)の出力：包括的核実験禁止条約
[自作Prompt] GPT-2の出力      ：「sigma」
[自作Prompt] GPT(1b)の出力    ：、このページ、このページ、このページ、
答え： START

# 事例4
入力：リンゴが木から落ちる様子を見て「万有引力」を発見したといわれる、17世紀のイギリスの科学者は誰?
[Prompt-tuning] GPT-2の出力  ：ニュートン
[Prompt-tuning] GPT(1b)の出力：ニュートン
[自作Prompt] GPT-2の出力      ：アインシュタイン
[自作Prompt] GPT(1b)の出力    ：、この答えの正解は、この答えの
答え：ニュートン （or アイザック・ニュートン）

import torch
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer

soft_prompt_path = './outputs/soft_prompt.pt'
model_name = "rinna/japanese-gpt2-medium"
n_prompt_tokens = 100

tokenizer = AutoTokenizer.from_pretrained(model_name)
config = AutoConfig.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, config=config)
soft_prompt = torch.nn.Embedding(n_prompt_tokens, config.hidden_size)
soft_prompt = torch.load(soft_prompt_path, map_location='cpu')

# 既存の単語に対するベクトルを抽出
wte = model.transformer.wte.weight
wte = torch.transpose(wte, 0, 1)

token_ids = []
for i in range(n_prompt_tokens):
    vector = soft_prompt.weight[i].unsqueeze(0)
    # ボキャブラリーから最も類似するベクトルを持つ単語を選択（内積を類似度とする）
    similarity = torch.mm(vector, wte)
    token_id = int(similarity.argmax())
    token_ids.append(token_id)

prompt = tokenizer.decode(token_ids)
print(prompt)