[converter] Trim jit output token_str (#3527)

deepjavalibrary · Nov 16, 2024 · 4ed3149 · 4ed3149
1 parent 6100647
commit 4ed3149
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/extensions/tokenizers/src/main/python/djl_converter/fill_mask_converter.py b/extensions/tokenizers/src/main/python/djl_converter/fill_mask_converter.py
@@ -46,7 +46,9 @@ def verify_jit_output(self, hf_pipeline, encoding, out):
                 text = text.replace("[MASK]", tokenizer.mask_token)
             pipeline_output = hf_pipeline(text)
 
-            if prediction not in [o["token_str"] for o in pipeline_output]:
+            if prediction not in [
+                    o["token_str"].strip() for o in pipeline_output
+            ]:
                 logging.error(f"Unexpected inference result: {prediction}")
                 return False, "Unexpected inference result"