[tokenizer] fixes token_classification import issues

extensions/tokenizers/src/main/python/djl_converter/huggingface_converter.py

@@ -17,6 +17,9 @@
 from argparse import Namespace
 
 import onnx
+from torch import nn
+from transformers.modeling_outputs import TokenClassifierOutput
+
 from djl_converter.safetensors_convert import convert_file
 import torch
 from huggingface_hub import hf_hub_download, HfApi
@@ -40,6 +43,28 @@ def __init__(self, config):
         self.config = config
 
 
+class ModelWrapper(nn.Module):
+
+    def __init__(self, model, include_types: bool) -> None:
+        super().__init__()
+        self.model = model
+        self.include_types = include_types
+
+    def forward(self,
+                input_ids: torch.Tensor,
+                attention_mask: torch.Tensor,
+                token_type_ids: torch.Tensor = None):
+        if self.include_types:
+            output = self.model(input_ids, attention_mask)
+        else:
+            output = self.model(input_ids, attention_mask, token_type_ids)
+        if isinstance(output, TokenClassifierOutput):
+            # TokenClassifierOutput may contains mix of Tensor and Tuple(Tensor)
+            return {"logits": output["logits"]}
+
+        return output
+
+
 class HuggingfaceConverter:
 
     def __init__(self):
@@ -279,11 +304,12 @@ def jit_trace_model(self, hf_pipeline, model_id: str, temp_dir: str,
         try:
             if include_types:
                 script_module = torch.jit.trace(
-                    hf_pipeline.model,
+                    ModelWrapper(hf_pipeline.model, include_types),
                     (input_ids, attention_mask, token_type_ids),
                     strict=False)
             else:
-                script_module = torch.jit.trace(hf_pipeline.model,
+                script_module = torch.jit.trace(ModelWrapper(
+                    hf_pipeline.model, include_types),
                                                 (input_ids, attention_mask),
                                                 strict=False)
 

extensions/tokenizers/src/main/python/djl_converter/token_classification_converter.py

@@ -61,10 +61,14 @@ def verify_jit_output(self, hf_pipeline, encoding, out):
                         or self.outputs[2] in entity):
                     return True, None
 
+        pipeline_output = hf_pipeline(self.inputs)
+        if len(pipeline_output) == 0:
+            logging.warning(f"Warning: pipeline output is empty")
+            return True, None
+
         if len(entities) == 0:
             return False, "TokenClassification returns with empty result"
 
-        pipeline_output = hf_pipeline(self.inputs)
         for e in pipeline_output:
             if e["word"] == entities[0]["word"]:
                 if e["entity"] == entities[0]["entity"]: