bug fix

model_partial.py

@@ -178,7 +178,7 @@ class PartialDD(nn.Module):
             # padding_mask: T.BoolTensor,
             *,
             eval_only: str = None,
-            subword_lengths: T.Tensor,
+            subword_lengths: T.Tensor = None,
             return_extra: bool = False
     ):
         # assert self._built and not self.training
@@ -195,7 +195,6 @@ class PartialDD(nn.Module):
                 word_ids,
                 char_ids,
                 _labels,
-                subword_lengths=subword_lengths,
             )
             out_shape = y_ctxt.shape[:-1]
         else:
@@ -328,7 +327,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)
+            output = self(*inputs, eval_only='ctxt')
 
             # output = np.argmax(T.softmax(output.detach(), dim=-1).cpu().numpy(), axis=-1)
             marks = output

predict.py

@@ -14,6 +14,7 @@ from torch.utils.data import DataLoader
 
 from diac_utils import HARAKAT_MAP, shakkel_char, diac_ids_of_line
 from model_partial import PartialDD
+from model_dd import DiacritizerD2
 from data_utils import DatasetUtils
 from dataloader import DataRetriever
 from segment import segment
@@ -44,7 +45,7 @@ class Predictor:
             if T.cuda.is_available() else 'cpu'
         )
     
-        self.model = PartialDD(config, d2=True)
+        self.model = DiacritizerD2(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)