pdd working now

.gitignore

@@ -1,4 +1,5 @@
 *.pyc
 *.pt
 *.vec
+*.pem
 .DS_Store

app.py

@@ -3,6 +3,7 @@ import yaml
 import gdown 
 import gradio as gr
 from predict import PredictTri
+from gradio import blocks 
 
 output_path = "tashkeela-d2.pt"
 if not os.path.exists(output_path):
@@ -20,18 +21,20 @@ with open("config.yaml", 'r', encoding="utf-8") as file:
 config["train"]["max-sent-len"] = config["predictor"]["window"]
 config["train"]["max-token-count"] = config["predictor"]["window"] * 3
 
-def diacritze(text):
-    print(text)
+def diacritze(text, do_partial):
     predictor = PredictTri(config, text)
-    diacritized_lines = predictor.predict_majority_vote()
-    return '\n'.join(diacritized_lines)
+    diacritized_lines = predictor.predict_partial(do_partial=do_partial)
+    return diacritized_lines
 
 with gr.Blocks() as demo:
     gr.Markdown(
     """
-    # Partial Diacritization
-    TODO: put paper links here
+    # Partial Diacritization: A Context-Contrastive Inference Approach
+    ## Authors: Muhammad ElNokrashy, Badr AlKhamissi
     """)
+
+    check_box = gr.Checkbox(label="Partial", info="Apply Partial Diacritics or Full Diacritics")
+
     input_txt = gr.Textbox(
         placeholder="اكتب هنا", 
         lines=5, 
@@ -50,7 +53,15 @@ with gr.Blocks() as demo:
     )
 
     btn = gr.Button(value="Shakkel")
-    btn.click(diacritze, inputs=input_txt, outputs=output_txt)
+    btn.click(diacritze, inputs=[input_txt, check_box], outputs=[output_txt])
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.queue().launch(
+        # share=False,
+        # debug=False,
+        # server_port=7860,
+        # server_name="0.0.0.0",
+        # ssl_verify=False,
+        # ssl_certfile="cert.pem",
+        # ssl_keyfile="key.pem"
+    )

model_partial.py

@@ -5,10 +5,11 @@ import numpy as np
 
 import torch as T
 from torch import nn
-from torch import functional as F
+from torch.nn import functional as F
 from diac_utils import flat_2_3head
 
 from model_dd import DiacritizerD2
+from model_dd import DatasetUtils
 
 class Readout(nn.Module):
     def __init__(
@@ -56,24 +57,27 @@ class PartialDiacOutput(NamedTuple):
     preds_hard: T.Tensor
     preds_ctxt_logit: T.Tensor
     preds_base_logit: T.Tensor
-    
 
 class PartialDD(nn.Module):
     def __init__(
             self,
             config: dict,
-            # feature_size: int,
-            # confidence_threshold: float,
-            d2=False
+            **kwargs
     ):
         super().__init__()
         self._built = False
         self.no_diac_id = 0
         self._dummy = nn.Parameter(T.ones(1, 1))
-
+        # with open('./configs/dd/config_d2.yaml', 'r', encoding='utf-8') as fin:
+        #     self.config_d2 = yaml.safe_load(fin)
+        # self.device = T.device('cuda' if T.cuda.is_available() else 'cpu')
         self.config = config
+        self._use_d2 = True
         self.sentence_diac = DiacritizerD2(self.config)
-
+     
+        # self.sentence_diac.to(self.device)
+        # self.build()
+        # self.word_diac = WordDD_LSTM(feature_size, num_classes=13, return_logits=False)
         self.eval()
 
     @property
@@ -114,6 +118,7 @@ class PartialDD(nn.Module):
         
         return toke_ids, char_ids, diac_ids, subword_lengths
 
+    T.jit.export
     def word_diac(
             self,
             toke_ids: T.Tensor,
@@ -169,6 +174,7 @@ class PartialDD(nn.Module):
         z = z.reshape(Nb, Tw, Tc, -1)
         return z
 
+    T.jit.ignore
     def forward(
             self,
             word_ids: T.Tensor,
@@ -178,8 +184,9 @@ class PartialDD(nn.Module):
             # padding_mask: T.BoolTensor,
             *,
             eval_only: str = None,
-            subword_lengths: T.Tensor = None,
-            return_extra: bool = False
+            subword_lengths: T.Tensor,
+            return_extra: bool = False,
+            do_partial: bool = False,
     ):
         # assert self._built and not self.training
         assert not self.training
@@ -195,6 +202,7 @@ class PartialDD(nn.Module):
                 word_ids,
                 char_ids,
                 _labels,
+                subword_lengths=subword_lengths,
             )
             out_shape = y_ctxt.shape[:-1]
         else:
@@ -219,6 +227,7 @@ class PartialDD(nn.Module):
         if eval_only == 'base':
             return y_base.argmax(-1)
 
+        #! TODO: Return the logits.
         ypred_ctxt = y_ctxt.argmax(-1)
         ypred_base = y_base.argmax(-1)
         #^ ypred: [b tw tc _]
@@ -226,7 +235,9 @@ class PartialDD(nn.Module):
         # Maybe for eval
         # ypred_ctxt[~((ypred_base == ground_truth) & (~padding_mask))] = self.no_diac_id
         # return ypred_ctxt
-        ypred_ctxt[(padding_mask) | (ypred_base == ypred_ctxt)] = self.no_diac_id
+        if do_partial:
+            ypred_ctxt[(padding_mask) | (ypred_base == ypred_ctxt)] = self.no_diac_id
+            
         if not return_extra:
             return ypred_ctxt
         else:
@@ -250,6 +261,7 @@ class PartialDD(nn.Module):
             dataloader,
             return_extra=False,
             eval_only: str = None,
+            do_partial=True,
     ):
         training = self.training
         self.eval()
