move inputs modification into forward

Author: eaidova

optimum/intel/openvino/modeling_decoder.py

@@ -380,7 +380,6 @@ def prepare_inputs(
         **kwargs,
     ) -> Dict:
         batch_size = input_ids.shape[0]
-        duplication_indices = None
         if self.config.model_type == "bloom":
             batch_size *= self.config.num_attention_heads
 
@@ -463,9 +462,7 @@ def prepare_inputs(
                 self.next_beam_idx if self.next_beam_idx is not None else np.arange(batch_size, dtype=int)
             )
 
-        if self._first_iter_beam_search:
-            inputs, duplication_indices = self._deduplicate_inputs(inputs)
-        return inputs, duplication_indices
+        return inputs
 
     def forward(
         self,
@@ -477,13 +474,16 @@ def forward(
     ) -> CausalLMOutputWithPast:
         self.compile()
 
-        inputs, duplication_idicies = self.prepare_inputs(
+        inputs = self.prepare_inputs(
             input_ids=input_ids,
             attention_mask=attention_mask,
             past_key_values=past_key_values,
             position_ids=position_ids,
             **kwargs,
         )
+
+        if self._first_iter_beam_search:
+            inputs, duplication_indices = self._deduplicate_inputs(inputs)
         # Run inference
         self.request.start_async(inputs, share_inputs=True)
         self.request.wait()
@@ -512,7 +512,7 @@ def forward(
                 past_key_values = None
 
         if self._first_iter_beam_search:
-            logits, past_key_values = self._expand_outputs_for_generation(duplication_idicies, logits, past_key_values)
+            logits, past_key_values = self._expand_outputs_for_generation(duplication_indices, logits, past_key_values)
             self._first_iter_beam_search = False
 
         return CausalLMOutputWithPast(logits=logits, past_key_values=past_key_values)

optimum/intel/openvino/quantization.py

@@ -688,7 +688,7 @@ def _prepare_builtin_dataset(self, quantization_config: OVWeightQuantizationConf
         nsamples = quantization_config.num_samples if quantization_config.num_samples else 128
         calibration_dataset = get_dataset(quantization_config.dataset, tokenizer, seqlen=32, nsamples=nsamples)
         calibration_dataset = prepare_dataset(calibration_dataset)
-        calibration_dataset = nncf.Dataset(calibration_dataset, lambda x: self.model.prepare_inputs(**x)[0])
+        calibration_dataset = nncf.Dataset(calibration_dataset, lambda x: self.model.prepare_inputs(**x))
 
         return calibration_dataset