# -*- coding: utf-8 -*-


"""
Pytorch models
__author__ = 'Jamie (krikit@naver.com)'
__copyright__ = 'No copyright. Just copyleft!'
"""


# pylint: disable=no-member
# pylint: disable=invalid-name


###########
# imports #
###########
import torch
import torch.nn as nn
from embedder import Embedder
from pos_models import PosTagger, FnnTagger, CnnTagger    # pylint: disable=unused-import


#############
# Ner Class #
#############
class Ner(nn.Module):
    """
    named entity recognizer pytorch model
    """
    def __init__(self, embedder, encoder, decoder):
        """
        * embedder (Embedder)
            [sentence_len, context_len] => [sentence_len, context_len, embed_dim]
        * encoder (nn.Module)
            [sentence_len, context_len, embed_dim] => [sentence_len, hidden_dim]
        * decoder (nn.Module)
            [sentence_len, hidden_dim] => [sentence_len, n_tags],
        """
        super().__init__()
        self.embedder = embedder
        self.encoder = encoder
        self.decoder = decoder

        assert isinstance(embedder, Embedder)
        assert isinstance(encoder, nn.Module)
        assert isinstance(decoder, nn.Module)

    def forward(self, sentence, gazet, pos, words): #pylint: disable=arguments-differ
        # [sentence_len, context_len] => [sentence_len, context_len, embed_dim]
        sentence_embed = self.embedder(sentence, gazet, pos, words)

        # [sentence_len, context_len, embed_dim] => [sentence_len, hidden_dim]
        hidden = self.encoder(sentence_embed)

        # [sentence_len, hidden_dim]   => [sentence_len, n_tags]
        predicted_tags = self.decoder(hidden)

        return predicted_tags

    def save(self, path):
        """
        모델을 저장하는 메소드
        :param  path:  경로
        """
        if torch.cuda.is_available():
            self.cpu()
        torch.save(self, str(path))
        if torch.cuda.is_available():
            self.cuda()

    @classmethod
    def load(cls, path):
        """
        저장된 모델을 로드하는 메소드
        :param  path:  경로
        :return:  모델 클래스 객체
        """
        model = torch.load(str(path))
        if torch.cuda.is_available():
            model.cuda()
        return model


#################
# Encoder Class #
#################
class Fnn5(nn.Module):
    """
    2-Layer Full-Connected Neural Networks
    """
    def __init__(self, context_len=21, in_dim=50, hidden_dim=500):
        super(Fnn5, self).__init__()

        self.context_len = context_len
        self.hidden_dim = hidden_dim
        self.out_dim = hidden_dim

        self.net = nn.Sequential(
            nn.Linear(context_len*in_dim, hidden_dim),
        )

    def forward(self, x):#pylint: disable=arguments-differ
        """
        Args:
            x: [sentence_len, context_len, in_dim]
        Return:
            x: [sentence_len, out_dim]
        """
        sentence_len = x.size(0)
        x = x.view(sentence_len, -1) # [sentence_len, context_len x in_dim]
        x = self.net(x) # [setence_len, out_dim]
        return x


