From 2c457fca9fc98e7c15c25c075d8c73105ab9a69c Mon Sep 17 00:00:00 2001
From: JingyaHuang <jingya@huggingface.co>
Date: Sun, 11 Feb 2024 23:15:26 +0000
Subject: [PATCH] improve tests

---
 optimum/commands/export/neuron.py   | 4 ++--
 optimum/exporters/neuron/convert.py | 2 +-
 tests/cli/test_export_cli.py        | 7 +++++--
 tests/exporters/test_export.py      | 1 -
 4 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/optimum/commands/export/neuron.py b/optimum/commands/export/neuron.py
index 00c1c2f78..43305aeeb 100644
--- a/optimum/commands/export/neuron.py
+++ b/optimum/commands/export/neuron.py
@@ -51,7 +51,7 @@ def parse_args_neuron(parser: "ArgumentParser"):
         type=str,
         choices=["transformers", "sentence_transformers"],
         default=None,
-        help=("The library on the model." " If not provided, will attempt to infer the local checkpoint's library."),
+        help=("The library on the model. If not provided, will attempt to infer the local checkpoint's library."),
     )
     optional_group.add_argument(
         "--subfolder",
@@ -76,7 +76,7 @@ def parse_args_neuron(parser: "ArgumentParser"):
     optional_group.add_argument(
         "--compiler_workdir",
         type=Path,
-        help="Path indicating the directory where to store intermediary files generated by Neuronx compiler.",
+        help="Path indicating the directory where to store intermediary files generated by Neuron compiler.",
     )
     optional_group.add_argument(
         "--disable-weights-neff-inline",
diff --git a/optimum/exporters/neuron/convert.py b/optimum/exporters/neuron/convert.py
index c6dc85671..fb0bcadb9 100644
--- a/optimum/exporters/neuron/convert.py
+++ b/optimum/exporters/neuron/convert.py
@@ -598,7 +598,7 @@ def export_neuron(
         output (`Path`):
             Directory to store the exported Neuron model.
         compiler_workdir (`Optional[Path]`, defaults to `None`):
-            The directory used by neuronx-cc, where you can find intermediary outputs (neff, weight, hlo...).
+            The directory used by neuron-cc, where you can find intermediary outputs (neff, weight, hlo...).
         inline_weights_to_neff (`bool`, defaults to `True`):
             Whether to inline the weights to the neff graph. If set to False, weights will be seperated from the neff.
         auto_cast (`Optional[str]`, defaults to `None`):
diff --git a/tests/cli/test_export_cli.py b/tests/cli/test_export_cli.py
index f72f32bbc..2bc38eaef 100644
--- a/tests/cli/test_export_cli.py
+++ b/tests/cli/test_export_cli.py
@@ -18,6 +18,7 @@
 import unittest
 
 from optimum.exporters.neuron.model_configs import *  # noqa: F403
+from optimum.neuron.utils import is_neuronx_available
 from optimum.neuron.utils.testing_utils import is_inferentia_test, requires_neuronx
 from optimum.utils import logging
 
@@ -109,12 +110,14 @@ def test_opt_level(self):
                     check=True,
                 )
 
-    @requires_neuronx
     def test_store_intemediary(self):
         model_id = "hf-internal-testing/tiny-random-BertModel"
         with tempfile.TemporaryDirectory() as tempdir:
             save_path = f"{tempdir}/neff"
-            neff_path = os.path.join(save_path, model_id.split("/")[-1], "graph.neff")
+            if is_neuronx_available():
+                neff_path = os.path.join(save_path, model_id.split("/")[-1], "graph.neff")
+            else:
+                neff_path = os.path.join(save_path, model_id.split("/")[-1], "32", "neff.json")
             subprocess.run(
                 [
                     "optimum-cli",
diff --git a/tests/exporters/test_export.py b/tests/exporters/test_export.py
index 80316db9a..a1b8c1ccd 100644
--- a/tests/exporters/test_export.py
+++ b/tests/exporters/test_export.py
@@ -170,7 +170,6 @@ def test_export(self, test_name, name, model_name, task, neuron_config_construct
         _get_models_to_test(EXPORT_MODELS_TINY, exclude_model_types=WEIGHTS_NEFF_SEPARATION_UNSUPPORTED_ARCH)
     )
     @is_inferentia_test
-    @requires_neuronx
     def test_export_separated_weights(self, test_name, name, model_name, task, neuron_config_constructor):
         self._neuronx_export(
             test_name, name, model_name, task, neuron_config_constructor, inline_weights_to_neff=False