(self, input_ids: torch.Tensor)
| 99 | print(f"BertModel: {n_params / 1e6:.1f}M parameters") |
| 100 | |
| 101 | def forward(self, input_ids: torch.Tensor) -> torch.Tensor: |
| 102 | B, T = input_ids.shape |
| 103 | positions = torch.arange(T, device=input_ids.device).unsqueeze(0) |
| 104 | token_types = torch.zeros_like(input_ids) |
| 105 | |
| 106 | x = self.word_embeddings(input_ids) + self.position_embeddings(positions) + self.token_type_embeddings(token_types) |
| 107 | x = self.embed_dropout(self.embed_norm(x)) |
| 108 | |
| 109 | for layer in self.layers: |
| 110 | x = layer(x) |
| 111 | |
| 112 | # Pooled output from [CLS] token |
| 113 | pooled = torch.tanh(self.pooler(x[:, 0])) |
| 114 | return pooled |
nothing calls this directly
no outgoing calls
no test coverage detected