class Cnn7(nn.Module):
    """
    ConvNet kernels=[2,3,4,5] + Fully-Connected
    """
    def __init__(self, in_dim=50, hidden_dim=500):
        """
        """
        super(Cnn7, self).__init__()

        self.in_dim = in_dim
        self.hidden_dim = hidden_dim
        self.out_dim = in_dim * 4

        self.conv2 = nn.Sequential(
            nn.Conv1d(in_dim, in_dim, kernel_size=2),    # 20
            nn.ReLU(),
            nn.MaxPool1d(kernel_size=2, ceil_mode=True),    # 10
            nn.Conv1d(in_dim, in_dim, kernel_size=2),    # 9
            nn.ReLU(),
            nn.MaxPool1d(kernel_size=2, ceil_mode=True),    # 5
            nn.Conv1d(in_dim, in_dim, kernel_size=2),    # 4
            nn.ReLU(),
            nn.MaxPool1d(kernel_size=2, ceil_mode=True),    # 2
            nn.Conv1d(in_dim, in_dim, kernel_size=2),    # 1
        )

        self.conv3 = nn.Sequential(
            nn.Conv1d(in_dim, in_dim, kernel_size=3, padding=1),    # 21
            nn.ReLU(),
            nn.MaxPool1d(kernel_size=2, ceil_mode=True),    # 11
            nn.Conv1d(in_dim, in_dim, kernel_size=3, padding=1),    # 11
            nn.ReLU(),
            nn.MaxPool1d(kernel_size=2, ceil_mode=True),    # 6
            nn.Conv1d(in_dim, in_dim, kernel_size=3, padding=1),    # 6
            nn.ReLU(),
            nn.MaxPool1d(kernel_size=2, ceil_mode=True),    # 3
            nn.Conv1d(in_dim, in_dim, kernel_size=3),    # 1
        )


        self.conv4 = nn.Sequential(
            nn.Conv1d(in_dim, in_dim, kernel_size=4, padding=1),    # 20
            nn.ReLU(),
            nn.MaxPool1d(kernel_size=2, ceil_mode=True),    # 10
            nn.Conv1d(in_dim, in_dim, kernel_size=4, padding=1),    # 9
            nn.ReLU(),
            nn.MaxPool1d(kernel_size=2, ceil_mode=True),    # 5
            nn.Conv1d(in_dim, in_dim, kernel_size=4, padding=1),    # 4
            nn.ReLU(),

            nn.Conv1d(in_dim, in_dim, kernel_size=4),    # 1
        )

        self.conv5 = nn.Sequential(
            nn.Conv1d(in_dim, in_dim, kernel_size=5, padding=2),    # 21
            nn.ReLU(),
            nn.MaxPool1d(kernel_size=2, ceil_mode=True),    # 11
            nn.Conv1d(in_dim, in_dim, kernel_size=5, padding=2),    # 11
            nn.ReLU(),
            nn.MaxPool1d(kernel_size=2, ceil_mode=True),    # 6
            nn.Conv1d(in_dim, in_dim, kernel_size=5, padding=2),    # 6
            nn.ReLU(),
            nn.MaxPool1d(kernel_size=2, ceil_mode=True),    # 3
            nn.Conv1d(in_dim, in_dim, kernel_size=5, padding=1),    # 1
        )

    def forward(self, x):   #pylint: disable=arguments-differ
        """
        Args:
            x: [sentence_length, context_len, in_dim]
        Return:
            x: [sentence_length, in_dim * 4]
        """

        # [sentence_length, in_dim, context_len]
        x = x.transpose(1, 2)

        conv2 = self.conv2(x).squeeze(-1) # [sentence_len, in_dim]
        conv3 = self.conv3(x).squeeze(-1) # [sentence_len, in_dim]
        conv4 = self.conv4(x).squeeze(-1) # [sentence_len, in_dim]
        conv5 = self.conv5(x).squeeze(-1) # [sentence_len, in_dim]

        # [sentence_len, in_dim * 4]
        out = torch.cat([conv2, conv3, conv4, conv5], dim=1)

        return out


