Language model embeddings. Arguments: hidden_size: hidden size vocab_size: vocabulary size max_sequence_length: maximum size of sequence. This is used for positional embedding
| 676 | |
| 677 | |
| 678 | class Embedding(paddle.nn.Layer): |
| 679 | """Language model embeddings. |
| 680 | |
| 681 | Arguments: |
| 682 | hidden_size: hidden size |
| 683 | vocab_size: vocabulary size |
| 684 | max_sequence_length: maximum size of sequence. This |
| 685 | is used for positional embedding |
| 686 | """ |
| 687 | |
| 688 | def __init__( |
| 689 | self, |
| 690 | hidden_size, |
| 691 | vocab_size, |
| 692 | max_sequence_length, |
| 693 | ): |
| 694 | super(Embedding, self).__init__() |
| 695 | self.hidden_size = hidden_size |
| 696 | self.vocab_size = vocab_size |
| 697 | self.max_sequence_length = max_sequence_length |
| 698 | |
| 699 | # Word embeddings. |
| 700 | self.word_embeddings = paddle.nn.Embedding(self.vocab_size, self.hidden_size) |
| 701 | self._word_embeddings_key = 'word_embeddings' |
| 702 | |
| 703 | # Position embedding. |
| 704 | self.position_embeddings = paddle.nn.Embedding(self.max_sequence_length, self.hidden_size) |
| 705 | self.position_embeddings = self.position_embeddings.to(dtype="float16") |
| 706 | self._position_embeddings_key = 'position_embeddings' |
| 707 | |
| 708 | def forward(self, input_ids, position_ids): |
| 709 | # Embeddings. |
| 710 | words_embeddings = self.word_embeddings(input_ids) |
| 711 | position_embeddings = self.position_embeddings(position_ids) |
| 712 | embeddings = words_embeddings + position_embeddings |
| 713 | |
| 714 | return embeddings |
| 715 | |
| 716 | def state_dict_for_save_checkpoint(self, destination=None, prefix='', |
| 717 | keep_vars=False): |
| 718 | """For easy load.""" |
| 719 | |
| 720 | state_dict_ = {} |
| 721 | state_dict_[self._word_embeddings_key] \ |
| 722 | = self.word_embeddings.state_dict(destination, prefix, keep_vars) |
| 723 | state_dict_[self._position_embeddings_key] \ |
| 724 | = self.position_embeddings.state_dict( |
| 725 | destination, prefix, keep_vars) |
| 726 | |
| 727 | return state_dict_ |
| 728 | |
| 729 | def set_state_dict(self, state_dict, use_structured_name=True): |
| 730 | """Customized load.""" |
| 731 | |
| 732 | # Word embedding. |
| 733 | if self._word_embeddings_key in state_dict: |
| 734 | state_dict_ = state_dict[self._word_embeddings_key] |
| 735 | else: |