predict.py

from typing import Iterable, Union, Tuple
from collections import Counter

import argparse
import os

import yaml
from pyarabic.araby import tokenize, strip_tatweel, strip_tashkeel
from tqdm import tqdm

import numpy as np
import torch as T
from torch.utils.data import DataLoader

from diac_utils import HARAKAT_MAP, shakkel_char, flat2_3head
from model_partial import PartialDD
from data_utils import DatasetUtils
from dataloader import DataRetriever
from dataloader_plm import DataRetriever as DataRetrieverPLM
from segment import segment

from partial_dd_metrics import (
    parse_data,
    load_data,
    make_mask_hard,
    make_mask_logits,
)

def apply_tashkeel(
        line: str,
        diacs: Union[np.ndarray, T.Tensor]
):
    line_w_diacs = ""
    ts, tw = diacs.shape
    diacs = diacs.flatten()
    diacs_h3 = flat2_3head(diacs)
    diacs_h3 = tuple(x.reshape(ts, tw) for x in diacs_h3)
    diac_char_idx = 0
    diac_word_idx = 0
    for ch in line:
        line_w_diacs += ch
        if ch == " ":
            diac_char_idx = 0
            diac_word_idx += 1
        else:
            tashkeel = (diacs_h3[0][diac_word_idx][diac_char_idx], diacs_h3[1][diac_word_idx][diac_char_idx], diacs_h3[2][diac_word_idx][diac_char_idx])
            diac_char_idx += 1
            line_w_diacs += shakkel_char(*tashkeel)
    return line_w_diacs
    
def diac_text(data, model_output_base, model_output_ctxt, selection_mode='contrastive-hard', threshold=0.1):

    mode = selection_mode
    if mode == 'contrastive-hard':
        # model_output_base = parse_data(data_base)[0]
        # model_output_ctxt = parse_data(data_ctxt)[0]
        # diacs = np.where(diacs_base != diacs_ctxt, diacs_ctxt, 0)
        diacritics = np.where(
            make_mask_hard(model_output_ctxt, model_output_base),
            model_output_ctxt.argmax(-1),
            0,
        ).astype(int)
    else:
        # model_output_base = parse_data(data_base, logits=True, side='base')[2]
        # model_output_ctxt = parse_data(data_ctxt, logits=True, side='ctxt')[2]
        diacritics = np.where(
            make_mask_logits(
                model_output_ctxt, model_output_base,
                version=mode, threshold=threshold,
            ),
            model_output_ctxt.argmax(-1),
            0,
        ).astype(int)
    #^ shape: [b, tc | ClassId]
    diacs_pred = model_output_base
    
    assert len(diacs_pred) == len(data)
    data = [
        ' '.join(tokenize(
            line.strip(),
            morphs=[strip_tashkeel, strip_tatweel]
        ))
        for line in data
    ]
    
    output = []
    for line, line_diacs in zip(
            tqdm(data),
            diacritics
    ):
        line = apply_tashkeel(line, line_diacs)
        output.append(line)

    return output

class Predictor:
    def __init__(self, config):

        self.data_utils = DatasetUtils(config)
        vocab_size = len(self.data_utils.letter_list)
        word_embeddings = self.data_utils.embeddings
        self.config = config 
       
        self.device = T.device(
            config['predictor'].get('device', 'cuda:0')
            if T.cuda.is_available() else 'cpu'
        )

        self.model = PartialDD(config)
        if config["model-name"] == "D2":
            self.model.sentence_diac.build(word_embeddings, vocab_size)
            state_dict = T.load(config["paths"]["load"], map_location=T.device(self.device))['state_dict']
        else:
            state_dict = T.load(config["paths"]["load-td2"], map_location=T.device(self.device))['state_dict']

        self.model.load_state_dict(state_dict, strict=False)
        self.model.to(self.device)
        self.model.eval()

    def create_dataloader(self, text, do_partial, do_hard_mask, threshold, model_name):
        self.threshold = threshold
        self.do_hard_mask = do_hard_mask

        stride = self.config["segment"]["stride"]
        window = self.config["segment"]["window"]
        min_window = self.config["segment"]["min-window"]
        if self.do_hard_mask or not do_partial:
            segments, mapping = segment([text], stride, window, min_window)

            mapping_lines = []
            for sent_idx, seg_idx, word_idx, char_idx in mapping:
                mapping_lines += [f"{sent_idx}, {seg_idx}, {word_idx}, {char_idx}"]

            self.mapping = self.data_utils.load_mapping_v3_from_list(mapping_lines)
            self.original_lines = [text]
            self.segments = segments
        else:
            segments = text.split('\n')

        self.segments = segments
        self.original_lines = text.split('\n')

        self.data_loader = DataLoader(
            DataRetriever(self.data_utils, segments) 
            if model_name == "D2" 
            else DataRetrieverPLM(segments, self.data_utils,
                is_test=True,
                tokenizer=self.model.tokenizer
            ),
            batch_size=self.config["predictor"].get("batch-size", 32),
            shuffle=False,
            num_workers=self.config['loader'].get('num-workers', 0),
        )
        
