你還弄不清xxxForCausalLM和xxxForConditionalGeneration嗎?

西西嘛呦發表於2023-04-26

Part1基本介紹

大語言模型目前一發不可收拾,在使用的時候經常會看到transformers庫的蹤影,其中xxxCausalLM和xxxForConditionalGeneration會經常出現在我們的視野中,接下來我們就來聊聊transformers庫中的一些基本任務。

這裡以三類模型為例:bert(自編碼)、gpt(自迴歸)、bart(編碼-解碼)

首先我們整體看下每個模型有什麼任務:

from ..bart.modeling_bart import (
    BartForCausalLM,
    BartForConditionalGeneration,
    BartForQuestionAnswering,
    BartForSequenceClassification,
    BartModel,
)
from ..bert.modeling_bert import (
    BertForMaskedLM,
    BertForMultipleChoice,
    BertForNextSentencePrediction,
    BertForPreTraining,
    BertForQuestionAnswering,
    BertForSequenceClassification,
    BertForTokenClassification,
    BertLMHeadModel,
    BertModel,
)
from ..gpt2.modeling_gpt2 import GPT2ForSequenceClassification, GPT2LMHeadModel, GPT2Model

1Bert

  • BertModel(BertPreTrainedModel):最原始的bert,可獲得句向量表示或者每個token的向量表示。

  • BertForPreTraining(BertPreTrainedModel):在BertModel的基礎上加了一個預訓練頭:

self.bert = BertModel(config)
self.cls = BertPreTrainingHeads(config)

