| 392 | ) |
| 393 | class MossModel(MossPreTrainedModel): |
| 394 | def __init__(self, config): |
| 395 | super().__init__(config) |
| 396 | |
| 397 | self.embed_dim = config.n_embd |
| 398 | self.vocab_size = config.vocab_size |
| 399 | self.wte = nn.Embedding(config.vocab_size, self.embed_dim) |
| 400 | self.drop = nn.Dropout(config.embd_pdrop) |
| 401 | self.h = nn.ModuleList([MossBlock(config) for _ in range(config.n_layer)]) |
| 402 | self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon) |
| 403 | self.rotary_dim = min(config.rotary_dim, config.n_ctx // config.num_attention_heads) |
| 404 | |
| 405 | self.gradient_checkpointing = False |
| 406 | |
| 407 | # Initialize weights and apply final processing |
| 408 | self.post_init() |
| 409 | |
| 410 | def get_input_embeddings(self): |
| 411 | return self.wte |