class Cnn8(nn.Module):
    """
    9-layer Conv NN + Batch Norm + Residual
    """
    def __init__(self, context_len=21, in_dim=64, hidden_dim=None):
        super(Cnn8, self).__init__()

        self.context_len = context_len

        # conv block 64
        self.conv_block1_1 = self.conv_block(in_dim, 2, False)
        self.conv_block1_2_1 = self.conv_block(in_dim, 1, False)
        self.conv_block1_2_2 = self.conv_block(in_dim, 1, True)
        self.pool1 = nn.MaxPool1d(kernel_size=2, padding=1, ceil_mode=True)

        # conv block 128
        self.conv_block2_1 = self.conv_block(in_dim*2, 2, False)
        self.conv_block2_2_1 = self.conv_block(in_dim*2, 1, False)
        self.conv_block2_2_2 = self.conv_block(in_dim*2, 1, True)
        self.pool2 = nn.MaxPool1d(kernel_size=2, padding=1, ceil_mode=True)

        # conv block 256
        self.conv_block3_1 = self.conv_block(in_dim*4, 2, False)
        self.conv_block3_2_1 = self.conv_block(in_dim*4, 1, False)
        self.conv_block3_2_2 = self.conv_block(in_dim*4, 1, True)
        self.pool3 = nn.MaxPool1d(kernel_size=2)

        # conv block 512
        self.conv_block4_1 = self.conv_block(in_dim*8, 2, False)
        self.conv_block4_2_1 = self.conv_block(in_dim*8, 1, False)
        self.conv_block4_2_2 = self.conv_block(in_dim*8, 1, True)
        self.pool4 = nn.MaxPool1d(kernel_size=3)

        self.out_dim = in_dim*16

    @classmethod
    def conv_block(cls, in_dim=64, depth=2, double=True):
        """
        Args:
            [batch_size, dim, length]
        Return:
            [batch_size, dim*2, length] if double=True
            [batch_size, dim, length] if double=False
        """
        out_dim = in_dim
        layers = []
        for i in range(depth):
            if double:
                if i == depth - 1:
                    out_dim = in_dim * 2
            layers.append(nn.Conv1d(in_dim, out_dim, kernel_size=3, padding=1))
            layers.append(nn.BatchNorm1d(out_dim))
            layers.append(nn.ReLU())
        return nn.Sequential(*layers)

    def forward(self, sentence):#pylint: disable=arguments-differ
        """
        Args:
            sentence: [sentence_len, context_len, embed_dim]
        Return:
            logit: [batch_size, out_dim]
        """
        # [sentence_len, embed_dim, context_len]
        x = sentence.transpose(1, 2)

        # conv block 64
        x = self.conv_block1_1(x) + x # [batch, in_dim, 21]
        x = self.conv_block1_2_1(x) + x # [batch, in_dim, 21]
        x = self.conv_block1_2_2(x) # [batch, in_dim*2, 21]
        x = self.pool1(x) # [batch, in_dim*2, 11]

        # conv block 128
        x = self.conv_block2_1(x) + x # [batch, in_dim*2, 11]
        x = self.conv_block2_2_1(x) + x # [batch, in_dim*2, 11]
        x = self.conv_block2_2_2(x) # [batch, in_dim*4, 11]
        x = self.pool2(x) # [batch, in_dim*4, 6]

        # conv block 256
        x = self.conv_block3_1(x) + x # [batch, in_dim*4, 6]
        x = self.conv_block3_2_1(x) + x # [batch, in_dim*4, 6]
        x = self.conv_block3_2_2(x) # [batch, in_dim*8, 6]
        x = self.pool3(x) # [batch, in_dim*8, 3]

        # conv block 512
        x = self.conv_block4_1(x) + x # [batch, in_dim*8, 3]
        x = self.conv_block4_2_1(x) + x # [batch, in_dim*8, 3]
        x = self.conv_block4_2_2(x) # [batch, in_dim*16, 3]
        x = self.pool4(x) # [batch_size, in_dim*16, 1]
        x = x.squeeze(-1) # [batch, in_dim*16]

        return x


class RnnEncoder(nn.Module):
    """
    RNN Encoder Module
    """
    def __init__(self, context_len=21, in_dim=1024, out_dim=1024,
                 num_layers=2, cell='gru'):
        super(RnnEncoder, self).__init__()

        self.hidden_dim = out_dim // 2

        if cell == 'gru':
            self.rnn = nn.GRU(
                input_size=in_dim,
                hidden_size=self.hidden_dim,
                num_layers=num_layers,
                dropout=0.5,
                bidirectional=True)

        if cell == 'lstm':
            self.rnn = nn.LSTM(
                input_size=in_dim,
                hidden_size=self.hidden_dim,
                num_layers=num_layers,
                dropout=0.5,
                bidirectional=True)

        elif cell == 'sru':
            from sru import SRU
            self.rnn = SRU(
                input_size=in_dim,
                hidden_size=self.hidden_dim,
                num_layers=num_layers,
                dropout=0.5,
                bidirectional=True)

    def forward(self, x):#pylint: disable=arguments-differ
        """
        Args:
            x: [sentence_len, context_len, input_size]
        Return:
            x: [sentence_len, hidden_size]
        """
        # input (seq_len, batch, input_size)
        # h_0 (num_layers * num_directions, batch, hidden_size)

        # output (seq_len, batch, hidden_size * num_directions)
        # h_n (num_layers * num_directions, batch, hidden_size)

        #   [sequence_len, context_len, input_size]
        # =>[sentence_len, context_len, hidden_size x 2]
        x, _ = self.rnn(x)

        # [sequence_len, hidden_size x 2]
        x = x[:, 10, :]

        return x