class BertPreTrainingHeads(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.predictions = BertLMPredictionHead(config)
        self.seq_relationship = nn.Linear(config.hidden_size, 2)

    def forward(self, sequence_output, pooled_output):
        prediction_scores = self.predictions(sequence_output)
        seq_relationship_score = self.seq_relationship(pooled_output)
        return prediction_scores, seq_relationship_score

class BertLMPredictionHead(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.transform = BertPredictionHeadTransform(config)

        # The output weights are the same as the input embeddings, but there is
        # an output-only bias for each token.
        self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        self.bias = nn.Parameter(torch.zeros(config.vocab_size))

        # Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings`
        self.decoder.bias = self.bias

    def forward(self, hidden_states):
        hidden_states = self.transform(hidden_states)
        hidden_states = self.decoder(hidden_states)
        return hidden_states

對應bert的兩個訓練任務:掩碼語言模型(MLM)和下一個句子預測(NSP)。
  • BertLMHeadModel(BertPreTrainedModel):MLM任務
self.bert = BertModel(config, add_pooling_layer=False)
self.cls = BertOnlyMLMHead(config)

class BertOnlyMLMHead(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.predictions = BertLMPredictionHead(config)

    def forward(self, sequence_output: torch.Tensor) -> torch.Tensor:
        prediction_scores = self.predictions(sequence_output)
        return prediction_scores

class BertLMPredictionHead(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.transform = BertPredictionHeadTransform(config)

        # The output weights are the same as the input embeddings, but there is
        # an output-only bias for each token.
        self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        self.bias = nn.Parameter(torch.zeros(config.vocab_size))

        # Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings`
        self.decoder.bias = self.bias

    def forward(self, hidden_states):
        hidden_states = self.transform(hidden_states)
        hidden_states = self.decoder(hidden_states)
        return hidden_states
  • BertForNextSentencePrediction(BertPreTrainedModel):NSP任務
self.bert = BertModel(config)
self.cls = BertOnlyNSPHead(config)

class BertOnlyNSPHead(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.seq_relationship = nn.Linear(config.hidden_size, 2)

    def forward(self, pooled_output):
        seq_relationship_score = self.seq_relationship(pooled_output)
        return seq_relationship_score
  • BertForSequenceClassification(BertPreTrainedModel):對句子進行分類
self.bert = BertModel(config)
classifier_dropout = (
    config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
)
self.dropout = nn.Dropout(classifier_dropout)
self.classifier = nn.Linear(config.hidden_size, config.num_labels)
  • BertForMultipleChoice(BertPreTrainedModel)::多項選擇
self.bert = BertModel(config)
classifier_dropout = (
        config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
    )
self.dropout = nn.Dropout(classifier_dropout)
self.classifier = nn.Linear(config.hidden_size, 1)
  • BertForTokenClassification(BertPreTrainedModel):對token進行分類,一般為命名實體識別任務
self.bert = BertModel(config, add_pooling_layer=False)
classifier_dropout = (
    config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
)
self.dropout = nn.Dropout(classifier_dropout)
self.classifier = nn.Linear(config.hidden_size, config.num_labels)
  • BertForQuestionAnswering(BertPreTrainedModel):QA任務,很多工都可以轉換為這種形式。即識別答案的開始位置和結束位置。
self.bert = BertModel(config, add_pooling_layer=False)
self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)

2GPT2

  • GPT2Model(GPT2PreTrainedModel):原始的GPT2模型,返回每個token的向量。

  • GPT2LMHeadModel(GPT2PreTrainedModel):進行語言模型任務。判斷每一個token的下一個token是什麼、

self.transformer = GPT2Model(config)
self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
  • GPT2DoubleHeadsModel(GPT2PreTrainedModel):除了語言模型任務外,額外定義了一個任務:多項選擇任務。比如一個問題有兩個回答,一個正確回答,一個錯誤回答,進行二分類任務判斷哪一個是正確回答。當然可以擴充套件到多個選項。
 config.num_labels = 1
 self.transformer = GPT2Model(config)
 self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
 self.multiple_choice_head = SequenceSummary(config)

這個要看個例子:

import torch
from transformers import GPT2Tokenizer, GPT2DoubleHeadsModel

tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
model = GPT2DoubleHeadsModel.from_pretrained('gpt2')

choices = [ "Bob likes candy ; what does Bob like ?  Bag <|endoftext|>",
                   "Bob likes candy ; what does Bob like ?  Burger <|endoftext|>",
                   "Bob likes candy ; what does Bob like ?  Candy <|endoftext|>",
                  "Bob likes candy ; what does Bob like ?  Apple <|endoftext|>"]

encoded_choices = [tokenizer.encode(s) for s in choices]

eos_token_location = [tokens.index(tokenizer.eos_token_id) for tokens in encoded_choices]
input_ids = torch.tensor(encoded_choices).unsqueeze(0
mc_token_ids = torch.tensor([eos_token_location]) 
print(input_ids.shape)
print(mc_token_ids.shape)
outputs = model(input_ids, mc_token_ids=mc_token_ids)
lm_prediction_scores, mc_prediction_scores = outputs[:2]
print(lm_prediction_scores.shape)
print(mc_prediction_scores)

"""
torch.Size([1, 4, 13])
torch.Size([1, 4])
torch.Size([1, 4, 13, 50257])
tensor([[-6.0075, -6.0649, -6.0657, -6.0585]], grad_fn=<SqueezeBackward1>)
"""

Confused by GPT2DoubleHeadsModel example · Issue #1794 · huggingface/transformers (github.com)

How to use GPT2DoubleHeadsModel? · Issue #3680 · huggingface/transformers (github.com)

  • GPT2ForSequenceClassification(GPT2PreTrainedModel):顯然,針對於文字分類任務
self.transformer = GPT2Model(config)
self.score = nn.Linear(config.n_embd, self.num_labels, bias=False)
  • GPT2ForTokenClassification(GPT2PreTrainedModel):針對於token分類(命名實體識別任務)
self.transformer = GPT2Model(config)
if hasattr(config, "classifier_dropout"and config.classifier_dropout is not None:
    classifier_dropout = config.classifier_dropout
elif hasattr(config, "hidden_dropout"and config.hidden_dropout is not None:
    classifier_dropout = config.hidden_dropout
else:
    classifier_dropout = 0.1
    self.dropout = nn.Dropout(classifier_dropout)
    self.classifier = nn.Linear(config.hidden_size, config.num_labels)

舉個例子:

import torch
from transformers import GPT2Tokenizer, GPT2DoubleHeadsModel, GPT2Model

tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
model = GPT2Model.from_pretrained('gpt2')

text = [
  "Bob likes candy ; what does Bob like ?  Bag <|endoftext|>"
  "Bob likes candy ; what does Bob like ?  Bag <|endoftext|>"
]

inputs = tokenizer(text, return_tensors="pt")
print(inputs)
print(tokenizer.decode(inputs["input_ids"][0]))

output = model(**inputs)
print(output[0].shape)

"""
{'input_ids': tensor([[18861,  7832, 18550,  2162,   644,   857,  5811,   588,  5633,   220,
         20127,   220, 50256],
        [18861,  7832, 18550,  2162,   644,   857,  5811,   588,  5633,   220,
         20127,   220, 50256]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])}
Bob likes candy ; what does Bob like?  Bag <|endoftext|>
torch.Size([2, 13, 768])
"""

3BART

  • BartModel(BartPretrainedModel):bart的原始模型,返回解碼器每個token的向量。當然還有其它可選的。

  • BartForConditionalGeneration(BartPretrainedModel):顧名思義,條件文字生成。

self.model = BartModel(config)
self.register_buffer("final_logits_bias", torch.zeros((1, self.model.shared.num_embeddings)))
self.lm_head = nn.Linear(config.d_model, self.model.shared.num_embeddings, bias=False)

輸入一般我們需要定義:input_ids(編碼器的輸入)、attention_mask (編碼器注意力)、decoder_input_ids(解碼器的輸入),target_attention_mask(解碼器注意力)輸出一般我們使用的有兩個 loss=masked_lm_loss和 logits=lm_logits。

  • BartForSequenceClassification(BartPretrainedModel):文字分類
self.model = BartModel(config)
self.classification_head = BartClassificationHead(
    config.d_model,
    config.d_model,
    config.num_labels,
    config.classifier_dropout,
)

class BartClassificationHead(nn.Module):
    """Head for sentence-level classification tasks."""

    def __init__(
        self,
        input_dim: int,
        inner_dim: int,
        num_classes: int,
        pooler_dropout: float,
    )
:

        super().__init__()
        self.dense = nn.Linear(input_dim, inner_dim)
        self.dropout = nn.Dropout(p=pooler_dropout)
        self.out_proj = nn.Linear(inner_dim, num_classes)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states = self.dropout(hidden_states)
        hidden_states = self.dense(hidden_states)
        hidden_states = torch.tanh(hidden_states)
        hidden_states = self.dropout(hidden_states)
        hidden_states = self.out_proj(hidden_states)
        return hidden_states

具體的獲取logits是這麼操作的:

hidden_states = outputs[0]  # last hidden state
# 找到eos_mask的位置
eos_mask = input_ids.eq(self.config.eos_token_id).to(hidden_states.device)

if len(torch.unique_consecutive(eos_mask.sum(1))) > 1:
    raise ValueError("All examples must have the same number of <eos> tokens.")
    sentence_representation = hidden_states[eos_mask, :].view(hidden_states.size(0), -1, hidden_states.size(-1))[
        :, -1, :
    ]
logits = self.classification_head(sentence_representation)


損失計算:

loss = None
if labels is not None:
    labels = labels.to(logits.device)
    if self.config.problem_type is None:
        if self.config.num_labels == 1:
            self.config.problem_type = "regression"
        elif self.config.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
            self.config.problem_type = "single_label_classification"
        else:
            self.config.problem_type = "multi_label_classification"

            if self.config.problem_type == "regression":
                loss_fct = MSELoss()
                if self.config.num_labels == 1:
                    loss = loss_fct(logits.squeeze(), labels.squeeze())
                else:
                    loss = loss_fct(logits, labels)
                elif self.config.problem_type == "single_label_classification":
                    loss_fct = CrossEntropyLoss()
                    loss = loss_fct(logits.view(-1, self.config.num_labels), labels.view(-1))
                elif self.config.problem_type == "multi_label_classification":
                    loss_fct = BCEWithLogitsLoss()
  • BartForQuestionAnswering(BartPretrainedModel):問答和之前GPT基本一致,只不過這裡的輸入到計算logits前的向量是解碼器的隱含層向量。
config.num_labels = 2
self.num_labels = config.num_labels

self.model = BartModel(config)
self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)

sequence_output = outputs[0]
logits = self.qa_outputs(sequence_output)
start_logits, end_logits = logits.split(1, dim=-1)
start_logits = start_logits.squeeze(-1).contiguous()
end_logits = end_logits.squeeze(-1).contiguous()

  • BartForCausalLM(BartPretrainedModel):語言模型任務,只使用bart的解碼器。
config = copy.deepcopy(config)
config.is_decoder = True
config.is_encoder_decoder = False
super().__init__(config)
self.model = BartDecoderWrapper(config)

self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

outputs = self.model.decoder(
            input_ids=input_ids,
            attention_mask=attention_mask,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_attention_mask,
            head_mask=head_mask,
            cross_attn_head_mask=cross_attn_head_mask,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
)

logits = self.lm_head(outputs[0])

>>> from transformers import AutoTokenizer, BartForCausalLM
>>> tokenizer = AutoTokenizer.from_pretrained("facebook/bart-base")
>>> model = BartForCausalLM.from_pretrained("facebook/bart-base", add_cross_attention=False)
>>> assert model.config.is_decoder, f"{model.__class__} has to be configured as a decoder."
>>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
>>> outputs = model(**inputs)

>>> logits = outputs.logits
>>> expected_shape = [1, inputs.input_ids.shape[-1], model.config.vocab_size]
>>> list(logits.shape) == expected_shape
True

Part2實操

接下來針對xxxCausalLM和xxxForConditionalGeneration,我們實際操作來更加深入的瞭解它們。首先需要安裝一些依賴:

pip install transformers==4.28.1
pip install evaluate
pip install datasets

4使用GPT2進行觀點評論生成

資料從這裡下載:https://raw.githubusercontent.com/SophonPlus/ChineseNlpCorpus/master/datasets/ChnSentiCorp_htl_all/ChnSentiCorp_htl_all.csv

直接上程式碼:

import torch
from tqdm import tqdm
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM
from transformers import (
    default_data_collator,
    get_linear_schedule_with_warmup,
)
from torch.utils.data import DataLoader

data_file = "./ChnSentiCorp_htl_all.csv"  # 資料檔案路徑,資料需要提前下載
max_length = 86
train_batch_size = 64
eval_batch_size = 64
num_epochs = 10
lr = 3e-4

# 載入資料集
dataset = load_dataset("csv", data_files=data_file)
dataset = dataset.filter(lambda x: x["review"is not None)
dataset = dataset["train"].train_test_split(0.2, seed=123)

model_name_or_path = "uer/gpt2-chinese-cluecorpussmall"
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
model = AutoModelForCausalLM.from_pretrained(model_name_or_path)


# example = {'label': 1, 'review': '早餐太差,無論去多少人,那邊也不加食品的。酒店應該重視一下這個問題了。房間本身很好。'}


def process(example):
    text = example["review"]
    # text = ["399真的很值得之前也住過別的差不多價位的酒店式公寓沒有這間好廚房很像廚房很大整個格局也都很舒服早上的早餐我訂的8點半的已經冷了。。。位置啊什麼還是很好的下次還會去服務也很周到"]
    batch_size = len(text)
    inputs = tokenizer(text, add_special_tokens=False, truncation=True, max_length=max_length)
    inputs["labels"] = []
    for i in range(batch_size):
        input_ids = inputs["input_ids"][i]
        if len(input_ids) + 1 <= max_length:
            inputs["input_ids"][i] = input_ids + [tokenizer.pad_token_id] + [0] * (max_length - len(input_ids) - 1)
            inputs["labels"].append(input_ids + [tokenizer.pad_token_id] + [-100] * (max_length - len(input_ids) - 1))
            inputs["attention_mask"][i] = [1] * len(input_ids) + [0] + [0] * (max_length - len(input_ids) - 1)
        else:
            inputs["input_ids"][i] = input_ids[:max_length - 1] + [tokenizer.pad_token_id]
            inputs["labels"].append(inputs["input_ids"][i])
            inputs["attention_mask"][i] = [1] * max_length

        inputs["token_type_ids"][i] = [0] * max_length
        # for k, v in inputs.items():
        #   print(k, len(v[0]))
        # assert len(inputs["labels"][i]) == len(inputs["input_ids"][i]) == len(inputs["token_type_ids"][i]) == len(inputs["attention_mask"][i]) == 86
    return inputs


# process(None)

train_dataset = dataset["train"].map(process, batched=True, num_proc=1, remove_columns=dataset["train"].column_names)
test_dataset = dataset["test"].map(process, batched=True, num_proc=1, remove_columns=dataset["test"].column_names)

train_dataloader = DataLoader(
    train_dataset, collate_fn=default_data_collator, shuffle=True, batch_size=train_batch_size, pin_memory=True
)

test_dataloader = DataLoader(
    test_dataset, collate_fn=default_data_collator, batch_size=eval_batch_size, pin_memory=True
)

# optimizer

optimizer = torch.optim.AdamW(model.parameters(), lr=lr)

# lr scheduler
lr_scheduler = get_linear_schedule_with_warmup(
    optimizer=optimizer,
    num_warmup_steps=0,
    num_training_steps=(len(train_dataloader) * num_epochs),
)

model.cuda()

from tqdm import tqdm

for epoch in range(num_epochs):
    model.train()
    total_loss = 0
    t = tqdm(train_dataloader)
    for step, batch in enumerate(t):
        for k, v in batch.items():
            batch[k] = v.cuda()
        outputs = model(
            input_ids=batch["input_ids"],
            token_type_ids=batch["token_type_ids"],
            attention_mask=batch["attention_mask"],
            labels=batch["labels"],
        )
        loss = outputs.loss
        t.set_description("loss:{:.6f}".format(loss.item()))
        total_loss += loss.detach().float()
        loss.backward()
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()

    train_epoch_loss = total_loss / len(train_dataloader)
    model.save_pretrained("gpt2-chinese/")
    print(f"epoch:{epoch+1}/{num_epochs} loss:{train_epoch_loss}")

訓練結果:

loss:2.416899100%|██████████| 98/98 [01:51<00:00,  1.14s/it]
epoch:1/10 loss:2.7781832218170166
loss:2.174688100%|██████████| 98/98 [01:54<00:00,  1.17s/it]
epoch:2/10 loss:2.3192219734191895
loss:2.123909100%|██████████| 98/98 [01:55<00:00,  1.17s/it]
epoch:3/10 loss:2.037835121154785
loss:1.785878100%|██████████| 98/98 [01:55<00:00,  1.18s/it]
epoch:4/10 loss:1.7687807083129883
loss:1.466153100%|██████████| 98/98 [01:55<00:00,  1.18s/it]
epoch:5/10 loss:1.524872064590454
loss:1.465316100%|██████████| 98/98 [01:54<00:00,  1.17s/it]
epoch:6/10 loss:1.3074666261672974
loss:1.150320100%|██████████| 98/98 [01:54<00:00,  1.16s/it]
epoch:7/10 loss:1.1217808723449707
loss:1.043044100%|██████████| 98/98 [01:53<00:00,  1.16s/it]
epoch:8/10 loss:0.9760875105857849
loss:0.790678100%|██████████| 98/98 [01:53<00:00,  1.16s/it]
epoch:9/10 loss:0.8597695827484131
loss:0.879025100%|██████████| 98/98 [01:53<00:00,  1.16s/it]
epoch:10/10 loss:0.790839433670044

可以這麼進行預測:

from transformers import AutoTokenizer, GPT2LMHeadModel, TextGenerationPipeline, AutoModelForCausalLM
from datasets import load_dataset
data_file = "./ChnSentiCorp_htl_all.csv" # 資料檔案路徑,資料需要提前下載
dataset = load_dataset("csv", data_files=data_file)
dataset = dataset.filter(lambda x: x["review"is not None)
dataset = dataset["train"].train_test_split(0.2, seed=123)

model_name_or_path = "uer/gpt2-chinese-cluecorpussmall"
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
model = AutoModelForCausalLM.from_pretrained("./gpt2-chinese/")
text_generator = TextGenerationPipeline(model, tokenizer)  

import random
examples = dataset["train"]
example = random.choice(examples)
text = example["review"]
print(text)
print(text[:10])
text_generator(text[:10], 
        max_length=100
        do_sample=False
        top_p=0.8
        repetition_penalty=10.0,
        temperature=0.95,
        eos_token_id=0
        )
"""
第一次住在這裡兒,我針對大家的意見,特別關注了一下,感覺如下吧!1、標準間雖然有點舊但很乾淨,被子蓋得很舒服,也很暖和,衛生間也蠻大的,因是在商業中心離很多還算很近。2、酒店服務還算可以,沒有像這裡說的那樣,入住時,退房時也挺快的,總的來說我很滿意。3、早餐也還可以,環境也不錯,有點江南的感覺;菜品種品也不少,挺可口。4、可能是在市或者離火車站的距離很近,稍微有點“熱鬧”,來找我辦事的人不方便停車,但還好這裡有地下停車場。總體來說,我感覺很不錯,值得推薦!!!
第一次住在這裡兒,我
[{'generated_text': '第一次住在這裡兒,我 感 覺 很 溫 馨 。 房 間 寬 敞 、 幹 淨 還 有 水 果 送 ( 每 人 10 元 ) ; 飯 菜 也 不 錯 ! 價 格 合 理 經 濟 實 惠 .'}]
"""

我們需要注意的幾點:

  • 不同模型使用的tokenizer是不一樣的,需要注意它們的區別,尤其是pad_token_id和eos_token_id。eos_token_id常常用於標識生成文字的結尾。

  • 有一些中文的生成預訓練模型使用的還是Bert的tokenizer,在進行token化的時候,透過指定add_special_tokens=False來避免新增[CLS]和[SEP]。

  • BertTokenizer的eos_token_id為None,這裡我們用[PAD]視為生成結束的符號,其索引為0.當然,你也可以設定它為詞表裡面的特殊符號,比如[SEP]。

  • 對於不需要計算損失的token,我們將其標籤設定為-100。

  • 我們的labels和input_ids為什麼是一樣的,不是說根據上一個詞生成下一個詞嗎?這是因為模型裡面幫我們處理了,見程式碼:

shift_logits = lm_logits[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous()
  • 進行預測有三種方式,控制文字生成的多樣性有很多引數可以選擇,具體剛興趣可參考最後面的連結。

5使用BART進行對聯生成

資料從這裡下載:https://www.modelscope.cn/datasets/minisnow/couplet_samll.git

直接看程式碼:

import json
import pandas as pd
import numpy as np
# import lawrouge
from transformers import BertTokenizer, BartForConditionalGeneration, Text2TextGenerationPipeline, pipeline
from datasets import load_dataset, Dataset
from transformers import default_data_collator
import torch
from tqdm import tqdm
from datasets import load_dataset
from transformers import (
    default_data_collator,
    get_linear_schedule_with_warmup,
)
from torch.utils.data import DataLoader


# ============================= 
# 載入資料
train_path = "couplet_samll/train.csv"
train_dataset = Dataset.from_csv(train_path)
test_path = "couplet_samll/test.csv"
test_dataset = Dataset.from_csv(test_path)

max_len = 24
train_batch_size = 64
eval_batch_size = 64
lr = 3e-4
num_epochs = 1

# 轉換為模型需要的格式
def tokenize_dataset(tokenizer, dataset, max_len):
  def convert_to_features(batch):
    text1 = batch["text1"]
    text2 = batch["text2"]
    inputs = tokenizer.batch_encode_plus(
      text1,
      max_length=max_len,
      padding="max_length",
      truncation=True,
    )
    targets = tokenizer.batch_encode_plus(
      text2,
      max_length=max_len,
      padding="max_length",
      truncation=True,
    )
    outputs = {
      "input_ids": inputs["input_ids"],
      "attention_mask": inputs["attention_mask"],
      "target_ids": targets["input_ids"],
      "target_attention_mask": targets["attention_mask"]
    }
    return outputs
  
  dataset = dataset.map(convert_to_features, batched=True)
  # Set the tensor type and the columns which the dataset should return
  columns = ['input_ids''target_ids''attention_mask''target_attention_mask']
  dataset.with_format(type='torch', columns=columns)
  dataset = dataset.rename_column('target_ids''labels')
  dataset = dataset.rename_column('target_attention_mask''decoder_attention_mask')
  dataset = dataset.remove_columns(['text1''text2'])
  return dataset

tokenizer = BertTokenizer.from_pretrained("fnlp/bart-base-chinese")

train_data = tokenize_dataset(tokenizer, train_dataset, max_len)
test_data = tokenize_dataset(tokenizer, test_dataset, max_len)

train_dataset = train_data
train_dataloader = DataLoader(
    train_dataset, collate_fn=default_data_collator, shuffle=True, batch_size=train_batch_size, pin_memory=True
)
test_dataset = test_data
test_dataloader = DataLoader(
    test_dataset, collate_fn=default_data_collator, batch_size=eval_batch_size, pin_memory=True
)

# optimizer
model = BartForConditionalGeneration.from_pretrained("fnlp/bart-base-chinese")
optimizer = torch.optim.AdamW(model.parameters(), lr=lr)

# lr scheduler
lr_scheduler = get_linear_schedule_with_warmup(
    optimizer=optimizer,
    num_warmup_steps=0,
    num_training_steps=(len(train_dataloader) * num_epochs),
)

model.cuda()

from tqdm import tqdm

for epoch in range(num_epochs):
    model.train()
    total_loss = 0
    t = tqdm(train_dataloader)
    for step, batch in enumerate(t):
        for k, v in batch.items():
            batch[k] = v.cuda()
        outputs = model(**batch)
        loss = outputs.loss
        t.set_description("loss:{:.6f}".format(loss.item()))
        total_loss += loss.detach().float()
        loss.backward()
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()

    train_epoch_loss = total_loss / len(train_dataloader)
    model.save_pretrained("bart-couplet/")
    tokenizer.save_pretrained("bart-couplet/")
    print(f"epoch:{epoch+1}/{num_epochs} loss:{train_epoch_loss}")

結果:

loss:1.593506100%|██████████| 4595/4595 [33:28<00:00,  2.29it/s]
epoch:1/1 loss:1.76453697681427

我們可以這麼預測:

from transformers import Text2TextGenerationPipeline
model_path = "bart-couplet"
# model_path = "fnlp/bart-base-chinese"
model = BartForConditionalGeneration.from_pretrained(model_path)
tokenizer = BertTokenizer.from_pretrained(model_path)
generator = Text2TextGenerationPipeline(model=model, tokenizer=tokenizer)

max_len = 24

test_path = "couplet_samll/test.csv"
test_data = pd.read_csv(test_path)
texts = test_data["text1"].values.tolist()
labels = test_data["text2"].values.tolist()

results = generator(texts, max_length=max_len, eos_token_id=102, pad_token_id=0, do_sample=True)
for text, label, res in zip(texts, labels, results):
  print("上聯:", text)
  print("真實下聯:", label)
  print("預測下聯:""".join(res["generated_text"].split(" ")))
  print("="*100)
    
"""
上聯: 幾幀山水關秋路
真實下聯: 無奈胭脂點絳唇
預測下聯: 天高雲淡月光明
====================================================================================================
上聯: 許多心事懶收拾
真實下聯: 大好青春莫撂荒
預測下聯: 何妨明月照寒窗
====================================================================================================
上聯: 誰同執手人間老
真實下聯: 自願並肩化外遊
預測下聯: 心中有夢月當頭
====================================================================================================
上聯: 畫地為牢封自步
真實下聯: 齊天大聖悟空行
預測下聯: 不妨一世好清閒
====================================================================================================
上聯: 布穀攜春臨五嶽
真實下聯: 流鶯送喜到千家
預測下聯: 萬家燈火慶豐年
====================================================================================================
上聯: 冤家宜解不宜結
真實下聯: 窮寇定殲必定追
預測下聯: 不因風雨誤春秋
====================================================================================================
上聯: 汪倫情義人間少
真實下聯: 法律條文格外繁
預測下聯: 一江春水向東流
====================================================================================================
上聯: 潑墨吟詩,銀髮人生添雅興
真實下聯: 手機簡訊,古稀老叟逐新潮
預測下聯: 春風得意,萬裡千帆逐浪高
====================================================================================================
上聯: 刊岫展屏山,雲凝罨畫
真實下聯: 平湖環鏡檻,波漾空明
預測下聯: 千年古邑,百花芳草淹春
====================================================================================================
上聯: 且向人間賒一醉
真實下聯: 直如島外泛孤舟
預測下聯: 春風得意樂逍遙
====================================================================================================
"""

需要注意的地方:

  • 這裡我們的輸入不再是單條文字,而是文字對。
  • 我們需要構造編碼器Input_ids,編碼器attention_mask,解碼器input_ids,解碼器attention_mask。
  • 這裡使用了一個技巧:取樣生成,設定do_sample=True。如果你嘗試設定它為False,你會發現生成的結果可能不是那麼好。
  • 同樣的這裡使用的還是Bert的tokenizer,這裡進行tokenizer的時候我們保留了bert的[CLS]和[SEP]。為了更直觀的理解,我們使用另一種更直接的方法來生成結果:
model = BartForConditionalGeneration.from_pretrained(model_path)
model = model.to("cuda")
model.eval()
inputs = tokenizer(
        texts,
        padding="max_length",
        truncation=True,
        max_length=max_len,
        return_tensors="pt",
    )
input_ids = inputs.input_ids.to(model.device)
attention_mask = inputs.attention_mask.to(model.device)
# 生成
outputs = model.generate(input_ids, 
              attention_mask=attention_mask, 
              max_length=max_len, 
              do_sample=True
              pad_token_id=0,
              eos_token_id=102)
# 將token轉換為文字
output_str = tokenizer.batch_decode(outputs, skip_special_tokens=False)
output_str = [s.replace(" ",""for s in output_str]
for text, label, pred in zip(texts, labels, output_str):
  print("上聯:", text)
  print("真實下聯:", label)
  print("預測下聯:", pred)
  print("="*100)

結果:

上聯: 幾幀山水關秋路
真實下聯: 無奈胭脂點絳唇
預測下聯: [SEP][CLS]春風送暖柳含煙[SEP][PAD][PAD][PAD][PAD][PAD]
====================================================================================================
上聯: 許多心事懶收拾
真實下聯: 大好青春莫撂荒
預測下聯: [SEP][CLS]無私奉獻為人民[SEP][PAD][PAD][PAD][PAD][PAD]
====================================================================================================
上聯: 誰同執手人間老
真實下聯: 自願並肩化外遊
預測下聯: [SEP][CLS]清風明月是知音[SEP][PAD][PAD][PAD][PAD][PAD]
====================================================================================================
上聯: 畫地為牢封自步
真實下聯: 齊天大聖悟空行
預測下聯: [SEP][CLS]月明何處不相逢[SEP][PAD][PAD][PAD][PAD][PAD]
====================================================================================================
上聯: 布穀攜春臨五嶽
真實下聯: 流鶯送喜到千家
預測下聯: [SEP][CLS]一壺老酒醉春風[SEP][PAD][PAD][PAD][PAD][PAD]
====================================================================================================
上聯: 冤家宜解不宜結
真實下聯: 窮寇定殲必定追
預測下聯: [SEP][CLS]風流人物不虛名[SEP][PAD][PAD][PAD][PAD][PAD]
====================================================================================================
上聯: 汪倫情義人間少
真實下聯: 法律條文格外繁
預測下聯: [SEP][CLS]萬裡江山萬裡春[SEP][PAD][PAD][PAD][PAD][PAD]
====================================================================================================
上聯: 潑墨吟詩,銀髮人生添雅興
真實下聯: 手機簡訊,古稀老叟逐新潮
預測下聯: [SEP][CLS]和諧社會,和諧和諧幸福家[SEP]
====================================================================================================
上聯: 刊岫展屏山,雲凝罨畫
真實下聯: 平湖環鏡檻,波漾空明
預測下聯: [SEP][CLS]天下無雙,人壽年豐[SEP][PAD][PAD][PAD]
====================================================================================================
上聯: 且向人間賒一醉
真實下聯: 直如島外泛孤舟
預測下聯: [SEP][CLS]不知何處有閒人[SEP][PAD][PAD][PAD][PAD][PAD]
====================================================================================================
  • 我們設定skip_special_tokens=False,在生成時不忽略特殊token。
  • 以"無奈胭脂點絳唇"為例。輸入[SEP],預測得到[CLS],輸入[SEP][CLS]得到正常的文字,最後以[SEP]結尾。因為我們的encoder_input_ids和decoder_input_ids都是加了特殊符號的。當然你可以不加或者自定義使用其它的特殊符號。

到這裡,你已經瞭解了transformers庫中自帶的模型及相關的一些任務了,特別是針對生成模型有了更深一層的瞭解,趕緊去試試吧。

最後附上相關的一些知識:

https://zhuanlan.zhihu.com/p/624845975

Part3參考

transformers.models.auto.modeling_auto — transformers 4.4.2 documentation (huggingface.co)

相關文章