(self, vocab_size: int, hidden_dim: int)
| 99 | """Tiny model with both LM logits and scalar rewards in one forward.""" |
| 100 | |
| 101 | def __init__(self, vocab_size: int, hidden_dim: int): |
| 102 | super().__init__() |
| 103 | self.embedding = torch.nn.Embedding(vocab_size, hidden_dim) |
| 104 | self.lm_head = torch.nn.Linear(hidden_dim, vocab_size) |
| 105 | self.reward_head = torch.nn.Linear(hidden_dim, 1) |
| 106 | self.use_cache_calls: list[bool | None] = [] |
| 107 | |
| 108 | def forward(self, input_ids, attention_mask=None, use_cache=None): |
| 109 | self.use_cache_calls.append(use_cache) |