utils.py

# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Any, Dict, List, Optional, Tuple, Union

import numpy as np
import paddle
import paddle.nn as nn

from paddlenlp.trainer import Trainer


class GLMTrainer(Trainer):
    def __init__(self, do_generation: bool, **kwargs):
        super().__init__(**kwargs)
        self.do_generation = do_generation

    def prediction_step(
        self,
        model: nn.Layer,
        inputs: Dict[str, Union[paddle.Tensor, Any]],
        prediction_loss_only: bool,
        ignore_keys: Optional[List[str]] = None,
    ) -> Tuple[Optional[paddle.Tensor], Optional[paddle.Tensor], Optional[paddle.Tensor]]:

        if not self.do_generation:
            return super().prediction_step(model, inputs, prediction_loss_only, ignore_keys)

        model.eval()
        with paddle.no_grad():
            tokens = model.generate(
                input_ids=inputs["input_ids"],
                position_ids=inputs["position_ids"],
                attention_mask=inputs["attention_mask"],
                decode_strategy="sampling",
                top_k=1,
                repetition_penalty=2.0,
                bos_token_id=self.tokenizer.sop_token_id,
                eos_token_id=self.tokenizer.eop_token_id,
                pad_token_id=self.tokenizer.pad_token_id,
            )[0]
            all_preds = []
            for pred_tokens in tokens:
                all_preds.append(pred_tokens[pred_tokens != self.tokenizer.pad_token_id].tolist())
            max_pred_length = max([len(x) for x in all_preds])
            for index, preds in enumerate(all_preds):
                all_preds[index] = preds + [-100] * (max_pred_length - len(preds))

            all_labels = []
            for label, mask in zip(inputs["labels"].numpy(), inputs["loss_mask"].numpy()):
                label = label[mask.astype("bool")]
                label = [x for x in label[label != self.tokenizer.pad_token_id]]
                all_labels.append(label)
            max_label_length = max([len(x) for x in all_labels])
            for index, labels in enumerate(all_labels):
                all_labels[index] = labels + [-100] * (max_label_length - len(labels))

        return (None, paddle.to_tensor(all_preds), paddle.to_tensor(all_labels))

    def log(self, logs: Dict[str, float], **kwargs) -> None:

        if self.state.epoch is not None:
            logs["epoch"] = round(self.state.epoch, 4)

        if "eval_loss" in logs:
            logs["eval_ppl"] = np.exp(logs["eval_loss"])
        output = {**logs, **{"step": self.state.global_step}}
        self.state.log_history.append(output)
        self.control = self.callback_handler.on_log(self.args, self.state, self.control, logs, **kwargs)