import torch
import torch.nn as nn
import torch.nn.functional as F
from BertModel import *

class BertModel(nn.Module):
    """Construct a 'BERT' model to rank relations for each step"""
    def __init__(self, config, device = None):
        super(BertModel, self).__init__()
        self.hidden_size = config.hidden_size
        self.embeddings = InputEmbeddings(config)
        self.encoder = EncoderStack(config)
        self.pooler = PoolerHead(config)
        self.classifier = nn.Linear(self.hidden_size, 2)
        self.device = device
        self.padding_id = 0

    def forward(self, sequence):
        
        sequence_mask = self.generate_attention_mask(sequence)
        sequence_emb = self.embeddings(sequence)
        sequence_enc = self.encoder(sequence_emb,
                                    sequence_mask,
                                    output_all_layers = False)[-1]
        sequence_pool = self.pooler(sequence_enc)

        logits = self.classifier(sequence_pool)
        return logits

    def generate_attention_mask(self, sequence):
        '''
        sequence --> (B, S), where B = batch_size, S = sequence length
        请对sequence进行掩码，对padding_id的位置赋值-1.e10，使其能在bert中使用
        输出掩码后的sequence
        attention_mask --> (B, S)

        '''

        # ==========
        # todo '''请补全代码'''
        # ==========
        pass