#################
# Decoder Class #
#################

class FCDecoder(nn.Module):
    """
    Fully-Connected Decoder
    """
    def __init__(self, in_dim, hidden_dim, n_tags):
        super(FCDecoder, self).__init__()
        self.net = nn.Sequential(
            nn.ReLU(),
            nn.Dropout(),
            nn.Linear(in_dim, n_tags)
        )

    def forward(self, x):#pylint: disable=arguments-differ
        """
        [sentence_len, in_dim] => [sentence_len, n_tags]
        """
        return self.net(x)

class RnnDecoder(nn.Module):
    """
    RNN-based Decoder
    """
    def __init__(self, in_dim=1024, hidden_dim=512, n_tags=11,
                 num_layers=2, cell='gru'):
        super(RnnDecoder, self).__init__()
        if cell == 'gru':
            self.rnn = nn.GRU(
                input_size=in_dim,
                hidden_size=hidden_dim,
                num_layers=num_layers,
                dropout=0.5,
                bidirectional=True)

        if cell == 'lstm':
            self.rnn = nn.LSTM(
                input_size=in_dim,
                hidden_size=hidden_dim,
                num_layers=num_layers,
                dropout=0.5,
                bidirectional=True)

        elif cell == 'sru':
            from sru import SRU
            self.rnn = SRU(
                input_size=in_dim,
                hidden_size=hidden_dim,
                num_layers=num_layers,
                dropout=0.5,
                bidirectional=True)

        self.out = nn.Sequential(
            nn.ReLU(),
            nn.Dropout(),
            nn.Linear(hidden_dim * 2, n_tags)
        )

    def forward(self, x):#pylint: disable=arguments-differ
        """
        [sentence_len, in_dim] => [sentence_len, n_tags]
        """
        # input (seq_len, batch, input_size)
        # h_0 (num_layers * num_directions, batch, hidden_size)

        # output (seq_len, batch, hidden_size * num_directions)
        # h_n (num_layers * num_directions, batch, hidden_size)

        # [sentence_len, batch=1, input_size]
        x = x.unsqueeze(1)

        # x: [sentence_len, batch=1, hidden_size x 2]
        # h_n: [num_layers * 2, batch=1, hidden_size]
        # c_n: [num_layers * 2, batch=1, hidden_size]
        x, _ = self.rnn(x)

        # [sequence_len, hidden_size x 2]
        x = x.squeeze(1)

        # [sequence_len, n_tags]
        x = self.out(x)

        return x

