| 11 | |
| 12 | |
| 13 | class CPMTrainer(Trainer): |
| 14 | def compute_loss(self, model, inputs, return_outputs=False): |
| 15 | if "labels" in inputs: |
| 16 | labels = inputs.pop("labels") |
| 17 | else: |
| 18 | labels = None |
| 19 | |
| 20 | if not self.args.use_lora: |
| 21 | outputs = self.model(data = inputs, use_cache=False) |
| 22 | else: |
| 23 | with self.model._enable_peft_forward_hooks(**inputs): |
| 24 | outputs = self.model.base_model(data = inputs, use_cache=False) |
| 25 | |
| 26 | if labels is not None: |
| 27 | # Flatten the tokens |
| 28 | loss_fct = nn.CrossEntropyLoss() |
| 29 | logits = outputs.logits.view(-1, |
| 30 | self.model.config.vocab_size).contiguous() |
| 31 | labels = labels.view(-1).long().contiguous() |
| 32 | # Enable model parallelism |
| 33 | labels = labels.to(logits.device) |
| 34 | loss = loss_fct(logits, labels) |
| 35 | else: |
| 36 | if isinstance(outputs, dict) and "loss" not in outputs: |
| 37 | raise ValueError( |
| 38 | "The model did not return a loss from the inputs, only the following keys: " |
| 39 | f"{','.join(outputs.keys())}. For reference, the inputs it received are {','.join(inputs.keys())}." |
| 40 | ) |
| 41 | # We don't use .loss here since the model may return tuples instead of ModelOutput. |
| 42 | loss = outputs["loss"] if isinstance(outputs, dict) else outputs[0] |
| 43 | |
| 44 | return (loss, outputs) if return_outputs else loss |
| 45 | |
| 46 | def prediction_step( |
| 47 | self, |
| 48 | model: nn.Module, |
| 49 | inputs: Dict[str, Union[torch.Tensor, Any]], |
| 50 | prediction_loss_only: bool, |
| 51 | ignore_keys: Optional[List[str]] = None, |
| 52 | ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]: |
| 53 | """ |
| 54 | Perform an evaluation step on `model` using `inputs`. |
| 55 | |
| 56 | Subclass and override to inject custom behavior. |
| 57 | |
| 58 | Args: |
| 59 | model (`nn.Module`): |
| 60 | The model to evaluate. |
| 61 | inputs (`Dict[str, Union[torch.Tensor, Any]]`): |
| 62 | The inputs and targets of the model. |
| 63 | |
| 64 | The dictionary will be unpacked before being fed to the model. Most models expect the targets under the |
| 65 | argument `labels`. Check your model's documentation for all accepted arguments. |
| 66 | prediction_loss_only (`bool`): |
| 67 | Whether or not to return the loss only. |
| 68 | ignore_keys (`List[str]`, *optional*): |
| 69 | A list of keys in the output of your model (if it is a dictionary) that should be ignored when |
| 70 | gathering predictions. |