model.py

import torch
import torch.nn as nn
from torch.autograd import Variable
import torch.nn.functional as F
import numpy as np
from torch.nn.functional import cross_entropy, softmax
from transformers import BertModel, BertConfig


class Eyettention(nn.Module):
    def __init__(self, cf):
        super(Eyettention, self).__init__()
        self.cf = cf
        self.window_width = 1  # D
        self.atten_type = cf["atten_type"]
        self.hidden_size = 128

        # Word-Sequence Encoder
        encoder_config = BertConfig.from_pretrained(self.cf["model_pretrained"])
        encoder_config.output_hidden_states = True
        # initiate Bert with pre-trained weights
        print("keeping Bert with pre-trained weights")
        self.encoder = BertModel.from_pretrained(self.cf["model_pretrained"], config=encoder_config)
        self.encoder.eval()
        # freeze the parameters in Bert model
        for param in self.encoder.parameters():
            param.requires_grad = False

        self.embedding_dropout = nn.Dropout(0.4)
        self.encoder_lstm = nn.LSTM(
            input_size=768,  # BERT embedding size
            hidden_size=int(self.hidden_size / 2),
            num_layers=8,
            batch_first=True,
            bidirectional=True,
            dropout=0.2,
        )

        # Fixation-Sequence Encoder
        self.position_embeddings = nn.Embedding(
            encoder_config.max_position_embeddings, encoder_config.hidden_size
        )
        self.LayerNorm = nn.LayerNorm(encoder_config.hidden_size, eps=encoder_config.layer_norm_eps)

        # The scanpath is generated in an autoregressive manner, the output of the previous timestep is fed to the input of the next time step.
        # So we use decoder cells and loop over all timesteps.
        # initialize eight decoder cells
        self.decoder_cell1 = nn.LSTMCell(
            768 + 2, self.hidden_size
        )  # first layer input size = #BERT embedding size + two fixation attributes:landing position and fixiation duration
        self.decoder_cell2 = nn.LSTMCell(self.hidden_size, self.hidden_size)
        self.decoder_cell3 = nn.LSTMCell(self.hidden_size, self.hidden_size)
        self.decoder_cell4 = nn.LSTMCell(self.hidden_size, self.hidden_size)
        self.decoder_cell5 = nn.LSTMCell(self.hidden_size, self.hidden_size)
        self.decoder_cell6 = nn.LSTMCell(self.hidden_size, self.hidden_size)
        self.decoder_cell7 = nn.LSTMCell(self.hidden_size, self.hidden_size)
        self.decoder_cell8 = nn.LSTMCell(self.hidden_size, self.hidden_size)
        self.dropout_LSTM = nn.Dropout(0.2)

        # Cross-Attention
        self.attn = nn.Linear(
            self.hidden_size, self.hidden_size + 1
        )  # +1 acount for the word length feature

        # Decoder
        # initialize five dense layers
        self.dropout_dense = nn.Dropout(0.2)
        self.decoder_dense = nn.Sequential(
            self.dropout_dense,
            nn.Linear(self.hidden_size * 2 + 1, 512),
            nn.ReLU(),
            self.dropout_dense,
            nn.Linear(512, 256),
            nn.ReLU(),
            self.dropout_dense,
            nn.Linear(256, 256),
            nn.ReLU(),
            self.dropout_dense,
            nn.Linear(256, 256),
            nn.ReLU(),
            nn.Linear(256, self.cf["max_sn_len"] * 2 - 3),  # number of output classes
        )

        # for scanpath generation
        self.softmax = nn.Softmax(dim=1)

    def pool_subword_to_word(self, subword_emb, word_ids_sn, target, pool_method="sum"):
        # batching computing
        # Pool bert token (subword) to word level
        if target == "sn":
            max_len = self.cf["max_sn_len"]  # CLS and SEP included
        elif target == "sp":
            max_len = self.cf["max_sp_len"] - 1  # do not account the 'SEP' token

        merged_word_emb = torch.empty(subword_emb.shape[0], 0, 768).to(subword_emb.device)
        for word_idx in range(max_len):
            word_mask = (word_ids_sn == word_idx).unsqueeze(2).repeat(1, 1, 768)
            # pooling method -> sum
            if pool_method == "sum":
                pooled_word_emb = torch.sum(subword_emb * word_mask, 1).unsqueeze(
                    1
                )  # [batch, 1, 768]
            elif pool_method == "mean":
                pooled_word_emb = torch.mean(subword_emb * word_mask, 1).unsqueeze(
                    1
                )  # [batch, 1, 768]
            merged_word_emb = torch.cat([merged_word_emb, pooled_word_emb], dim=1)

        mask_word = torch.sum(merged_word_emb, 2).bool()
        return merged_word_emb, mask_word

    def encode(self, sn_emd, sn_mask, word_ids_sn, sn_word_len):
        # Word-Sequence Encoder
        outputs = self.encoder(input_ids=sn_emd, attention_mask=sn_mask)
        hidden_rep_orig, pooled_rep = outputs[0], outputs[1]
        if word_ids_sn != None:
            # Pool bert subword to word level for english corpus
            merged_word_emb, sn_mask_word = self.pool_subword_to_word(
                hidden_rep_orig, word_ids_sn, target="sn", pool_method="sum"
            )
        else:  # no pooling for Chinese corpus
            merged_word_emb, sn_mask_word = hidden_rep_orig, None

        hidden_rep = self.embedding_dropout(merged_word_emb)
        x, (hn, hc) = self.encoder_lstm(hidden_rep, None)

        # concatenate with the word length feature
        x = torch.cat((x, sn_word_len[:, :, None]), dim=2)
        return x, sn_mask_word

    def cross_attention(self, ht, hs, sn_mask, cur_word_index):
        # General Attention:
        # score(ht,hs) = (ht^T)(Wa)hs
        # hs is the output from word-Sequence Encoder
        # ht is the previous hidden state from Fixation-Sequence Encoder
        # self.attn(o): [batch, step, units]
        attn_prod = torch.matmul(
            self.attn(ht.unsqueeze(1)), hs.permute(0, 2, 1)
        )  # [batch, 1, step]
        if self.atten_type == "global":  # global attention
            attn_prod += (~sn_mask).unsqueeze(1) * -1e9
            att_weight = softmax(attn_prod, dim=2)  # [batch, 1, step]

        else:  # local attention
            # current fixated word index
            aligned_position = cur_word_index
            # Get window borders
            left = torch.where(
                aligned_position - self.window_width >= 0, aligned_position - self.window_width, 0
            )
            right = torch.where(
                aligned_position + self.window_width <= self.cf["max_sn_len"] - 1,
                aligned_position + self.window_width,
                self.cf["max_sn_len"] - 1,
            )

            # exclude padding tokens
            # only consider words in the window
            sen_seq = (
                torch.arange(self.cf["max_sn_len"])[None, :]
                .expand(sn_mask.shape[0], self.cf["max_sn_len"])
                .to(sn_mask.device)
            )
            outside_win_mask = (sen_seq < left.unsqueeze(1)) + (sen_seq > right.unsqueeze(1))

            attn_prod += (~sn_mask + outside_win_mask).unsqueeze(1) * -1e9
            att_weight = softmax(attn_prod, dim=2)  # [batch, 1, step]

            if self.atten_type == "local-g":  # local attention with Gaussian Kernel
                gauss = lambda s: torch.exp(
                    -torch.square(s - aligned_position.unsqueeze(1))
                    / (2 * torch.square(torch.tensor(self.window_width / 2)))
                )
                gauss_factor = gauss(sen_seq)
                att_weight = att_weight * gauss_factor.unsqueeze(1)

        return att_weight

    def decode(self, sp_emd, sn_mask, sp_pos, enc_out, sp_fix_dur, sp_landing_pos, word_ids_sp):
        # Fixation-Sequence Encoder + Decoder
        # Initialize hidden state and cell state with zeros,
        hn = torch.zeros(8, sp_emd.shape[0], self.hidden_size).to(sp_emd.device)
        hc = torch.zeros(8, sp_emd.shape[0], self.hidden_size).to(sp_emd.device)
        hx, cx = hn[0, :, :], hc[0, :, :]
        hx2, cx2 = hn[1, :, :], hc[1, :, :]
        hx3, cx3 = hn[2, :, :], hc[2, :, :]
        hx4, cx4 = hn[3, :, :], hc[3, :, :]
        hx5, cx5 = hn[4, :, :], hc[4, :, :]
        hx6, cx6 = hn[5, :, :], hc[5, :, :]
        hx7, cx7 = hn[6, :, :], hc[6, :, :]
        hx8, cx8 = hn[7, :, :], hc[7, :, :]

        dec_emb_in = self.encoder.embeddings.word_embeddings(sp_emd[:, :-1])
        if word_ids_sp is not None:
            # Pool bert subword to word level for English corpus
            sp_merged_word_emd, sp_mask_word = self.pool_subword_to_word(
                dec_emb_in, word_ids_sp[:, :-1], target="sp", pool_method="sum"
            )
        else:  # no pooling for Chinese corpus
            sp_merged_word_emd, sp_mask_word = dec_emb_in, None

        # add positional embeddings
        position_embeddings = self.position_embeddings(sp_pos[:, :-1])
        dec_emb_in = sp_merged_word_emd + position_embeddings
        dec_emb_in = self.LayerNorm(dec_emb_in)
        dec_emb_in = dec_emb_in.permute(1, 0, 2)  # [step, n, emb_dim]
        dec_emb_in = self.embedding_dropout(dec_emb_in)

        # concatenate two additional gaze features
        if sp_landing_pos is not None:
            dec_emb_in = torch.cat((dec_emb_in, sp_landing_pos.permute(1, 0)[:-1, :, None]), dim=2)

        if sp_fix_dur is not None:
            dec_emb_in = torch.cat((dec_emb_in, sp_fix_dur.permute(1, 0)[:-1, :, None]), dim=2)

        # Predict output for each time step in turn
        output = []
        # save attention scores for visualization
        atten_weights_batch = torch.empty(sp_emd.shape[0], 0, self.cf["max_sn_len"]).to(
            sp_emd.device
        )
        for i in range(dec_emb_in.shape[0]):
            hx, cx = self.decoder_cell1(dec_emb_in[i], (hx, cx))  # [batch, units]
            hx2, cx2 = self.decoder_cell2(self.dropout_LSTM(hx), (hx2, cx2))
            hx3, cx3 = self.decoder_cell3(self.dropout_LSTM(hx2), (hx3, cx3))
            hx4, cx4 = self.decoder_cell4(self.dropout_LSTM(hx3), (hx4, cx4))
            hx5, cx5 = self.decoder_cell5(self.dropout_LSTM(hx4), (hx5, cx5))
            hx6, cx6 = self.decoder_cell6(self.dropout_LSTM(hx5), (hx6, cx6))
            hx7, cx7 = self.decoder_cell7(self.dropout_LSTM(hx6), (hx7, cx7))
            hx8, cx8 = self.decoder_cell8(self.dropout_LSTM(hx7), (hx8, cx8))

            att_weight = self.cross_attention(
                ht=hx8, hs=enc_out, sn_mask=sn_mask, cur_word_index=sp_pos[:, i]
            )
            atten_weights_batch = torch.cat([atten_weights_batch, att_weight], dim=1)

            context = torch.matmul(att_weight, enc_out)  # [batch, 1, units]

            # Decoder
            hc = torch.cat([context.squeeze(1), hx8], dim=1)  # [batch, units *2]
            result = self.decoder_dense(hc)  # [batch, dec_o_dim]
            output.append(result)

        output = torch.stack(output, dim=0)  # [step, batch, dec_o_dim]
        # output = F.softmax(output, dim=2) # cross entropy in pytorch includes softmax
        return output.permute(1, 0, 2), atten_weights_batch  # [batch, step, dec_o_dim]

    def forward(
        self,
        sn_emd,
        sn_mask,
        sp_emd,
        sp_pos,
        word_ids_sn,
        word_ids_sp,
        sp_fix_dur,
        sp_landing_pos,
        sn_word_len,
    ):
        x, sn_mask_word = self.encode(
            sn_emd, sn_mask, word_ids_sn, sn_word_len
        )  # [batch, step, units], [batch, units]

        if sn_mask_word is None:  # for Chinese dataset without token pooling
            sn_mask = torch.Tensor.bool(sn_mask)
            pred, atten_weights = self.decode(
                sp_emd, sn_mask, sp_pos, x, sp_fix_dur, sp_landing_pos, word_ids_sp
            )  # [batch, step, dec_o_dim]

        else:  # for English dataset with token pooling
            pred, atten_weights = self.decode(
                sp_emd, sn_mask_word, sp_pos, x, sp_fix_dur, sp_landing_pos, word_ids_sp
            )  # [batch, step, dec_o_dim]

        return pred, atten_weights

    def scanpath_generation(
        self, sn_emd, sn_mask, word_ids_sn, sn_word_len, le, max_pred_len=60, previous_scanpath=None
    ):

        if max_pred_len <= 0:
            raise ValueError("max_pred_len must be positive.")

        prev_scanpath_len = 0

        if previous_scanpath is not None:
            if (
                len(previous_scanpath) == 0 or previous_scanpath[0] != 0
            ):  # prepend with 0 if no 0 at the start of previous scanpaths
                previous_scanpath = [0] + previous_scanpath

            prev_scanpath_len = len(previous_scanpath)

            previous_scanpath = torch.as_tensor(
                previous_scanpath, dtype=torch.long, device=sn_emd.device
            )
            if previous_scanpath.ndim == 1:
                previous_scanpath = previous_scanpath.unsqueeze(0)
            # if input includes CLS: [0, 1, 2, ...]
            previous_scanpath = previous_scanpath[:, 1:]

        # compute the scan path generated from the model when the first CLS taken is given
        enc_out, sn_mask_word = self.encode(sn_emd, sn_mask, word_ids_sn, sn_word_len)
        if sn_mask_word is None:
            sn_mask = torch.Tensor.bool(sn_mask)
        else:
            sn_mask = sn_mask_word
        sn_len = torch.sum(sn_mask, axis=1) - 2

        # decode
        # Initialize hidden state and cell state with zeros,
        hn = torch.zeros(8, sn_emd.shape[0], self.hidden_size).to(sn_emd.device)
        hc = torch.zeros(8, sn_emd.shape[0], self.hidden_size).to(sn_emd.device)
        hx, cx = hn[0, :, :], hc[0, :, :]
        hx2, cx2 = hn[1, :, :], hc[1, :, :]
        hx3, cx3 = hn[2, :, :], hc[2, :, :]
        hx4, cx4 = hn[3, :, :], hc[3, :, :]
        hx5, cx5 = hn[4, :, :], hc[4, :, :]
        hx6, cx6 = hn[5, :, :], hc[5, :, :]
        hx7, cx7 = hn[6, :, :], hc[6, :, :]
        hx8, cx8 = hn[7, :, :], hc[7, :, :]

        # use CLS token (101) as start token
        dec_in_start = (torch.ones(sn_mask.shape[0]) * 101).long().to(sn_mask.device)
        dec_emb_in = self.encoder.embeddings.word_embeddings(dec_in_start)  # [batch, emb_dim]
        # dec_in_start = sp_emd[:, 0]
        # dec_emb_in = self.encoder.embeddings.word_embeddings(dec_in_start) # [batch, emb_dim]

        # add positional embeddings
        start_pos = torch.zeros(sn_mask.shape[0]).to(sn_mask.device)
        position_embeddings = self.position_embeddings(start_pos.long())
        dec_emb_in = dec_emb_in + position_embeddings
        dec_emb_in = self.LayerNorm(dec_emb_in)

        # concatenate two additional gaze features, which are set to zeros for CLS token
        dec_in = torch.cat(
            (dec_emb_in, torch.zeros(dec_emb_in.shape[0], 2).to(sn_emd.device)), dim=1
        )

        # generate fixation one by one in an autoregressive way
        output = []
        density_prediction = []
        pred_counter = 0
        # output.append(sp_pos[:, pred_counter])
        output.append(start_pos.long())
        for p in range(prev_scanpath_len + max_pred_len - 1):
            hx, cx = self.decoder_cell1(dec_in, (hx, cx))  # [batch, units]
            hx2, cx2 = self.decoder_cell2(self.dropout_LSTM(hx), (hx2, cx2))
            hx3, cx3 = self.decoder_cell3(self.dropout_LSTM(hx2), (hx3, cx3))
            hx4, cx4 = self.decoder_cell4(self.dropout_LSTM(hx3), (hx4, cx4))
            hx5, cx5 = self.decoder_cell5(self.dropout_LSTM(hx4), (hx5, cx5))
            hx6, cx6 = self.decoder_cell6(self.dropout_LSTM(hx5), (hx6, cx6))
            hx7, cx7 = self.decoder_cell7(self.dropout_LSTM(hx6), (hx7, cx7))
            hx8, cx8 = self.decoder_cell8(self.dropout_LSTM(hx7), (hx8, cx8))

            att_weight = self.cross_attention(
                ht=hx8, hs=enc_out, sn_mask=sn_mask, cur_word_index=output[-1]
            )

            context = torch.matmul(att_weight, enc_out)  # [batch, 1, units]
            hc = torch.cat([context.squeeze(1), hx8], dim=1)  # [batch, units *2]

            result = self.decoder_dense(hc)  # [batch, dec_o_dim]
            result = self.softmax(result)  # [batch, dec_o_dim]
            density_prediction.append(result)

            # we can either take argmax or sampling from the output distribution,
            # we do sampling in the paper
            # pred_indx = result.argmax(dim=1)
            # sampling next fixation location according to the distribution
            if previous_scanpath is not None and p < previous_scanpath.shape[1]:
                # replay observed fixation
                pred_pos = previous_scanpath[:, p].clone()
            else:
                pred_indx = torch.multinomial(result, 1)
                pred_class = [le.classes_[pred_indx[i]] for i in torch.arange(result.shape[0])]
                pred_class = torch.from_numpy(np.array(pred_class)).to(sn_emd.device)
                # predict fixation word index = last fixation word index + predicted saccade range
                pred_pos = output[-1] + pred_class

            # larger than sentence max length -- set to sentence length+1, i.e. token <'SEP'>
            # prepare the input to the next timstep
            input_ids = []
            for i in range(pred_pos.shape[0]):
                if pred_pos[i] > sn_len[i]:
                    pred_pos[i] = sn_len[i] + 1
                elif pred_pos[i] < 1:
                    pred_pos[i] = 1

                if word_ids_sn is not None:
                    input_ids.append(sn_emd[i, word_ids_sn[i, :] == pred_pos[i]])
                else:
                    input_ids.append(sn_emd[i, pred_pos[i]])
            output.append(pred_pos)

            # prepare next timestamp input token
            pred_counter += 1
            if word_ids_sn is not None:
                # merge tokens
                dec_emb_in = torch.empty(0, 768).to(sn_emd.device)
                for id in input_ids:
                    dec_emb_in = torch.cat(
                        [
                            dec_emb_in,
                            torch.sum(self.encoder.embeddings.word_embeddings(id), axis=0)[None, :],
                        ],
                        dim=0,
                    )

            else:
                input_ids = torch.stack(input_ids)
                dec_emb_in = self.encoder.embeddings.word_embeddings(input_ids)  # [batch, emb_dim]
            # add positional embeddings
            position_embeddings = self.position_embeddings(output[-1])
            dec_emb_in = dec_emb_in + position_embeddings
            dec_emb_in = self.LayerNorm(dec_emb_in)
            # concatenate two additional gaze features
            dec_in = torch.cat(
                (dec_emb_in, torch.zeros(dec_emb_in.shape[0], 2).to(sn_emd.device)), dim=1
            )

        output = torch.stack(output, dim=0)  # [step, batch]
        return output.permute(1, 0), density_prediction  # [batch, step]


