add more comments

Author: xadupre

_doc/conf.py

@@ -146,11 +146,12 @@ def linkcode_resolve(domain, info):
     ("py:class", "transformers.cache_utils.EncoderDecoderCache"),
     ("py:class", "transformers.cache_utils.HybridCache"),
     ("py:class", "transformers.cache_utils.MambaCache"),
-    ("py:class", "transformers.models.mamba.modeling_mamba.MambaCache"),
     ("py:class", "transformers.cache_utils.SlidingWindowCache"),
     ("py:class", "transformers.cache_utils.StaticCache"),
     ("py:class", "transformers.configuration_utils.PretrainedConfig"),
+    ("py:class", "transformers.configuration_utils.PreTrainedConfig"),
     ("py:class", "transformers.modeling_outputs.BaseModelOutput"),
+    ("py:class", "transformers.models.mamba.modeling_mamba.MambaCache"),
     ("py:class", "transformers.models.phi3.modeling_phi3.Phi3RotaryEmbedding"),
     ("py:func", "torch.export._draft_export.draft_export"),
     ("py:func", "torch._export.tools.report_exportability"),

_doc/examples/plot_export_tiny_llm_dim01.py

@@ -34,6 +34,15 @@
 data = get_untrained_model_with_inputs("arnir0/Tiny-LLM", add_second_input=True)
 model, dynamic_shapes = data["model"], data["dynamic_shapes"]
 
+# %%
+# The trained model can be obtained with:
+#
+# .. code-block:: python
+#
+#   MODEL_NAME = "arnir0/Tiny-LLM"
+#   tokenizer = transformers.AutoTokenizer.from_pretrained(MODEL_NAME)
+#   model = transformers.AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+
 input_sets = {k: v for k, v in data.items() if k.startswith("inputs")}
 
 for k, v in input_sets.items():
@@ -207,6 +216,6 @@ def validation(ep, input_sets, expected):
 
 # %%
 # If you have any error, then look at example
-# :ref:`l-plot-tiny-llm-export-cache_patched`.
+# :ref:`l-plot-tiny-llm-export-patched`.
 
 doc.plot_legend("Tiny-LLM\nexport with\ndimension in {0,1}", "torch.export.export", "tomato")

_doc/examples/plot_export_tiny_llm_dim01_onnx.py

@@ -34,6 +34,15 @@
 data = get_untrained_model_with_inputs("arnir0/Tiny-LLM", add_second_input=True)
 model, dynamic_shapes = data["model"], data["dynamic_shapes"]
 
+# %%
+# The trained model can be obtained with:
+#
+# .. code-block:: python
+#
+#   MODEL_NAME = "arnir0/Tiny-LLM"
+#   tokenizer = transformers.AutoTokenizer.from_pretrained(MODEL_NAME)
+#   model = transformers.AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+
 input_sets = {k: v for k, v in data.items() if k.startswith("inputs")}
 
 for k, v in input_sets.items():
@@ -221,6 +230,6 @@ def validation(ep, input_sets, expected, catch_exception=True):
 
 # %%
 # If you have any error, then look at example
-# :ref:`l-plot-tiny-llm-export-cache_patched`.
+# :ref:`l-plot-tiny-llm-export-patched`.
 
 doc.plot_legend("Tiny-LLM\nexport with\ndimension in {0,1}", "torch.onnx.export", "tomato")

_doc/examples/plot_export_tiny_llm_dim01_onnx_custom.py

@@ -33,6 +33,15 @@
 data = get_untrained_model_with_inputs("arnir0/Tiny-LLM", add_second_input=True)
 model, dynamic_shapes = data["model"], data["dynamic_shapes"]
 
+# %%
+# The trained model can be obtained with:
+#
+# .. code-block:: python
+#
+#   MODEL_NAME = "arnir0/Tiny-LLM"
+#   tokenizer = transformers.AutoTokenizer.from_pretrained(MODEL_NAME)
+#   model = transformers.AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+
 input_sets = {k: v for k, v in data.items() if k.startswith("inputs")}
 
 for k, v in input_sets.items():
@@ -219,6 +228,6 @@ def validation(onx, input_sets, expected, catch_exception=True):
 
 # %%
 # If you have any error, then look at example
-# :ref:`l-plot-tiny-llm-export-cache_patched`.
+# :ref:`l-plot-tiny-llm-export-patched`.
 
 doc.plot_legend("Tiny-LLM\nexport with\ndimension in {0,1}", "to_onnx", "tomato")