(self, args: ModelArgs)
| 99 | # https://github.com/karpathy/nanoGPT. |
| 100 | class Transformer(nn.Module): |
| 101 | def __init__(self, args: ModelArgs): |
| 102 | super().__init__() |
| 103 | assert args.vocab_size is not None |
| 104 | assert args.max_seq_len is not None |
| 105 | self.model_args = args |
| 106 | self.max_seq_len = args.max_seq_len |
| 107 | self.tok_embeddings = nn.Embedding(args.vocab_size, args.dim) |
| 108 | self.pos_embeddings = nn.Embedding(args.max_seq_len, args.dim) |
| 109 | self.dropout = nn.Dropout(args.dropout_p) |
| 110 | self.layers = nn.ModuleList() |
| 111 | for _ in range(args.n_layers): |
| 112 | self.layers.append(TransformerBlock(args)) |
| 113 | self.norm = nn.LayerNorm(args.dim) |
| 114 | self.output = nn.Linear(args.dim, args.vocab_size, bias=False) |
| 115 | |
| 116 | def forward(self, tokens): |
| 117 | _bsz, seq_len = tokens.size() |
nothing calls this directly
no test coverage detected