Update unit test

neuralmagic · Nov 8, 2023 · 0c618a6 · 0c618a6
1 parent bf1b0cf
commit 0c618a6
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tests/deepsparse/transformers/pipelines/test_text_generation.py b/tests/deepsparse/transformers/pipelines/test_text_generation.py
@@ -114,10 +114,10 @@ def test_pipeline_for_ppl_eval(pipeline, prompt):
         max_length=1,
     )
     assert hasattr(predictions, "generations")
+    assert hasattr(predictions, "input_tokens")
     assert hasattr(predictions.generations[0], "score")
-    assert hasattr(predictions.generations[0], "input_tokens")
-    assert "input_ids" in predictions.generations[0].input_tokens
-    assert "attention_mask" in predictions.generations[0].input_tokens
+    assert "input_ids" in predictions.input_tokens
+    assert "attention_mask" in predictions.input_tokens
 
 
 def test_streaming_mode_returns_generator(pipeline, prompt):