| 19 | |
| 20 | |
| 21 | class FakeGenerationReferenceModel(torch.nn.Module): |
| 22 | def __init__(self, logits: torch.Tensor): |
| 23 | super().__init__() |
| 24 | self.register_buffer("fixed_logits", logits) |
| 25 | self.use_cache_calls: list[bool | None] = [] |
| 26 | |
| 27 | def forward(self, input_ids, attention_mask=None, use_cache=None): |
| 28 | del attention_mask |
| 29 | self.use_cache_calls.append(use_cache) |
| 30 | batch, seq_len = input_ids.shape |
| 31 | key = torch.empty(batch, 1, seq_len, 4, device=input_ids.device) |
| 32 | value = torch.empty_like(key) |
| 33 | return SimpleNamespace( |
| 34 | logits=self.fixed_logits[: input_ids.shape[0], : input_ids.shape[1]], |
| 35 | past_key_values=((key, value),) if use_cache else None, |
| 36 | ) |
| 37 | |
| 38 | |
| 39 | def _inputs() -> StatelessForwardInputs: |
no outgoing calls