Fix 2 tests in FillMaskPipelineTests (#27889)

* fix * fix * fix * fix * fix --------- Co-authored-by: ydshieh <[email protected]>
huggingface · Dec 8, 2023 · e366937 · e366937
1 parent 79e7655
commit e366937
Showing 1 changed file with 13 additions and 4 deletions.
diff --git a/tests/pipelines/test_pipelines_fill_mask.py b/tests/pipelines/test_pipelines_fill_mask.py
@@ -216,15 +216,24 @@ def run_large_test(self, unmasker):
             ],
         )
 
+        dummy_str = "Lorem ipsum dolor sit amet, consectetur adipiscing elit," * 100
         outputs = unmasker(
-            "My name is <mask>" + "Lorem ipsum dolor sit amet, consectetur adipiscing elit," * 100,
+            "My name is <mask>" + dummy_str,
             tokenizer_kwargs={"truncation": True},
         )
+        simplified = nested_simplify(outputs, decimals=4)
         self.assertEqual(
-            nested_simplify(outputs, decimals=6),
+            [{"sequence": x["sequence"][:100]} for x in simplified],
             [
-                {"sequence": "My name is grouped", "score": 2.2e-05, "token": 38015, "token_str": " grouped"},
-                {"sequence": "My name is accuser", "score": 2.1e-05, "token": 25506, "token_str": " accuser"},
+                {"sequence": f"My name is,{dummy_str}"[:100]},
+                {"sequence": f"My name is:,{dummy_str}"[:100]},
+            ],
+        )
+        self.assertEqual(
+            [{k: x[k] for k in x if k != "sequence"} for x in simplified],
+            [
+                {"score": 0.2819, "token": 6, "token_str": ","},
+                {"score": 0.0954, "token": 46686, "token_str": ":,"},
             ],
         )