from torch import nn
from transformers import BertForTokenClassification


class BertPunctuation(nn.Module):
    def __init__(
        self,
        num_labels,
        pretrained: str = 'hfl/chinese-macbert-base',
    ):
        super().__init__()
        self.model = BertForTokenClassification.from_pretrained(
            pretrained,
            num_labels = num_labels,
            output_hidden_states = True,
            output_attentions = False,
        )

    def forward(self, batch):
        input_ids, att_masks = batch[0], batch[1]
        labels = batch[2] if len(batch)>2 else None
        y = self.model(
            input_ids = input_ids,
            attention_mask = att_masks,
            labels = labels,
        )
        return y
