(self, ids, mask=None, encoder_states=None, encoder_mask=None)
| 354 | self.apply(init_weights) |
| 355 | |
| 356 | def forward(self, ids, mask=None, encoder_states=None, encoder_mask=None): |
| 357 | b, s = ids.size() |
| 358 | |
| 359 | # causal mask |
| 360 | if mask is None: |
| 361 | mask = torch.tril(torch.ones(1, s, s).to(ids.device)) |
| 362 | elif mask.ndim == 2: |
| 363 | mask = torch.tril(mask.unsqueeze(1).expand(-1, s, -1)) |
| 364 | |
| 365 | # layers |
| 366 | x = self.token_embedding(ids) |
| 367 | x = self.dropout(x) |
| 368 | e = self.pos_embedding(x.size(1), |
| 369 | x.size(1)) if self.shared_pos else None |
| 370 | for block in self.blocks: |
| 371 | x = block(x, mask, encoder_states, encoder_mask, pos_bias=e) |
| 372 | x = self.norm(x) |
| 373 | x = self.dropout(x) |
| 374 | return x |
| 375 | |
| 376 | |
| 377 | class T5Model(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected