Skip to content

transformer_encoder

Transformer Encoder Model.

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

TransformerEncoderModel

Bases: Module, Module

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

Source code in meds_torch/models/components/transformer_encoder.py 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
class TransformerEncoderModel(torch.nn.Module, Module):
    """Wrapper of Encoder Transformer for use in MEDS with triplet token embeddings."""

    def __init__(self, cfg: DictConfig):
        super().__init__()
        self.model = cfg.model

        if cfg.token_emb:
            self.model.token_emb = cfg.token_emb

    def forward(self, batch):
        input_data, mask = batch[INPUT_ENCODER_TOKENS_KEY], batch[INPUT_ENCODER_MASK_KEY]
        output, embeddings = self.model(input_data, mask=mask, return_logits_and_embeddings=True)
        batch[BACKBONE_TOKENS_KEY] = output
        batch[BACKBONE_EMBEDDINGS_KEY] = embeddings
        return batch