baseline.py

import torch
import numpy as np
import torch.nn as nn
from dataset import data_generator
from torch.optim import Adam
import time
import math
from utils import *
from monotonic_gru import MonotonicGru

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


class Baseline(nn.Module):
    def __init__(self, mapper, embedding_dimensions, data_dir=None, use_loc_ind=False, shared_layer_dim=512, hidden_dim=512, dropout=0.25, model='nemo', num_gru_layers=1, pre_trained=False):
        super().__init__()
        # model can be 'gru', 'm_gru'
        self.use_loc_ind = use_loc_ind
        self.embedding_dimensions = embedding_dimensions
        self.company_embed = nn.Embedding(
            max(mapper['companies'].values())+1, embedding_dimensions['companies'])
        self.school_embed = nn.Embedding(
            max(mapper['schools'].values())+1, embedding_dimensions['schools'])
        self.degree_embed = nn.Embedding(
            max(mapper['degrees'].values())+1, embedding_dimensions['degrees'])
        self.time_embed = nn.Linear(1, embedding_dimensions['times'])
        self.interval_embed = nn.Linear(1, embedding_dimensions['intervals'])
        if use_loc_ind:
            self.locality_embed = nn.Embedding(
                max(mapper['localities'].values())+1, embedding_dimensions['locality'])
            self.industry_embed = nn.Embedding(
                max(mapper['industries'].values())+1, embedding_dimensions['industry'])

        self.majors_embed = nn.Embedding(
            max(mapper['majors'].values())+1, embedding_dimensions['majors'])
        self.titles_embed = nn.Embedding(
            max(mapper['occupations'].values())+1, embedding_dimensions['occupations'])

        temp_dim = embedding_dimensions['industry'] + \
            embedding_dimensions['locality']
        self.total_input_dimension = sum(
            embedding_dimensions.values()) - temp_dim-embedding_dimensions['skills']
        self.total_fixed_dimension = temp_dim if use_loc_ind else 0
        self.num_gru_layers = num_gru_layers

        # Model parameter
        self.model = model
        if self.model == 'nemo':
            self.rnn = nn.GRU(input_size=self.total_input_dimension, hidden_size=hidden_dim,
                              num_layers=num_gru_layers, batch_first=False, bidirectional=False)
        elif self.model in {'nss'}:
            self.rnn = MonotonicGru(input_size=self.total_input_dimension,
                                    hidden_size=hidden_dim, num_layers=num_gru_layers, batch_first=False)

        # Final Parameters
        self.intermediate_layers = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(hidden_dim+self.total_fixed_dimension, shared_layer_dim),
            nn.Dropout(dropout),
            nn.ReLU()
        )
        self.company_final_layers = nn.Sequential(
            nn.Linear(shared_layer_dim, 256),
            nn.Dropout(dropout),
            nn.ReLU(),
            nn.Linear(256, embedding_dimensions['companies'])
        )
        self.title_final_layers = nn.Sequential(
            nn.Linear(shared_layer_dim, 256),
            nn.Dropout(dropout),
            nn.ReLU(),
            nn.Linear(256, embedding_dimensions['occupations'])
        )
        if self.model == 'nemo':
            self.skill_final_layers = nn.Sequential(
                nn.Dropout(dropout),
                nn.Linear(hidden_dim+self.total_fixed_dimension,
                          embedding_dimensions['skills'])
            )
            self.skill_embed = nn.Linear(
                embedding_dimensions['skills'], max(mapper['skills'].values())+1)
        elif self.model == 'nss':
            self.skill_final_layers = nn.Sequential(
                nn.Dropout(dropout),
                nn.Linear(hidden_dim+self.total_fixed_dimension,
                          embedding_dimensions['skills'])
            )
            self.skill_embed = nn.Linear(
                embedding_dimensions['skills'], max(mapper['skills'].values())+1)

        self.num_skill_label = max(mapper['skills'].values())+1
        if pre_trained:
            self.company_embed.weight.data = torch.load(
                data_dir+'company_pretrained.pt')
            self.titles_embed.weight.data = torch.load(
                data_dir+'title_pretrained.pt')
            self.skill_embed.weight.data = torch.load(
                data_dir+'skill_pretrained.pt')
            self.skill_embed.bias.data = torch.zeros(
                max(mapper['skills'].values())+1)

    def forward(self, batch, negative_batch):
        # unpack data
        L, T, B, job_mask, edu_mask, locality, industry, times, intervals, schools, degrees, majors,  companies, titles, skill_label = tuple(
            i.to(device) if isinstance(i, torch.Tensor) else i for i in batch)
        # setup mask
        company_mask = companies > 0  # T, B
        title_mask = titles > 0  # T, B
        total_time_steps = job_mask.sum()

        if self.use_loc_ind is False:
            del locality, industry
        # unpack negative data
        negative_companies, negative_titles, company_qt, title_qt = tuple(
            i.to(device) if isinstance(i, torch.Tensor) else i for i in negative_batch)

        # Majors
        majors_embed = self.majors_embed(majors)  # T, B, 20
        # Occupation
        titles_embed = self.titles_embed(titles)  # T,B, 20
        negative_titles_embed = self.titles_embed(
            negative_titles)  # T, B, N, 20

        all_titles_embed = torch.cat([titles_embed.unsqueeze(
            2), negative_titles_embed], dim=2)  # T, B, N+1, 20

        # companies
        companies_embed = self.company_embed(companies)  # T, B, 20
        negative_companies_embed = self.company_embed(
            negative_companies)  # T, B, N , 20
        all_companies_embed = torch.cat([companies_embed.unsqueeze(
            2), negative_companies_embed], dim=2)  # T, B, N+1, 20

        # Embed
        if self.use_loc_ind:
            # locality
            locality_embed = self.locality_embed(locality).expand(T, B, -1)
            industry_embed = self.industry_embed(industry).expand(T, B, -1)
            #locality_embed, industry_embed
            input_fixed_embeddings = torch.cat(
                [locality_embed, industry_embed], dim=-1)  # T, B, H1
        # schools
        schools_embed = self.school_embed(schools)
        degrees_embed = self.degree_embed(degrees)
        times = (times.float().unsqueeze(2)-180.)/30.
        times_embed = self.time_embed(times)
        intervals_ = (intervals.float().unsqueeze(2) - 5.)/5.
        intervals_embed = self.interval_embed(intervals_)

        # input Embedding
        input_edu_embeddings = torch.cat(
            [degrees_embed, schools_embed, majors_embed], dim=-1)*edu_mask.unsqueeze(2)  # T, B, H1
        input_job_embeddings = torch.cat(
            [companies_embed, titles_embed], dim=-1)*job_mask.unsqueeze(2)  # T, B, H1
        input_embeddings = torch.cat(
            [input_edu_embeddings, input_job_embeddings, intervals_embed, times_embed], dim=-1)

        # Run RNN #T,B, H and B,H
        hidden_states, last_states = self.run_rnn(
            input_embeddings, self.rnn, L, B, T, self.model)

        # Add fixed info
        if self.use_loc_ind:
            augmented_hidden_states = torch.cat(
                [hidden_states, input_fixed_embeddings], dim=-1)
            augmented_last_states = torch.cat(
                [last_states, input_fixed_embeddings[0, :, :]], dim=-1)
        else:
            augmented_hidden_states = hidden_states
            augmented_last_states = last_states

        shared_output = self.intermediate_layers(
            augmented_hidden_states)  # T, B, 100

        # predict companies (T, B, 100)
        company_shared_layer = self.company_final_layers(
            shared_output)  # T, B , 100
        # loss
        company_loss = (restricted_sigmoid(company_shared_layer,
                                           all_companies_embed, company_qt, company_mask)).sum()
        # predict title (T, B, 100)
        title_shared_layer = self.title_final_layers(
            shared_output)  # T, B , 100
        # loss
        title_loss = (restricted_sigmoid(title_shared_layer,
                                         all_titles_embed, title_qt, title_mask)).sum()

        # predict skills (B, 100)
        skill_shared_layer = self.skill_embed(
            self.skill_final_layers(augmented_last_states))  # B, number_skills
        # loss
        skill_loss_fn = torch.nn.BCEWithLogitsLoss(reduction='none')
        skill_loss = (skill_loss_fn(skill_shared_layer, skill_label)).mean(
            dim=-1).sum()  # B, S

        total_loss = (company_loss+title_loss)/total_time_steps+skill_loss/B

        return total_loss

    def run_rnn(self, data_embed, rnn, data_length, batch_size, max_len, model):
        if model == 'nemo':
            packed_data = torch.nn.utils.rnn.pack_padded_sequence(
                data_embed, lengths=(data_length), batch_first=False, enforce_sorted=False)
            packed_output, ht = rnn(packed_data)
            ht = ht.permute(1, 2, 0).contiguous().reshape(batch_size, -1)
            output, input_sizes = torch.nn.utils.rnn.pad_packed_sequence(
                packed_output, batch_first=False)
        elif model in {'nss'}:
            output, ht = rnn(data_embed, data_length)
        init_state = torch.zeros(1, batch_size, output.shape[-1]).to(device)
        # append init states
        output = torch.cat([init_state, output[:-1, :, :]], dim=0)
        return output, ht