address PR review

Author: Kotomi-Du

onnxruntime/core/providers/openvino/ov_interface.cc

@@ -365,7 +365,7 @@ StatefulOVInferRequest::StatefulOVInferRequest(ov::InferRequest infer_request, s
   // check if there is input_ids tensors and if the tensor type is int64,
   // because logic prefill_use_full_chat_history is only for specific inputs and data type
   auto input_ids_opt = FindTensor("input_ids");
-  if (gpu_or_npu && input_ids_opt.has_value() && input_ids_opt->get_element_type() != ov::element::i64) {
+  if (gpu_or_npu && input_ids_opt.has_value() && input_ids_opt->get_element_type() == ov::element::i64) {
     prefill_use_full_chat_history = true;
   }
 }

onnxruntime/core/providers/openvino/ov_stateful_patch_utils.cc

@@ -59,6 +59,17 @@ bool ModelHasInputOutputNames(std::shared_ptr<ov::Model> model, const std::strin
   return false;
 }
 
+std::string GetInputOutputName(std::shared_ptr<ov::Model> ov_model,
+                            const std::vector<std::string>& candidate_names) {
+  for (const auto& name : candidate_names) {
+    if (ModelHasInputOutputNames(ov_model, name)) {
+      return name;
+    }
+  }
+  // Return the first candidate as default if none are found
+  return candidate_names.empty() ? "" : candidate_names[0];
+}
+
 void FuseCacheReorder(std::shared_ptr<ov::Model> ov_model,
                       std::vector<std::string>& not_kv_inputs,
                       const std::vector<std::string>& key_value_input_names,
@@ -67,18 +78,15 @@ void FuseCacheReorder(std::shared_ptr<ov::Model> ov_model,
     throw std::runtime_error("Model already has fused cache");
   }
 
-  std::string main_input_name = "inputs_embeds";
-  if (ModelHasInputOutputNames(ov_model, "input_ids")) {
-    main_input_name = "input_ids";
-  }
-
-  if (ModelHasInputOutputNames(ov_model, "input_hidden_states")) {
-    main_input_name = "input_hidden_states";
-  }
+    // Define input name candidates in priority order
+  const std::vector<std::string> input_name_candidates = {
+    "inputs_embeds",                           // Default fallback
+    "input_ids",                               // Most common
+    "input_hidden_states",                     // Alternative
+    "/model/embed_tokens/Gather_output_0"      // Specific model type
+  };
 
-  if (ModelHasInputOutputNames(ov_model, "/model/embed_tokens/Gather_output_0")) {
-    main_input_name = "/model/embed_tokens/Gather_output_0";
-  }
+  std::string main_input_name = GetInputOutputName(ov_model, input_name_candidates);
 
   auto input_batch = ov_model->input(main_input_name).get_partial_shape()[0];
 
@@ -131,7 +139,7 @@ void PatchStatefulDecoder(std::shared_ptr<ov::Model> model) {
   std::vector<std::string> not_kv_inputs;
   const auto& params = model->get_parameters();
   bool found = false;
-  for (auto i = 0; i < params.size(); i++) {
+  for (size_t i = 0; i < params.size(); i++) {
     auto param_name = params.at(i)->output(0).get_any_name();
     if (param_name.find("key_values") != std::string::npos) {
          key_value_input_names.push_back(param_name);