models/instruction_prompting/incontext.py

import torch
from tqdm import tqdm
from typing import Optional, Tuple
from turtle import forward
from torch.nn import CrossEntropyLoss
from transformers import AutoModelForCausalLM
from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions
from transformers.models.gpt2.modeling_gpt2 import GPT2LMHeadModel, GPT2Model


class GPT2ForInContextClassification(GPT2LMHeadModel):

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None, # input token id
        past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        token_type_ids: Optional[torch.LongTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        label_masks: Optional[torch.LongTensor] = None, # mask=1 means it should be calculated loss
        options :Optional[list] = None, # 如果是分类任务，则可以添加候选label
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
    ):
        assert len(input_ids.shape) == 3 and input_ids.shape[1] == len(options) # [n, option_size, len]
        batch_size = input_ids.shape[0]
        option_size = input_ids.shape[1]
        input_ids = input_ids.view(-1, input_ids.shape[1], input_ids.shape[2]) # [n*option_size, len]
        attention_mask = attention_mask.view(-1, input_ids.shape[1], input_ids.shape[2]) if attention_mask is not None else None # [n*option_size, len]
        token_type_ids = token_type_ids.view(-1, input_ids.shape[1], input_ids.shape[2]) if token_type_ids is not None else None# [n*option_size, len]
        # labels = labels.view(-1, input_ids.shape[1], input_ids.shape[2]) # [n*option_size, len]

        transformer_outputs = self.transformer(
            input_ids,
            past_key_values=past_key_values,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        hidden_states = transformer_outputs[0] # [n*option_size, len, hidden_size]
        lm_logits = self.lm_head(hidden_states) # [n*option_size, len, vocab_size]
        lm_logits = lm_logits.view(batch_size, option_size, input_ids.shape[-1], -1) # [n, option_size, len, vocab_size]

        # print("len(input_ids)=", len(input_ids[0]))
        # print("input_ids[-1]=", input_ids[0][-1])
        print("lm_logits.shape=", lm_logits.shape)

        losses = list()
        if labels is not None:
            for label, lm_logit in zip(labels, lm_logits):
                # label: [option_size, len]
                # lm_logit: [option_size, len, vocab_size]
                shift_logits = lm_logit[..., :-1, :].contiguous()
                # print("shift_logits.shape=", shift_logits.shape)
                shift_labels = label[..., 1:].contiguous()
                # print("shift_labels=", shift_labels)
                # print("shift_labels.shape=", shift_labels.shape)
                # Flatten the tokens
                loss_fct = CrossEntropyLoss()
                print("shift_logits.shape=", shift_logits.shape)
                print("shift_labels.shape=", shift_labels.shape)
                loss = [loss_fct(shift_logit.view(-1, shift_logit.size(-1)), shift_label.view(-1)) for shift_logit, shift_label in zip(shift_logits, shift_labels)]
                loss = torch.stack(loss)
                # print("loss=", loss)
                if label_masks is not None:
                    loss = loss.view(lm_logits.size(0), lm_logits.size(1)) * label_masks # [option_size, len]
                    loss = torch.sum(loss, axis=1) / torch.sum(label_mask, axis=1) # [option_size]
                losses.append(loss)
        losses = torch.stack(losses) # [n, option_size]
        # 将各个option的loss视为logit，loss越小，对应的概率应越大
        loss_logits = torch.softmax(-losses, -1) # [n, option_size]
        print("losses.shape=", losses.shape)
        print("loss_logits.shape=", loss_logits.shape)

        if not return_dict:
            output = (lm_logits,) + transformer_outputs[1:]
            return ((loss,) + output) if loss is not None else output

        return CausalLMOutputWithCrossAttentions(
            loss=losses,
            logits=loss_logits,
            past_key_values=transformer_outputs.past_key_values,
            hidden_states=transformer_outputs.hidden_states,
            attentions=transformer_outputs.attentions,
            cross_attentions=transformer_outputs.cross_attentions,
        )

if __name__ == "__main__":
    from transformers import GPT2Tokenizer
    tokenizer = GPT2Tokenizer.from_pretrained("/Users/wangjianing/Desktop/开源代码与数据模型/模型/gpt2")
    model = GPT2ForInContextClassification.from_pretrained("/Users/wangjianing/Desktop/开源代码与数据模型/模型/gpt2")
    # input_text = "The capital city of China is Beijing. The capital city of Japan is Tokyo. The capital city of America"
    input_text1 = "What are follows emotions? \n\n Input: The book is very nice.\n Output: Great. \n\n Input: I never eat chocolate!\n Output: Bad. \n\n Input: This film is not wonderful.\n Output: Great"
    input_text2 = "What are follows emotions? \n\n Input: The book is very nice.\n Output: Great. \n\n Input: I never eat chocolate!\n Output: Bad. \n\n Input: This film is not wonderful.\n Output: Bad"
    # input_text = "This film is wonderful.\n Great."
    # input_text = "Mr. Chen was born in Shanghai. Obama was born in US. Jinping Xi was born in China."
    tokenizer.pad_token = tokenizer.eos_token
    inputs = tokenizer(
        [input_text1, input_text2], return_tensors="pt",
        max_length=60,
        padding="max_length")
    inputs["input_ids"] = inputs["input_ids"].view(-1, inputs["input_ids"].shape[0], inputs["input_ids"].shape[1])
    # inputs["token_type_ids"] = inputs["token_type_ids"].view(-1, inputs["input_ids"].shape[0], inputs["input_ids"].shape[1])
    inputs["attention_mask"] = inputs["attention_mask"].view(-1, inputs["input_ids"].shape[0], inputs["input_ids"].shape[1])
    inputs["labels"] = inputs["input_ids"]
    inputs["options"] = torch.Tensor([[0, 1], [0, 1]]).long()
    print(inputs["input_ids"].shape)
    label_mask = torch.zeros([1, 2, inputs["input_ids"].shape[2]])
    # print(label_mask)
    label_mask[0][0][20] = 1
    label_mask[0][1][20] = 1
    print(label_mask)
    output = model(**inputs, return_dict=True)
    # print(output["last_hidden_state"])
    # print(output["last_hidden_state"].size())
    # print(output["logits"])
    # print(output["logits"].size())
    losses, logits = output["loss"], output["logits"]
    print("loss=", losses)
    print("logits=", logits)
    # gen_output = model.generate(**inputs, max_length=60)
    # for i in range(len(gen_output)):
    #     gen_result = tokenizer.decode(gen_output[i])
    #     print("gen_result=", gen_result[len(inputs["input_ids"]):])