| 594 | |
| 595 | |
| 596 | class TelechatModel(TelechatPreTrainedModel): |
| 597 | def __init__(self, config: TelechatConfig): |
| 598 | super().__init__(config) |
| 599 | |
| 600 | self.embed_dim = config.hidden_size |
| 601 | self.num_heads = config.n_head |
| 602 | self.config = config |
| 603 | self.word_embeddings = nn.Embedding(config.vocab_size, self.embed_dim) |
| 604 | if self.config.embed_layernorm: |
| 605 | self.word_embeddings_layernorm = MixedFusedRMSNorm(self.embed_dim, eps=config.layer_norm_epsilon) |
| 606 | |
| 607 | self.h = nn.ModuleList([TelechatBlock(config, _) for _ in range(config.num_hidden_layers)]) |
| 608 | self.ln_f = MixedFusedRMSNorm(self.embed_dim, eps=config.layer_norm_epsilon) |
| 609 | self.gradient_checkpointing = False |
| 610 | self.post_init() |
| 611 | |
| 612 | def get_input_embeddings(self): |
| 613 | return self.word_embeddings |
| 614 | |
| 615 | def _prepare_attn_mask( |
| 616 | self, attention_mask: torch.Tensor, input_shape: Tuple[int, int], past_key_values_length: int |
| 617 | ) -> torch.BoolTensor: |
| 618 | combined_attention_mask = None |
| 619 | device = attention_mask.device |
| 620 | _, src_length = input_shape |
| 621 | |
| 622 | if src_length > 1: |
| 623 | combined_attention_mask = _make_causal_mask( |
| 624 | input_shape, device=device, past_key_values_length=past_key_values_length |
| 625 | ) |
| 626 | expanded_attn_mask = _expand_mask(attention_mask, tgt_length=src_length) |
| 627 | combined_attention_mask = ( |
| 628 | expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask | combined_attention_mask |
| 629 | ) |
| 630 | |
| 631 | return combined_attention_mask |
| 632 | |
| 633 | def set_input_embeddings(self, new_embeddings: torch.Tensor): |
| 634 | self.word_embeddings = new_embeddings |
| 635 | |
| 636 | def forward( |
| 637 | self, |
| 638 | input_ids: Optional[torch.LongTensor] = None, |
| 639 | past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None, |
| 640 | attention_mask: Optional[torch.Tensor] = None, |
| 641 | inputs_embeds: Optional[torch.LongTensor] = None, |
| 642 | use_cache: Optional[bool] = None, |
| 643 | output_attentions: Optional[bool] = None, |
| 644 | output_hidden_states: Optional[bool] = None, |
| 645 | return_dict: Optional[bool] = None, |
| 646 | **deprecated_arguments, |
| 647 | ) -> Union[Tuple[torch.Tensor, ...], BaseModelOutputWithPastAndCrossAttentions]: |
| 648 | |
| 649 | output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions |
| 650 | output_hidden_states = ( |
| 651 | output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states |
| 652 | ) |
| 653 | use_cache = use_cache if use_cache is not None else self.config.use_cache |