Transformer language model. Arguments: transformer_hparams: transformer hyperparameters attention_mask_func: a function that takes `unmaksed-attention-scores` with size [b, np, s, s] and an `attention-mask` and will apply the masking. The function should
| 817 | |
| 818 | |
| 819 | class TransformerLanguageModel(torch.nn.Module): |
| 820 | """Transformer language model. |
| 821 | |
| 822 | Arguments: |
| 823 | transformer_hparams: transformer hyperparameters |
| 824 | attention_mask_func: a function that takes `unmaksed-attention-scores` |
| 825 | with size [b, np, s, s] and an `attention-mask` and will apply |
| 826 | the masking. The function should return a masked score of the |
| 827 | same size [b, np, s, s]. |
| 828 | masked-attention-scores = attention_mask_func( |
| 829 | unmaksed-attention-scores, attention-mask) |
| 830 | vocab_size: vocabulary size |
| 831 | max_sequence_length: maximum size of sequence. This |
| 832 | is used for positional embedding |
| 833 | """ |
| 834 | |
| 835 | def __init__( |
| 836 | self, |
| 837 | hidden_size, |
| 838 | num_layers, |
| 839 | num_attention_heads, |
| 840 | padded_vocab_size, |
| 841 | max_position_embeddings, |
| 842 | ): |
| 843 | super(TransformerLanguageModel, self).__init__() |
| 844 | self.hidden_size = hidden_size |
| 845 | self.num_layers = num_layers |
| 846 | self.num_attention_heads = num_attention_heads |
| 847 | self.padded_vocab_size = padded_vocab_size |
| 848 | self.max_position_embeddings = max_position_embeddings |
| 849 | |
| 850 | # Embeddings |
| 851 | self.embedding = Embedding(self.hidden_size, |
| 852 | self.padded_vocab_size, |
| 853 | self.max_position_embeddings) |
| 854 | self._embedding_key = 'embedding' |
| 855 | |
| 856 | # Query embeddings |
| 857 | self.topQueryEmbedding = QueryEmbedding(self.hidden_size, |
| 858 | self.padded_vocab_size, |
| 859 | self.max_position_embeddings) |
| 860 | self._topQueryEmbedding_key = 'topQueryEmbedding' |
| 861 | |
| 862 | # Transformer |
| 863 | self.transformer = Transformer(self.hidden_size, |
| 864 | self.num_attention_heads, |
| 865 | self.num_layers) |
| 866 | self._transformer_key = 'transformer' |
| 867 | |
| 868 | def forward( |
| 869 | self, |
| 870 | input_ids, |
| 871 | position_ids, |
| 872 | attention_mask, |
| 873 | layer_past=None, |
| 874 | get_key_value=False, |
| 875 | prompt_length=None, |
| 876 | context_length=None, |