find names pattern

Author: xadupre

CHANGELOGS.rst

@@ -4,8 +4,9 @@ Change Logs
 0.8.0
 +++++
 
-* :pr:`276`: implements onnx_generate which implements method generate for an onnx model,
-  changes the serialization for all caches to reorder the model outputs (key_1, value_1, key_2, ...)
+* :pr:`278`: implements ``onnx_generate_with_genai``
+* :pr:`277`: changes the serialization for all caches to reorder the model outputs (key_1, value_1, key_2, ...)
+* :pr:`276`: implements ``onnx_generate`` which implements method generate for an onnx model,
 * :pr:`275`: fixes function ``patched_vmap``
 
 0.7.16

_unittests/ut_helpers/test_model_builder_helper.py

@@ -11,6 +11,7 @@
     import_model_builder,
     create_model_builder,
     save_model_builder,
+    find_names_pattern,
 )
 from onnx_diagnostic.torch_models.hghub import get_untrained_model_with_inputs
 from onnx_diagnostic.helpers.rt_helper import make_feeds
@@ -63,6 +64,11 @@ def test_model_builder_id(self):
                 raise unittest.SkipTest("batch_size must be 1 when sequence_length > 1")
         self.assertEqualAny(expected, got)
 
+    def test_find_names_pattern(self):
+        pats = ["past_key_values_key_0", "past_key_values_key_1"]
+        self.assertEqual("past_key_values_key_%d", find_names_pattern(pats))
+        self.assertEqual("past_key_values_key_%d", find_names_pattern(pats[:1]))
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

_unittests/ut_helpers/test_rt_helper.py

@@ -138,8 +138,6 @@ def test_onnx_generate(self):
             model, input_ids[:1], 2, max_new_tokens=10, session=session
         )
         self.assertEqualArray(input_ids[:1], expected[:, :n_inputs])
-        print("******", res)
-        print("******", expected)
         self.assertEqual(expected.dtype, torch.int64)
         self.assertEqual(expected.shape, (1, 13))
         self.assertEqualArray(expected, res)

onnx_diagnostic/helpers/model_builder_helper.py

@@ -1,10 +1,11 @@
 import copy
 import importlib.util
 import os
+import re
 import requests
 import sys
 from pathlib import Path
-from typing import Any, Dict, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 from urllib.parse import urlparse
 from onnx import ModelProto, TensorProto, load as load_model
 
@@ -340,6 +341,26 @@ def _post(onnx_model):
     return onnx_model
 
 
+def find_names_pattern(names: List[str]) -> str:
+    """
+    Finds a repeatable patterns in a list of names.
+    It tries to locate the figures.
+
+    .. runpython::
+        :showcode:
+
+        from onnx_diagnostic.helpers.model_builder_helper import find_names_pattern
+        pattern = find_names_pattern(["past_key_values_key_0", "past_key_values_key_1"])
+        print(pattern)
+    """
+    patterns = [re.sub(r"(\d+)", r"%d", t) for t in names]
+    unique = set(patterns)
+    assert (
+        len(unique) == 1
+    ), f"Unable to guess a pattern from {names} which led to the unique patterns {unique}"
+    return patterns[0]
+
+
 def make_genai_config(
     config,
     onnx_filename: str,
@@ -398,8 +419,17 @@ def make_genai_config(
                 "filename": onnx_filename,
                 "head_size": shape[-1],
                 "hidden_size": config.hidden_size,
-                "inputs": input_names,
-                "outputs": output_names,
+                "inputs": {
+                    "input_ids": input_names[0],
+                    "attention_mask": input_names[1],
+                    "past_key_names": find_names_pattern(input_names[2::2]),
+                    "past_value_names": find_names_pattern(input_names[3::2]),
+                },
+                "outputs": {
+                    "logits": output_names[0],
+                    "present_key_names": find_names_pattern(output_names[1::2]),
+                    "present_value_names": find_names_pattern(output_names[2::2]),
+                },
                 "num_attention_heads": config.num_attention_heads,
                 "num_hidden_layers": len(past_key_values) // 2,
                 "num_key_value_heads": shape[1],