[((1783, 1808), 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), '()\n', (1806, 1808), False, 'import torch\n'), ((1880, 1905), 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), '()\n', (1903, 1905), False, 'import torch\n'), ((2126, 2151), 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), '()\n', (2149, 2151), False, 'import torch\n'), ((6216, 6262), 'torch.cat', 'torch.cat', (['[conv2, conv3, conv4, conv5]'], {'dim': '(1)'}), '([conv2, conv3, conv4, conv5], dim=1)\n', (6225, 6262), False, 'import torch\n'), ((6751, 6805), 'torch.nn.MaxPool1d', 'nn.MaxPool1d', ([], {'kernel_size': '(2)', 'padding': '(1)', 'ceil_mode': '(True)'}), '(kernel_size=2, padding=1, ceil_mode=True)\n', (6763, 6805), True, 'import torch.nn as nn\n'), ((7051, 7105), 'torch.nn.MaxPool1d', 'nn.MaxPool1d', ([], {'kernel_size': '(2)', 'padding': '(1)', 'ceil_mode': '(True)'}), '(kernel_size=2, padding=1, ceil_mode=True)\n', (7063, 7105), True, 'import torch.nn as nn\n'), ((7351, 7378), 'torch.nn.MaxPool1d', 'nn.MaxPool1d', ([], {'kernel_size': '(2)'}), '(kernel_size=2)\n', (7363, 7378), True, 'import torch.nn as nn\n'), ((7624, 7651), 'torch.nn.MaxPool1d', 'nn.MaxPool1d', ([], {'kernel_size': '(3)'}), '(kernel_size=3)\n', (7636, 7651), True, 'import torch.nn as nn\n'), ((8321, 8343), 'torch.nn.Sequential', 'nn.Sequential', (['*layers'], {}), '(*layers)\n', (8334, 8343), True, 'import torch.nn as nn\n'), ((2599, 2642), 'torch.nn.Linear', 'nn.Linear', (['(context_len * in_dim)', 'hidden_dim'], {}), '(context_len * in_dim, hidden_dim)\n', (2608, 2642), True, 'import torch.nn as nn\n'), ((3382, 3422), 'torch.nn.Conv1d', 'nn.Conv1d', (['in_dim', 'in_dim'], {'kernel_size': '(2)'}), '(in_dim, in_dim, kernel_size=2)\n', (3391, 3422), True, 'import torch.nn as nn\n'), ((3444, 3453), 'torch.nn.ReLU', 'nn.ReLU', ([], {}), '()\n', (3451, 3453), True, 'import torch.nn as nn\n'), ((3467, 3510), 'torch.nn.MaxPool1d', 'nn.MaxPool1d', ([], {'kernel_size': '(2)', 'ceil_mode': '(True)'}), '(kernel_size=2, ceil_mode=True)\n', (3479, 3510), True, 'import torch.nn as nn\n'), ((3532, 3572), 'torch.nn.Conv1d', 'nn.Conv1d', (['in_dim', 'in_dim'], {'kernel_size': '(2)'}), '(in_dim, in_dim, kernel_size=2)\n', (3541, 3572), True, 'import torch.nn as nn\n'), ((3593, 3602), 'torch.nn.ReLU', 'nn.ReLU', ([], {}), '()\n', (3600, 3602), True, 'import torch.nn as nn\n'), ((3616, 3659), 'torch.nn.MaxPool1d', 'nn.MaxPool1d', ([], {'kernel_size': '(2)', 'ceil_mode': '(True)'}), '(kernel_size=2, ceil_mode=True)\n', (3628, 3659), True, 'import torch.nn as nn\n'), ((3680, 3720), 'torch.nn.Conv1d', 'nn.Conv1d', (['in_dim', 'in_dim'], {'kernel_size': '(2)'}), '(in_dim, in_dim, kernel_size=2)\n', (3689, 3720), True, 'import torch.nn as nn\n'), ((3741, 3750), 'torch.nn.ReLU', 'nn.ReLU', ([], {}), '()\n', (3748, 3750), True, 'import torch.nn as nn\n'), ((3764, 3807), 'torch.nn.MaxPool1d', 'nn.MaxPool1d', ([], {'kernel_size': '(2)', 'ceil_mode': '(True)'}), '(kernel_size=2, ceil_mode=True)\n', (3776, 3807), True, 'import torch.nn as nn\n'), ((3828, 3868), 'torch.nn.Conv1d', 'nn.Conv1d', (['in_dim', 'in_dim'], {'kernel_size': '(2)'}), '(in_dim, in_dim, kernel_size=2)\n', (3837, 3868), True, 'import torch.nn as nn\n'), ((3936, 3987), 'torch.nn.Conv1d', 'nn.Conv1d', (['in_dim', 'in_dim'], {'kernel_size': '(3)', 'padding': '(1)'}), '(in_dim, in_dim, kernel_size=3, padding=1)\n', (3945, 3987), True, 'import torch.nn as nn\n'), ((4009, 4018), 'torch.nn.ReLU', 'nn.ReLU', ([], {}), '()\n', (4016, 4018), True, 'import torch.nn as nn\n'), ((4032, 4075), 'torch.nn.MaxPool1d', 'nn.MaxPool1d', ([], {'kernel_size': '(2)', 'ceil_mode': '(True)'}), '(kernel_size=2, ceil_mode=True)\n', (4044, 4075), True, 'import torch.nn as nn\n'), ((4097, 4148), 'torch.nn.Conv1d', 'nn.Conv1d', (['in_dim', 'in_dim'], {'kernel_size': '(3)', 'padding': '(1)'}), '(in_dim, in_dim, kernel_size=3, padding=1)\n', (4106, 4148), True, 'import torch.nn as nn\n'), ((4170, 4179), 'torch.nn.ReLU', 'nn.ReLU', ([], {}), '()\n', (4177, 4179), True, 'import torch.nn as nn\n'), ((4193, 4236), 'torch.nn.MaxPool1d', 'nn.MaxPool1d', ([], {'kernel_size': '(2)', 'ceil_mode': '(True)'}), '(kernel_size=2, ceil_mode=True)\n', (4205, 4236), True, 'import torch.nn as nn\n'), ((4257, 4308), 'torch.nn.Conv1d', 'nn.Conv1d', (['in_dim', 'in_dim'], {'kernel_size': '(3)', 'padding': '(1)'}), '(in_dim, in_dim, kernel_size=3, padding=1)\n', (4266, 4308), True, 'import torch.nn as nn\n'), ((4329, 4338), 'torch.nn.ReLU', 'nn.ReLU', ([], {}), '()\n', (4336, 4338), True, 'import torch.nn as nn\n'), ((4352, 4395), 'torch.nn.MaxPool1d', 'nn.MaxPool1d', ([], {'kernel_size': '(2)', 'ceil_mode': '(True)'}), '(kernel_size=2, ceil_mode=True)\n', (4364, 4395), True, 'import torch.nn as nn\n'), ((4416, 4456), 'torch.nn.Conv1d', 'nn.Conv1d', (['in_dim', 'in_dim'], {'kernel_size': '(3)'}), '(in_dim, in_dim, kernel_size=3)\n', (4425, 4456), True, 'import torch.nn as nn\n'), ((4525, 4576), 'torch.nn.Conv1d', 'nn.Conv1d', (['in_dim', 'in_dim'], {'kernel_size': '(4)', 'padding': '(1)'}), '(in_dim, in_dim, kernel_size=4, padding=1)\n', (4534, 4576), True, 'import torch.nn as nn\n'), ((4598, 4607), 'torch.nn.ReLU', 'nn.ReLU', ([], {}), '()\n', (4605, 4607), True, 'import torch.nn as nn\n'), ((4621, 4664), 'torch.nn.MaxPool1d', 'nn.MaxPool1d', ([], {'kernel_size': '(2)', 'ceil_mode': '(True)'}), '(kernel_size=2, ceil_mode=True)\n', (4633, 4664), True, 'import torch.nn as nn\n'), ((4686, 4737), 'torch.nn.Conv1d', 'nn.Conv1d', (['in_dim', 'in_dim'], {'kernel_size': '(4)', 'padding': '(1)'}), '(in_dim, in_dim, kernel_size=4, padding=1)\n', (4695, 4737), True, 'import torch.nn as nn\n'), ((4758, 4767), 'torch.nn.ReLU', 'nn.ReLU', ([], {}), '()\n', (4765, 4767), True, 'import torch.nn as nn\n'), ((4781, 4824), 'torch.nn.MaxPool1d', 'nn.MaxPool1d', ([], {'kernel_size': '(2)', 'ceil_mode': '(True)'}), '(kernel_size=2, ceil_mode=True)\n', (4793, 4824), True, 'import torch.nn as nn\n'), ((4845, 4896), 'torch.nn.Conv1d', 'nn.Conv1d', (['in_dim', 'in_dim'], {'kernel_size': '(4)', 'padding': '(1)'}), '(in_dim, in_dim, kernel_size=4, padding=1)\n', (4854, 4896), True, 'import torch.nn as nn\n'), ((4917, 4926), 'torch.nn.ReLU', 'nn.ReLU', ([], {}), '()\n', (4924, 4926), True, 'import torch.nn as nn\n'), ((4941, 4981), 'torch.nn.Conv1d', 'nn.Conv1d', (['in_dim', 'in_dim'], {'kernel_size': '(4)'}), '(in_dim, in_dim, kernel_size=4)\n', (4950, 4981), True, 'import torch.nn as nn\n'), ((5049, 5100), 'torch.nn.Conv1d', 'nn.Conv1d', (['in_dim', 'in_dim'], {'kernel_size': '(5)', 'padding': '(2)'}), '(in_dim, in_dim, kernel_size=5, padding=2)\n', (5058, 5100), True, 'import torch.nn as nn\n'), ((5122, 5131), 'torch.nn.ReLU', 'nn.ReLU', ([], {}), '()\n', (5129, 5131), True, 'import torch.nn as nn\n'), ((5145, 5188), 'torch.nn.MaxPool1d', 'nn.MaxPool1d', ([], {'kernel_size': '(2)', 'ceil_mode': '(True)'}), '(kernel_size=2, ceil_mode=True)\n', (5157, 5188), True, 'import torch.nn as nn\n'), ((5210, 5261), 'torch.nn.Conv1d', 'nn.Conv1d', (['in_dim', 'in_dim'], {'kernel_size': '(5)', 'padding': '(2)'}), '(in_dim, in_dim, kernel_size=5, padding=2)\n', (5219, 5261), True, 'import torch.nn as nn\n'), ((5283, 5292), 'torch.nn.ReLU', 'nn.ReLU', ([], {}), '()\n', (5290, 5292), True, 'import torch.nn as nn\n'), ((5306, 5349), 'torch.nn.MaxPool1d', 'nn.MaxPool1d', ([], {'kernel_size': '(2)', 'ceil_mode': '(True)'}), '(kernel_size=2, ceil_mode=True)\n', (5318, 5349), True, 'import torch.nn as nn\n'), ((5370, 5421), 'torch.nn.Conv1d', 'nn.Conv1d', (['in_dim', 'in_dim'], {'kernel_size': '(5)', 'padding': '(2)'}), '(in_dim, in_dim, kernel_size=5, padding=2)\n', (5379, 5421), True, 'import torch.nn as nn\n'), ((5442, 5451), 'torch.nn.ReLU', 'nn.ReLU', ([], {}), '()\n', (5449, 5451), True, 'import torch.nn as nn\n'), ((5465, 5508), 'torch.nn.MaxPool1d', 'nn.MaxPool1d', ([], {'kernel_size': '(2)', 'ceil_mode': '(True)'}), '(kernel_size=2, ceil_mode=True)\n', (5477, 5508), True, 'import torch.nn as nn\n'), ((5529, 5580), 'torch.nn.Conv1d', 'nn.Conv1d', (['in_dim', 'in_dim'], {'kernel_size': '(5)', 'padding': '(1)'}), '(in_dim, in_dim, kernel_size=5, padding=1)\n', (5538, 5580), True, 'import torch.nn as nn\n'), ((10072, 10187), 'torch.nn.GRU', 'nn.GRU', ([], {'input_size': 'in_dim', 'hidden_size': 'self.hidden_dim', 'num_layers': 'num_layers', 'dropout': '(0.5)', 'bidirectional': '(True)'}), '(input_size=in_dim, hidden_size=self.hidden_dim, num_layers=\n    num_layers, dropout=0.5, bidirectional=True)\n', (10078, 10187), True, 'import torch.nn as nn\n'), ((10315, 10431), 'torch.nn.LSTM', 'nn.LSTM', ([], {'input_size': 'in_dim', 'hidden_size': 'self.hidden_dim', 'num_layers': 'num_layers', 'dropout': '(0.5)', 'bidirectional': '(True)'}), '(input_size=in_dim, hidden_size=self.hidden_dim, num_layers=\n    num_layers, dropout=0.5, bidirectional=True)\n', (10322, 10431), True, 'import torch.nn as nn\n'), ((11722, 11731), 'torch.nn.ReLU', 'nn.ReLU', ([], {}), '()\n', (11729, 11731), True, 'import torch.nn as nn\n'), ((11745, 11757), 'torch.nn.Dropout', 'nn.Dropout', ([], {}), '()\n', (11755, 11757), True, 'import torch.nn as nn\n'), ((11771, 11796), 'torch.nn.Linear', 'nn.Linear', (['in_dim', 'n_tags'], {}), '(in_dim, n_tags)\n', (11780, 11796), True, 'import torch.nn as nn\n'), ((12242, 12351), 'torch.nn.GRU', 'nn.GRU', ([], {'input_size': 'in_dim', 'hidden_size': 'hidden_dim', 'num_layers': 'num_layers', 'dropout': '(0.5)', 'bidirectional': '(True)'}), '(input_size=in_dim, hidden_size=hidden_dim, num_layers=num_layers,\n    dropout=0.5, bidirectional=True)\n', (12248, 12351), True, 'import torch.nn as nn\n'), ((12480, 12590), 'torch.nn.LSTM', 'nn.LSTM', ([], {'input_size': 'in_dim', 'hidden_size': 'hidden_dim', 'num_layers': 'num_layers', 'dropout': '(0.5)', 'bidirectional': '(True)'}), '(input_size=in_dim, hidden_size=hidden_dim, num_layers=num_layers,\n    dropout=0.5, bidirectional=True)\n', (12487, 12590), True, 'import torch.nn as nn\n'), ((12983, 12992), 'torch.nn.ReLU', 'nn.ReLU', ([], {}), '()\n', (12990, 12992), True, 'import torch.nn as nn\n'), ((13006, 13018), 'torch.nn.Dropout', 'nn.Dropout', ([], {}), '()\n', (13016, 13018), True, 'import torch.nn as nn\n'), ((13032, 13065), 'torch.nn.Linear', 'nn.Linear', (['(hidden_dim * 2)', 'n_tags'], {}), '(hidden_dim * 2, n_tags)\n', (13041, 13065), True, 'import torch.nn as nn\n'), ((8164, 8216), 'torch.nn.Conv1d', 'nn.Conv1d', (['in_dim', 'out_dim'], {'kernel_size': '(3)', 'padding': '(1)'}), '(in_dim, out_dim, kernel_size=3, padding=1)\n', (8173, 8216), True, 'import torch.nn as nn\n'), ((8244, 8267), 'torch.nn.BatchNorm1d', 'nn.BatchNorm1d', (['out_dim'], {}), '(out_dim)\n', (8258, 8267), True, 'import torch.nn as nn\n'), ((8295, 8304), 'torch.nn.ReLU', 'nn.ReLU', ([], {}), '()\n', (8302, 8304), True, 'import torch.nn as nn\n'), ((10592, 10703), 'sru.SRU', 'SRU', ([], {'input_size': 'in_dim', 'hidden_size': 'self.hidden_dim', 'num_layers': 'num_layers', 'dropout': '(0.5)', 'bidirectional': '(True)'}), '(input_size=in_dim, hidden_size=self.hidden_dim, num_layers=num_layers,\n    dropout=0.5, bidirectional=True)\n', (10595, 10703), False, 'from sru import SRU\n'), ((12752, 12858), 'sru.SRU', 'SRU', ([], {'input_size': 'in_dim', 'hidden_size': 'hidden_dim', 'num_layers': 'num_layers', 'dropout': '(0.5)', 'bidirectional': '(True)'}), '(input_size=in_dim, hidden_size=hidden_dim, num_layers=num_layers,\n    dropout=0.5, bidirectional=True)\n', (12755, 12858), False, 'from sru import SRU\n')]