| 3 | import numpy as np |
| 4 | |
| 5 | class GPT1Embedding(nn.Module): |
| 6 | def __init__( |
| 7 | self, |
| 8 | vocabulary_size, |
| 9 | embedding_size, |
| 10 | sequence_length, |
| 11 | _tokens_embedding_weight=None, |
| 12 | _positional_embedding_weight=None, |
| 13 | ): |
| 14 | |
| 15 | super(GPT1Embedding, self).__init__() |
| 16 | self.vocabulary_size = vocabulary_size |
| 17 | self.embedding_size = embedding_size |
| 18 | self.sequence_length = sequence_length |
| 19 | |
| 20 | self.tokens = nn.Embedding( |
| 21 | vocabulary_size, |
| 22 | embedding_size, |
| 23 | padding_idx=0, |
| 24 | _weight=_tokens_embedding_weight, |
| 25 | ) |
| 26 | self.position = nn.Embedding( |
| 27 | sequence_length, embedding_size, _weight=_positional_embedding_weight |
| 28 | ) |
| 29 | |
| 30 | def forward(self, tokens, positions): |
| 31 | """Embedding module for GPT-1. |
| 32 | |
| 33 | This module combines token and positional embeddings, to return the |
| 34 | embeddings from a sequence of input tokens and positions. |
| 35 | See Lecture 06, slides 30-32. |
| 36 | |
| 37 | Parameters |
| 38 | ---------- |
| 39 | tokens (`torch.LongTensor` of shape `(batch_size, sequence_length)`) |
| 40 | The input tensor containing the token sequences. All the tokens |
| 41 | must be integers in [0, vocabulary_size). |
| 42 | |
| 43 | positions (`torch.LongTensor` of shape `(batch_size, sequence_length)`) |
| 44 | The tensor containing the position indices in the sequence. All |
| 45 | the positions must be integers in [0, sequence_length) |
| 46 | |
| 47 | Returns |
| 48 | ------- |
| 49 | embeddings (`torch.FloatTensor` of shape `(batch_size, sequence_length, embedding_size)`) |
| 50 | The tensor containing the embeddings. For example, `embeddings[0, 2]` |
| 51 | is the embedding vector for the token in 3rd position (index 2) |
| 52 | of the 1st sequence in the batch (index 0). |
| 53 | """ |
| 54 | if torch.any(positions >= self.sequence_length): |
| 55 | raise RuntimeError( |
| 56 | "Some position indices are larger than the " "maximum sequence length." |
| 57 | ) |
| 58 | |
| 59 | if torch.any(tokens >= self.vocabulary_size): |
| 60 | raise RuntimeError( |
| 61 | "Some tokens are larger than the size of " "the vocabulary." |
| 62 | ) |