Text Encoder.

text_encoder.py

# Encoding text using gated-recurrent unit
class TextEncoderTinyGRU(nn.Module):
    def __init__(self, vocab_size, d_word=64, d_model=128):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, d_word)
        self.gru = nn.GRU(d_word, d_model, batch_first=True)
        self.ln = nn.LayerNorm(d_model)

    def forward(self, token_ids):
        x = self.embed(token_ids)
        _, h_last = self.gru(x)
        x = h_last[0]
        x = self.ln(x)
        return x