| 739 | |
| 740 | |
| 741 | class TelechatForCausalLM(TelechatPreTrainedModel): |
| 742 | # _tied_weights_keys = ["lm_head.weight"] |
| 743 | _keys_to_ignore_on_load_missing = [r"lm_head.weight"] |
| 744 | |
| 745 | def __init__(self, config: TelechatConfig): |
| 746 | super().__init__(config) |
| 747 | self.transformer = TelechatModel(config) |
| 748 | self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) |
| 749 | self.post_init() |
| 750 | |
| 751 | def get_output_embeddings(self): |
| 752 | return self.lm_head |
| 753 | |
| 754 | def set_output_embeddings(self, new_embeddings: torch.Tensor): |
| 755 | self.lm_head = new_embeddings |
| 756 | |
| 757 | def prepare_inputs_for_generation( |
| 758 | self, |
| 759 | input_ids: torch.LongTensor, |
| 760 | past_key_values: Optional[torch.Tensor] = None, |
| 761 | attention_mask: Optional[torch.Tensor] = None, |
| 762 | inputs_embeds: Optional[torch.Tensor] = None, |
| 763 | **kwargs, |
| 764 | ) -> dict: |
| 765 | if past_key_values: |
| 766 | input_ids = input_ids[:, -1].unsqueeze(-1) |
| 767 | if inputs_embeds is not None and past_key_values is None: |
| 768 | model_inputs = {"inputs_embeds": inputs_embeds} |
| 769 | else: |
| 770 | model_inputs = {"input_ids": input_ids} |
| 771 | |
| 772 | model_inputs.update( |
| 773 | { |
| 774 | "past_key_values": past_key_values, |
| 775 | "use_cache": kwargs.get("use_cache"), |
| 776 | "attention_mask": attention_mask, |
| 777 | } |
| 778 | ) |
| 779 | return model_inputs |
| 780 | |
| 781 | def forward( |
| 782 | self, |
| 783 | input_ids: Optional[torch.LongTensor] = None, |
| 784 | past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None, |
| 785 | attention_mask: Optional[torch.Tensor] = None, |
| 786 | inputs_embeds: Optional[torch.Tensor] = None, |
| 787 | labels: Optional[torch.Tensor] = None, |
| 788 | use_cache: Optional[bool] = None, |
| 789 | output_attentions: Optional[bool] = None, |
| 790 | output_hidden_states: Optional[bool] = None, |
| 791 | return_dict: Optional[bool] = None, |
| 792 | **deprecated_arguments, |
| 793 | ) -> Union[Tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]: |
| 794 | |
| 795 | return_dict = return_dict if return_dict is not None else self.config.use_return_dict |
| 796 | |
| 797 | transformer_outputs = self.transformer( |
| 798 | input_ids, |
nothing calls this directly
no outgoing calls
no test coverage detected