Skip to content

eic_encoder

EicEncoder

Bases: Module, Module

Placeholder encoder, does nothing as the x-transformer model handles the embedding of tokens.

Source code in meds_torch/input_encoder/eic_encoder.py 
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class EicEncoder(nn.Module, Module):
    """Placeholder encoder, does nothing as the x-transformer model handles the embedding of tokens."""

    def __init__(self, cfg: DictConfig):
        super().__init__()
        self.cfg = cfg
        if self.cfg.model_type == "eic_lstm":
            self.code_embedder = nn.Embedding(cfg.vocab_size, cfg.token_dim)

    def forward(self, batch):
        batch[INPUT_ENCODER_MASK_KEY] = batch["mask"]
        if self.cfg.model_type == "eic_lstm":
            embedded_codes = self.code_embedder(batch["code"])
            batch[INPUT_ENCODER_TOKENS_KEY] = embedded_codes
        else:
            batch[INPUT_ENCODER_TOKENS_KEY] = batch["code"]
        return batch