| 243 | class MomModel(MomPreTrainedModel): |
| 244 | |
| 245 | def __init__(self, config: MomConfig): |
| 246 | super().__init__(config) |
| 247 | self.padding_idx = config.pad_token_id |
| 248 | self.vocab_size = config.vocab_size |
| 249 | |
| 250 | self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) |
| 251 | self.layers = nn.ModuleList([MomBlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]) |
| 252 | self.norm = RMSNorm(config.hidden_size, eps=config.norm_eps) |
| 253 | |
| 254 | self.gradient_checkpointing = False |
| 255 | |
| 256 | self.post_init() |
| 257 | |
| 258 | def get_input_embeddings(self): |
| 259 | return self.embeddings |