@@ -261,10 +273,11 @@ class PartialDD(nn.Module):
             'diacs':   [],
             'y_ctxt':  [],
             'y_base':  [],
+            'subword_lengths': [],
         }
         print("> Predicting...")
         # breakpoint()
-        for i_batch, (inputs, _, subword_lengths) in enumerate(tqdm(dataloader)):
+        for i_batch, (inputs, _) in enumerate(tqdm(dataloader)):
             # if i_batch > 10:
             #     break
             #^ inputs: [toke_ids, char_ids, diac_ids]
@@ -282,15 +295,19 @@ class PartialDD(nn.Module):
                     subword_lengths=subword_lengths,
                     return_extra=return_extra,
                     eval_only=eval_only,
+                    do_partial=do_partial,
                 )
 
             # output = np.argmax(T.softmax(output.detach(), dim=-1).cpu().numpy(), axis=-1)
             if return_extra:
                 assert isinstance(output, PartialDiacOutput)
                 marks = output.preds_hard
+                if eval_only == 'recalibrated':
+                    marks = (output.preds_ctxt_logit + output.preds_base_logit).argmax(-1)
                 preds['diacs'].extend(list(marks.detach().cpu().numpy()))
                 preds['y_ctxt'].extend(list(output.preds_ctxt_logit.detach().cpu().numpy()))
                 preds['y_base'].extend(list(output.preds_base_logit.detach().cpu().numpy()))
+                preds['subword_lengths'].extend(list(subword_lengths.detach().cpu().numpy()))
             else:
                 assert isinstance(output, T.Tensor)
                 marks = output
@@ -312,9 +329,10 @@ class PartialDD(nn.Module):
                 np.array(preds["shadda"]),
             ),
             'other': ( # Would be empty when !return_extra
-                preds['y_ctxt'],
-                preds['y_base'],
-                preds['diacs'],
+                np.array(preds['y_ctxt']),
+                np.array(preds['y_base']),
+                np.array(preds['diacs']),
+                np.array(preds['subword_lengths']),
             )
         }
 
@@ -327,7 +345,7 @@ class PartialDD(nn.Module):
         for inputs, _ in tqdm(dataloader, total=len(dataloader)):
             inputs[0] = inputs[0].to(self.device)
             inputs[1] = inputs[1].to(self.device)
