import inspect import warnings import torch from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss from transformers import PreTrainedModel from transformers.modeling_outputs import SequenceClassifierOutput from .tuners import LoRAModel, PrefixEncoder, PromptEmbedding, PromptEncoder from .utils import PETConfig, PETType, TaskType, shift_tokens_right class PETModel(torch.nn.Module): def __init__(self, model, pet_config: PETConfig): super().__init__() self.pet_config = pet_config self.base_model = model if pet_config.pet_type != PETType.LORA: self._setup_prompt_encoder() else: self.base_model = LoRAModel(pet_config, model) def _setup_prompt_encoder(self): num_transformer_submodules = 0 transformer_backbone = None for name, module in self.base_model.named_children(): if isinstance(module, PreTrainedModel): # Make sure to freeze Tranformers model for param in module.parameters(): param.requires_grad = False if transformer_backbone is None: transformer_backbone = module self.transformer_backbone_name = name num_transformer_submodules += 1 self.pet_config.num_transformer_submodules = 2 if self.pet_config.task_type == TaskType.SEQ_2_SEQ_LM else 1 for named_param, value in list(transformer_backbone.named_parameters()): if value.shape[0] == self.base_model.config.vocab_size: self.word_embeddings = transformer_backbone.get_submodule(named_param.replace(".weight", "")) break if self.pet_config.pet_type == PETType.PROMPT_TUNING: prompt_encoder = PromptEmbedding(self.pet_config, self.word_embeddings) elif self.pet_config.pet_type == PETType.P_TUNING: prompt_encoder = PromptEncoder(self.pet_config) elif self.pet_config.pet_type == PETType.PREFIX_TUNING: prompt_encoder = PrefixEncoder(self.pet_config) else: raise ValueError("Not supported") self.prompt_encoder = prompt_encoder self.prompt_tokens = torch.arange( self.pet_config.num_virtual_tokens * self.pet_config.num_transformer_submodules ).long() def get_prompt(self, batch_size): prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(batch_size, -1).to(self.base_model.device) if self.pet_config.pet_type == PETType.PREFIX_TUNING: prompt_tokens = prompt_tokens[:, : self.pet_config.num_virtual_tokens] if self.pet_config.inference_mode: past_key_values = self.prompt_encoder.embedding.weight.repeat(batch_size, 1, 1) else: past_key_values = self.prompt_encoder(prompt_tokens) past_key_values = past_key_values.view( batch_size, self.pet_config.num_virtual_tokens, self.pet_config.num_layers * 2, self.pet_config.num_attention_heads, self.pet_config.token_dim // self.pet_config.num_attention_heads, ) if self.pet_config.num_transformer_submodules == 2: past_key_values = torch.cat([past_key_values, past_key_values], dim=2) past_key_values = past_key_values.permute([2, 0, 3, 1, 4]).split( self.pet_config.num_transformer_submodules * 2 ) if self.pet_config.postprocess_past_key_value_function is not None: post_process_fn = self.pet_config.postprocess_past_key_value_function past_key_values = post_process_fn(past_key_values) return past_key_values else: if self.pet_config.inference_mode: prompts = self.prompt_encoder.embedding.weight.repeat(batch_size, 1, 1) else: prompts = self.prompt_encoder(prompt_tokens) return prompts def print_trainable_parameters(self): trainable_params = 0 all_param = 0 for _, param in self.named_parameters(): all_param += param.numel() if param.requires_grad: trainable_params += param.numel() print( f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}" ) class PETModelForSequenceClassification(PETModel): def __init__(self, model, pet_config: PETConfig): super().__init__(model, pet_config) self.config = self.base_model.config for name, module in self.base_model.named_children(): if isinstance(module, torch.nn.Linear): self.cls_layer_name = name break def forward( self, input_ids=None, attention_mask=None, inputs_embeds=None, labels=None, output_attentions=None, output_hidden_states=None, return_dict=None, **kwargs, ): return_dict = return_dict if return_dict is not None else self.config.use_return_dict if self.pet_config.pet_type == PETType.LORA: return self.base_model( input_ids=input_ids, attention_mask=attention_mask, inputs_embeds=inputs_embeds, labels=labels, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, **kwargs, ) batch_size = input_ids.shape[0] if attention_mask is not None: # concat prompt attention mask prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to( self.base_model.device ) attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1) if kwargs.get("position_ids", None) is not None: warnings.warn("Position ids are not supported for parameter efficient tuning. Ignoring position ids.") kwargs["position_ids"] = None kwargs.update( { "attention_mask": attention_mask, "labels": labels, "output_attentions": output_attentions, "output_hidden_states": output_hidden_states, "return_dict": return_dict, } ) if self.pet_config.pet_type == PETType.PREFIX_TUNING: return self.prefix_tuning_forward(input_ids=input_ids, **kwargs) else: if kwargs.get("token_type_ids", None) is not None: kwargs["token_type_ids"] = torch.cat( ( torch.zeros(batch_size, self.pet_config.num_virtual_tokens).to(self.base_model.device), kwargs["token_type_ids"], ), dim=1, ).long() if inputs_embeds is None: inputs_embeds = self.word_embeddings(input_ids) prompts = self.get_prompt(batch_size=batch_size) inputs_embeds = torch.cat((prompts, inputs_embeds), dim=1) return self.base_model(inputs_embeds=inputs_embeds, **kwargs) def prefix_tuning_forward( self, input_ids=None, attention_mask=None, inputs_embeds=None, labels=None, output_attentions=None, output_hidden_states=None, return_dict=None, **kwargs, ): batch_size = input_ids.shape[0] past_key_values = self.get_prompt(batch_size) fwd_params = list(inspect.signature(self.base_model.forward).parameters.keys()) kwargs.update( { "input_ids": input_ids, "attention_mask": attention_mask, "inputs_embeds": inputs_embeds, "output_attentions": output_attentions, "output_hidden_states": output_hidden_states, "return_dict": return_dict, "past_key_values": past_key_values, } ) if "past_key_values" in fwd_params: return self.base_model(labels=labels, **kwargs) else: transformer_backbone_name = self.base_model.get_submodule(self.transformer_backbone_name) fwd_params = list(inspect.signature(transformer_backbone_name.forward).parameters.keys()) if "past_key_values" not in fwd_params: raise ValueError("Model does not support past key values which are required for prefix tuning.") outputs = transformer_backbone_name(**kwargs) pooled_output = outputs[1] if len(outputs) > 1 else outputs[0] if "dropout" in [name for name, _ in list(self.base_model.named_children())]: pooled_output = self.base_model.dropout(pooled_output) logits = self.base_model.get_submodule(self.cls_layer_name)(pooled_output) loss = None if labels is not None: if self.config.problem_type is None: if self.base_model.num_labels == 1: self.config.problem_type = "regression" elif self.base_model.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.base_model.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.base_model.num_labels), labels.view(-1)) elif self.config.problem_type == "multi_label_classification": loss_fct = BCEWithLogitsLoss() loss = loss_fct(logits, labels) if not return_dict: output = (logits,) + outputs[2:] return ((loss,) + output) if loss is not None else output return SequenceClassifierOutput( loss=loss, logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions, ) class PETModelForCausalLM(PETModel): def __init__(self, model, pet_config: PETConfig): super().__init__(model, pet_config) self.config = self.base_model.config def forward( self, input_ids=None, attention_mask=None, inputs_embeds=None, labels=None, output_attentions=None, output_hidden_states=None, return_dict=None, **kwargs, ): if self.pet_config.pet_type == PETType.LORA: return self.base_model( input_ids=input_ids, attention_mask=attention_mask, inputs_embeds=inputs_embeds, labels=labels, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, **kwargs, ) batch_size = input_ids.shape[0] if attention_mask is not None: # concat prompt attention mask prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to( self.base_model.device ) attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1) if kwargs.get("position_ids", None) is not None: warnings.warn("Position ids are not supported for parameter efficient tuning. Ignoring position ids.") kwargs["position_ids"] = None if kwargs.get("token_type_ids", None) is not None: warnings.warn("Token type ids are not supported for parameter efficient tuning. Ignoring token type ids") kwargs["token_type_ids"] = None kwargs.update( { "attention_mask": attention_mask, "labels": labels, "output_attentions": output_attentions, "output_hidden_states": output_hidden_states, "return_dict": return_dict, } ) if self.pet_config.pet_type == PETType.PREFIX_TUNING: past_key_values = self.get_prompt(batch_size) return self.base_model(input_ids=input_ids, past_key_values=past_key_values, **kwargs) else: if inputs_embeds is None: inputs_embeds = self.word_embeddings(input_ids) # concat prompt labels if labels is not None: prefix_labels = torch.full((batch_size, self.pet_config.num_virtual_tokens), -100).to( self.base_model.device ) kwargs["labels"] = torch.cat((prefix_labels, labels), dim=1) prompts = self.get_prompt(batch_size=batch_size) inputs_embeds = torch.cat((prompts, inputs_embeds), dim=1) return self.base_model(inputs_embeds=inputs_embeds, **kwargs) class PETModelForSeq2SeqLM(PETModel): def __init__(self, model, pet_config: PETConfig): super().__init__(model, pet_config) self.config = self.base_model.config def forward( self, input_ids=None, attention_mask=None, inputs_embeds=None, decoder_input_ids=None, decoder_attention_mask=None, decoder_inputs_embeds=None, labels=None, output_attentions=None, output_hidden_states=None, return_dict=None, **kwargs, ): if self.pet_config.pet_type == PETType.LORA: return self.base_model( input_ids=input_ids, attention_mask=attention_mask, inputs_embeds=inputs_embeds, decoder_input_ids=decoder_input_ids, decoder_attention_mask=decoder_attention_mask, decoder_inputs_embeds=decoder_inputs_embeds, labels=labels, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, **kwargs, ) batch_size = input_ids.shape[0] if decoder_attention_mask is not None: # concat prompt attention mask prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to( self.base_model.device ) decoder_attention_mask = torch.cat((prefix_attention_mask, decoder_attention_mask), dim=1) if kwargs.get("position_ids", None) is not None: warnings.warn("Position ids are not supported for parameter efficient tuning. Ignoring position ids.") kwargs["position_ids"] = None if kwargs.get("token_type_ids", None) is not None: warnings.warn("Token type ids are not supported for parameter efficient tuning. Ignoring token type ids") kwargs["token_type_ids"] = None kwargs.update( { "attention_mask": attention_mask, "decoder_attention_mask": decoder_attention_mask, "labels": labels, "output_attentions": output_attentions, "output_hidden_states": output_hidden_states, "return_dict": return_dict, } ) if self.pet_config.pet_type == PETType.PREFIX_TUNING: past_key_values = self.get_prompt(batch_size) return self.base_model( input_ids=input_ids, decoder_input_ids=decoder_input_ids, past_key_values=past_key_values, **kwargs ) else: if inputs_embeds is None: inputs_embeds = self.word_embeddings(input_ids) if decoder_inputs_embeds is None and decoder_input_ids is None: decoder_input_ids = shift_tokens_right( labels, self.config.pad_token_id, self.config.decoder_start_token_id ) decoder_inputs_embeds = self.word_embeddings(decoder_input_ids) if attention_mask is not None: # concat prompt attention mask prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to( self.base_model.device ) kwargs["attention_mask"] = torch.cat((prefix_attention_mask, attention_mask), dim=1) # concat prompt labels if labels is not None: prefix_labels = torch.full((batch_size, self.pet_config.num_virtual_tokens), -100).to( self.base_model.device ) kwargs["labels"] = torch.cat((prefix_labels, labels), dim=1) prompts = self.get_prompt(batch_size=batch_size) inputs_embeds = torch.cat((prompts[:, : self.pet_config.num_virtual_tokens], inputs_embeds), dim=1) decoder_inputs_embeds = torch.cat( (prompts[:, self.pet_config.num_virtual_tokens :], decoder_inputs_embeds), dim=1 ) return self.base_model(inputs_embeds=inputs_embeds, decoder_inputs_embeds=decoder_inputs_embeds, **kwargs)