class AI(nn.Module):
    def __init__(self):
        super().__init__()
        self.embed = nn.Embedding(len(wti), 128)
        self.lstm = nn.LSTM(128, 128, batch_first=True)
        self.attention = nn.MultiheadAttention(embed_dim=128, num_heads=4, batch_first=True)
        self.layernorm_after = nn.LayerNorm(128)
        self.linear = nn.Linear(128, len(itw))
    def forward(self, x):
        x = self.embed(x)
        x_origin = x
        x, _ = self.lstm(x)
        x_origin = x
        x, _ = self.attention(x, x, x)
        x+= x_origin
        x = self.layernorm_after(x)
        x = self.linear(x)
        x = x[:, -1, :]
        return x