-            output = self(*inputs, eval_only='ctxt')
+            output = self(*inputs)
 
             # output = np.argmax(T.softmax(output.detach(), dim=-1).cpu().numpy(), axis=-1)
             marks = output
@@ -344,4 +362,4 @@ class PartialDD(nn.Module):
             np.array(preds['haraka']),
             np.array(preds["tanween"]),
             np.array(preds["shadda"]),
-        )
+        )

partial_dd_metrics.py

@@ -0,0 +1,329 @@
+from typing import NamedTuple
+from argparse import ArgumentParser
+
+from tqdm import tqdm
+import logging
+
+import numpy as np
+import torch as T
+from torch.nn import functional as F
+
+import diac_utils as du
+
+_x = [
+    'a'
+]
+
+# logging.setLevel(logging.INFO)
+logger = logging.getLogger(__file__)
+logger.setLevel(logging.INFO)
+
+def logln(*texts: str):
+    # logger.info(' '.join(texts))
+    print(*texts)
+
+# Relative improvement:
+#     T.mean((pred_c.argmax('c') == gt) - (pred_m.argmax('c') == gt))
+# Coverage Confidence:
+#     pred_c.argmax('c')[pred_c.argmax('c') != pred_m.argmax('c')].mean()
+
+class PartialDiacMetrics(NamedTuple):
+    diff_total: float
+    worse_total: float
+    diff_relative: float
+    der_total: float
+    selectivity: float
+    hidden_der: float
+    partial_der: float
+    reader_error: float
+
+def load_data(path: str):
+    if path.endswith('.txt'):
+        with open(path, 'r', encoding='utf-8') as fin:
+            return fin.readlines()
+    else:
+        return T.load(path)
+
+def parse_data(
+        data,
+        logits: bool = False,
+        side=None,
+):
+    if logits:
+        ld = data['line_data']
+        diac_logits = T.tensor(ld[f'diac_logits_{side}'])
+        # diac_pred: T.Tensor = ld['diac_pred']
+        diac_pred: T.Tensor = diac_logits.argmax(dim=-1)
+        diac_gt  : T.Tensor = ld['diac_gt']
+        # diac_logits = (ld['diac_logits_ctxt'], ld['diac_logits_base'])
+        return diac_pred, diac_gt, diac_logits
+    if isinstance(data, dict):
+        ld = data.get('line_data_fix', data['line_data'])
+        if side is None:
+            diac_pred: T.Tensor = ld['diac_pred']
+        else:
+            diac_pred: T.Tensor = ld[f'diac_logits_{side}'].argmax(axis=-1)
+        diac_gt  : T.Tensor = ld['diac_gt']
+        return diac_pred, diac_gt
+    elif isinstance(data, list):
+        data_indices = [
+            du.diac_ids_of_line(du.strip_tatweel(du.normalize_spaces(line)))
+            for line in data
+        ]
+        max_len = max(map(len, data_indices))
+        out = np.full((len(data), max_len), fill_value=du.DIAC_PAD_IDX)
+        for i_line, line_indices in enumerate(data_indices):
+            out[i_line][:len(line_indices)] = line_indices
+        return out, None
+    elif isinstance(data, (T.Tensor, np.ndarray)):
+        return data, None
+    else:
+        raise NotImplementedError
+
+def make_mask_hard(
+        pred_c: T.Tensor,
+        pred_m: T.Tensor,
+):  
+    selection  = (pred_c != pred_m)
+    return selection
+
+def make_mask_logits(
+        pred_c: T.Tensor,
+        pred_m: T.Tensor,
+        threshold: float = 0.1,
+        version: str = '2',
+) -> T.BoolTensor:
+    logger.warning(f"{version=}, {threshold=}")
+    pred_c = T.softmax(T.tensor(pred_c), dim=-1)
+    pred_m = T.softmax(T.tensor(pred_m), dim=-1)
+    # pred_i = pred_c.argmax(dim=-1)
+    if version == 'hard':
+        selection = pred_c.argmax(-1) != pred_m.argmax(-1)
+    elif version == '0':
+        selection = pred_c.max(dim=-1).values > pred_m.max(dim=-1).values
+        selection = selection & (pred_m.max(dim=-1).values > threshold)
+    elif version == '1':
+        pred_c_conf = pred_c.max(dim=-1).values
+        pred_m_conf = pred_m.max(dim=-1).values
+        selection = (pred_c_conf - pred_m_conf) > threshold
+    elif version == '1.1':
+        pred_c_conf = pred_c.max(dim=-1).values
+        pred_m_conf = pred_m.max(dim=-1).values
+        selection = (pred_c_conf - pred_m_conf).abs() > threshold
+    elif version.startswith('2'):
+        if version == '2':
+            max_c = pred_c.argmax(dim=-1, keepdims=True)
+            selection = T.gather(pred_c - pred_m, dim=-1, index=max_c) > threshold
+        elif version == '2.1':
+            max_c = pred_m.argmax(dim=-1, keepdims=True)
+            selection = T.gather(pred_c - pred_m, dim=-1, index=max_c) > threshold
+        elif version == '2.abs':
+            max_c = pred_c.argmax(dim=-1, keepdims=True)
+            selection = T.gather(pred_c - pred_m, dim=-1, index=max_c).abs() > threshold
+        elif version == '2.1.abs':
+            max_c = pred_m.argmax(dim=-1, keepdims=True)
+            selection = T.gather(pred_c - pred_m, dim=-1, index=max_c).abs() > threshold
+    elif version == '3':
+        selection = (pred_c - pred_m).max(dim=-1).values > threshold
+    elif version == '4':
+        selection_hard   = (pred_c.argmax(-1) != pred_m.argmax(-1))
+        # selection_logits = (pred_c.max(-1).values - pred_m.max(-1).values) > threshold
+        selection_logits = T.gather(pred_c - pred_m, dim=-1, index=pred_c.argmax(-1, keepdims=True)) > threshold
+        selection = selection_hard & selection_logits.squeeze()
+    # selection  = (pred_c != pred_m)
+    return selection.squeeze()
+
+def analysis_summary(
+        pred_c      : T.LongTensor,
+        pred_m      : T.LongTensor,
+        labels      : T.LongTensor,
+        padding_mask: T.BoolTensor,
+        *,
+        selection   : T.Tensor = None,
+        random: bool = False,
+        logits: tuple = None
+):
+    #^ pred_c: [b tw tc | ClassId]
+    #^ pred_m: [b tw tc | ClassId]
+    #^ labels: [b tw tc | ClassId]
+    padding_mask = T.tensor(padding_mask)
+    # padding_mask[:, 200:] = False
+    nonpad_mask = ~padding_mask
+    num_chars = nonpad_mask.sum()
+    
+    if logits is not None:
+        logits = tuple(map(T.tensor, logits))
+        # pred_c = (logits[0] + logits[1]).argmax(-1)
+        pred_c = (T.softmax(logits[0], dim=-1) + T.softmax(logits[1], dim=-1)).argmax(-1)
+    pred_c = T.tensor(pred_c)[nonpad_mask]
+    pred_m = T.tensor(pred_m)[nonpad_mask]
+    labels = T.tensor(labels)[nonpad_mask]
+    #^ : [(b * tw * tc) | ClassId]
+
+    ctxt_match = (pred_c == labels).float()
+    base_match = (pred_m == labels).float()
+    
+    selection = T.tensor(selection)[nonpad_mask]
+    if random:
+        selection  = pred_c.new_empty(pred_c.shape).bernoulli_(p=selection.float().mean()).to(bool)
+    unselected = ~selection
+
+    assert num_chars > 0
+    assert selection.sum() > 0
+    base_accuracy = base_match[unselected].sum() / unselected.sum()
+    ctxt_accuracy = ctxt_match[selection].sum() / selection.sum()
+    correct_total = ctxt_match.sum() / num_chars
+    der_total = 1 - correct_total
+    
+    cmp = (ctxt_match - base_match)[selection]
+    diff = T.sum(cmp)
+    diff_total = diff / num_chars
+    diff_relative = diff / selection.sum()
+    
+    selectivity    = selection.sum() / num_chars
+    worse_total = base_match[selection].sum() / num_chars
+    
+    hidden_der   = 1.0 - base_accuracy
+    partial_der  = 1.0 - ctxt_accuracy
+    reader_error = selectivity * partial_der + (1 - selectivity) * hidden_der
+    
+    return PartialDiacMetrics(
+        diff_total      = round(diff_total.item() * 100, 2),
+        worse_total     = round(worse_total.item() * 100, 2),
+        diff_relative   = round(diff_relative.item() * 100, 2),
+        der_total       = round(der_total.item() * 100, 2),
+        selectivity     = round(selectivity.item() * 100, 2),
+        hidden_der      = round(hidden_der.item() * 100, 2),
+        partial_der     = round(partial_der.item() * 100, 2),
+        reader_error    = round(reader_error.item() * 100, 2)
+    )
+
+
+def relative_improvement_soft(
+        pred_c      : T.Tensor,
+        pred_m      : T.Tensor,
+        labels      : T.LongTensor,
+        padding_mask: T.Tensor,
+):
+    #^ pred_c: [b tw tc Classes="15"]
+    #^ pred_m: [b tw tc Classes="15"]
+    padding_mask = T.tensor(padding_mask)
+    nonpad_mask = 1 - padding_mask.float()
+    num_chars = nonpad_mask.sum()
+    
+    pred_c = T.tensor(pred_c)[~padding_mask]
+    pred_m = T.tensor(pred_m)[~padding_mask]
+    #^ : [(b * tw * tc), Classes]
+    labels = T.tensor(labels)[~padding_mask]
+    #^ : [(b * tw * tc) | ClassId]
+
+    ctxt_match = T.gather(pred_c, dim=1, index=labels)
+    base_match = T.gather(pred_m, dim=1, index=labels)
+    selection  = (pred_c.argmax(-1) != pred_m.argmax(-1))
+
+    better = T.sum(ctxt_match - base_match) / num_chars
+    selectivity = selection.sum() / num_chars
+    worse = base_match[selection].sum() / num_chars
+    return better, worse, selectivity
+
+def relative_improvement_masked_soft(
+        pred_c: T.Tensor,
+        pred_m: T.Tensor,
+        ground_truth: T.LongTensor,
+        padding_mask: T.Tensor,
+):
+    raise NotImplementedError
+    #^ pred_c: [b tw tc "13"]
+    #^ pred_m: [b tw tc "13"]
+    #^ ground_truth: [b tw tc ClassId]
+    nonpad_mask = 1 - padding_mask
+
+    selection_mask = pred_c.argmax(3) != pred_m.argmax(3)
+    #^ selection_mask: [b tw tc]
+    probs = F.softmax(pred_c.clone(), dim=-1)
+    probs_gt = T.gather(probs, dim=-1, index=ground_truth.unsqueeze(-1)).squeeze(-1)
+    #^ probs_gt: [b tw tc]
+    result = probs_gt[selection_mask & nonpad_mask].mean()
+    return result
+
+def coverage_confidence(
+        pred_c: T.Tensor,
+        pred_m: T.Tensor,
+        padding_mask: T.Tensor,
+        # selection_mask: T.Tensor,
+):
+    raise NotImplementedError
+    #^ pred_c:         [b tw tc "13"]
+    #^ pred_m:         [b tw tc "13"]
+    #^ selection_mask: [b tw tc (bool)]
+    pred_c_id = pred_c.argmax(3)
+    pred_m_id = pred_m.argmax(3)
+    selected = pred_c_id[pred_c_id != pred_m_id]
+    nonpad_mask = 1 - padding_mask
+    result = selected.sum() / nonpad_mask.sum()
+    return result
+
+def cli():
+    parser = ArgumentParser('Compare diacritics from base/ctxt systems with partial diac metrics.')
+    parser.add_argument('-m', '--model-output-base', help="Path to tensor.pt dump files of base diacs.")
+    parser.add_argument('-c', '--model-output-ctxt', help="Path to tensor.pt dump files of ctxt diacs.")
+    parser.add_argument('--gt', default=None, help="Path to tensor.pt for gt only.")
+    parser.add_argument('--mode', choices=['hard', 'logits'], default='hard')
+    args = parser.parse_args()
+    
+    model_output_base = parse_data(
+        load_data(args.model_output_base),
+        # logits=args.mode == 'logits',
+        logits=True,
+        side='base',
+    )
+    model_output_ctxt = parse_data(
+        load_data(args.model_output_ctxt),
+        # logits=args.mode == 'logits',
+        logits=True,
+        side='ctxt',
+    )
+    #^ shape: [b, tc] -> ClassId
+    diacs_pred = model_output_base
+    
+    logln(f"{model_output_base[0].shape=} , {model_output_ctxt[0].shape=}")
+    
+    assert len(model_output_base[0]) == len(model_output_ctxt[0])
+    
+    # for diacs_base, diacs_ctxt in zip(
+    #         tqdm(model_output_base, dynamic_cols=True),
+    #         model_output_ctxt
+    # ):
+    #     diacs = np.where(diacs_base != diacs_ctxt, diacs_ctxt, 0)[diacs_ctxt != -1] #< Ignore padding
+
+    xc = model_output_ctxt
+    xm = model_output_base
+    # if args.mode == 'logits':
+    # elif args.mode == 'hard':
+    #     xc = model_output_ctxt
+    #     xm = model_output_base
+    # if args.gt is not None:
+    #     ground_truth = parse_data(load_data(args.gt))[1]
+    if xm[1] is not None:
+        ground_truth = xm[1]
+    elif xc[1] is not None:
+        ground_truth = xc[1]
+    assert ground_truth is not None
+
+    if args.mode == 'hard':
+        selection = make_mask_hard(xc[0], xm[0])
+    elif args.mode == 'logits':
+        selection = make_mask_logits(xc[2], xm[2])
+
+    metrics = analysis_summary(
+        xc[0], xm[0], ground_truth, ground_truth == -1,
+        selection=selection,
+        logits=(xc[2], xm[2])
+    )
+    logln("Actual Totals:", metrics)
+    metrics = analysis_summary(
+        xc[0], xm[0], ground_truth, ground_truth == -1, random=True,
+        selection=selection,
+        logits=(xc[2], xm[2])
+    )
+    logln("Random Marked Chars:", metrics)

