from .layers import *


class KoSBERT(nn.Module):
    def __init__(self,
                 bert_weight: str = "kykim/bert-kor-base") -> None:
        super(KoSBERT, self).__init__()
        self.bert = KoBERT(bert_weight=bert_weight)
        self.pooling = Pooling()

    def forward(self,
                input_ids: torch.Tensor,
                attention_mask: torch.Tensor,
                token_type_ids: torch.Tensor) -> torch.Tensor:

        bert_output = self.bert(input_ids=input_ids,
                                attention_mask=attention_mask,
                                token_type_ids=token_type_ids).last_hidden_state
        features = {'token_embeddings': bert_output, 'attention_mask': attention_mask}
        return self.pooling(features)
