from torchtext import data
import torch
from torch import nn


class RNN(nn.Module):
    def __init__(self, num_embeddings,
                 embedding_dim=50,
                 hidden_size=50,
                 output_size=1,
                 num_layers=1,
                 dropout=0.2):
        super().__init__()
        self.emb = nn.Embedding(num_embeddings, embedding_dim,
                                padding_idx=0)
        self.lstm = nn.LSTM(embedding_dim,
                            hidden_size, num_layers,
                            batch_first=True, dropout=dropout)
        self.linear = nn.Linear(hidden_size, output_size)

    def forward(self, x, h0=None):
        x = self.emb(x)
        x, h = self.lstm(x, h0)
        x = x[:, -1, :]
        x = self.linear(x)
        return x


TEXT = data.Field(sequential=True, lower=True, batch_first=True)
LABELS = data.Field(sequential=False, batch_first=True, use_vocab=False)

train, val, test = data.TabularDataset.splits(
    path='ch06', train='train2.txt',
    validation='valid2.txt', test='test2.txt', format='tsv',
    fields=[('TEXT', TEXT), ('LABEL', LABELS)])

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
train_iter, val_iter, test_iter = data.Iterator.splits(
    (train, val, test), batch_sizes=(64, 64, 64), device=device, repeat=False, sort=False)

TEXT.build_vocab(train, min_freq=2)
LABELS.build_vocab(train)
model = RNN(len(TEXT.vocab.stoi) + 1, num_layers=2, output_size=4)

for epoch in range(1):
    model.train()
    for batch in train_iter:
        x, y = batch.TEXT, batch.LABEL
        y_pred = model(x)
        print(y_pred)
        print(y_pred.shape)
pythonを用いて、ID番号で表現された単語列$$\boldsymbol{x} = (x_1, x_2, \dots, x_T)$$がある。ただし、$$T$$は単語列の長さ、$$x_t \in \mathbb{R}^{V}$$は単語のID番号のone-hot表記である（$$V$$は単語の総数である）。再帰型ニューラルネットワーク（RNN: Recurrent Neural Network）を用い、単語列$$\boldsymbol{x}$$からカテゴリ$$y$$を予測するモデルとして、次式を実装せよ。

$$
\overrightarrow{h}_0 = 0, \\
\overrightarrow{h}_t = {\rm \overrightarrow{RNN}}(\mathrm{emb}(x_t), \overrightarrow{h}_{t-1}), \\
y = {\rm softmax}(W^{(yh)} \overrightarrow{h}_T + b^{(y)})
$$

ただし、$$\mathrm{emb}(x) \in \mathbb{R}^{d_w}$$は単語埋め込み（単語のone-hot表記から単語ベクトルに変換する関数）、$$\overrightarrow{h}_t \in \mathbb{R}^{d_h}$$は時刻$$t$$の隠れ状態ベクトル、$${\rm \overrightarrow{RNN}}(x,h)$$は入力$$x$$と前時刻の隠れ状態$$h$$から次状態を計算するRNNユニット、$$W^{(yh)} \in \mathbb{R}^{L \times d_h}$$は隠れ状態ベクトルからカテゴリを予測するための行列、$$b^{(y)} \in \mathbb{R}^{L}$$はバイアス項である（$$d_w, d_h, L$$はそれぞれ、単語埋め込みの次元数、隠れ状態ベクトルの次元数、ラベル数である）。RNNユニット$${\rm \overrightarrow{RNN}}(x,h)$$には様々な構成が考えられるが、典型例として次式が挙げられる。

$$
{\rm \overrightarrow{RNN}}(x,h) = g(W^{(hx)} x + W^{(hh)}h + b^{(h)})
$$

ただし、$$W^{(hx)} \in \mathbb{R}^{d_h \times d_w}、W^{(hh)} \in \mathbb{R}^{d_h \times d_h}, b^{(h)} \in \mathbb{R}^{d_h}$$はRNNユニットのパラメータ、$$g$$は活性化関数（例えば$$\tanh$$やReLUなど）である。

なお、この問題ではパラメータの学習を行わず、ランダムに初期化されたパラメータで$$y$$を計算するだけでよい。次元数などのハイパーパラメータは、$$d_w = 300, d_h=50$$など、適当な値に設定せよ（以降の問題でも同様である）。