XLMRobertaModel with no pooler and no LM head.
| 78 | |
| 79 | |
| 80 | class XLMRoberta(nn.Module): |
| 81 | """ |
| 82 | XLMRobertaModel with no pooler and no LM head. |
| 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 | """ |
| 124 | ids: [B, L] of torch.LongTensor. |
| 125 | """ |
| 126 | b, s = ids.shape |
| 127 | mask = ids.ne(self.pad_id).long() |
| 128 | |
| 129 | # embeddings |
| 130 | x = self.token_embedding(ids) + \ |
| 131 | self.type_embedding(torch.zeros_like(ids)) + \ |
| 132 | self.pos_embedding(self.pad_id + torch.cumsum(mask, dim=1) * mask) |
| 133 | if self.post_norm: |
| 134 | x = self.norm(x) |
| 135 | x = self.dropout(x) |
| 136 | |
| 137 | # blocks |