predict.py

@@ -5,7 +5,7 @@ import argparse
 import os
 
 import yaml
-from pyarabic.araby import tokenize, strip_tatweel
+from pyarabic.araby import tokenize, strip_tatweel, strip_tashkeel
 from tqdm import tqdm
 
 import numpy as np
@@ -19,6 +19,69 @@ from data_utils import DatasetUtils
 from dataloader import DataRetriever
 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 = ""
+    diacs_h3 = DatasetUtils.flat2_3head(diacs)
+    for ch, tashkeel in zip(line, zip(*diacs_h3)):
+        line_w_diacs += ch
+        line_w_diacs += DatasetUtils.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 '\n'.join(output)
+
 class Predictor:
     def __init__(self, config, text):
 
@@ -45,8 +108,8 @@ class Predictor:
             if T.cuda.is_available() else 'cpu'
         )
     
-        self.model = DiacritizerD2(config)
-        self.model.build(word_embeddings, vocab_size)
+        self.model = PartialDD(config)
+        self.model.sentence_diac.build(word_embeddings, vocab_size)
         state_dict = T.load(config["paths"]["load"], map_location=T.device(self.device))['state_dict']
         self.model.load_state_dict(state_dict)
         self.model.to(self.device)
@@ -82,6 +145,13 @@ class PredictTri(Predictor):
         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):
+        outputs = self.model.predict_partial(self.data_loader, return_extra=True, eval_only='both', do_partial=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)
+        return '\n'.join(diac_lines)
 
     def predict_majority_vote_context_contrastive(self, overwrite_cache=False):
         assert isinstance(self.model, PartialDD)
@@ -89,7 +159,7 @@ class PredictTri(Predictor):
             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='ctxt')
+            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")
@@ -107,6 +177,7 @@ class PredictTri(Predictor):
                 # 'logits': segment_outputs['logits'],
             }
         }
+
         return diacritized_lines, extra_out
 
     def coalesce_votes_by_majority(