onnx_generate_with_genai

Author: xadupre

_doc/technical/plot_generate.py

@@ -94,6 +94,9 @@
 # %%
 # Custom method generate
 # ======================
+#
+# Let's implement a simple function replicating when method
+# ``generate`` does.
 
 
 def simple_generate_with_cache(

_unittests/ut_helpers/test_rt_helper.py

@@ -3,12 +3,17 @@
 import torch
 from onnx_diagnostic.ext_test_case import (
     ExtTestCase,
+    has_onnxruntime_genai,
     hide_stdout,
     requires_transformers,
     requires_torch,
 )
 from onnx_diagnostic.helpers import max_diff, flatten_object
-from onnx_diagnostic.helpers.rt_helper import onnx_generate, make_empty_cache
+from onnx_diagnostic.helpers.rt_helper import (
+    onnx_generate,
+    onnx_generate_with_genai,
+    make_empty_cache,
+)
 from onnx_diagnostic.helpers.torch_helper import torch_deepcopy
 from onnx_diagnostic.helpers.ort_session import InferenceSessionForTorch
 from onnx_diagnostic.torch_models.hghub import get_untrained_model_with_inputs
@@ -101,6 +106,7 @@ def test_onnx_generate(self):
         print("-- test_onnx_generate: get model")
         data = get_untrained_model_with_inputs(mid)
         model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
+        configuration = data["configuration"]
         del inputs["position_ids"]
         del ds["position_ids"]
         input_ids = inputs["input_ids"]
@@ -118,25 +124,38 @@ def test_onnx_generate(self):
                 exporter="custom",
             )
 
-            print("-- test_onnx_generate: generate")
-            res, session = onnx_generate(
-                model_name, input_ids[:1], 2, max_new_tokens=10, return_session=True
-            )
-            n_inputs = input_ids.shape[1]
-            self.assertEqualArray(input_ids[:1], res[:, :n_inputs])
-            self.assertEqual(res.dtype, torch.int64)
-            self.assertEqual(res.shape, (1, 13))
-            print("-- test_onnx_generate: done")
-            # expected = model.generate(input_ids[:1], max_new_tokens=10)
-            expected = self.simple_generate_with_cache(
-                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)
+        print("-- test_onnx_generate: generate")
+        res, session = onnx_generate(
+            model_name, input_ids[:1], 2, max_new_tokens=10, return_session=True
+        )
+        n_inputs = input_ids.shape[1]
+        self.assertEqualArray(input_ids[:1], res[:, :n_inputs])
+        self.assertEqual(res.dtype, torch.int64)
+        self.assertEqual(res.shape, (1, 13))
+        print("-- test_onnx_generate: done")
+        # expected = model.generate(input_ids[:1], max_new_tokens=10)
+        expected = self.simple_generate_with_cache(
+            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)
+
+        if not has_onnxruntime_genai():
+            raise unittest.SkipTest("onnxruntime_genai is missing")
+
+        res, session = onnx_generate_with_genai(
+            model_name,
+            input_ids[:1],
+            max_new_tokens=10,
+            return_session=True,
+            transformers_config=configuration,
+        )
+        self.assertNotEmpty(session)
+        self.assertEqualArray(expected, res)
 
 
 if __name__ == "__main__":

onnx_diagnostic/ext_test_case.py

@@ -630,6 +630,17 @@ def has_onnxruntime_training(push_back_batch: bool = False):
     return True
 
 
+def has_onnxruntime_genai():
+    """Tells if onnxruntime_genai is installed."""
+    try:
+        import onnxruntime_genai  # noqa: F401
+
+        return True
+    except ImportError:
+        # onnxruntime not training
+        return False
+
+
 def requires_onnxruntime_training(
     push_back_batch: bool = False, ortmodule: bool = False, msg: str = ""
 ) -> Callable:

onnx_diagnostic/helpers/model_builder_helper.py

