| 24 | |
| 25 | class FakeReferenceModel(torch.nn.Module): |
| 26 | def __init__(self, logits: torch.Tensor): |
| 27 | super().__init__() |
| 28 | self.register_buffer("fixed_logits", logits) |
| 29 | self.config = SimpleNamespace(use_cache=True, _attn_implementation="eager") |
| 30 | self.generation_config = SimpleNamespace(use_cache=True) |
| 31 | self.use_cache_calls: list[bool | None] = [] |
| 32 | self.config_use_cache_calls: list[bool | None] = [] |
| 33 | self.generation_config_use_cache_calls: list[bool | None] = [] |
| 34 | self.attn_implementation_calls: list[str | None] = [] |
| 35 | |
| 36 | def forward(self, input_ids, attention_mask=None, use_cache=None): |
| 37 | del attention_mask |