fix: add a flag for return_word_ids

aligning with NGramTokenizer

Author: meilame-tayebjee

torchTextClassifiers/tokenizers/base.py

@@ -122,7 +122,10 @@ def __init__(
         self.output_dim = output_dim  # constant context size for all batch
 
     def tokenize(
-        self, text: Union[str, List[str]], return_offsets_mapping: Optional[bool] = False
+        self,
+        text: Union[str, List[str]],
+        return_offsets_mapping: Optional[bool] = False,
+        return_word_ids: Optional[bool] = False,
     ) -> list:
         if not self.trained:
             raise RuntimeError("Tokenizer must be trained before tokenization.")
@@ -142,11 +145,16 @@ def tokenize(
 
         encoded_text = tokenize_output["input_ids"]
 
+        if return_word_ids:
+            word_ids = np.array([tokenize_output.word_ids(i) for i in range(len(encoded_text))])
+        else:
+            word_ids = None
+
         return TokenizerOutput(
             input_ids=encoded_text,
             attention_mask=tokenize_output["attention_mask"],
             offset_mapping=tokenize_output.get("offset_mapping", None),
-            word_ids=np.array([tokenize_output.word_ids(i) for i in range(len(encoded_text))]),
+            word_ids=word_ids,
         )
 
     @classmethod

torchTextClassifiers/tokenizers/ngram.py

@@ -282,7 +282,7 @@ def __init__(
             self.subword_cache = None
 
         self.vocab_size = 3 + self.nwords + self.num_tokens
-        print("brrrrr ", self.vocab_size)
+
         super().__init__(
             vocab_size=self.vocab_size, padding_idx=self.pad_token_id, output_dim=output_dim
         )

torchTextClassifiers/torchTextClassifiers.py

@@ -460,6 +460,7 @@ def predict(
 
         if explain:
             return_offsets_mapping = True  # to be passed to the tokenizer
+            return_word_ids = True
             if self.pytorch_model.text_embedder is None:
                 raise RuntimeError(
                     "Explainability is not supported when the tokenizer outputs vectorized text directly. Please use a tokenizer that outputs token IDs."
@@ -474,6 +475,7 @@ def predict(
                 )  # initialize a Captum layer gradient integrator
         else:
             return_offsets_mapping = False
+            return_word_ids = False
 
         X_test = self._check_X(X_test)
         text = X_test["text"]
@@ -482,7 +484,9 @@ def predict(
         self.pytorch_model.eval().cpu()
 
         tokenize_output = self.tokenizer.tokenize(
-            text.tolist(), return_offsets_mapping=return_offsets_mapping
+            text.tolist(),
+            return_offsets_mapping=return_offsets_mapping,
+            return_word_ids=return_word_ids,
         )
 
         if not isinstance(tokenize_output, TokenizerOutput):