Skip to content

Commit 9dec8f9

Browse files
authored
test: optimum - add unit tests for configs (#3220)
1 parent 9d7f6de commit 9dec8f9

1 file changed

Lines changed: 51 additions & 0 deletions

File tree

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
# SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
2+
#
3+
# SPDX-License-Identifier: Apache-2.0
4+
5+
import pytest
6+
from optimum.onnxruntime.configuration import OptimizationConfig, QuantizationConfig
7+
8+
from haystack_integrations.components.embedders.optimum.optimization import (
9+
OptimumEmbedderOptimizationConfig,
10+
OptimumEmbedderOptimizationMode,
11+
)
12+
from haystack_integrations.components.embedders.optimum.quantization import (
13+
OptimumEmbedderQuantizationConfig,
14+
OptimumEmbedderQuantizationMode,
15+
)
16+
17+
18+
class TestOptimumEmbedderQuantizationConfig:
19+
@pytest.mark.parametrize(
20+
"mode",
21+
[
22+
OptimumEmbedderQuantizationMode.ARM64,
23+
OptimumEmbedderQuantizationMode.AVX2,
24+
OptimumEmbedderQuantizationMode.AVX512,
25+
OptimumEmbedderQuantizationMode.AVX512_VNNI,
26+
],
27+
)
28+
def test_to_optimum_config_returns_quantization_config_for_each_mode(self, mode):
29+
config = OptimumEmbedderQuantizationConfig(mode=mode, per_channel=True)
30+
optimum_config = config.to_optimum_config()
31+
32+
assert isinstance(optimum_config, QuantizationConfig)
33+
assert optimum_config.is_static is False
34+
assert optimum_config.per_channel is True
35+
36+
37+
class TestOptimumEmbedderOptimizationConfig:
38+
@pytest.mark.parametrize(
39+
"mode",
40+
[
41+
OptimumEmbedderOptimizationMode.O1,
42+
OptimumEmbedderOptimizationMode.O2,
43+
OptimumEmbedderOptimizationMode.O3,
44+
OptimumEmbedderOptimizationMode.O4,
45+
],
46+
)
47+
def test_to_optimum_config_returns_optimization_config_for_each_mode(self, mode):
48+
config = OptimumEmbedderOptimizationConfig(mode=mode, for_gpu=False)
49+
optimum_config = config.to_optimum_config()
50+
51+
assert isinstance(optimum_config, OptimizationConfig)

0 commit comments

Comments
 (0)