Ensure backward compatibility of ORTModel (#933)

* fix

* Update optimum/onnxruntime/base.py

Co-authored-by: Michael Benayoun <mickbenayoun@gmail.com>

---------

Co-authored-by: Michael Benayoun <mickbenayoun@gmail.com>

optimum/onnxruntime/base.py

@@ -22,7 +22,11 @@
 from onnxruntime import InferenceSession
 
 from ..utils import NormalizedConfigManager
-from .utils import get_ordered_input_names
+from ..utils.logging import warn_once
+from .utils import get_ordered_input_names, logging
+
+
+logger = logging.get_logger(__name__)
 
 
 if TYPE_CHECKING:
@@ -402,6 +406,10 @@ def __init__(
         else:
             self.num_pkv = 4
 
+        self.legacy_outputs = any(
+            "encoder" in output_name and "present" in output_name for output_name in self.output_names
+        )
+
     def compute_past_key_values_output_shapes(
         self, input_ids, encoder_hidden_states, past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
     ) -> Dict[str, int]:
@@ -569,10 +577,21 @@ def filter_out_output(output_name):
                 out_past_key_values[i : i + self.num_pkv] for i in range(0, len(out_past_key_values), self.num_pkv)
             )
         else:
-            # grab the cross attention key/values from the inputs
-            out_past_key_values = tuple(
-                out_past_key_values[i : i + self.num_pkv] + past_key_values[2 * i + 2 : 2 * i + 2 + self.num_pkv]
-                for i in range(0, len(out_past_key_values), self.num_pkv)
-            )
+            if self.legacy_outputs is True:
+                msg = (
+                    "For the decoder with past, using ONNX models outputting cross attention past key values"
+                    " is deprecated and the support will be removed in optimum 2.0. We recommend exporting again the model"
+                    " with optimum>=1.7.3."
+                )
+                warn_once(logger, msg=msg)
+                out_past_key_values = tuple(
+                    out_past_key_values[i : i + self.num_pkv] for i in range(0, len(out_past_key_values), self.num_pkv)
+                )
+            else:
+                # Grab the cross attention key/values from the inputs
+                out_past_key_values = tuple(
+                    out_past_key_values[i : i + self.num_pkv] + past_key_values[2 * i + 2 : 2 * i + 2 + self.num_pkv]
+                    for i in range(0, len(out_past_key_values), self.num_pkv)
+                )
 
         return Seq2SeqLMOutput(loss=loss, logits=logits, past_key_values=out_past_key_values)

optimum/utils/logging.py

@@ -19,6 +19,7 @@
 import os
 import sys
 import threading
+from functools import lru_cache
 from logging import (
     CRITICAL,  # NOQA
     DEBUG,  # NOQA
@@ -29,7 +30,11 @@
     WARN,  # NOQA
     WARNING,  # NOQA
 )
-from typing import Optional
+from typing import TYPE_CHECKING, Optional
+
+
+if TYPE_CHECKING:
+    from logging import Logger
 
 
 _lock = threading.Lock()
@@ -262,3 +267,8 @@ def reset_format() -> None:
 
     for handler in handlers:
         handler.setFormatter(None)
+
+
+@lru_cache(None)
+def warn_once(logger: "Logger", msg: str):
+    logger.warning(msg)

tests/onnxruntime/test_modeling.py

@@ -1983,6 +1983,15 @@ class ORTModelForCausalLMIntegrationTest(ORTModelTestMixin):
     GENERATION_LENGTH = 100
     SPEEDUP_CACHE = 1.1
 
+    def test_inference_old_onnx_model(self):
+        model = ORTModelForCausalLM.from_pretrained("optimum/gpt2")
+
+        tokenizer = get_preprocessor("optimum/gpt2")
+        text = "This is a sample output"
+        tokens = tokenizer(text, return_tensors="pt")
+
+        model.generate(**tokens)
+
     def test_load_model_from_hub_onnx(self):
         model = ORTModelForCausalLM.from_pretrained("fxmarty/onnx-tiny-random-gpt2-without-merge")
 
@@ -3064,6 +3073,15 @@ class ORTModelForSeq2SeqLMIntegrationTest(ORTModelTestMixin):
     GENERATION_LENGTH = 100
     SPEEDUP_CACHE = 1.1
 
+    def test_inference_old_onnx_model(self):
+        model = ORTModelForSeq2SeqLM.from_pretrained("optimum/t5-small")
+
+        tokenizer = get_preprocessor("optimum/t5-small")
+        text = "This is a sample output"
+        tokens = tokenizer(text, return_tensors="pt")
+
+        model.generate(**tokens)
+
     def test_load_vanilla_transformers_which_is_not_supported(self):
         with self.assertRaises(Exception) as context:
             _ = ORTModelForSeq2SeqLM.from_pretrained(MODEL_NAMES["bert"], from_transformers=True)