BatchEncoding wrapper for custom tokenizer output

harim_plus.py

@@ -7,9 +7,12 @@ import torch.nn.functional as F
 from transformers import (AutoModelForSeq2SeqLM,
                         AutoTokenizer,
                         PreTrainedTokenizer,
-                        PreTrainedTokenizerFast)
+                        PreTrainedTokenizerFast,
+                        )
+from transformers.tokenization_utils_base import BatchEncoding # for custom tokenizer other than huggingface
+
 import pandas as pd
-from tqdm import tqdm
+from tqdm import tqdme
 
 from typing import List, Dict, Union
 from collections import defaultdict
@@ -201,8 +204,15 @@ class Harimplus_Scorer:
                 emp_in = self._prep_input( mini_e_, src_or_tgt='src' )
 
 
-            tgt_mask = tgt_in.attention_mask
-
+            tgt_mask = tgt_in.attention_mask # torch.Tensor
+            # if not tokenizer loaded from huggingface, this might cause some problem (.to(device))
+            if not isinstance(src_in, BatchEncoding):
+                src_in = BatchEncoding(src_in)
+            if not isinstance(emp_in, BatchEncoding):
+                emp_in = BatchEncoding(emp_in)
+            if not isinstance(tgt_in, BatchEncoding):
+                tgt_in = BatchEncoding(tgt_in)
+                
             src_in = src_in.to(self._device)
             emp_in = emp_in.to(self._device)
             tgt_in = tgt_in.to(self._device)