optimum v2

.github/workflows/build_pr_documentation.yml

@@ -31,6 +31,8 @@ jobs:
           sudo mv ./doc-build/optimum.amd ./doc-build/optimum-amd
           echo ${{ env.COMMIT_SHA }} > ./commit_sha
           echo ${{ env.PR_NUMBER }} > ./pr_number
+        env:
+          COMMIT_SHA_SUBPACKAGE: ${{ github.event.pull_request.head.sha }}
 
       - uses: actions/upload-artifact@v4
         with:

optimum/amd/brevitas/quantizer.py

@@ -13,7 +13,7 @@
 from brevitas_examples.llm.llm_quant.equalize import apply_weight_equalization
 from tqdm import tqdm
 
-from optimum.exporters import TasksManager
+from optimum.exporters.tasks import TasksManager
 from optimum.quantization_base import OptimumQuantizer
 from transformers.utils.fx import symbolic_trace
 

optimum/amd/ryzenai/modeling.py

@@ -17,8 +17,8 @@
 from onnx import shape_inference
 from onnx.tools import update_model_dims
 
-from optimum.exporters import TasksManager
 from optimum.exporters.onnx import main_export
+from optimum.exporters.tasks import TasksManager
 from optimum.modeling_base import FROM_PRETRAINED_START_DOCSTRING, OptimizedModel
 from optimum.onnx.utils import _get_external_data_paths
 from optimum.utils.save_utils import maybe_load_preprocessors

optimum/amd/ryzenai/pipelines/__init__.py

@@ -4,7 +4,7 @@
 
 from typing import TYPE_CHECKING, Any, Optional, Union
 
-from optimum.exporters import TasksManager
+from optimum.exporters.tasks import TasksManager
 from transformers import ImageClassificationPipeline, Pipeline, PretrainedConfig
 from transformers import pipeline as transformers_pipeline
 from transformers.image_processing_utils import BaseImageProcessor

setup.py

@@ -14,7 +14,7 @@
     assert False, "Error: Could not open '%s' due %s\n" % (filepath, error)
 
 # ORT 1.16 is not compatible: https://github.com/Xilinx/Vitis-AI/issues/1343
-INSTALL_REQUIRE = ["optimum", "transformers>=4.38", "onnx", "onnxruntime-extensions"]
+INSTALL_REQUIRE = ["optimum<2.0.0", "transformers>=4.38", "onnx", "onnxruntime-extensions"]
 
 # TODO: unpin pytest once https://github.com/huggingface/transformers/pull/29154 is merged & released
 TESTS_REQUIRE = [

tests/brevitas/test_onnx_export.py

@@ -15,12 +15,12 @@
 from testing_utils import SUPPORTED_MODELS_TINY, VALIDATE_EXPORT_ON_SHAPES, get_quantized_model
 
 from optimum.amd.brevitas.export import find_and_insert_matmulinteger
-from optimum.exporters import TasksManager
 from optimum.exporters.onnx import (
     export_models,
     validate_models_outputs,
 )
 from optimum.exporters.onnx.utils import _get_submodels_and_onnx_configs
+from optimum.exporters.tasks import TasksManager
 from optimum.utils import DEFAULT_DUMMY_SHAPES
 from optimum.utils.testing_utils import grid_parameters
 from transformers.modeling_utils import get_parameter_dtype