class PredictTri(Predictor):
    def __init__(self, config):
        super().__init__(config)
        self.diacritics = {
            "FATHA": 1,
            "KASRA": 2,
            "DAMMA": 3,
            "SUKUN": 4
        }
        self.votes: Union[Counter[int], Counter[bool]] = Counter()

    def count_votes(
            self,
            things: Union[Iterable[int], Iterable[bool]]
    ):
        self.votes.clear()
        self.votes.update(things)
        return self.votes.most_common(1)[0][0]

    def predict_majority_vote(self):
        y_gen_diac, y_gen_tanween, y_gen_shadda = self.model.predict(self.data_loader)
        diacritized_lines, _ = self.coalesce_votes_by_majority(y_gen_diac, y_gen_tanween, y_gen_shadda)
        return diacritized_lines
    
    def predict_partial(self, do_partial, lines):
        outputs = self.model.predict_partial(self.data_loader, return_extra=True, eval_only='both', do_partial=do_partial)

        if self.do_hard_mask or not do_partial:
            y_gen_diac, y_gen_tanween, y_gen_shadda = outputs['diacritics']
            diac_lines, _ = self.coalesce_votes_by_majority(y_gen_diac, y_gen_tanween, y_gen_shadda)
        else:
            diac_lines = diac_text(lines, outputs["other"][1], outputs["other"][0], selection_mode='1', threshold=self.threshold)

        return '\n'.join(diac_lines)

    def predict_majority_vote_context_contrastive(self, overwrite_cache=False):
        assert isinstance(self.model, PartialDD)
        if not os.path.exists("dataset/cache/y_gen_diac.npy") or overwrite_cache:
            if not os.path.exists("dataset/cache"):
                os.mkdir("dataset/cache")
            # segment_outputs = self.model.predict_partial(self.data_loader, return_extra=True)
            segment_outputs = self.model.predict_partial(self.data_loader, return_extra=False, eval_only='both')
            T.save(segment_outputs, "dataset/cache/cache.pt")
        else:
            segment_outputs = T.load("dataset/cache/cache.pt")
        
        y_gen_diac, y_gen_tanween, y_gen_shadda = segment_outputs['diacritics']
        diacritized_lines, extra_for_lines = self.coalesce_votes_by_majority(
            y_gen_diac, y_gen_tanween, y_gen_shadda,
        )
        extra_out = {
            'line_data': {
                **extra_for_lines,
            },
            'segment_data': {
                **segment_outputs,
                # 'logits': segment_outputs['logits'],
            }
        }

        return diacritized_lines, extra_out

    def coalesce_votes_by_majority(
            self,
            y_gen_diac: np.ndarray,
            y_gen_tanween: np.ndarray,
            y_gen_shadda: np.ndarray,
    ):
        prepped_lines_og = [' '.join(tokenize(strip_tatweel(line))) for line in self.original_lines]
        max_line_chars = max(len(line) for line in prepped_lines_og)
        diacritics_pred = np.full((len(self.original_lines), max_line_chars), fill_value=-1, dtype=int)

        count_processed_sents = 0
        do_break = False
        diacritized_lines = []
        for sent_idx, line in enumerate(tqdm(prepped_lines_og)):
            count_processed_sents = sent_idx + 1
            line = line.strip()
            diacritized_line = ""
            for char_idx, char in enumerate(line):
                diacritized_line += char
                char_vote_diacritic = []
                # ? This is the voting part
                if sent_idx not in self.mapping:
                    continue
                
                mapping_s_i = self.mapping[sent_idx]
                for seg_idx in mapping_s_i:
                    if self.data_utils.debug and seg_idx >= 256:
                        do_break = True
                        break

                    mapping_g_i = mapping_s_i[seg_idx]
                    for t_idx in mapping_g_i:
                    
                        mapping_t_i = mapping_g_i[t_idx]
                        if char_idx in mapping_t_i:
                            c_idx = mapping_t_i.index(char_idx)
                            output_idx = np.s_[seg_idx, t_idx, c_idx]
                            diac_h3 = (y_gen_diac[output_idx], y_gen_tanween[output_idx], y_gen_shadda[output_idx])
                            diac_char_i = HARAKAT_MAP.index(diac_h3)
                            if c_idx < 13 and diac_char_i != 0:
                                char_vote_diacritic.append(diac_char_i)
                
                if do_break:
                    break
                if len(char_vote_diacritic) > 0:
                    char_mv_diac = self.count_votes(char_vote_diacritic)
                    diacritized_line += shakkel_char(*HARAKAT_MAP[char_mv_diac])
                    diacritics_pred[sent_idx, char_idx] = char_mv_diac
                else:
                    diacritics_pred[sent_idx, char_idx] = 0
            if do_break:
                break
            
            diacritized_lines += [diacritized_line.strip()]

        print(f'[INFO] Cutting stats from {len(diacritics_pred)} to {count_processed_sents}')
        extra = {
            'diac_pred': diacritics_pred[:count_processed_sents],
        }
        return diacritized_lines, extra