Skip to content

lstm

Transformer Encoder Model.

Uses transformer encoder to get contextualized representations of all tokens and we take the CLS token representation as the embedding.

LstmModel

Bases: Module, Module

Wrapper of Encoder Transformer for use in MEDS with triplet token embeddings.

Source code in meds_torch/models/components/lstm.py 
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
class LstmModel(torch.nn.Module, Module):
    """Wrapper of Encoder Transformer for use in MEDS with triplet token embeddings."""

    def __init__(self, cfg: DictConfig):
        super().__init__()
        dropout = cfg.dropout
        self.cfg = cfg
        self.model = nn.LSTM(
            cfg.token_dim,
            cfg.token_dim,
            num_layers=cfg.n_layers,
            batch_first=True,
            dropout=dropout,
        )

    def forward(self, batch):
        input_data, mask = batch[INPUT_ENCODER_TOKENS_KEY], batch[INPUT_ENCODER_MASK_KEY]
        # pass tokens and attention mask to the lstm
        output = self.model(input_data)[0]
        # extract the representation token's embedding
        batch[BACKBONE_TOKENS_KEY] = output
        batch[BACKBONE_EMBEDDINGS_KEY] = get_last_token(output, ~mask)
        return batch