<reponame>hamed-sadeghi-layer6/transformer-kk-mimic<filename>Transformer_classifier.py
# -*- coding: utf-8 -*-
"""
Created on Tue Oct  9 16:46:24 2018

@author: hamed
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformer.Models import Encoder
#from transformer.SubLayers import PositionwiseFeedForward


class ffn_compressed(nn.Module):
    ''' A two-feed-forward-layer module '''

    def __init__(self, d_in, d_hid, d_out, dropout=0.1):
        super(ffn_compressed, self).__init__()
        
#        self.w_1 = nn.Conv1d(d_in, d_hid, 1) # position-wise
#        self.w_2 = nn.Conv1d(d_hid, d_out, 1) # position-wise
        self.w_1 = nn.Conv1d(d_in, d_out, 1)
        self.w_2 = nn.Conv1d(d_out, d_out, 1)
        
        self.layer_norm = nn.LayerNorm(d_out)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        output = x.transpose(1, 2)
        
        output = self.w_2(F.relu(self.w_1(output)))
        output = output.transpose(1, 2)
        output = self.dropout(output)
        output = self.layer_norm(output)
        return output
        
        
class model(nn.Module):
        
    def __init__(self, d_src_vec,            
                 len_seq,
                 d_emb_vec,
                 n_layers,
                 n_head, d_k,
                 d_v,
                 d_inner, dropout):
        
        super(model, self).__init__()
        self.d_src_vec = d_src_vec
        self.d_emb_vec = d_emb_vec
        self.len_seq = len_seq
        self.n_layers= n_layers
        self.n_head = n_head
        self.dropout = dropout
        self.d_inner = d_inner
        
        self.ffn = ffn_compressed(d_in=self.d_src_vec, d_hid=self.d_inner,
                                  d_out=self.d_emb_vec)
        
        self.encoder = Encoder(len_seq=self.len_seq, d_word_vec=self.d_emb_vec,
            n_layers=self.n_layers, n_head=self.n_head, d_k=self.d_emb_vec//self.n_head,
            d_v=self.d_emb_vec//self.n_head, d_inner=self.d_inner,
            dropout=self.dropout) 
            
        #Fully connected. Seems to have a lot of params
#        self.FC1 = nn.Linear(self.d_emb_vec * self.len_seq , 64)   
#        self.FC2 = nn.Linear(64, 8)
#        self.FC3 = nn.Linear(8, 2)
            
#        #Average pooling over features
#        self.avg_pooling = nn.AvgPool1d(d_emb_vec-1, stride=1)  #d_emb_vec-1: To have 2 classes
#        self.FC = nn.Linear(len_seq * 2, 2)  #2: binary classification
#        self.softmax = nn.Softmax(dim=-1)
        
        #Average pooling over sequence
        self.avg_pooling = nn.AvgPool1d(len_seq, stride=1)  #self.len_seq: To have 1 averaged token
        self.FC = nn.Linear(d_emb_vec, 2)  #2: binary classification
        self.softmax = nn.Softmax(dim=-1)
        
    def forward(self, x, x_pos):
        
        x = self.ffn(x)       
        x = self.encoder(x, x_pos, return_attns=False)
        
        #Fully connected
#        x =  x.view(x.shape[0], x.shape[1]*x.shape[2])
#        x = self.FC1(x); x = self.FC2(x); x = self.FC3(x)
        
#        #Average pooling over features
#        x = self.avg_pooling(x)
##        print("size after avg pooling = ", x.shape)
##        x = torch.squeeze(x)  #To get rid of 1-dimensional feature and results in [batch_size, len_seq] size
#        x = x.view(-1, self.len_seq*2)
##        print('shape after squeeze = ', x.shape)
#        x = self.FC(x)
        
        #Average pooling over sequence
        x = x.transpose(2,1)
        x = self.avg_pooling(x)
        x = x.transpose(2,1)
        x = torch.squeeze(x)  #To get rid of 1-dimensional feature and results in [batch_size, feature_size] size
        x = self.FC(x)
        
        return self.softmax(x)
    
       
# -*- coding: utf-8 -*-
"""
Created on Tue Oct  9 16:46:24 2018

@author: hamed
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformer.Models import Encoder
#from transformer.SubLayers import PositionwiseFeedForward


class ffn_compressed(nn.Module):
    ''' A two-feed-forward-layer module '''

    def __init__(self, d_in, d_hid, d_out, dropout=0.1):
        super(ffn_compressed, self).__init__()
        
#        self.w_1 = nn.Conv1d(d_in, d_hid, 1) # position-wise
#        self.w_2 = nn.Conv1d(d_hid, d_out, 1) # position-wise
        self.w_1 = nn.Conv1d(d_in, d_out, 1)
        self.w_2 = nn.Conv1d(d_out, d_out, 1)
        
        self.layer_norm = nn.LayerNorm(d_out)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        output = x.transpose(1, 2)
        
        output = self.w_2(F.relu(self.w_1(output)))
        output = output.transpose(1, 2)
        output = self.dropout(output)
        output = self.layer_norm(output)
        return output
        
        
