Upload model

config.json

@@ -0,0 +1,52 @@
+{
+  "activation_function": "gelu_new",
+  "architectures": [
+    "GPT2PrefixTuningWithLMHeadModel"
+  ],
+  "attn_pdrop": 0.1,
+  "auto_map": {
+    "AutoConfig": "gpt2.GPT2PrefixTuningConfig",
+    "AutoModelForCausalLM": "gpt2.GPT2PrefixTuningWithLMHeadModel"
+  },
+  "bos_token_id": 50256,
+  "embd_pdrop": 0.1,
+  "eos_token_id": 50256,
+  "initializer_range": 0.02,
+  "is_flat": false,
+  "layer_norm_epsilon": 1e-05,
+  "model_type": "gpt2",
+  "n_ctx": 1024,
+  "n_embd": 1024,
+  "n_head": 16,
+  "n_inner": null,
+  "n_layer": 24,
+  "n_positions": 1024,
+  "n_special": 0,
+  "objective_type": "sentence",
+  "pad_token_id": 50257,
+  "plm_name_or_path": "gpt2-medium",
+  "predict_special_tokens": true,
+  "prefix_dropout_prob": 0.0,
+  "prefix_hidden_size": 512,
+  "prefix_len": 5,
+  "reorder_and_upcast_attn": false,
+  "resid_pdrop": 0.1,
+  "scale_attn_by_inverse_layer_idx": false,
+  "scale_attn_weights": true,
+  "summary_activation": null,
+  "summary_first_dropout": 0.1,
+  "summary_proj_to_labels": true,
+  "summary_type": "cls_index",
+  "summary_use_proj": true,
+  "task_specific_params": {
+    "text-generation": {
+      "do_sample": true,
+      "max_length": 50
+    }
+  },
+  "torch_dtype": "float32",
+  "transformers_version": "4.26.0",
+  "use_cache": true,
+  "use_layer_dep": false,
+  "vocab_size": 50258
+}

generation_config.json

@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 50256,
+  "eos_token_id": 50256,
+  "pad_token_id": 50257,
+  "transformers_version": "4.26.0"
+}

gpt2.py

@@ -0,0 +1,209 @@
+import torch
+from torch import nn
+from transformers import PretrainedConfig, AutoConfig
+from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions
+from transformers.models.gpt2.modeling_gpt2 import GPT2PreTrainedModel, GPT2LMHeadModel
+
+from src.utils.prefix import PrefixEncoder
+
+class GPT2PrefixTuningConfig(PretrainedConfig):
+    attribute_map = {
+        "hidden_size": "n_embd",
+        "max_position_embeddings": "n_positions",
+        "num_attention_heads": "n_head",
+        "num_hidden_layers": "n_layer",
+    }
+    model_type = "gpt2"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(self, 
+        plm_name_or_path='gpt2-medium',
+        prefix_len=5,
+        prefix_dropout_prob=0.0,
+        prefix_hidden_size=512,
+        is_flat=False,
+        pad_token_id=50257,
+        objective_type='sentence',
+        use_layer_dep=False,
+        **kwargs):
+        super().__init__(**kwargs)
+        self.plm_name_or_path = plm_name_or_path
+        self.prefix_len = prefix_len
+        self.prefix_dropout_prob = prefix_dropout_prob
+        self.prefix_hidden_size = prefix_hidden_size
+        self.is_flat = is_flat
+        plm_config = AutoConfig.from_pretrained(plm_name_or_path).to_dict()
+        del plm_config['_name_or_path']
+        self.update(plm_config)
+        self.pad_token_id = pad_token_id
+        self.vocab_size = self.pad_token_id + 1 
+        self.objective_type = objective_type # or 'sentence' or 'token' which is the classical objective
+        self.use_layer_dep = use_layer_dep
+
+class GPT2PrefixTuningWithLMHeadModel(GPT2PreTrainedModel):
+    def __init__(self, config, pretrained_model=None):
+        super().__init__(config)
+        print(config)
+        if pretrained_model is None:
+            self.pretrained_model = GPT2LMHeadModel.from_pretrained(config.plm_name_or_path, pad_token_id=config.pad_token_id)
+            self.pretrained_model.resize_token_embeddings(config.vocab_size)
+        else:
+            self.pretrained_model = pretrained_model
+
+        for param in self.pretrained_model.parameters():
+            param.requires_grad = False
+
+        self.prefix_len = config.prefix_len
+        self.prefix_encoder = PrefixEncoder(config)
+
+    def train(self, mode=True):
+        super().train(mode)
+        self.pretrained_model.eval()
+    
+    def get_input_embeddings(self) -> nn.Module:
+        return self.pretrained_model.get_input_embeddings()
+
+    def get_output_embeddings(self):
+        return self.pretrained_model.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.pretrained_model.set_output_embeddings(new_embeddings=new_embeddings)
+    
+    def get_input_embeddings(self):
+        return self.pretrained_model.get_input_embeddings()
+
+    def set_input_embeddings(self, new_embeddings):
+        self.pretrained_model.set_input_embeddings(new_embeddings=new_embeddings)
+
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=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_key_values:
+            input_ids = input_ids[:, -1].unsqueeze(-1)
+            if token_type_ids is not None:
+                token_type_ids = token_type_ids[:, -1].unsqueeze(-1)
+        
+        batch_size = input_ids.shape[0]
+        attention_mask = kwargs.get("attention_mask", None)
+        position_ids = kwargs.get("position_ids", None)
+
+        if attention_mask is not None:
+            prefix_attention_mask = torch.ones(batch_size, self.prefix_len).to(input_ids.device)
+            attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1)
+
+        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_key_values:
+                position_ids = position_ids[:, -1].unsqueeze(-1)
+        else:
+            position_ids = None
+        
+        if past_key_values is None:
+            past_key_values = self.prefix_encoder(batch_size=batch_size)
+            if position_ids is not None:
+                position_ids = position_ids[:, self.prefix_len:]
+        
+        return {
+            "input_ids": input_ids,
+            "past_key_values": past_key_values,
+            "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,
+        past_key_values=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        labels=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):  
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if past_key_values is not None and self.training:
+            raise ValueError("past_key_value is dedicated to prefix tokens in this implementation. Please don't use it for anything else.")
+        
+        if past_key_values is None:
+            batch_size = input_ids.shape[0]
+            past_key_values = self.prefix_encoder(batch_size=batch_size)
+            if attention_mask is not None:
+                prefix_attention_mask = torch.ones(batch_size, self.prefix_len).to(input_ids.device)
+                attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1)
+
+        labels_for_plm = None if self.config.objective_type == 'sentence' else labels
+
+        position_ids = None if not self.training and input_ids.shape[1] == 1 else position_ids
+        if position_ids is not None:
+            position_ids = position_ids.contiguous()
+
+        transformer_outputs = self.pretrained_model(
+            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,
+            labels=labels_for_plm,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        
+        if labels_for_plm is None:
+            lm_logits = transformer_outputs.logits if return_dict else transformer_outputs[0]
+
+            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()
+                loss_fct = nn.CrossEntropyLoss(reduction='none')
+                batch_size, seqlen, _ = shift_logits.shape
+                loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
+                loss = loss.view(batch_size, seqlen).sum(dim=-1)
+                loss = loss.mean()
+
+            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,
+            )
+        else:
+            return transformer_outputs
+
+    @staticmethod
+    def _reorder_cache(past, beam_idx):
+        """
+        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
+        )

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ecf2cdfb5b0a3726e5eebf5461a65ba9adc54ae6ca427a918f7dda0624d6c3ec
+size 1547553167