| 60 | |
| 61 | class BertEmbeddings(nn.Module): |
| 62 | def __init__(self, config): |
| 63 | super().__init__() |
| 64 | self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size) |
| 65 | self.token_type_embeddings = nn.Embedding( |
| 66 | config.type_vocab_size, config.hidden_size |
| 67 | ) |
| 68 | self.position_embeddings = nn.Embedding( |
| 69 | config.max_position_embeddings, config.hidden_size |
| 70 | ) |
| 71 | self.norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) |
| 72 | |
| 73 | def __call__( |
| 74 | self, input_ids: mx.array, token_type_ids: mx.array = None |