class model(nn.Module):
        
    def __init__(self, d_src_vec,            
                 len_seq,
                 d_emb_vec,
                 n_layers,
                 n_head, d_k,
                 d_v,
                 d_inner, dropout):
        
        super(model, self).__init__()
        self.d_src_vec = d_src_vec
        self.d_emb_vec = d_emb_vec
        self.len_seq = len_seq
        self.n_layers= n_layers
        self.n_head = n_head
        self.dropout = dropout
        self.d_inner = d_inner
        
        self.ffn = ffn_compressed(d_in=self.d_src_vec, d_hid=self.d_inner,
                                  d_out=self.d_emb_vec)
        
        self.encoder = Encoder(len_seq=self.len_seq, d_word_vec=self.d_emb_vec,
            n_layers=self.n_layers, n_head=self.n_head, d_k=self.d_emb_vec//self.n_head,
            d_v=self.d_emb_vec//self.n_head, d_inner=self.d_inner,
            dropout=self.dropout) 
            
        #Fully connected. Seems to have a lot of params
#        self.FC1 = nn.Linear(self.d_emb_vec * self.len_seq , 64)   
#        self.FC2 = nn.Linear(64, 8)
#        self.FC3 = nn.Linear(8, 2)
            
#        #Average pooling over features
#        self.avg_pooling = nn.AvgPool1d(d_emb_vec-1, stride=1)  #d_emb_vec-1: To have 2 classes
#        self.FC = nn.Linear(len_seq * 2, 2)  #2: binary classification
#        self.softmax = nn.Softmax(dim=-1)
        
        #Average pooling over sequence
        self.avg_pooling = nn.AvgPool1d(len_seq, stride=1)  #self.len_seq: To have 1 averaged token
        self.FC = nn.Linear(d_emb_vec, 2)  #2: binary classification
        self.softmax = nn.Softmax(dim=-1)
        
    def forward(self, x, x_pos):
        
        x = self.ffn(x)       
        x = self.encoder(x, x_pos, return_attns=False)
        
        #Fully connected
#        x =  x.view(x.shape[0], x.shape[1]*x.shape[2])
#        x = self.FC1(x); x = self.FC2(x); x = self.FC3(x)
        
#        #Average pooling over features
#        x = self.avg_pooling(x)
##        print("size after avg pooling = ", x.shape)
##        x = torch.squeeze(x)  #To get rid of 1-dimensional feature and results in [batch_size, len_seq] size
#        x = x.view(-1, self.len_seq*2)
##        print('shape after squeeze = ', x.shape)
#        x = self.FC(x)
        
        #Average pooling over sequence
        x = x.transpose(2,1)
        x = self.avg_pooling(x)
        x = x.transpose(2,1)
        x = torch.squeeze(x)  #To get rid of 1-dimensional feature and results in [batch_size, feature_size] size
        x = self.FC(x)
        
        return self.softmax(x)
# -*- coding: utf-8 -*- Created on Tue Oct  9 16:46:24 2018  @author: hamed #from transformer.SubLayers import PositionwiseFeedForward A two-feed-forward-layer module #        self.w_1 = nn.Conv1d(d_in, d_hid, 1) # position-wise #        self.w_2 = nn.Conv1d(d_hid, d_out, 1) # position-wise #Fully connected. Seems to have a lot of params #        self.FC1 = nn.Linear(self.d_emb_vec * self.len_seq , 64) #        self.FC2 = nn.Linear(64, 8) #        self.FC3 = nn.Linear(8, 2) #        #Average pooling over features #        self.avg_pooling = nn.AvgPool1d(d_emb_vec-1, stride=1)  #d_emb_vec-1: To have 2 classes #        self.FC = nn.Linear(len_seq * 2, 2)  #2: binary classification #        self.softmax = nn.Softmax(dim=-1) #Average pooling over sequence #self.len_seq: To have 1 averaged token #2: binary classification #Fully connected #        x =  x.view(x.shape[0], x.shape[1]*x.shape[2]) #        x = self.FC1(x); x = self.FC2(x); x = self.FC3(x) #        #Average pooling over features #        x = self.avg_pooling(x) ##        print("size after avg pooling = ", x.shape) ##        x = torch.squeeze(x)  #To get rid of 1-dimensional feature and results in [batch_size, len_seq] size #        x = x.view(-1, self.len_seq*2) ##        print('shape after squeeze = ', x.shape) #        x = self.FC(x) #Average pooling over sequence #To get rid of 1-dimensional feature and results in [batch_size, feature_size] size