@@ -1,11 +1,12 @@
+import copy
 import importlib.util
 import os
 import requests
 import sys
 from pathlib import Path
-from typing import Any, Optional, Union
+from typing import Any, Dict, Optional, Union
 from urllib.parse import urlparse
-from onnx import ModelProto, TensorProto
+from onnx import ModelProto, TensorProto, load as load_model
 
 CACHE_SUBDIR = "onnx-diagnostic"
 
@@ -337,3 +338,102 @@ def _post(onnx_model):
     # onnx_model.make_genai_config(hf_name, extra_kwargs, output_dir)
     # onnx_model.save_processing(hf_name, extra_kwargs, output_dir)
     return onnx_model
+
+
+def make_genai_config(
+    config,
+    onnx_filename: str,
+) -> Dict:
+    """
+    Creates genai config file for a model.
+
+    :param config: configuration from transformers
+    :param onnx_filename: onnx configuration
+    :return: configuration
+    """
+    onx = load_model(onnx_filename, load_external_data=False)
+    config = copy.deepcopy(config)
+    defaults = {
+        "bos_token_id": None,
+        "do_sample": False,
+        "eos_token_id": None,
+        "pad_token_id": None,
+        "temperature": 1.0,
+        "top_k": 50,
+        "top_p": 1.0,
+    }
+    for key, default_val in defaults.items():
+        if not hasattr(config, key):
+            setattr(config, key, default_val)
+
+    bos_token_id = (
+        config.bos_token_id
+        if hasattr(config, "bos_token_id") and config.bos_token_id is not None
+        else 1
+    )
+    eos_token_id = config.eos_token_id
+    pad_token_id = (
+        config.pad_token_id
+        if hasattr(config, "pad_token_id") and config.pad_token_id is not None
+        else (
+            config.eos_token_id[0]
+            if isinstance(config.eos_token_id, list)
+            else config.eos_token_id
+        )
+    )
+    input_names = [i.name for i in onx.graph.input]
+    output_names = [i.name for i in onx.graph.output]
+    past_key_values = [s for s in input_names if s.startswith("past_key_value")]
+    first = [i for i in onx.graph.input if i.name == past_key_values[0]][0]  # noqa: RUF015
+    shape = tuple(d.dim_value or d.dim_param for d in first.type.tensor_type.shape.dim)
+    return {
+        "model": {
+            "bos_token_id": bos_token_id,
+            "context_length": config.max_position_embeddings,
+            "decoder": {
+                "session_options": {
+                    "log_id": "onnxruntime-genai",
+                    "provider_options": [],
+                },
+                "filename": onnx_filename,
+                "head_size": shape[-1],
+                "hidden_size": config.hidden_size,
+                "inputs": input_names,
+                "outputs": output_names,
+                "num_attention_heads": config.num_attention_heads,
+                "num_hidden_layers": len(past_key_values) // 2,
+                "num_key_value_heads": shape[1],
+            },
+            "eos_token_id": eos_token_id,
+            "pad_token_id": pad_token_id,
+            # "type": self.model_type[ : self.model_type.find("For")
+            # if "For" in self.model_type else len(self.model_type)].lower(),
+            "vocab_size": config.vocab_size,
+        },
+        "search": {
+            "diversity_penalty": (
+                config.diversity_penalty if hasattr(config, "diversity_penalty") else 0.0
+            ),
+            "do_sample": config.do_sample if hasattr(config, "do_sample") else False,
+            "early_stopping": True,
+            "length_penalty": (
+                config.length_penalty if hasattr(config, "length_penalty") else 1.0
+            ),
+            "max_length": config.max_position_embeddings,
+            "min_length": 0,
+            "no_repeat_ngram_size": (
+                config.no_repeat_ngram_size if hasattr(config, "no_repeat_ngram_size") else 0
+            ),
+            "num_beams": config.num_beams if hasattr(config, "num_beams") else 1,
+            "num_return_sequences": (
+                config.num_return_sequences if hasattr(config, "num_return_sequences") else 1
+            ),
+            "past_present_share_buffer": False,
+            "repetition_penalty": (
+                config.repetition_penalty if hasattr(config, "repetition_penalty") else 1.0
+            ),
+            "temperature": config.temperature if hasattr(config, "temperature") else 1.0,
+            "top_k": config.top_k if hasattr(config, "top_k") else 50,
+            "top_p": config.top_p if hasattr(config, "top_p") else 1.0,
+        },
+    }

