From fa051e4ccbfdcda6c4c5b05fffb33879c8219dcc Mon Sep 17 00:00:00 2001 From: Nikita Savelyev Date: Thu, 25 Jul 2024 21:40:47 +0200 Subject: [PATCH] Also adopt changes from #672 --- llm_bench/python/llm_bench_utils/conversion_utils/helpers.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/llm_bench/python/llm_bench_utils/conversion_utils/helpers.py b/llm_bench/python/llm_bench_utils/conversion_utils/helpers.py index 48ce10eabe..387b0ece8d 100644 --- a/llm_bench/python/llm_bench_utils/conversion_utils/helpers.py +++ b/llm_bench/python/llm_bench_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