# -*- 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

[(66, 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), False, 'import torch\n'), (69, 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), False, 'import torch\n'), (80, 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), False, 'import torch\n'), (199, 'torch.cat', 'torch.cat', (['[conv2, conv3, conv4, conv5]'], {'dim': '(1)'}), False, 'import torch\n'), (217, 'torch.nn.MaxPool1d', 'nn.MaxPool1d', ([], {'kernel_size': '(2)', 'padding': '(1)', 'ceil_mode': '(True)'}), True, 'import torch.nn as nn\n'), (223, 'torch.nn.MaxPool1d', 'nn.MaxPool1d', ([], {'kernel_size': '(2)', 'padding': '(1)', 'ceil_mode': '(True)'}), True, 'import torch.nn as nn\n'), (229, 'torch.nn.MaxPool1d', 'nn.MaxPool1d', ([], {'kernel_size': '(2)'}), True, 'import torch.nn as nn\n'), (235, 'torch.nn.MaxPool1d', 'nn.MaxPool1d', ([], {'kernel_size': '(3)'}), True, 'import torch.nn as nn\n'), (257, 'torch.nn.Sequential', 'nn.Sequential', (['*layers'], {}), True, 'import torch.nn as nn\n'), (100, 'torch.nn.Linear', 'nn.Linear', (['(context_len * in_dim)', 'hidden_dim'], {}), True, 'import torch.nn as nn\n'), (130, 'torch.nn.Conv1d', 'nn.Conv1d', (['in_dim', 'in_dim'], {'kernel_size': '(2)'}), True, 'import torch.nn as nn\n'), (131, 'torch.nn.ReLU', 'nn.ReLU', ([], {}), True, 'import torch.nn as nn\n'), (132, 'torch.nn.MaxPool1d', 'nn.MaxPool1d', ([], {'kernel_size': '(2)', 'ceil_mode': '(True)'}), True, 'import torch.nn as nn\n'), (133, 'torch.nn.Conv1d', 'nn.Conv1d', (['in_dim', 'in_dim'], {'kernel_size': '(2)'}), True, 'import torch.nn as nn\n'), (134, 'torch.nn.ReLU', 'nn.ReLU', ([], {}), True, 'import torch.nn as nn\n'), (135, 'torch.nn.MaxPool1d', 'nn.MaxPool1d', ([], {'kernel_size': '(2)', 'ceil_mode': '(True)'}), True, 'import torch.nn as nn\n'), (136, 'torch.nn.Conv1d', 'nn.Conv1d', (['in_dim', 'in_dim'], {'kernel_size': '(2)'}), True, 'import torch.nn as nn\n'), (137, 'torch.nn.ReLU', 'nn.ReLU', ([], {}), True, 'import torch.nn as nn\n'), (138, 'torch.nn.MaxPool1d', 'nn.MaxPool1d', ([], {'kernel_size': '(2)', 'ceil_mode': '(True)'}), True, 'import torch.nn as nn\n'), (139, 'torch.nn.Conv1d', 'nn.Conv1d', (['in_dim', 'in_dim'], {'kernel_size': '(2)'}), True, 'import torch.nn as nn\n'), (143, 'torch.nn.Conv1d', 'nn.Conv1d', (['in_dim', 'in_dim'], {'kernel_size': '(3)', 'padding': '(1)'}), True, 'import torch.nn as nn\n'), (144, 'torch.nn.ReLU', 'nn.ReLU', ([], {}), True, 'import torch.nn as nn\n'), (145, 'torch.nn.MaxPool1d', 'nn.MaxPool1d', ([], {'kernel_size': '(2)', 'ceil_mode': '(True)'}), True, 'import torch.nn as nn\n'), (146, 'torch.nn.Conv1d', 'nn.Conv1d', (['in_dim', 'in_dim'], {'kernel_size': '(3)', 'padding': '(1)'}), True, 'import torch.nn as nn\n'), (147, 'torch.nn.ReLU', 'nn.ReLU', ([], {}), True, 'import torch.nn as nn\n'), (148, 'torch.nn.MaxPool1d', 'nn.MaxPool1d', ([], {'kernel_size': '(2)', 'ceil_mode': '(True)'}), True, 'import torch.nn as nn\n'), (149, 'torch.nn.Conv1d', 'nn.Conv1d', (['in_dim', 'in_dim'], {'kernel_size': '(3)', 'padding': '(1)'}), True, 'import torch.nn as nn\n'), (150, 'torch.nn.ReLU', 'nn.ReLU', ([], {}), True, 'import torch.nn as nn\n'), (151, 'torch.nn.MaxPool1d', 'nn.MaxPool1d', ([], {'kernel_size': '(2)', 'ceil_mode': '(True)'}), True, 'import torch.nn as nn\n'), (152, 'torch.nn.Conv1d', 'nn.Conv1d', (['in_dim', 'in_dim'], {'kernel_size': '(3)'}), True, 'import torch.nn as nn\n'), (157, 'torch.nn.Conv1d', 'nn.Conv1d', (['in_dim', 'in_dim'], {'kernel_size': '(4)', 'padding': '(1)'}), True, 'import torch.nn as nn\n'), (158, 'torch.nn.ReLU', 'nn.ReLU', ([], {}), True, 'import torch.nn as nn\n'), (159, 'torch.nn.MaxPool1d', 'nn.MaxPool1d', ([], {'kernel_size': '(2)', 'ceil_mode': '(True)'}), True, 'import torch.nn as nn\n'), (160, 'torch.nn.Conv1d', 'nn.Conv1d', (['in_dim', 'in_dim'], {'kernel_size': '(4)', 'padding': '(1)'}), True, 'import torch.nn as nn\n'), (161, 'torch.nn.ReLU', 'nn.ReLU', ([], {}), True, 'import torch.nn as nn\n'), (162, 'torch.nn.MaxPool1d', 'nn.MaxPool1d', ([], {'kernel_size': '(2)', 'ceil_mode': '(True)'}), True, 'import torch.nn as nn\n'), (163, 'torch.nn.Conv1d', 'nn.Conv1d', (['in_dim', 'in_dim'], {'kernel_size': '(4)', 'padding': '(1)'}), True, 'import torch.nn as nn\n'), (164, 'torch.nn.ReLU', 'nn.ReLU', ([], {}), True, 'import torch.nn as nn\n'), (166, 'torch.nn.Conv1d', 'nn.Conv1d', (['in_dim', 'in_dim'], {'kernel_size': '(4)'}), True, 'import torch.nn as nn\n'), (170, 'torch.nn.Conv1d', 'nn.Conv1d', (['in_dim', 'in_dim'], {'kernel_size': '(5)', 'padding': '(2)'}), True, 'import torch.nn as nn\n'), (171, 'torch.nn.ReLU', 'nn.ReLU', ([], {}), True, 'import torch.nn as nn\n'), (172, 'torch.nn.MaxPool1d', 'nn.MaxPool1d', ([], {'kernel_size': '(2)', 'ceil_mode': '(True)'}), True, 'import torch.nn as nn\n'), (173, 'torch.nn.Conv1d', 'nn.Conv1d', (['in_dim', 'in_dim'], {'kernel_size': '(5)', 'padding': '(2)'}), True, 'import torch.nn as nn\n'), (174, 'torch.nn.ReLU', 'nn.ReLU', ([], {}), True, 'import torch.nn as nn\n'), (175, 'torch.nn.MaxPool1d', 'nn.MaxPool1d', ([], {'kernel_size': '(2)', 'ceil_mode': '(True)'}), True, 'import torch.nn as nn\n'), (176, 'torch.nn.Conv1d', 'nn.Conv1d', (['in_dim', 'in_dim'], {'kernel_size': '(5)', 'padding': '(2)'}), True, 'import torch.nn as nn\n'), (177, 'torch.nn.ReLU', 'nn.ReLU', ([], {}), True, 'import torch.nn as nn\n'), (178, 'torch.nn.MaxPool1d', 'nn.MaxPool1d', ([], {'kernel_size': '(2)', 'ceil_mode': '(True)'}), True, 'import torch.nn as nn\n'), (179, 'torch.nn.Conv1d', 'nn.Conv1d', (['in_dim', 'in_dim'], {'kernel_size': '(5)', 'padding': '(1)'}), True, 'import torch.nn as nn\n'), (308, 'torch.nn.GRU', 'nn.GRU', ([], {'input_size': 'in_dim', 'hidden_size': 'self.hidden_dim', 'num_layers': 'num_layers', 'dropout': '(0.5)', 'bidirectional': '(True)'}), True, 'import torch.nn as nn\n'), (316, 'torch.nn.LSTM', 'nn.LSTM', ([], {'input_size': 'in_dim', 'hidden_size': 'self.hidden_dim', 'num_layers': 'num_layers', 'dropout': '(0.5)', 'bidirectional': '(True)'}), True, 'import torch.nn as nn\n'), (366, 'torch.nn.ReLU', 'nn.ReLU', ([], {}), True, 'import torch.nn as nn\n'), (367, 'torch.nn.Dropout', 'nn.Dropout', ([], {}), True, 'import torch.nn as nn\n'), (368, 'torch.nn.Linear', 'nn.Linear', (['in_dim', 'n_tags'], {}), True, 'import torch.nn as nn\n'), (385, 'torch.nn.GRU', 'nn.GRU', ([], {'input_size': 'in_dim', 'hidden_size': 'hidden_dim', 'num_layers': 'num_layers', 'dropout': '(0.5)', 'bidirectional': '(True)'}), True, 'import torch.nn as nn\n'), (393, 'torch.nn.LSTM', 'nn.LSTM', ([], {'input_size': 'in_dim', 'hidden_size': 'hidden_dim', 'num_layers': 'num_layers', 'dropout': '(0.5)', 'bidirectional': '(True)'}), True, 'import torch.nn as nn\n'), (410, 'torch.nn.ReLU', 'nn.ReLU', ([], {}), True, 'import torch.nn as nn\n'), (411, 'torch.nn.Dropout', 'nn.Dropout', ([], {}), True, 'import torch.nn as nn\n'), (412, 'torch.nn.Linear', 'nn.Linear', (['(hidden_dim * 2)', 'n_tags'], {}), True, 'import torch.nn as nn\n'), (254, 'torch.nn.Conv1d', 'nn.Conv1d', (['in_dim', 'out_dim'], {'kernel_size': '(3)', 'padding': '(1)'}), True, 'import torch.nn as nn\n'), (255, 'torch.nn.BatchNorm1d', 'nn.BatchNorm1d', (['out_dim'], {}), True, 'import torch.nn as nn\n'), (256, 'torch.nn.ReLU', 'nn.ReLU', ([], {}), True, 'import torch.nn as nn\n'), (325, 'sru.SRU', 'SRU', ([], {'input_size': 'in_dim', 'hidden_size': 'self.hidden_dim', 'num_layers': 'num_layers', 'dropout': '(0.5)', 'bidirectional': '(True)'}), False, 'from sru import SRU\n'), (402, 'sru.SRU', 'SRU', ([], {'input_size': 'in_dim', 'hidden_size': 'hidden_dim', 'num_layers': 'num_layers', 'dropout': '(0.5)', 'bidirectional': '(True)'}), False, 'from sru import SRU\n')]