onnx_diagnostic/helpers/rt_helper.py

@@ -1,4 +1,6 @@
-from typing import Any, Dict, List, Tuple, Union
+import json
+import os
+from typing import Any, Dict, List, Optional, Tuple, Union
 import numpy as np
 import onnx
 import torch
@@ -224,3 +226,73 @@ def onnx_generate(
     if return_session:
         return input_ids, session
     return input_ids
+
+
+def onnx_generate_with_genai(
+    model_or_path: Union[onnx.ModelProto, str, InferenceSessionForTorch],
+    input_ids: torch.Tensor,
+    max_new_tokens=100,
+    return_session: bool = False,
+    transformers_config: Optional[Any] = None,
+) -> Union[torch.Tensor, Tuple[torch.Tensor, InferenceSessionForTorch]]:
+    """
+    Uses :epkg:`onnxruntime-genai` to implement a simple method ``generate``
+    for an ONNX model. The function does not expect any ``position_ids`` as input.
+
+    :param model_or_path: model or loaded model
+    :param input_ids: input tokens
+    :param eos_token_ids: token representing the end of an answer
+    :param max_new_tokens: stops after this number of generated tokens
+    :param return_session: returns the instance of class
+        :class:`InferenceSessionForTorch
+        <onnx_diagnostic.helpers.ort_session.InferenceSessionForTorch>`
+        created if necessary
+    :param transformers_config: write configuration
+        if missing and if this configuration is provided
+    :return: input tokens concatenated with new tokens
+    """
+    import onnxruntime_genai as og
+
+    if not isinstance(model_or_path, og.Model):
+        from .model_builder_helper import make_genai_config
+
+        assert isinstance(
+            model_or_path, str
+        ), f"Only a filename is allowed for model_or_path but type is {type(model_or_path)}"
+        folder = os.path.dirname(model_or_path)
+        assert os.path.exists(folder), f"Folder {folder!r} does not exists."
+        assert os.path.exists(model_or_path), f"Folder {model_or_path!r} does not exists."
+        config_file = os.path.join(folder, "genai_config.json")
+        if not os.path.exists(config_file):
+            if not transformers_config:
+                raise FileNotFoundError(
+                    f"Folder {model_or_path!r} does not contain 'genai_config.json'."
+                )
+            config = make_genai_config(transformers_config, model_or_path)
+            with open(config_file, "w") as f:
+                json.dump(config, f, indent=4)
+
+        config = og.Config(os.path.dirname(config_file))
+        if input_ids.is_cuda:
+            config.clear_providers()
+            config.append_provider("cuda")
+        session = og.Model(config)
+    else:
+        session = model_or_path
+
+    params = og.GeneratorParams(session)
+    params.set_search_options(max_new_tokens=max_new_tokens, batch_size=input_ids.shape[0])
+    generator = og.Generator(session, params)
+
+    # First call: prefill
+    cats = [input_ids]
+    generator.append_tokens(input_ids)
+    while not generator.is_done():
+        generator.generate_next_token()
+        new_token = generator.get_next_tokens()[0]
+        cats.append(new_token)
+
+    input_ids = torch.cat(cats, dim=-1)
+    if return_session:
+        return input_ids, session
+    return input_ids

requirements-dev.txt

@@ -7,6 +7,7 @@ huggingface_hub
 matplotlib
 onnx-array-api>=0.3.1
 onnx
+onnxruntime-genai
 onnxscript
 openpyxl
 packaging