[prompt] Fix bug in Template (#4456)

* [trainer] fix unknown variable bug in prompt trainer

* [prompt] fix bug in template

Author: LemonNoel

paddlenlp/prompt/prompt_model.py

@@ -24,7 +24,7 @@
     SequenceClassifierOutput,
 )
 from .prompt_utils import signature
-from .template import Template
+from .template import PrefixTemplate, Template
 from .verbalizer import Verbalizer
 
 
@@ -55,6 +55,9 @@ def __init__(
         self.forward_keys = signature(self.plm.forward)
         self._mask_token_id = self.template.tokenizer.mask_token_id
         self._pad_token_id = self.template.tokenizer.pad_token_id
+        if isinstance(self.template, PrefixTemplate):
+            self.plm = self.template.process_model(self.plm)
+            self.forward_keys.append("past_key_values")
 
     def forward(
         self,
@@ -82,6 +85,7 @@ def forward(
             **kwargs,
         }
         input_dict = self.template.process_batch(input_dict)
+        input_dict = {**input_dict, **kwargs}
         model_inputs = {k: input_dict[k] for k in input_dict if k in self.forward_keys}
         if "masked_positions" in model_inputs:
             model_inputs.pop("masked_positions")

paddlenlp/prompt/prompt_trainer.py

@@ -198,15 +198,18 @@ def create_optimizer(self, lr_scheduler=None):
                 else:
                     params = plm_parameters
             else:
-                args = self.init_num_steps(self.args, len(self.train_dataset))
+                if self.args.max_steps > 0:
+                    max_steps = self.args.max_steps
+                else:
+                    raise ValueError("Please use `max_steps` to set the maximum training steps.")
                 warmup = (
-                    args.warmup_steps if args.warmup_steps > 0 else int(args.warmup_ratio * args.num_training_steps)
+                    self.args.warmup_steps if self.args.warmup_steps > 0 else int(self.args.warmup_ratio * max_steps)
                 )
                 self.lr_scheduler = get_scheduler(
-                    args.lr_scheduler_type,
+                    self.args.lr_scheduler_type,
                     learning_rate=self.args.ppt_learning_rate,
                     num_warmup_steps=warmup,
-                    num_training_steps=args.num_training_steps,
+                    num_training_steps=max_steps,
                 )
                 lr = self.lr_scheduler
                 params = ppt_parameters

paddlenlp/prompt/prompt_utils.py

@@ -18,11 +18,13 @@
 
 import inspect
 from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import numpy as np
 import paddle
+from paddle import Tensor
 
+from ..transformers.model_outputs import MaskedLMOutput, SequenceClassifierOutput
 from ..transformers.tokenizer_utils_base import PaddingStrategy, PretrainedTokenizerBase
 
 
@@ -114,3 +116,95 @@ def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]:
                     continue
                 batch[key] = self._convert_to_tensors(values)
         return batch
+
+
+def sequence_classification_forward_with_past_key_values(
+    self,
+    input_ids: Optional[Tensor] = None,
+    token_type_ids: Optional[Tensor] = None,
+    position_ids: Optional[Tensor] = None,
+    attention_mask: Optional[Tensor] = None,
+    inputs_embeds: Optional[Tensor] = None,
+    labels: Optional[Tensor] = None,
+    output_hidden_states: Optional[bool] = None,
+    output_attentions: Optional[bool] = None,
+    return_dict: Optional[bool] = None,
+    past_key_values: Optional[Tuple[Tuple[Tensor]]] = None,
+):
+    outputs = self.ernie(
+        input_ids,
+        token_type_ids=token_type_ids,
+        position_ids=position_ids,
+        attention_mask=attention_mask,
+        inputs_embeds=inputs_embeds,
+        past_key_values=past_key_values,
+        output_attentions=output_attentions,
+        output_hidden_states=output_hidden_states,
+        return_dict=True,
+    )
+    pooled_output = outputs[1]
+
+    pooled_output = self.dropout(pooled_output)
+    logits = self.classifier(pooled_output)
+
+    loss = None
+    if labels is not None:
+        if self.num_labels == 1:
+            loss_fct = paddle.nn.MSELoss()
+            loss = loss_fct(logits, labels)
+        elif labels.dtype == paddle.int64 or labels.dtype == paddle.int32:
+            loss_fct = paddle.nn.CrossEntropyLoss()
+            loss = loss_fct(logits.reshape((-1, self.num_labels)), labels.reshape((-1,)))
+        else:
+            loss_fct = paddle.nn.BCEWithLogitsLoss()
+            loss = loss_fct(logits, labels)
+
+    return SequenceClassifierOutput(
+        loss=loss,
+        logits=logits,
+        hidden_states=outputs.hidden_states,
+        attentions=outputs.attentions,
+    )
+
+
+def masked_lm_forward_with_past_key_values(
+    self,
+    input_ids: Optional[Tensor] = None,
+    token_type_ids: Optional[Tensor] = None,
+    position_ids: Optional[Tensor] = None,
+    attention_mask: Optional[Tensor] = None,
+    masked_positions: Optional[Tensor] = None,
+    inputs_embeds: Optional[Tensor] = None,
+    labels: Optional[Tensor] = None,
+    output_hidden_states: Optional[bool] = None,
+    output_attentions: Optional[bool] = None,
+    return_dict: Optional[bool] = None,
+    past_key_values: Optional[Tuple[Tuple[Tensor]]] = None,
+):
+    outputs = self.ernie(
+        input_ids,
+        token_type_ids=token_type_ids,
+        position_ids=position_ids,
+        attention_mask=attention_mask,
+        inputs_embeds=inputs_embeds,
+        past_key_values=past_key_values,
+        output_attentions=output_attentions,
+        output_hidden_states=output_hidden_states,
+        return_dict=True,
+    )
+    sequence_output = outputs[0]
+    prediction_scores = self.cls(sequence_output, masked_positions=masked_positions)
+
+    masked_lm_loss = None
+    if labels is not None:
+        loss_fct = paddle.nn.CrossEntropyLoss()
+        masked_lm_loss = loss_fct(
+            prediction_scores.reshape((-1, paddle.shape(prediction_scores)[-1])), labels.reshape((-1,))
+        )
+
+    return MaskedLMOutput(
+        loss=masked_lm_loss,
+        logits=prediction_scores,
+        hidden_states=outputs.hidden_states,
+        attentions=outputs.attentions,
+    )

paddlenlp/prompt/template.py

@@ -21,6 +21,7 @@
 import re
 import traceback
 from abc import abstractmethod
+from functools import partial
 from typing import Any, Dict, List, Optional
 
 import numpy as np
@@ -32,6 +33,10 @@
 from paddlenlp.utils.log import logger
 
 from .prompt_tokenizer import MLMPromptTokenizer
+from .prompt_utils import (
+    masked_lm_forward_with_past_key_values,
+    sequence_classification_forward_with_past_key_values,
+)
 
 __all__ = ["Template", "ManualTemplate", "SoftTemplate", "PrefixTemplate", "AutoTemplate", "UTCTemplate"]
 
@@ -263,8 +268,10 @@ def save(self, save_path):
         if not os.path.exists(save_path):
             os.makedirs(save_path, exist_ok=True)
         template_config_file = os.path.join(save_path, TEMPLATE_CONFIG_FILE)
+        template_class = self.__class__.__name__
         with open(template_config_file, "w", encoding="utf-8") as fp:
-            fp.write(json.dumps(self._prompt, ensure_ascii=False))
+            fp.write(json.dumps(self._prompt, ensure_ascii=False) + "\n")
+            fp.write(json.dumps({"class": template_class}, ensure_ascii=False) + "\n")
         template_param_file = os.path.join(save_path, TEMPLATE_PARAMETER_FILE)
         template_state_dict = self.state_dict()
         if len(template_state_dict) > 0:
@@ -709,36 +716,54 @@ def parse_soft_prompt(self):
                 raise ValueError("Keyword `prefix` should locate at the beginning of template.")
             part["soft"] = part["prefix"]
             part.pop("prefix")
+            if "encoder" not in part:
+                part["encoder"] = "mlp"
             prompt[index] = part
 
         self._prompt = prompt
         return super(PrefixTemplate, self).parse_soft_prompt()
 
+    def process_model(self, model):
+        if model.__class__.__name__.endswith("ForSequenceClassification"):
+            model.forward = partial(sequence_classification_forward_with_past_key_values, self=model)
+        elif model.__class__.__name__.endswith("ForMaskedLM"):
+            model.forward = partial(masked_lm_forward_with_past_key_values, self=model)
+        return model
+
     def process_batch(self, input_dict: Dict[str, Tensor]) -> Dict[str, Tensor]:
         word_embeds = self.word_embeddings(input_dict["input_ids"])
+        batch_size, _ = input_dict["soft_token_ids"].shape
+
+        soft_token_ids = paddle.masked_select(input_dict["soft_token_ids"], input_dict["soft_token_ids"] > 0)
+        soft_token_ids = soft_token_ids.reshape([batch_size, -1])
+        _, soft_len = soft_token_ids.shape
+
+        token_type_ids = paddle.masked_select(input_dict["token_type_ids"], input_dict["soft_token_ids"] == 0)
+        input_dict["token_type_ids"] = token_type_ids.reshape([batch_size, -1])
+        position_ids = paddle.masked_select(input_dict["position_ids"], input_dict["soft_token_ids"] == 0)
+        input_dict["position_ids"] = position_ids.reshape([batch_size, -1])
+        if "masked_position" in input_dict and input_dict["masked_positions"] is not None:
+            input_dict["masked_positions"] = input_dict["masked_positions"] - soft_len
+        input_dict["inputs_embeds"] = paddle.concat(
+            [word_embeds[:, 0, :].unsqueeze(1), word_embeds[:, soft_len + 1 :, :]], axis=1
+        )
+
         if "attention_mask" not in input_dict or input_dict["attention_mask"] is None:
             pad_token_id = self.tokenizer.pad_token_id
             attention_mask = paddle.unsqueeze(
                 (input_dict["input_ids"] == pad_token_id).astype("float32") * -1e4, axis=[1, 2]
             )
             input_dict["attention_mask"] = attention_mask
         input_dict["input_ids"] = None
-
-        batch_size, _ = input_dict["soft_token_ids"].shape
-        soft_token_ids = paddle.masked_select(input_dict["soft_token_ids"], input_dict["soft_token_ids"] > 0)
-        soft_token_ids = soft_token_ids.reshape([batch_size, -1])
-        _, soft_len = soft_token_ids.shape
-
-        input_dict["inputs_embeds"] = word_embeds[:, soft_len:, :]
+        input_dict.pop("soft_token_ids")
+        input_dict.pop("encoder_ids")
 
         soft_embeds = self.soft_embeddings(soft_token_ids)
-        for encoder_id in range(1, len(self.encoder_list)):
-            to_encode = paddle.where(input_dict["encoder_ids"] == encoder_id)
-            encoded = self.encoder_list[encoder_id](to_encode)
-            soft_embeds = paddle.where(input_dict["encoder_ids"] == encoder_id, encoded, soft_embeds)
+        soft_embeds = self.encoder_list[1](soft_embeds)
         soft_embeds = soft_embeds.reshape(
             [batch_size, soft_len, self.n_layer * 2, self.n_heads, self.embed_size // self.n_heads]
         )
+
         soft_embeds = self.dropout(soft_embeds)
         soft_embeds = paddle.transpose(soft_embeds, perm=[2, 0, 3, 1, 4])
         soft_embeds = paddle.split(soft_embeds, num_or_sections=self.n_layer)
@@ -776,6 +801,7 @@ def create_from(
         model: PretrainedModel = None,
         soft_embeddings: Tensor = None,
         prefix_dropout: float = 0.1,
+        template_class: str = None,
     ):
         # Default template if not defined.
         if prompt is None:
@@ -791,12 +817,20 @@ def create_from(
             if "mask" not in template_keywords:
                 prompt = prompt + [{"mask": None}]
 
+        if template_class is None:
+            if "prefix" in template_keywords:
+                template_class = "PrefixTemplate"
+            elif "soft" in template_keywords or "soft_id" in template_keywords:
+                template_class = "SoftTemplate"
+            else:
+                template_class = "ManualTemplate"
+
         # Choose Template according to template keywords.
-        if "prefix" in template_keywords:
+        if template_class == "PrefixTemplate":
             return PrefixTemplate(
                 prompt=prompt, tokenizer=tokenizer, max_length=max_length, model=model, prefix_dropout=prefix_dropout
             )
-        elif "soft" in template_keywords or "soft_id" in template_keywords:
+        elif template_class == "SoftTemplate":
             word_embeddings = model.get_input_embeddings()
             return SoftTemplate(
                 prompt=prompt,
@@ -805,10 +839,12 @@ def create_from(
                 word_embeddings=word_embeddings,
                 soft_embeddings=soft_embeddings,
             )
-        elif "options" in template_keywords:
+        elif template_class == "UTCTemplate":
             return UTCTemplate(tokenizer=tokenizer, max_length=max_length)
-        else:
+        elif template_class == "ManualTemplate":
             return ManualTemplate(prompt=prompt, tokenizer=tokenizer, max_length=max_length)
+        else:
+            raise ValueError(f"Unknown template: {template_class}.")
 
     @classmethod
     def load_from(
@@ -818,9 +854,15 @@ def load_from(
         if not os.path.isfile(template_config_file):
             raise ValueError("{} not found under {}".format(TEMPLATE_CONFIG_FILE, data_path))
         with open(template_config_file, "r") as fp:
-            prompt = json.loads(fp.readline().strip())
-        # TODO (Huijuan): Load all configs from data_path.
-        template = cls.create_from(prompt=prompt, tokenizer=tokenizer, max_length=max_length, model=model)
+            config = [x.strip() for x in fp]
+            prompt = json.loads(config[0])
+            if len(config) > 1:
+                template_class = json.loads(config[1])
+            else:
+                template_class = None  # Compatible with previous versions
+        template = cls.create_from(
+            prompt=prompt, tokenizer=tokenizer, max_length=max_length, model=model, template_class=template_class
+        )
         template_param_file = os.path.join(data_path, TEMPLATE_PARAMETER_FILE)
         if os.path.isfile(template_param_file):
             template.set_state_dict(paddle.load(template_param_file))
@@ -834,10 +876,14 @@ class UTCTemplate(Template):
 
     template_special_tokens = ["text", "hard", "sep", "cls", "options"]
 
-    def __init__(self, tokenizer: PretrainedTokenizer, max_length: int):
+    def __init__(self, tokenizer: PretrainedTokenizer, max_length: int, prompt: str = None):
         prompt = (
-            "{'options': 'choices', 'add_omask': True, 'position': 0, 'token_type': 1}"
-            "{'sep': None, 'token_type': 0, 'position': 0}{'text': 'text_a'}{'sep': None, 'token_type': 1}{'text': 'text_b'}"
+            (
+                "{'options': 'choices', 'add_omask': True, 'position': 0, 'token_type': 1}"
+                "{'sep': None, 'token_type': 0, 'position': 0}{'text': 'text_a'}{'sep': None, 'token_type': 1}{'text': 'text_b'}"
+            )
+            if prompt is None
+            else prompt
         )
         super(UTCTemplate, self).__init__(prompt, tokenizer, max_length)
         self.max_position_id = self.tokenizer.model_max_length - 1