| 859 | """ |
| 860 | |
| 861 | def __init__( |
| 862 | self, |
| 863 | hidden_size, |
| 864 | vocab_size, |
| 865 | max_sequence_length, |
| 866 | ): |
| 867 | super(QueryEmbedding, self).__init__() |
| 868 | |
| 869 | self.hidden_size = hidden_size |
| 870 | self.vocab_size = vocab_size |
| 871 | self.max_sequence_length = max_sequence_length |
| 872 | |
| 873 | # Top query position embedding (serial). |
| 874 | self.top_query_embeddings = torch.nn.Embedding(self.max_sequence_length, self.hidden_size) |
| 875 | self.top_query_embeddings = self.top_query_embeddings.half() |
| 876 | self._top_query_embeddings_key = 'top_query_embeddings' |
| 877 | |
| 878 | def forward(self, position_ids): |
| 879 | # Embeddings. |