(self,
vocab_size=250002,
max_seq_len=514,
type_size=1,
pad_id=1,
dim=1024,
num_heads=16,
num_layers=24,
post_norm=True,
dropout=0.1,
eps=1e-5)
| 83 | """ |
| 84 | |
| 85 | def __init__(self, |
| 86 | vocab_size=250002, |
| 87 | max_seq_len=514, |
| 88 | type_size=1, |
| 89 | pad_id=1, |
| 90 | dim=1024, |
| 91 | num_heads=16, |
| 92 | num_layers=24, |
| 93 | post_norm=True, |
| 94 | dropout=0.1, |
| 95 | eps=1e-5): |
| 96 | super().__init__() |
| 97 | self.vocab_size = vocab_size |
| 98 | self.max_seq_len = max_seq_len |
| 99 | self.type_size = type_size |
| 100 | self.pad_id = pad_id |
| 101 | self.dim = dim |
| 102 | self.num_heads = num_heads |
| 103 | self.num_layers = num_layers |
| 104 | self.post_norm = post_norm |
| 105 | self.eps = eps |
| 106 | |
| 107 | # embeddings |
| 108 | self.token_embedding = nn.Embedding(vocab_size, dim, padding_idx=pad_id) |
| 109 | self.type_embedding = nn.Embedding(type_size, dim) |
| 110 | self.pos_embedding = nn.Embedding(max_seq_len, dim, padding_idx=pad_id) |
| 111 | self.dropout = nn.Dropout(dropout) |
| 112 | |
| 113 | # blocks |
| 114 | self.blocks = nn.ModuleList([ |
| 115 | AttentionBlock(dim, num_heads, post_norm, dropout, eps) |
| 116 | for _ in range(num_layers) |
| 117 | ]) |
| 118 | |
| 119 | # norm layer |
| 120 | self.norm = nn.LayerNorm(dim, eps=eps) |
| 121 | |
| 122 | def forward(self, ids): |
| 123 | """ |
nothing calls this directly
no test coverage detected