fix

openvinotoolkit · Nov 15, 2024 · a4f86ac · a4f86ac
1 parent 45ec25a
commit a4f86ac
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/tests/python_tests/test_cache_optimizations.py b/tests/python_tests/test_cache_optimizations.py
@@ -219,7 +219,7 @@ def test_unoptimized_generation_longbench(phi3_converted_model, subset):
                 )
                 for i, output in enumerate(ans_batch, start=p_idx-len(batch)+1):
                     context_len = answers[i]["context_len"]
-                    pred = post_process_pred(output.m_generation_ids, subset, model_name)
+                    pred = post_process_pred(output.m_generation_ids[0], subset, model_name)
                     answers[i]["pred"] = pred
 
                 batch.clear()

diff --git a/tests/python_tests/utils_longbench.py b/tests/python_tests/utils_longbench.py
@@ -249,7 +249,7 @@ def preprocess_prompt(tokenizer, data_sample, subset, model_name):
     return prompt, context_len
 
 
-def post_process_pred(subset, model_name):
+def post_process_pred(pred, subset, model_name):
     if subset in ["samsum", "qsum", "hotpotqa", "qasper"] and "Llama-3-8B" in model_name:
         pred = pred[:pred.find("assistant")]
     elif subset == "samsum":