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 r
| 909 | |
| 910 | |
| 911 | class TransformerLanguageModel(torch.nn.Module): |
| 912 | """Transformer language model. |
| 913 | Arguments: |
| 914 | transformer_hparams: transformer hyperparameters |
| 915 | attention_mask_func: a function that takes `unmaksed-attention-scores` |
| 916 | with size [b, np, s, s] and an `attention-mask` and will apply |
| 917 | the masking. The function should return a masked score of the |
| 918 | same size [b, np, s, s]. |
| 919 | masked-attention-scores = attention_mask_func( |
| 920 | unmaksed-attention-scores, attention-mask) |
| 921 | vocab_size: vocabulary size |
| 922 | max_sequence_length: maximum size of sequence. This |
| 923 | is used for positional embedding |
| 924 | """ |
| 925 | |
| 926 | def __init__( |
| 927 | self, |
| 928 | hidden_size, |
| 929 | num_layers, |
| 930 | num_attention_heads, |
| 931 | padded_vocab_size, |
| 932 | max_position_embeddings, |
| 933 | ): |
| 934 | super(TransformerLanguageModel, self).__init__() |
| 935 | self.hidden_size = hidden_size |
| 936 | self.num_layers = num_layers |
| 937 | self.num_attention_heads = num_attention_heads |
| 938 | self.padded_vocab_size = padded_vocab_size |
| 939 | self.max_position_embeddings = max_position_embeddings |
| 940 | |
| 941 | # Embeddings |
| 942 | self.embedding = Embedding(self.hidden_size, |
| 943 | self.padded_vocab_size, |
| 944 | self.max_position_embeddings) |
| 945 | self._embedding_key = 'embedding' |
| 946 | |
| 947 | # Query embeddings |
| 948 | self.topQueryEmbedding = QueryEmbedding(self.hidden_size, |
| 949 | self.padded_vocab_size, |
| 950 | self.max_position_embeddings) |
| 951 | self._topQueryEmbedding_key = 'topQueryEmbedding' |
| 952 | |
| 953 | # Transformer |
| 954 | self.transformer = Transformer(self.hidden_size, |
| 955 | self.num_attention_heads, |
| 956 | self.num_layers) |
| 957 | self._transformer_key = 'transformer' |
| 958 | |
| 959 | def forward( |
| 960 | self, |
| 961 | input_ids, |
| 962 | position_ids, |
| 963 | attention_mask, |
| 964 | layer_past=None, |
| 965 | get_key_value=False, |
| 966 | prompt_length=None, |
| 967 | context_length=None, |
| 968 | ): |