Fix max token slicing

openvinotoolkit · Jul 24, 2024 · 6f82e2b · 6f82e2b
1 parent 5d21486
commit 6f82e2b
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/llm_bench/python/utils/conversion_utils/helpers.py b/llm_bench/python/utils/conversion_utils/helpers.py
@@ -96,8 +96,8 @@ def save_tokenizer(tokenizer, out_dir):
 
 def transform_fn(item, item_name, input_shapes, tokenizer, config, max_tokens=127):
     tokenized_text = tokenizer(item[item_name], return_tensors="np")
-    input_ids = tokenized_text["input_ids"][:max_tokens]
-    attention_mask = tokenized_text["attention_mask"][:max_tokens]
+    input_ids = tokenized_text["input_ids"][:, :max_tokens]
+    attention_mask = tokenized_text["attention_mask"][:, :max_tokens]
 
     inputs = {}
     inputs["input_ids"] = input_ids