a few fixes

Author: xadupre

_doc/examples/plot_export_with_modelbuilder.py

@@ -6,9 +6,11 @@
 
 """
 
+import sys
 import os
 import pandas
-from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
 from onnx_diagnostic import doc
 from onnx_diagnostic.investigate.input_observer import InputObserver
 from onnx_diagnostic.helpers.rt_helper import onnx_generate
@@ -28,10 +30,11 @@ def generate_text(
     top_k=50,
     top_p=0.95,
     do_sample=True,
+    device="cpu",
 ):
     inputs = tokenizer(prompt, return_tensors="pt")
-    input_ids = inputs["input_ids"]
-    attention_mask = inputs["attention_mask"]
+    input_ids = inputs["input_ids"].to(device)
+    attention_mask = inputs["attention_mask"].to(device)
 
     outputs = model.generate(
         input_ids=input_ids,
@@ -47,58 +50,71 @@ def generate_text(
     return generated_text
 
 
+# %%
+# filename for the model
+MODEL_NAME = sys.argv[1] if sys.argv and len(sys.argv) > 1 else "arnir0/Tiny-LLM"
+cache_dir = "dump_modelbuilder"
+os.makedirs(cache_dir, exist_ok=True)
+name = MODEL_NAME.replace("/", "_")
+filename = os.path.join(cache_dir, f"plot_export_with_modelbuilder_{name}.onnx")
+
+
 # %%
 # Creating the model
-print("-- creating...")
-MODEL_NAME = "arnir0/Tiny-LLM"
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+if not os.path.exists(filename):
+    print(f"-- creating... on {device} into {filename!r}")
+    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.bfloat16)
+    model = model.to(device)
+    config = model.config
+else:
+    config = AutoConfig.from_pretrained(MODEL_NAME)
 
 
 # %%
 # Capturing inputs/outputs to infer dynamic shapes and arguments
 print("-- capturing...")
 prompt = "Continue: it rains, what should I do?"
-observer = InputObserver()
-with register_additional_serialization_functions(patch_transformers=True), observer(model):
-    generate_text(prompt, model, tokenizer)
+if not os.path.exists(filename):
+    observer = InputObserver()
+    with register_additional_serialization_functions(patch_transformers=True), observer(model):
+        generate_text(prompt, model, tokenizer, device=device)
 
 
 # %%
 # Exporting.
-print("-- exporting...")
-observer.remove_inputs(["cache_position", "logits_to_keep", "position_ids"])
-ds = observer.infer_dynamic_shapes(set_batch_dimension_for=True)
-kwargs = observer.infer_arguments()
-
-cache_dir = "dump_modelbuilder"
-os.makedirs(cache_dir, exist_ok=True)
-filename = os.path.join(cache_dir, "plot_export_with_modelbuilder.onnx")
-with torch_export_patches(patch_transformers=True):
-    to_onnx(
-        model,
-        filename=filename,
-        kwargs=kwargs,
-        dynamic_shapes=ds,
-        exporter="modelbuilder",
-    )
-
-data = observer.check_discrepancies(filename, progress_bar=True)
-print(pandas.DataFrame(data))
+if not os.path.exists(filename):
+    print("-- exporting...")
+    observer.remove_inputs(["cache_position", "logits_to_keep", "position_ids"])
+    ds = observer.infer_dynamic_shapes(set_batch_dimension_for=True)
+    kwargs = observer.infer_arguments()
+
+    with torch_export_patches(patch_transformers=True):
+        to_onnx(
+            model,
+            filename=filename,
+            kwargs=kwargs,
+            dynamic_shapes=ds,
+            exporter="modelbuilder",
+        )
+
+    data = observer.check_discrepancies(filename, progress_bar=True)
+    print(pandas.DataFrame(data))
 
 # %%
 # ONNX Prompt
 # +++++++++++
 print("-- ONNX prompts...")
 inputs = tokenizer(prompt, return_tensors="pt")
-input_ids = inputs["input_ids"]
-attention_mask = inputs["attention_mask"]
+input_ids = inputs["input_ids"].to(device)
+attention_mask = inputs["attention_mask"].to(device)
 
 onnx_tokens = onnx_generate(
     filename,
     input_ids=input_ids,
     attention_mask=attention_mask,
-    eos_token_id=model.config.eos_token_id,
+    eos_token_id=config.eos_token_id,
     max_new_tokens=50,
 )
 onnx_generated_text = tokenizer.decode(onnx_tokens, skip_special_tokens=True)
@@ -108,4 +124,5 @@ def generate_text(
 print("-----------------")
 
 # %%
-doc.save_fig(doc.plot_dot(filename), f"{filename}.png", dpi=400)
+if os.stat(filename).st_size < 2**14:
+    doc.save_fig(doc.plot_dot(filename), f"{filename}.png", dpi=400)

_unittests/ut_investigate/test_input_observer.py

@@ -1196,7 +1196,6 @@ def forward(self, a, *args, **kwargs):
             )
             torch.export.export(model, args, kwargs=kwargs, dynamic_shapes=ds)
 
-
     def test_remove_inputs_kwargs(self):
         """Test that remove_inputs removes a kwarg from the observer info."""
 

onnx_diagnostic/export/api.py

@@ -320,12 +320,7 @@ def to_onnx(
         temp_filename = os.path.join(os.path.dirname(filename), "model.onnx")
         # renaming
         onx = onnx.load(temp_filename, load_external_data=True)
-        onnx.save(
-            onx,
-            filename,
-            save_as_external_data=True,
-            location=f"{os.path.splitext(filename[0])}.data",
-        )
+        onnx.save(onx, filename, save_as_external_data=True)
         return onx
 
     raise ValueError(f"Unknown exporter={exporter!r}")