(self, input_ids, attention_mask=None)
| 158 | self.lm_head.weight = self.token_embedding.weight |
| 159 | |
| 160 | def forward(self, input_ids, attention_mask=None): |
| 161 | |
| 162 | batch_size, seq_len = input_ids.size() |
| 163 | device = input_ids.device |
| 164 | |
| 165 | position_ids = torch.arange(0, seq_len, dtype=torch.long, device=device) |
| 166 | position_ids = position_ids.unsqueeze(0).expand(batch_size, seq_len) |
| 167 | |
| 168 | token_embeds = self.token_embedding(input_ids) |
| 169 | pos_embeds = self.position_embedding(position_ids) |
| 170 | hidden_states = token_embeds + pos_embeds |
| 171 | hidden_states = self.dropout(hidden_states) |
| 172 | |
| 173 | if attention_mask is not None: |
| 174 | attention_mask = attention_mask.unsqueeze(1).unsqueeze(2) |
| 175 | attention_mask = attention_mask.to(dtype=hidden_states.dtype) |
| 176 | attention_mask = (1.0 - attention_mask) * -10000.0 |
| 177 | |
| 178 | for _ in range(self.num_iterations): |
| 179 | for block in self.blocks: |
| 180 | hidden_states = block(hidden_states, attention_mask=attention_mask)[0] |
| 181 | |
| 182 | hidden_states = self.ln_f(hidden_states) |
| 183 | logits = self.lm_head(hidden_states) |
| 184 | |
| 185 | return type("Output", (object,), {"logits": logits}) |
nothing calls this directly
no outgoing calls
no test coverage detected