add a test

Author: titaiwangms

_unittests/ut_tasks/test_tasks_text_generation.py

@@ -0,0 +1,39 @@
+import unittest
+from onnx_diagnostic.ext_test_case import (
+    ExtTestCase,
+    hide_stdout,
+    requires_transformers,
+    requires_torch,
+)
+from onnx_diagnostic.torch_models.validate import validate_model
+
+
+class TestTasksMaskGeneration(ExtTestCase):
+    @hide_stdout()
+    @requires_transformers("4.53")
+    @requires_torch("2.7.99")
+    def test_text_generation(self):
+        mid = "microsoft/phi-2"
+        summary, data = validate_model(
+            mid,
+            do_run=True,
+            verbose=10,
+            exporter="onnx-dynamo",
+            dump_folder="dump_test/microsoft_phi-2",
+            inputs2=True,
+            patch=True,
+        )
+        self.assertIsInstance(summary, dict)
+        # token generation
+        self.assertLess(summary["disc_onnx_ort_run_abs"], 3e-2)
+        # prompt processing
+        self.assertLess(summary["disc_onnx_ort_run2_abs"], 3e-2)
+        # multi-turn conversation
+        self.assertLess(summary["disc_onnx_ort_run3_abs"], 3e-2)
+        self.assertIsInstance(data, dict)
+        onnx_filename = data["onnx_filename"]
+        self.assertExists(onnx_filename)
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

onnx_diagnostic/tasks/text_generation.py

@@ -278,6 +278,24 @@ def get_inputs(
             add_second_input=0,
             **kwargs,
         )["inputs"]
+        # multi-turn conversation
+        # prompt-processing -> token-generation(loop output) ->
+        # prompt-processing from the loop output
+        res["inputs3"] = get_inputs(
+            model=model,
+            config=config,
+            dummy_max_token_id=dummy_max_token_id,
+            num_hidden_layers=num_hidden_layers,
+            batch_size=1,
+            past_sequence_length=32,
+            sequence_length=8,
+            dynamic_rope=dynamic_rope,
+            num_key_value_heads=num_key_value_heads,
+            head_dim=head_dim,
+            cls_cache=cls_cache,
+            add_second_input=0,
+            **kwargs,
+        )["inputs"]
     return res
 
 

onnx_diagnostic/torch_models/validate.py

@@ -573,6 +573,7 @@ def validate_model(
         if verbose:
             print(f"[validate_model] new inputs: {string_type(data['inputs'])}")
             print(f"[validate_model] new dynamic_hapes: {string_type(data['dynamic_shapes'])}")
+        # NOTE: The dynamic_shapes is always the same across inputs sets
         if inputs2:
             assert (
                 "inputs2" in data
@@ -583,6 +584,14 @@ def validate_model(
                 model=data["model"],
                 dynamic_shapes=data["dynamic_shapes"],
             )
+            # NOTE: text-generation tests 3rd inputs for multi-turn conversation
+            if "inputs3" in data:
+                data["inputs3"], _ = filter_inputs(
+                    data["inputs3"],
+                    drop_names=drop_inputs,
+                    model=data["model"],
+                    dynamic_shapes=data["dynamic_shapes"],
+                )
 
     if not empty(dtype):
         if isinstance(dtype, str):
@@ -594,6 +603,8 @@ def validate_model(
         summary["model_dtype"] = str(dtype)
         if "inputs2" in data:
             data["inputs2"] = to_any(data["inputs2"], dtype)  # type: ignore
+        if "inputs3" in data:
+            data["inputs3"] = to_any(data["inputs3"], dtype)  # type: ignore
 
     if not empty(device):
         if verbose:
@@ -603,6 +614,8 @@ def validate_model(
         summary["model_device"] = str(device)
         if "inputs2" in data:
             data["inputs2"] = to_any(data["inputs2"], device)  # type: ignore
+        if "inputs3" in data:
+            data["inputs3"] = to_any(data["inputs3"], device)  # type: ignore
 
     for k in ["task", "size", "n_weights"]:
         summary[f"model_{k.replace('_','')}"] = data[k]
@@ -638,10 +651,12 @@ def validate_model(
         _validate_do_run_model(
             data, summary, "inputs", "run", "run_expected", verbose, repeat, warmup, quiet
         )
-        if inputs2:
-            _validate_do_run_model(
-                data, summary, "inputs2", "run2", "run_expected2", verbose, 1, 0, quiet
-            )
+        _validate_do_run_model(
+            data, summary, "inputs2", "run2", "run_expected2", verbose, 1, 0, quiet
+        )
+        _validate_do_run_model(
+            data, summary, "inputs3", "run3", "run_expected3", verbose, 1, 0, quiet
+        )
 
     if exporter:
         print(
@@ -899,6 +914,10 @@ def _validate_do_run_model(
     if verbose:
         print(f"[validate_model] -- run the model inputs={key!r}...")
         print(f"[validate_model] {key}={string_type(data[key], with_shape=True)}")
+    if key not in data:
+        if verbose:
+            print(f"[validate_model] input; {key!r} not defined, skip.")
+        return
     # We make a copy of the input just in case the model modifies them inplace
     hash_inputs = string_type(data[key], with_shape=True)
     inputs = torch_deepcopy(data[key])
@@ -1329,6 +1348,9 @@ def _mk(key, flavour=flavour):
     keys = [("inputs", "run_expected", "")]
     if inputs2:
         keys.append(("inputs2", "run_expected2", "2"))
+        # text-generation tests multi-turn conversation as 3rd inputs
+        if "inputs3" in data:
+            keys.append(("inputs3", "run_expected3", "3"))
     for k_input, k_expected, suffix in keys:
         # make_feeds
         if verbose: