model_plm.py

from typing import List, Iterator, cast

import copy
import numpy as np

import torch as T
from torch import nn
from torch.nn import functional as F
from transformers import BertConfig, BertModel
from transformers import AutoTokenizer, AutoModel, AutoConfig
from transformers import PreTrainedModel
from transformers.modeling_outputs import BaseModelOutputWithPoolingAndCrossAttentions

class Diacritizer(nn.Module):
    def __init__(
            self,
            config,
            device=None,
            load_pretrained=True
    ) -> None:
        super().__init__()
        self._dummy = nn.Parameter(T.ones(1))

        if 'modeling' in config:
            config = config['modeling']
        self.config = config

        model_name = config.get('base_model', "CAMeL-Lab/bert-base-arabic-camelbert-mix-ner")
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)

        if load_pretrained:
            self.token_model: BertModel = AutoModel.from_pretrained(model_name)
        else:
            marbert_config = AutoConfig.from_pretrained(model_name)
            self.token_model = AutoModel.from_config(marbert_config)

        self.num_classes  = 15
        self.diac_model_config = BertConfig(**config['diac_model_config'])
        self.token_model_config: BertConfig = self.token_model.config

        self.char_embs      = nn.Embedding(config["num-chars"], embedding_dim=config["char-embed-dim"])
        self.diac_emb_model = self.build_diac_model(self.token_model)

        self.down_project_token_embeds_deep = None
        self.down_project_token_embeds = None
        if 'token_hidden_size' in config:
            if config['token_hidden_size'] == 'auto':
                down_proj_size = self.diac_emb_model.config.hidden_size
            else:
                down_proj_size = config['token_hidden_size']
            if config.get('deep-down-proj', False):
                self.down_project_token_embeds_deep = nn.Sequential(
                    nn.Linear(
                        self.token_model_config.hidden_size + config["char-embed-dim"],
                        down_proj_size * 4,
                        bias=False,
                    ),
                    nn.Tanh(),
                    nn.Linear(
                        down_proj_size * 4,
                        down_proj_size,
                        bias=False,
                    )
                )
            # else:
                self.down_project_token_embeds = nn.Linear(
        self.token_model_config.hidden_size + config["char-embed-dim"],
                    down_proj_size,
                    bias=False,
                )

        # assert self.down_project_token_embeds_deep is None or self.down_project_token_embeds is None
        classifier_feature_size = self.diac_model_config.hidden_size
        if config.get('deep-cls', False):
            # classifier_feature_size = 512
            self.final_feature_transform = nn.Linear(
                self.diac_model_config.hidden_size
                    + self.token_model_config.hidden_size,
                #^ diac_features + [residual from token_model]
                out_features=classifier_feature_size,
                bias=False
            )
        else:
            self.final_feature_transform = None

        self.feature_layer_norm = nn.LayerNorm(classifier_feature_size)
        self.classifier = nn.Linear(classifier_feature_size, self.num_classes, bias=True)

        self.trim_model_(config)

        self.dropout = nn.Dropout(config['dropout'])
        self.sent_dropout_p = config['sentence_dropout']
        self.closs = F.cross_entropy

    def build_diac_model(self, token_model=None):
        if self.config.get('pre-init-diac-model', False):
            model = copy.deepcopy(self.token_model)
            model.pooler = None
            model.embeddings.word_embeddings = None

            num_layers = self.config.get('keep-token-model-layers', None)
            model.encoder.layer = nn.ModuleList(
                list(model.encoder.layer[num_layers:num_layers*2])
            )

            model.encoder.config.num_hidden_layers = num_layers
        else:
            model = BertModel(self.diac_model_config)
        return model

    def trim_model_(self, config):
        self.token_model.pooler = None
        self.diac_emb_model.pooler = None
        # self.diac_emb_model.embeddings = None
        self.diac_emb_model.embeddings.word_embeddings = None

        num_token_model_kept_layers = config.get('keep-token-model-layers', None)
        if num_token_model_kept_layers is not None:
            self.token_model.encoder.layer = nn.ModuleList(
                list(self.token_model.encoder.layer[:num_token_model_kept_layers])
            )
            self.token_model.encoder.config.num_hidden_layers = num_token_model_kept_layers

        if not config.get('full-finetune', False):
            for param in self.token_model.parameters():
                param.requires_grad = False
            finetune_last_layers = config.get('num-finetune-last-layers', 4)
            if finetune_last_layers > 0:
                unfrozen_layers = self.token_model.encoder.layer[-finetune_last_layers:]
                for layer in unfrozen_layers:
                    for param in layer.parameters():
                        param.requires_grad = True

    def get_grouped_params(self):
        downstream_params: Iterator[nn.Parameter] = cast(
            Iterator,
            (param
                for module in (self.diac_emb_model, self.classifier, self.char_embs)
                for param in module.parameters())
        )
        pg = {
            'pretrained': self.token_model.parameters(),
            'downstream': downstream_params,
        }
        return pg

    @property
    def device(self):
        return self._dummy.device

    def step(self, xt, yt, mask=None, subword_lengths: T.Tensor=None):
        # ^ word_x, char_x, diac_x are Indices
        # ^ xt             : self.preprocess((word_x, char_x, diac_x)),
        # ^ yt             : T.tensor(diac_y, dtype=T.long),
        # ^ subword_lengths: T.tensor(subword_lengths, dtype=T.long)
        #< Move char_x, diac_x to device because they're small and trainable
        xt[0], xt[1], yt, subword_lengths = self._slim_batch_size(xt[0], xt[1], yt, subword_lengths)
        xt[0] = xt[0].to(self.device)
        xt[1] = xt[1].to(self.device)
        # xt[2] = xt[2].to(self.device)

        yt = yt.to(self.device)
        #^ yt: [b tw tc]

        Nb, Tword, Tchar = xt[1].shape
        if Tword * Tchar < 500:
            diac = self(*xt, subword_lengths)
            loss = self.closs(diac.view(-1, self.num_classes), yt.view(-1), reduction='sum')
        else:
            num_chunks = Tword * Tchar / 300
            loss = 0
            for i in range(round(num_chunks+0.5)):
                _slice = slice(i*300, (i+1)*300)
                chunk = self._slice_batch(xt, _slice)
                diac = self(*chunk, subword_lengths[_slice])
                chunk_loss = self.closs(diac.view(-1, self.num_classes), yt.view(-1), reduction='sum')
                loss = loss + chunk_loss

        return loss

    def _slice_batch(self, xt: List[T.Tensor], _slice):
        return [xt[0][_slice], xt[1][_slice], xt[2][_slice]]

    def _slim_batch_size(
            self,
            tx: T.Tensor,
            cx: T.Tensor,
            yt: T.Tensor,
            subword_lengths: T.Tensor
    ):
        #^ tx : [b tt]
        #^ cx : [b tw tc]
        #^ yt : [b tw tc]
        token_nonpad_mask = tx.ne(self.tokenizer.pad_token_id)
        Ttoken = token_nonpad_mask.sum(1).max()
        tx = tx[:, :Ttoken]

        char_nonpad_mask = cx.ne(0)
        Tword = char_nonpad_mask.any(2).sum(1).max()
        Tchar = char_nonpad_mask.sum(2).max()
        cx = cx[:, :Tword, :Tchar]
        yt = yt[:, :Tword, :Tchar]
        subword_lengths = subword_lengths[:, :Tword]

        return tx, cx, yt, subword_lengths

    def token_dropout(self, toke_x):
        #^ toke_x : [b tw]
        if self.training:
            q = 1.0 - self.sent_dropout_p
            sdo = T.bernoulli(T.full(toke_x.shape, q))
            toke_x[sdo == 0] = self.tokenizer.pad_token_id
        return toke_x

    def sentence_dropout(self, word_embs: T.Tensor):
        #^ word_embs : [b tw dwe]
        if self.training:
            q = 1.0 - self.sent_dropout_p
            sdo = T.bernoulli(T.full(word_embs.shape[:2], q))
            sdo = sdo.detach().unsqueeze(-1).to(word_embs)
            word_embs = word_embs * sdo
            # toke_x[sdo == 0] = self.tokenizer.pad_token_id
        return word_embs

    def embed_tokens(self, input_ids: T.Tensor, attention_mask: T.Tensor):
        y: BaseModelOutputWithPoolingAndCrossAttentions
        y = self.token_model(input_ids, attention_mask=attention_mask)
        z = y.last_hidden_state
        return z

    def forward(
            self,
            toke_x          : T.Tensor,
            char_x          : T.Tensor,
            diac_x          : T.Tensor,
            subword_lengths : T.Tensor,
    ):
        #^ toke_x : [b tt]
        #^ char_x : [b tw tc]
        #^ diac_x/labels : [b tw tc]
        #^ subword_lengths : [b, tw]
        # !TODO Use `subword_lengths` to aggregate subword embeddings first before ...
        # ... passing concatenated contextual embedding to chars in diac_model

        token_nonpad_mask = toke_x.ne(self.tokenizer.pad_token_id)
        char_nonpad_mask = char_x.ne(0)

        Nb, Tw, Tc = char_x.shape
        # assert Tw == Tw_0 and Tc == Tc_0, f"{Tw=} {Tw_0=}, {Tc=} {Tc_0=}"

        # toke_x = self.token_dropout(toke_x)
        token_embs = self.embed_tokens(toke_x, attention_mask=token_nonpad_mask)
        # token_embs = self.sentence_dropout(token_embs)
        #? Strip BOS,EOS
        token_embs = token_embs[:, 1:-1, ...]

        sent_word_strides = subword_lengths.cumsum(1)
        sent_enc: T.Tensor = T.zeros(Nb, Tw, token_embs.shape[-1]).to(token_embs)
        for i_b in range(Nb):
            token_embs_ib = token_embs[i_b]
            start_iw = 0
            for i_word, end_iw in enumerate(sent_word_strides[i_b]):
                if end_iw == start_iw: break
                word_emb = token_embs_ib[start_iw : end_iw].sum(0) / (end_iw - start_iw)
                sent_enc[i_b, i_word] = word_emb
                start_iw = end_iw
        #^ sent_enc: [b tw dwe]

        char_x_flat = char_x.reshape(Nb*Tw, Tc)
        char_nonpad_mask = char_x_flat.gt(0)
        # ^ char_nonpad_mask [b*tw tc]

        char_x_flat = char_x_flat * char_nonpad_mask

        cembs = self.char_embs(char_x_flat)
    
        #^ cembs: [b*tw tc dce]
        wembs = sent_enc.unsqueeze(-2).expand(Nb, Tw, Tc, -1).view(Nb*Tw, Tc, -1)
        #^ wembs: [b tw dwe] => [b tw _ dwe] => [b*tw tc dwe]
        cw_embs = T.cat([cembs, wembs], dim=-1)
        #^ char_embs : [b*tw tc dcw] ; dcw = dc + dwe
        cw_embs = self.dropout(cw_embs)

        cw_embs_ = cw_embs
        if self.down_project_token_embeds is not None:
            cw_embs_ = self.down_project_token_embeds(cw_embs)
        if self.down_project_token_embeds_deep is not None:
            cw_embs_ = cw_embs_ + self.down_project_token_embeds_deep(cw_embs)
        cw_embs = cw_embs_

        diac_enc: BaseModelOutputWithPoolingAndCrossAttentions
        diac_enc = self.diac_emb_model(inputs_embeds=cw_embs, attention_mask=char_nonpad_mask)
        diac_emb = diac_enc.last_hidden_state
        diac_emb = self.dropout(diac_emb)
        #^ diac_emb: [b*tw tc dce]
        diac_emb = diac_emb.view(Nb, Tw, Tc, -1)

        sent_residual = sent_enc.unsqueeze(2).expand(-1, -1, Tc, -1)
        final_feature = T.cat([sent_residual, diac_emb], dim=-1)
        if self.final_feature_transform is not None:
            final_feature = self.final_feature_transform(final_feature)
            final_feature = F.tanh(final_feature)
            final_feature = self.dropout(final_feature)
        else:
            final_feature = diac_emb

        # final_feature = self.feature_layer_norm(final_feature)
        diac_out = self.classifier(final_feature)
        # if T.isnan(diac_out).any():
        #     breakpoint()
        return diac_out

    def predict(self, dataloader):
        from tqdm import tqdm
        import diac_utils as du
        training = self.training
        self.eval()

        preds = {'haraka': [], 'shadda': [], 'tanween': []}
        print("> Predicting...")
        for inputs, _, subword_lengths in tqdm(dataloader, total=len(dataloader)):
            inputs[0] = inputs[0].to(self.device)
            inputs[1] = inputs[1].to(self.device)
            output = self(*inputs, subword_lengths).detach()

            marks = np.argmax(T.softmax(output, dim=-1).cpu().numpy(), axis=-1)
            #^ [b ts tw]

            haraka, tanween, shadda = du.flat_2_3head(marks)

            preds['haraka'].extend(haraka)
            preds['tanween'].extend(tanween)
            preds['shadda'].extend(shadda)

        self.train(training)
        return (
            np.array(preds['haraka']),
            np.array(preds["tanween"]),
            np.array(preds["shadda"]),
        )

if __name__ == "__main__":
    model = Diacritizer({
        "num-chars": 36,
        "hidden_size": 768,
        "char-embed-dim": 32,
        "dropout": 0.25,
        "sentence_dropout": 0.2,
        "diac_model_config": {
            "num_layers": 4,
            "hidden_size": 768 + 32,
            "intermediate_size": (768 + 32) * 4,
        },
    }, load_pretrained=False)

    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)

    print(model)
    print(f"{trainable_params:,}/{total_params:,} Trainable Parameters")