| 391 | MOSS_START_DOCSTRING, |
| 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 |
| 412 | |
| 413 | def set_input_embeddings(self, new_embeddings): |
| 414 | self.wte = new_embeddings |
| 415 | |
| 416 | @add_start_docstrings_to_model_forward(MOSS_INPUTS_DOCSTRING.format("batch_size, sequence_length")) |
| 417 | @add_code_sample_docstrings( |
| 418 | checkpoint=_CHECKPOINT_FOR_DOC, |
| 419 | output_type=BaseModelOutputWithPast, |
| 420 | config_class=_CONFIG_FOR_DOC, |
| 421 | ) |
| 422 | def forward( |
| 423 | self, |
| 424 | input_ids: Optional[torch.LongTensor] = None, |
| 425 | past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None, |
| 426 | attention_mask: Optional[torch.FloatTensor] = None, |
| 427 | token_type_ids: Optional[torch.LongTensor] = None, |
| 428 | position_ids: Optional[torch.LongTensor] = None, |
| 429 | head_mask: Optional[torch.FloatTensor] = None, |
| 430 | inputs_embeds: Optional[torch.FloatTensor] = None, |
| 431 | use_cache: Optional[bool] = None, |
| 432 | output_attentions: Optional[bool] = None, |
| 433 | output_hidden_states: Optional[bool] = None, |
| 434 | return_dict: Optional[bool] = None, |
| 435 | ) -> Union[Tuple, BaseModelOutputWithPast]: |
| 436 | output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions |
| 437 | output_hidden_states = ( |
| 438 | output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states |
| 439 | ) |
| 440 | use_cache = use_cache if use_cache is not None else self.config.use_cache |
| 441 | return_dict = return_dict if return_dict is not None else self.config.use_return_dict |
| 442 | |
| 443 | if input_ids is not None and inputs_embeds is not None: |
| 444 | raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") |
| 445 | elif input_ids is not None: |
| 446 | input_shape = input_ids.size() |
| 447 | input_ids = input_ids.view(-1, input_shape[-1]) |
| 448 | batch_size = input_ids.shape[0] |
| 449 | elif inputs_embeds is not None: |
| 450 | input_shape = inputs_embeds.size()[:-1] |