File size: 8,335 Bytes
9bbba62 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 | """
Test different backends (PyTorch, ONNX, OpenVINO) for the Transformer class with transformer_task="fill-mask" and SparseEncoder.
This module tests loading and using models with different inference backends.
"""
from __future__ import annotations
import gc
import json
import os
import tempfile
from contextlib import nullcontext
from pathlib import Path
import pytest
from packaging.version import Version, parse
from tests.sparse_encoder.utils import sparse_allclose
from tests.utils import is_ci
try:
from optimum.intel import OVModelForMaskedLM
from optimum.intel.version import __version__ as optimum_intel_version
from optimum.onnxruntime import ORTModelForMaskedLM
from optimum.version import __version__ as optimum_version
except ImportError:
pytest.skip("OpenVINO and ONNX backends are not available", allow_module_level=True)
from sentence_transformers.sparse_encoder import SparseEncoder
if is_ci():
pytest.skip("Skip test in CI to try and avoid 429 Client Error", allow_module_level=True)
## Testing exporting:
@pytest.mark.parametrize(
["backend", "expected_auto_model_class"],
[
("onnx", ORTModelForMaskedLM),
("openvino", OVModelForMaskedLM),
],
)
@pytest.mark.parametrize(
"model_kwargs", [{}, {"file_name": "wrong_file_name"}]
) # <- Using a file_name is fine when exporting
def test_backend_export(backend, expected_auto_model_class, model_kwargs) -> None:
model = SparseEncoder("sparse-encoder-testing/splade-bert-tiny-nq", backend=backend, model_kwargs=model_kwargs)
assert model.get_backend() == backend
assert isinstance(model[0].auto_model, expected_auto_model_class)
embedding = model.encode("Hello, World!")
assert embedding.shape == (model.get_embedding_dimension(),)
def test_backend_no_export_crash():
# Prior to optimum v1.25.0, ONNX Crashes when it can't export & the model repo/path doesn't contain an exported model
# Since then, it auto-updates export to True
with pytest.raises(OSError) if parse(optimum_version) < Version("1.25.0") else nullcontext():
model = SparseEncoder(
"sparse-encoder-testing/splade-bert-tiny-nq", backend="onnx", model_kwargs={"export": False}
)
assert isinstance(model[0].auto_model, ORTModelForMaskedLM)
# OpenVINO will forcibly override the export=False if the model repo/path doesn't contain an exported model
# But only starting from optimum-intel=v1.19.0
with pytest.raises(OSError) if parse(optimum_intel_version) < Version("1.19.0") else nullcontext():
model = SparseEncoder(
"sparse-encoder-testing/splade-bert-tiny-nq",
backend="openvino",
model_kwargs={"export": False},
)
assert isinstance(model[0].auto_model, OVModelForMaskedLM)
## Testing loading exported models:
@pytest.mark.parametrize(
["backend", "model_id"],
[
("onnx", "sparse-encoder-testing/splade-bert-tiny-nq-onnx"),
("openvino", "sparse-encoder-testing/splade-bert-tiny-nq-openvino"),
],
)
@pytest.mark.parametrize(
["model_kwargs", "exception"],
[
[{}, False],
[{"file_name": "wrong_file_name", "export": True}, False], # Using a file_name is fine when exporting
[{"file_name": "wrong_file_name", "export": False}, True], # ... but fails when not exporting
],
)
def test_backend_load(backend, model_id, model_kwargs, exception) -> None:
if exception:
with pytest.raises((OSError, RuntimeError)):
SparseEncoder(model_id, backend=backend, model_kwargs=model_kwargs)
else:
model = SparseEncoder(model_id, backend=backend, model_kwargs=model_kwargs)
assert model.get_backend() == backend
tokens = model.encode("Hello, World!")
assert len(tokens) > 0
def test_onnx_provider_crash() -> None:
with pytest.raises(ValueError):
SparseEncoder(
"sparse-encoder-testing/splade-bert-tiny-nq-onnx",
backend="onnx",
model_kwargs={"provider": "incorrect_provider"},
)
def test_openvino_provider() -> None:
model = SparseEncoder(
"sparse-encoder-testing/splade-bert-tiny-nq-openvino",
backend="openvino",
model_kwargs={"ov_config": {"INFERENCE_PRECISION_HINT": "precision_1"}},
)
assert model[0].auto_model.ov_config == {
"INFERENCE_PRECISION_HINT": "precision_1",
"PERFORMANCE_HINT": "LATENCY",
}
with tempfile.TemporaryDirectory() as temp_dir:
ov_config_path = os.path.join(temp_dir, "ov_config.json")
with open(ov_config_path, "w") as ov_config_file:
json.dump({"INFERENCE_PRECISION_HINT": "precision_2"}, ov_config_file)
model = SparseEncoder(
"sparse-encoder-testing/splade-bert-tiny-nq-openvino",
backend="openvino",
model_kwargs={"ov_config": ov_config_path},
)
assert model[0].auto_model.ov_config == {
"INFERENCE_PRECISION_HINT": "precision_2",
"PERFORMANCE_HINT": "LATENCY",
}
def test_incorrect_backend() -> None:
with pytest.raises(ValueError):
SparseEncoder("sparse-encoder-testing/splade-bert-tiny-nq", backend="incorrect_backend")
def test_openvino_backend() -> None:
model_id = "sparse-encoder-testing/splade-bert-tiny-nq"
# Test that OpenVINO output is close to PyTorch output
pytorch_model = SparseEncoder(model_id)
openvino_model = SparseEncoder(
model_id,
backend="openvino",
model_kwargs={"ov_config": {"INFERENCE_PRECISION_HINT": "f32"}},
)
pytorch_result = pytorch_model.encode(["Hello there!"], save_to_cpu=True)
openvino_result = openvino_model.encode(["Hello there!"])
assert sparse_allclose(openvino_result, pytorch_result, atol=0.00001), "OpenVINO and Pytorch outputs are not close"
with tempfile.TemporaryDirectory() as tmpdirname:
# Test that loading with ov_config file works as expected
config_file = str(Path(tmpdirname) / "ov_config.json")
with open(Path(config_file), "w") as f:
f.write('{"NUM_STREAMS" : "2"}')
openvino_model_with_config = SparseEncoder(
model_id,
backend="openvino",
model_kwargs={"ov_config": config_file},
)
# The transformers model is an Optimum model with an OpenVINO inference request property
assert openvino_model_with_config[0].auto_model.request.get_property("NUM_STREAMS") == 2
# Test that saving and loading local OpenVINO models works as expected
openvino_model_with_config.save_pretrained(tmpdirname)
local_openvino_model = SparseEncoder(
tmpdirname, backend="openvino", model_kwargs={"ov_config": {"INFERENCE_PRECISION_HINT": "f32"}}
)
local_openvino_result = local_openvino_model.encode(["Hello there!"])
assert sparse_allclose(local_openvino_result, openvino_result), (
"OpenVINO saved model output differs from in-memory converted model"
)
del local_openvino_model
gc.collect()
def test_export_false_subfolder() -> None:
model_id = "sparse-encoder-testing/splade-bert-tiny-nq-openvino"
def from_pretrained_decorator(method):
def decorator(*args, **kwargs):
assert not kwargs["export"]
assert kwargs["subfolder"] == "openvino"
assert kwargs["file_name"] == "openvino_model.xml"
return method(*args, **kwargs)
return decorator
OVModelForMaskedLM.from_pretrained = from_pretrained_decorator(OVModelForMaskedLM.from_pretrained)
SparseEncoder(model_id, backend="openvino", model_kwargs={"export": False})
def test_export_set_nested_filename() -> None:
model_id = "sparse-encoder-testing/splade-bert-tiny-nq-openvino"
def from_pretrained_decorator(method):
def decorator(*args, **kwargs):
assert kwargs["subfolder"] == "openvino"
assert kwargs["file_name"] == "openvino_model.xml"
return method(*args, **kwargs)
return decorator
OVModelForMaskedLM.from_pretrained = from_pretrained_decorator(OVModelForMaskedLM.from_pretrained)
SparseEncoder(model_id, backend="openvino", model_kwargs={"file_name": "openvino/openvino_model.xml"})
|