(self, model, inputs, return_outputs=False)
| 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, |
no outgoing calls
no test coverage detected