(self, input_ids: torch.Tensor, *, logit_bias: float = 0.0)
| 25 | self.register_buffer("base_logits", logits.clone()) |
| 26 | |
| 27 | def forward(self, input_ids: torch.Tensor, *, logit_bias: float = 0.0): |
| 28 | del input_ids |
| 29 | logits = self.base_logits * self.weight + float(logit_bias) |
| 30 | if self.output_kind == "tensor": |
| 31 | return logits |
| 32 | if self.output_kind == "mapping": |
| 33 | return {"logits": logits} |
| 34 | if self.output_kind == "object": |
| 35 | return ObjectOutput(logits) |
| 36 | if self.output_kind == "tuple": |
| 37 | return (logits, {"hidden_states": None}) |
| 38 | if self.output_kind == "loss_tuple": |
| 39 | return (logits.sum(), logits, {"hidden_states": None}) |
| 40 | raise AssertionError(f"unknown output kind: {self.output_kind}") |
| 41 | |
| 42 | |
| 43 | @pytest.mark.parametrize("output_kind", ("tensor", "mapping", "object", "tuple")) |
nothing calls this directly
no test coverage detected