allow optionally fill attention mask in forward (#457)

Author: eaidova

optimum/intel/openvino/modeling_decoder.py

@@ -352,7 +352,10 @@ def forward(
             input_ids = input_ids[:, -1:]
 
         inputs = {}
+        past_len = 0
         if past_key_values is not None:
+            seq_len_dim = 1 if self.model.input(self.key_value_input_names[0]).get_partial_shape()[1].is_dynamic else 2
+            past_len = past_key_values[0][0].shape[seq_len_dim]
             if self._pkv_precision == Type.bf16:
                 # numpy does not support bf16, pretending f16, should change to bf16
                 past_key_values = tuple(
@@ -387,8 +390,13 @@ def forward(
         inputs["input_ids"] = np.array(input_ids)
 
         # Add the attention_mask inputs when needed
-        if "attention_mask" in self.input_names and attention_mask is not None:
-            inputs["attention_mask"] = np.array(attention_mask)
+        if "attention_mask" in self.input_names:
+            if attention_mask is not None:
+                inputs["attention_mask"] = np.array(attention_mask)
+            else:
+                inputs["attention_mask"] = np.ones(
+                    (input_ids.shape[0], input_ids.shape[1] + past_len), dtype=inputs["input_ids"].dtype
+                )
 
         # Run inference
         self.request.start_async(inputs, shared_memory=True)

tests/openvino/test_modeling.py

@@ -573,6 +573,29 @@ def test_auto_device_loading(self):
             del model
             gc.collect()
 
+    def test_default_filling_attention_mask(self):
+        model_id = MODEL_NAMES["gpt2"]
+        model_with_cache = OVModelForCausalLM.from_pretrained(model_id, export=True, use_cache=True)
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        tokenizer.pad_token = tokenizer.eos_token
+        texts = ["this is a simple input"]
+        tokens = tokenizer(texts, return_tensors="pt")
+        self.assertTrue("attention_mask" in model_with_cache.input_names)
+        outs = model_with_cache(**tokens)
+        attention_mask = tokens.pop("attention_mask")
+        outs_without_attn_mask = model_with_cache(**tokens)
+        self.assertTrue(torch.allclose(outs.logits, outs_without_attn_mask.logits))
+        input_ids = torch.argmax(outs.logits, dim=2)
+        past_key_values = outs.past_key_values
+        attention_mask = torch.ones((input_ids.shape[0], tokens.input_ids.shape[1] + 1), dtype=torch.long)
+        outs_step2 = model_with_cache(
+            input_ids=input_ids, attention_mask=attention_mask, past_key_values=past_key_values
+        )
+        outs_without_attn_mask_step2 = model_with_cache(input_ids=input_ids, past_key_values=past_key_values)
+        self.assertTrue(torch.allclose(outs_step2.logits, outs_without_attn_mask_step2.logits))
+        del model_with_cache
+        gc.collect()
+
 
 class OVModelForMaskedLMIntegrationTest(unittest.TestCase):
     SUPPORTED_ARCHITECTURES = (