| import torch |
| from torch import nn |
|
|
| from src.rinalmo.model.modules import Transformer, MaskedLanguageModelHead, TokenDropout |
|
|
| class RiNALMo(nn.Module): |
| def __init__(self, config): |
| super().__init__() |
|
|
| self.config = config |
| self.pad_tkn_idx = self.config.model.embedding['padding_idx'] |
|
|
| self.embedding = nn.Embedding(**self.config.model['embedding']) |
| self.transformer = Transformer(**self.config.model['transformer']) |
|
|
| self.lm_mask_head = MaskedLanguageModelHead(**self.config.model['lm_mask_head']) |
|
|
| self.token_dropout = TokenDropout(**self.config.model['token_dropout']) |
|
|
| def forward(self, tokens, need_attn_weights=False): |
| pad_mask = tokens.eq(self.pad_tkn_idx) |
| x = self.embedding(tokens) |
| x = self.token_dropout(x, tokens) |
|
|
| with torch.cuda.amp.autocast(dtype=torch.float16): |
| if self.config.model.transformer.use_flash_attn: |
| representation, attn_weights = self.transformer( |
| x, |
| key_padding_mask=torch.logical_not(pad_mask) if pad_mask is not None else None, |
| need_attn_weights=need_attn_weights |
| ) |
| else: |
| representation, attn_weights = self.transformer(x, key_padding_mask=pad_mask, need_attn_weights=need_attn_weights) |
| x = self.lm_mask_head(representation) |
|
|
| result = {"logits": x, "representation": representation} |
| if need_attn_weights: |
| result["attentions"] = torch.stack(attn_weights, dim=1) |
| |
| return result |
|
|
