Upload 6 files

models/instruction_prompting/causal_incontext.py

@@ -0,0 +1,193 @@
+import sys
+sys.path.append("./")
+sys.path.append("../")
+sys.path.append("../../")
+sys.path.append("../../../")
+import torch
+import torch.nn as nn
+import transformers
+import torch.nn.functional as F
+from typing import Optional, Tuple, Union
+from torch.nn import CrossEntropyLoss
+from transformers import AutoModelForCausalLM
+from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions
+from transformers.models.gpt2.modeling_gpt2 import GPT2PreTrainedModel, GPT2Model, GPT2LMHeadModel
+from tools.runner_utils.log_util import logging
+from tools.model_utils.parameter_freeze import ParameterFreeze
+
+logger = logging.getLogger(__name__)
+freezer = ParameterFreeze()
+
+"""
+Function: Use Causal LM to prompt for cls
+Notes:
+- For classification, the model only calculate the loss at the position of label, the other position is set as -100
+- During inference, generate result at the last position.
+"""
+
+
+class GPT2ForInContextLearning(GPT2PreTrainedModel):
+    _keys_to_ignore_on_load_missing = [r"attn.masked_bias", r"attn.bias", r"lm_head.weight"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.transformer = GPT2Model(config)
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+
+        # if self.config.use_freezing:
+        #     self.transformer = freezer.freeze_lm(self.transformer)
+
+        # Model parallel
+        self.model_parallel = False
+        self.device_map = None
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def prepare_inputs_for_generation(self, input_ids, past=None, **kwargs):
+        token_type_ids = kwargs.get("token_type_ids", None)
+        # only last token for inputs_ids if past is defined in kwargs
+        if past:
+            input_ids = input_ids[:, -1].unsqueeze(-1)
+            if token_type_ids is not None:
+                token_type_ids = token_type_ids[:, -1].unsqueeze(-1)
+
+        attention_mask = kwargs.get("attention_mask", None)
+        position_ids = kwargs.get("position_ids", None)
+
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past:
+                position_ids = position_ids[:, -1].unsqueeze(-1)
+        else:
+            position_ids = None
+        return {
+            "input_ids": input_ids,
+            "past_key_values": past,
+            "use_cache": kwargs.get("use_cache"),
+            "position_ids": position_ids,
+            "attention_mask": attention_mask,
+            "token_type_ids": token_type_ids,
+        }
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        token_type_ids: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        head_mask: Optional[torch.FloatTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        encoder_attention_mask: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, CausalLMOutputWithCrossAttentions]:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
+            `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
+            are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        transformer_outputs = self.transformer(
+            input_ids,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        hidden_states = transformer_outputs[0]
+
+        # Set device for model parallelism
+        if self.model_parallel:
+            torch.cuda.set_device(self.transformer.first_device)
+            hidden_states = hidden_states.to(self.lm_head.weight.device)
+
+        lm_logits = self.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = lm_logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # print("shift_labels=", shift_labels)
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
+
+        if not return_dict:
+            output = (lm_logits,) + transformer_outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+
+        return CausalLMOutputWithCrossAttentions(
+            loss=loss,
+            logits=lm_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,
+        )
+
+
+    @staticmethod
+    def _reorder_cache(past: Tuple[Tuple[torch.Tensor]], beam_idx: torch.Tensor) -> Tuple[Tuple[torch.Tensor]]:
+        """
+        This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or
+        [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct
+        beam_idx at every generation step.
+        """
+        return tuple(
+            tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past)
+            for layer_past in past
+        )
+
+
+
+
+if __name__ == "__main__":
+    from transformers import GPT2Tokenizer
+    tokenizer = GPT2Tokenizer.from_pretrained("/Users/wangjianing/Desktop/开源代码与数据模型/模型/gpt2")
+    model = GPT2ForInContextLearning.from_pretrained("/Users/wangjianing/Desktop/开源代码与数据模型/模型/gpt2")
+
+    # In-Context Learning for classification
+    # input_text = "The capital city of China is Beijing. \n\n The capital city of Japan is Tokyo. \n\n The capital city of America is"
+    input_text = "What are follows emotions? \n\n Input: The book is very nice.\n Output: Great. \n\n Input: I never eat chocolate!\n Output:"
+    # input_text = "This film is wonderful.\n Great."
+    tokenizer.pad_token = tokenizer.eos_token
+    inputs = tokenizer(input_text, return_tensors="pt")
+    input_len = inputs["input_ids"].shape[-1]
+    gen_output = model.generate(**inputs, max_length=input_len + 10)
+    gen_result = tokenizer.decode(gen_output[0])
+    print("classification result:\n", gen_result)
+
+    # In-Context Learning for generation
+    input_text = "Please tell me what is the transformer? "
+    # input_text = "This film is wonderful.\n Great."
+    tokenizer.pad_token = tokenizer.eos_token
+    inputs = tokenizer(input_text, return_tensors="pt")
+    input_len = inputs["input_ids"].shape[-1]
+    gen_output = model.generate(**inputs, max_length=input_len + 60)
+    gen_result = tokenizer.decode(gen_output[0])
+    print("generation result:\n", gen_result)

models/instruction_prompting/incontext.py

@@ -0,0 +1,127 @@
+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"]):])

models/instruction_prompting/test.py

@@ -0,0 +1,30 @@
+from transformers import AutoModelForCausalLM
+from transformers import GPT2Tokenizer
+from transformers.models.gpt2.modeling_gpt2 import GPT2Model, GPT2LMHeadModel
+
+if __name__ == "__main__":
+    gpt2_tokenizer: GPT2Tokenizer = GPT2Tokenizer.from_pretrained("/Users/wangjianing/Desktop/开源代码与数据模型/模型/gpt2")
+    # gpt2_model = GPT2LMHeadModel.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_text = "What are follows emotions? \n\n The book is very nice.\n great. \n\n I never eat chocolate!\n bad. \n\n This film is wonderful.\n Great"
+    # # input_text = "Mr. Chen was born in Shanghai. Obama was born in US. Trump was born in"
+    # inputs = gpt2_tokenizer(input_text, return_tensors="pt")
+    # print(inputs)
+    # output = gpt2_model(**inputs)
+    # # print(output["last_hidden_state"])
+    # # print(output["last_hidden_state"].size())
+    # print(output["logits"])
+    # print(output["logits"].size())
+    # gen_output = gpt2_model.generate(**inputs, max_length=60)
+    # # gen_result = gpt2_tokenizer.convert_ids_to_tokens(gen_output[0])
+    # gen_result = gpt2_tokenizer.decode(gen_output[0])
+    # print(gen_result)
+
+
+    gpt2_tokenizer(
+                [["What are follows emotions?", "What are follows emotions?"], ["What are follows emotions?"]],
+                truncation=True,
+                max_length=30,
+                padding="max_length",
+                return_offsets_mapping=True
+            )

models/reinforcement_learning/actor.py

@@ -0,0 +1,103 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2023/5/6 3:53 p.m.
+# @Author  : JianingWang
+# @File    : actor.py
+
+from typing import Optional, Tuple, Union
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers import AutoModelForCausalLM, AutoConfig
+from models.basic_modules.generation import generate
+
+
+"""
+Actor model.
+"""
+class Actor(nn.Module):
+    """
+    Actor model base class.
+
+    Args:
+        model (nn.Module): Actor Model.
+    """
+
+    def __init__(self, model: nn.Module) -> None:
+        self.model = model
+    
+    def log_probs_from_logits(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
+        log_probs = F.log_softmax(logits, dim=-1)
+        log_probs_labels = log_probs.gather(dim=-1, index=labels.unsqueeze(-1))
+        return log_probs_labels.squeeze(-1)
+
+    """
+    For generative model, needs generate function.
+    """
+    @torch.no_grad()
+    def generate(
+        self,
+        input_ids: torch.Tensor,
+        return_action_mask: bool = True,
+        **kwargs
+    ) -> Union[Tuple[torch.LongTensor, torch.LongTensor], Tuple[torch.LongTensor, torch.LongTensor, torch.BoolTensor]]:
+        sequences = generate(self.model, input_ids, **kwargs)
+        attention_mask = None
+        pad_token_id = kwargs.get('pad_token_id', None)
+        if pad_token_id is not None:
+            attention_mask = sequences.not_equal(pad_token_id).to(dtype=torch.long, device=sequences.device)
+        if not return_action_mask:
+            return sequences, attention_mask, None
+        input_len = input_ids.size(1)
+        eos_token_id = kwargs.get('eos_token_id', None)
+        if eos_token_id is None:
+            action_mask = torch.ones_like(sequences, dtype=torch.bool)
+        else:
+            # left padding may be applied, only mask action
+            action_mask = (sequences[:, input_len:] == eos_token_id).cumsum(dim=-1) == 0
+            action_mask = F.pad(action_mask, (1 + input_len, -1), value=True)    # include eos token and input
+        action_mask[:, :input_len] = False
+        action_mask = action_mask[:, 1:]
+        return sequences, attention_mask, action_mask[:, -(sequences.size(1) - input_len):]
+
+    def forward(self,
+                sequences: torch.LongTensor,
+                num_actions: int,
+                attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """Returns action log probs
+        """
+        output = self.model(sequences, attention_mask=attention_mask)
+        logits = output['logits']
+        log_probs = self.log_probs_from_logits(logits[:, :-1, :], sequences[:, 1:])
+        return log_probs[:, -num_actions:]
+
+    def get_base_model(self):
+        return self.model
+
+
+"""
+Causal LM as a actor, e.g., GPT-2, OPT, BLOOM, etc.
+"""
+class CausalActor(Actor):
+    """
+    Causal LM Actor model.
+
+    Args:
+        pretrained (str): Pretrained model name or path.
+        config (AutoConfig): Model config.
+        checkpoint (bool): Enable gradient checkpointing.
+    """
+
+    def __init__(self,
+                 pretrained: str = None,
+                 config: Optional[AutoConfig] = None,
+                 checkpoint: bool = False) -> None:
+        if pretrained is not None:
+            model = AutoModelForCausalLM.from_pretrained(pretrained)
+        elif config is not None:
+            model = AutoModelForCausalLM(config)
+        else:
+            model = AutoModelForCausalLM(AutoConfig())
+        if checkpoint:
+            model.gradient_checkpointing_enable()
+        super().__init__(model)

models/reinforcement_learning/critic.py

@@ -0,0 +1,96 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2023/5/6 4:12 p.m.
+# @Author  : JianingWang
+# @File    : critic.py
+
+from typing import Optional, Tuple, Union
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers import AutoModel, AutoConfig
+from models.basic_modules.generation import generate
+
+
+def masked_mean(tensor: torch.Tensor, mask: torch.Tensor, dim: int = 1) -> torch.Tensor:
+    tensor = tensor * mask
+    tensor = tensor.sum(dim=dim)
+    mask_sum = mask.sum(dim=dim)
+    mean = tensor / (mask_sum + 1e-8)
+    return mean
+
+
+"""
+Critic model.
+"""
+class Critic(nn.Module):
+    """
+    Critic model base class.
+
+    Args:
+        model (nn.Module): Critic model.
+        value_head (nn.Module): Value head to get value.
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        value_head: nn.Module,
+        use_action_mask: bool = False,
+    ) -> None:
+
+        self.model = model
+        self.value_head = value_head # critic layer for predict value function
+        self.use_action_mask = use_action_mask
+
+    def forward(self,
+                sequences: torch.LongTensor,
+                action_mask: Optional[torch.Tensor] = None,
+                attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        outputs = self.model(sequences, attention_mask=attention_mask)
+        last_hidden_states = outputs['last_hidden_state']
+
+        values = self.value_head(last_hidden_states).squeeze(-1)
+
+        if action_mask is not None and self.use_action_mask:
+            num_actions = action_mask.size(1)
+            prompt_mask = attention_mask[:, :-num_actions]
+            values = values[:, :-num_actions]
+            value = masked_mean(values, prompt_mask, dim=1)
+            return value
+
+        values = values[:, :-1]
+        value = values.mean(dim=1)
+        return value
+
+
+"""
+Auto Model for Critic
+"""
+class AutoModelCritic(Critic):
+    """
+    AutoModel Critic model.
+
+    Args:
+        pretrained (str): Pretrained model name or path.
+        config (AutoConfig): Model config.
+        checkpoint (bool): Enable gradient checkpointing.
+    """
+
+    def __init__(self,
+                 pretrained: Optional[str] = None,
+                 config: Optional[AutoConfig] = None,
+                 checkpoint: bool = False,
+                 lora_rank: int = 0,
+                 lora_train_bias: str = 'none',
+                 **kwargs) -> None:
+        if pretrained is not None:
+            model = AutoModel.from_pretrained(pretrained)
+        elif config is not None:
+            model = AutoModel(config)
+        else:
+            model = AutoModel(AutoConfig())
+        if checkpoint:
+            model.gradient_checkpointing_enable()
+        value_head = nn.Linear(model.config.word_embed_proj_dim, 1)
+        super().__init__(model, value_head, **kwargs)

models/reinforcement_learning/reward_model.py

@@ -0,0 +1,131 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2023/5/6 4:29 p.m.
+# @Author  : JianingWang
+# @File    : reward_model.py
+
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+from transformers import AutoModel, AutoConfig
+from loss.rl_loss import LogSigLoss, LogExpLoss
+from transformers.models.roberta.modeling_roberta import RobertaPreTrainedModel, RobertaModel
+from transformers.models.gpt2.modeling_gpt2 import GPT2PreTrainedModel, GPT2Model
+
+"""
+RoERTa for Reward Model
+"""
+class RobertaForReward(RobertaPreTrainedModel):
+    """
+    Reward model base class.
+
+    Args:
+        model (nn.Module): Reward model.
+        value_head (nn.Module): Value head to get reward score.
+    """
+
+    def __init__(self, config) -> None:
+        super().__init__(config)
+        self.config = config
+        self.roberta = RobertaModel(config)
+        self.value_head = nn.Linear(self.config.n_embd, 1)
+        self.init_weights()
+
+    def forward(
+            self, 
+            chosen_sequences: torch.LongTensor, 
+            chosen_attention_mask: Optional[torch.Tensor],
+            rejected_sequences: Optional[torch.LongTensor] = None,
+            rejected_attention_mask: Optional[torch.Tensor] = None,
+        ) -> torch.Tensor:
+        # obtain reward value of chosen sequence
+        chosen_outputs = self.roberta(chosen_sequences, attention_mask=chosen_attention_mask)
+        chosen_last_hidden_states = chosen_outputs['last_hidden_state']
+        chosen_values = self.value_head(chosen_last_hidden_states)[:, :-1]
+        chosen_values = chosen_values.mean(dim=1).squeeze(1)    # ensure shape is (B)
+
+        return_dict = {
+            "chosen_values": chosen_values,
+        }
+        # if has rejected, obtain reward of rejected sequence, and calculate the loss
+        if rejected_sequences is not None:
+            rejected_outputs = self.roberta(rejected_sequences, attention_mask=rejected_attention_mask)
+            rejected_last_hidden_states = rejected_outputs['last_hidden_state']
+            rejected_values = self.value_head(rejected_last_hidden_states)[:, :-1]
+            rejected_values = rejected_values.mean(dim=1).squeeze(1)    # ensure shape is (B)
+            return_dict["rejected_values"] = rejected_values
+            
+            loss_fn = LogSigLoss()
+            loss = loss_fn(chosen_values, rejected_values)
+            
+            return_dict["loss"] = loss
+        
+        return return_dict
+
+
+"""
+GPT2 for Reward Model
+"""
+class GPT2ForReward(GPT2PreTrainedModel):
+    _keys_to_ignore_on_load_missing = [r"attn.masked_bias", r"attn.bias", r"lm_head.weight"]
+    """
+    Reward model base class.
+
+    Args:
+        model (nn.Module): Reward model.
+        value_head (nn.Module): Value head to get reward score.
+    """
+
+    def __init__(self, config) -> None:
+        super().__init__(config)
+        self.config = config
+        self.transformer = GPT2Model(config)
+        self.value_head = nn.Linear(self.config.n_embd, 1)
+
+        # Model parallel
+        self.model_parallel = False
+        self.device_map = None
+
+        self.post_init()
+
+    def forward(
+            self, 
+            chosen_sequences: torch.LongTensor, 
+            chosen_attention_mask: Optional[torch.Tensor],
+            rejected_sequences: Optional[torch.LongTensor] = None,
+            rejected_attention_mask: Optional[torch.Tensor] = None,
+        ) -> torch.Tensor:
+        # obtain reward value of chosen sequence
+        chosen_outputs = self.transformer(chosen_sequences, attention_mask=chosen_attention_mask)
+        chosen_last_hidden_states = chosen_outputs['last_hidden_state']
+        chosen_values = self.value_head(chosen_last_hidden_states)[:, :-1]
+        chosen_values = chosen_values.mean(dim=1).squeeze(1)    # ensure shape is (B)
+
+        return_dict = {
+            "chosen_values": chosen_values,
+        }
+        # if has rejected, obtain reward of rejected sequence, and calculate the loss
+        if rejected_sequences is not None:
+            rejected_outputs = self.transformer(rejected_sequences, attention_mask=rejected_attention_mask)
+            rejected_last_hidden_states = rejected_outputs['last_hidden_state']
+            rejected_values = self.value_head(rejected_last_hidden_states)[:, :-1]
+            rejected_values = rejected_values.mean(dim=1).squeeze(1)    # ensure shape is (B)
+            return_dict["rejected_values"] = rejected_values
+            loss_fn = LogSigLoss()
+            loss = loss_fn(chosen_values, rejected_values)
+            
+            return_dict["loss"] = loss
+        
+        return return_dict
+    
+    @staticmethod
+    def _reorder_cache(past: Tuple[Tuple[torch.Tensor]], beam_idx: torch.Tensor) -> Tuple[Tuple[torch.Tensor]]:
+        """
+        This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or
+        [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct
+        beam_idx at every generation step.
+        """
+        return tuple(
+            tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past)
+            for layer_past in past
+        )