From a16ea64f7854e8c9189ab894e7036bf7fbbf979f Mon Sep 17 00:00:00 2001 From: VictorSanh Date: Wed, 16 Aug 2023 21:04:12 +0000 Subject: [PATCH] fix lol --- server/text_generation_server/models/idefics_causal_lm.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/server/text_generation_server/models/idefics_causal_lm.py b/server/text_generation_server/models/idefics_causal_lm.py index f79caa9d..faad63ba 100644 --- a/server/text_generation_server/models/idefics_causal_lm.py +++ b/server/text_generation_server/models/idefics_causal_lm.py @@ -156,10 +156,11 @@ class IdeficsCausalLMBatch(Batch): ) # Copy tokenizer attention_mask into fully allocated attention_mask attention_mask[:, :max_input_length] = tokenized_inputs["attention_mask"] - # Do the same for image_attention_mask - I CHANGED THINGS HERE - mostly testing for now + # Do the same for image_attention_mask image_attention_mask = input_ids.new_zeros( (pb.size, max_input_length + padding_right_offset, tokenized_inputs["pixel_values"].size(1)) ) + image_attention_mask[:, :max_input_length, :] = tokenized_inputs["image_attention_mask"] position_ids = tokenized_inputs["attention_mask"].long().cumsum(-1) - 1