class Eyettention_readerID(nn.Module):
    def __init__(self, cf):
        super(Eyettention_readerID, self).__init__()
        self.cf = cf
        self.window_width = 1  # D
        self.atten_type = cf["atten_type"]
        self.hidden_size = 128
        self.sub_emb_size = cf["subid_emb_size"]

        # Word-Sequence Encoder
        encoder_config = BertConfig.from_pretrained(self.cf["model_pretrained"])
        encoder_config.output_hidden_states = True
        # initiate Bert with pre-trained weights
        print("keeping Bert with pre-trained weights")
        self.encoder = BertModel.from_pretrained(self.cf["model_pretrained"], config=encoder_config)
        self.encoder.eval()
        # freeze the parameters in Bert model
        for param in self.encoder.parameters():
            param.requires_grad = False

        self.embedding_dropout = nn.Dropout(0.4)
        self.encoder_lstm = nn.LSTM(
            input_size=768,  # BERT embedding size
            hidden_size=int(self.hidden_size / 2),
            num_layers=8,
            batch_first=True,
            bidirectional=True,
            dropout=0.2,
        )

        # Fixation-Sequence Encoder
        self.position_embeddings = nn.Embedding(
            encoder_config.max_position_embeddings, encoder_config.hidden_size
        )
        self.LayerNorm = nn.LayerNorm(encoder_config.hidden_size, eps=encoder_config.layer_norm_eps)
        # for reader-specific setting
        self.sub_embeddings = nn.Embedding(400, self.sub_emb_size)

        # The scanpath is generated in an autoregressive manner, the output of the previous timestep is fed to the input of the next time step.
        # So we use decoder cells and loop over all timesteps.
        # initialize eight decoder cells
        self.decoder_cell1 = nn.LSTMCell(
            768 + 2 + self.sub_emb_size, self.hidden_size
        )  # first layer input size = #BERT embedding size + two fixation attributes:landing position and fixiation duration
        self.decoder_cell2 = nn.LSTMCell(self.hidden_size, self.hidden_size)
        self.decoder_cell3 = nn.LSTMCell(self.hidden_size, self.hidden_size)
        self.decoder_cell4 = nn.LSTMCell(self.hidden_size, self.hidden_size)
        self.decoder_cell5 = nn.LSTMCell(self.hidden_size, self.hidden_size)
        self.decoder_cell6 = nn.LSTMCell(self.hidden_size, self.hidden_size)
        self.decoder_cell7 = nn.LSTMCell(self.hidden_size, self.hidden_size)
        self.decoder_cell8 = nn.LSTMCell(self.hidden_size, self.hidden_size)
        self.dropout_LSTM = nn.Dropout(0.2)

        # Cross-Attention
        self.attn = nn.Linear(
            self.hidden_size, self.hidden_size + 1
        )  # +1 acoount for the word length feature

        # Decoder
        # initialize five dense layers
        self.dropout_dense = nn.Dropout(0.2)
        self.decoder_dense = nn.Sequential(
            self.dropout_dense,
            nn.Linear(self.hidden_size * 2 + 1, 512),
            nn.ReLU(),
            self.dropout_dense,
            nn.Linear(512, 256),
            nn.ReLU(),
            self.dropout_dense,
            nn.Linear(256, 256),
            nn.ReLU(),
            self.dropout_dense,
            nn.Linear(256, 256),
            nn.ReLU(),
            nn.Linear(256, self.cf["max_sn_len"] * 2 - 3),  # number of output classes
        )

        # for scanpath generation
        self.softmax = nn.Softmax(dim=1)

    def pool_subword_to_word(self, subword_emb, word_ids_sn, target, pool_method="sum"):
        # batching computing
        # Pool bert token (subword) to word level
        if target == "sn":
            max_len = self.cf["max_sn_len"]  # CLS and SEP included
        elif target == "sp":
            max_len = self.cf["max_sp_len"] - 1  # do not account the 'SEP' token

        merged_word_emb = torch.empty(subword_emb.shape[0], 0, 768).to(subword_emb.device)
        for word_idx in range(max_len):
            word_mask = (word_ids_sn == word_idx).unsqueeze(2).repeat(1, 1, 768)
            # pooling method -> sum
            if pool_method == "sum":
                pooled_word_emb = torch.sum(subword_emb * word_mask, 1).unsqueeze(
                    1
                )  # [batch, 1, 768]
            elif pool_method == "mean":
                pooled_word_emb = torch.mean(subword_emb * word_mask, 1).unsqueeze(
                    1
                )  # [batch, 1, 768]
            merged_word_emb = torch.cat([merged_word_emb, pooled_word_emb], dim=1)

        mask_word = torch.sum(merged_word_emb, 2).bool()
        return merged_word_emb, mask_word

    def encode(self, sn_emd, sn_mask, word_ids_sn, sn_word_len):
        # Word-Sequence Encoder
        outputs = self.encoder(input_ids=sn_emd, attention_mask=sn_mask)
        hidden_rep_orig, pooled_rep = outputs[0], outputs[1]
        if word_ids_sn != None:
            # Pool bert subword to word level for english corpus
            merged_word_emb, sn_mask_word = self.pool_subword_to_word(
                hidden_rep_orig, word_ids_sn, target="sn", pool_method="sum"
            )
        else:  # no pooling for Chinese corpus
            merged_word_emb, sn_mask_word = hidden_rep_orig, None

        hidden_rep = self.embedding_dropout(merged_word_emb)
        x, (hn, hc) = self.encoder_lstm(hidden_rep, None)

        # concatenate with the word length feature
        x = torch.cat((x, sn_word_len[:, :, None]), dim=2)
        return x, sn_mask_word

    def cross_attention(self, ht, hs, sn_mask, cur_word_index):
        # General Attention:
        # score(ht,hs) = (ht^T)(Wa)hs
        # hs is the output from word-Sequence Encoder
        # ht is the previous hidden state from Fixation-Sequence Encoder
        # self.attn(o): [batch, step, units]
        attn_prod = torch.matmul(
            self.attn(ht.unsqueeze(1)), hs.permute(0, 2, 1)
        )  # [batch, 1, step]
        if self.atten_type == "global":  # global attention
            attn_prod += (~sn_mask).unsqueeze(1) * -1e9
            att_weight = softmax(attn_prod, dim=2)  # [batch, 1, step]

        else:  # local attention
            # current fixated word index
            aligned_position = cur_word_index
            # Get window borders
            left = torch.where(
                aligned_position - self.window_width >= 0, aligned_position - self.window_width, 0
            )
            right = torch.where(
                aligned_position + self.window_width <= self.cf["max_sn_len"] - 1,
                aligned_position + self.window_width,
                self.cf["max_sn_len"] - 1,
            )

            # exclude padding tokens
            # only consider words in the window
            sen_seq = (
                torch.arange(self.cf["max_sn_len"])[None, :]
                .expand(sn_mask.shape[0], self.cf["max_sn_len"])
                .to(sn_mask.device)
            )
            outside_win_mask = (sen_seq < left.unsqueeze(1)) + (sen_seq > right.unsqueeze(1))

            attn_prod += (~sn_mask + outside_win_mask).unsqueeze(1) * -1e9
            att_weight = softmax(attn_prod, dim=2)  # [batch, 1, step]

            if self.atten_type == "local-g":  # local attention with Gaussian Kernel
                gauss = lambda s: torch.exp(
                    -torch.square(s - aligned_position.unsqueeze(1))
                    / (2 * torch.square(torch.tensor(self.window_width / 2)))
                )
                gauss_factor = gauss(sen_seq)
                att_weight = att_weight * gauss_factor.unsqueeze(1)

        return att_weight

    def decode(
        self, sp_emd, sn_mask, sp_pos, enc_out, sp_fix_dur, sp_landing_pos, word_ids_sp, sub_id
    ):
        # Fixation-Sequence Encoder + Decoder
        # Initialize hidden state and cell state with zeros
        hn = torch.zeros(8, sp_emd.shape[0], self.hidden_size).to(sp_emd.device)
        hc = torch.zeros(8, sp_emd.shape[0], self.hidden_size).to(sp_emd.device)
        hx, cx = hn[0, :, :], hc[0, :, :]
        hx2, cx2 = hn[1, :, :], hc[1, :, :]
        hx3, cx3 = hn[2, :, :], hc[2, :, :]
        hx4, cx4 = hn[3, :, :], hc[3, :, :]
        hx5, cx5 = hn[4, :, :], hc[4, :, :]
        hx6, cx6 = hn[5, :, :], hc[5, :, :]
        hx7, cx7 = hn[6, :, :], hc[6, :, :]
        hx8, cx8 = hn[7, :, :], hc[7, :, :]

        dec_emb_in = self.encoder.embeddings.word_embeddings(sp_emd[:, :-1])
        if word_ids_sp is not None:
            # Pool bert subword to word level for English corpus
            sp_merged_word_emd, sp_mask_word = self.pool_subword_to_word(
                dec_emb_in, word_ids_sp[:, :-1], target="sp", pool_method="sum"
            )
        else:  # no pooling for Chinese
            sp_merged_word_emd, sp_mask_word = dec_emb_in, None

        # add positional embeddings and layer normalization
        position_embeddings = self.position_embeddings(sp_pos[:, :-1])
        dec_emb_in = sp_merged_word_emd + position_embeddings
        dec_emb_in = self.LayerNorm(dec_emb_in)
        dec_emb_in = dec_emb_in.permute(1, 0, 2)  # [step, n, emb_dim]
        dec_emb_in = self.embedding_dropout(dec_emb_in)

        # concatenate two additional gaze features
        if sp_landing_pos is not None:
            dec_emb_in = torch.cat((dec_emb_in, sp_landing_pos.permute(1, 0)[:-1, :, None]), dim=2)

        if sp_fix_dur is not None:
            dec_emb_in = torch.cat((dec_emb_in, sp_fix_dur.permute(1, 0)[:-1, :, None]), dim=2)

        # concatenate subject id for Eyettention_reader setting
        if sub_id is not None:
            dec_emb_in = torch.cat(
                (dec_emb_in, self.sub_embeddings(sub_id).repeat(dec_emb_in.shape[0], 1, 1)), dim=2
            )

        # Predict output for each time step in turn
        output = []
        # save attention scores for visualization
        atten_weights_batch = torch.empty(sp_emd.shape[0], 0, self.cf["max_sn_len"]).to(
            sp_emd.device
        )
        for i in range(dec_emb_in.shape[0]):
            hx, cx = self.decoder_cell1(dec_emb_in[i], (hx, cx))  # [batch, units]
            hx2, cx2 = self.decoder_cell2(self.dropout_LSTM(hx), (hx2, cx2))
            hx3, cx3 = self.decoder_cell3(self.dropout_LSTM(hx2), (hx3, cx3))
            hx4, cx4 = self.decoder_cell4(self.dropout_LSTM(hx3), (hx4, cx4))
            hx5, cx5 = self.decoder_cell5(self.dropout_LSTM(hx4), (hx5, cx5))
            hx6, cx6 = self.decoder_cell6(self.dropout_LSTM(hx5), (hx6, cx6))
            hx7, cx7 = self.decoder_cell7(self.dropout_LSTM(hx6), (hx7, cx7))
            hx8, cx8 = self.decoder_cell8(self.dropout_LSTM(hx7), (hx8, cx8))

            att_weight = self.cross_attention(
                ht=hx8, hs=enc_out, sn_mask=sn_mask, cur_word_index=sp_pos[:, i]
            )
            atten_weights_batch = torch.cat([atten_weights_batch, att_weight], dim=1)

            context = torch.matmul(att_weight, enc_out)  # [batch, 1, units]

            # Decoder
            hc = torch.cat([context.squeeze(1), hx8], dim=1)  # [batch, units *2]
            result = self.decoder_dense(hc)  # [batch, dec_o_dim]
            output.append(result)

        output = torch.stack(output, dim=0)  # [step, batch, dec_o_dim]
        # output = F.softmax(output, dim=2) # cross entropy in pytorch includes softmax
        return output.permute(1, 0, 2), atten_weights_batch  # [batch, step, dec_o_dim]

    def forward(
        self,
        sn_emd,
        sn_mask,
        sp_emd,
        sp_pos,
        word_ids_sn,
        word_ids_sp,
        sp_fix_dur,
        sp_landing_pos,
        sn_word_len,
        sub_id,
    ):
        x, sn_mask_word = self.encode(
            sn_emd, sn_mask, word_ids_sn, sn_word_len
        )  # [batch, step, units], [batch, units]

        if sn_mask_word is None:  # for Chinese dataset without token pooling
            sn_mask = torch.Tensor.bool(sn_mask)
            pred, atten_weights = self.decode(
                sp_emd, sn_mask, sp_pos, x, sp_fix_dur, sp_landing_pos, word_ids_sp, sub_id
            )  # [batch, step, dec_o_dim]

        else:  # for English dataset with token pooling
            pred, atten_weights = self.decode(
                sp_emd, sn_mask_word, sp_pos, x, sp_fix_dur, sp_landing_pos, word_ids_sp, sub_id
            )  # [batch, step, dec_o_dim]
        return pred, atten_weights