Fix deprecated optimum interface (#17112)

The `latest_model_name` argument to create an {self.__class__.__name__}
is deprecated since optimum 1.6.0. Replace it with `model_name`
This commit is contained in:
PeixuanZuo 2023-08-16 12:33:36 +08:00 committed by GitHub
parent 6b29837ed2
commit ebcd9b5cae
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

@ -29,8 +29,13 @@ from datasets import load_dataset
from evaluate import evaluator
from optimum.onnxruntime import ORTModelForQuestionAnswering
from optimum.onnxruntime.modeling_ort import ORTModel
from optimum.version import __version__ as optimum_version
from packaging import version as version_check
from transformers import AutoTokenizer, pipeline
if version_check.parse(optimum_version) < version_check.parse("1.6.0"):
raise ImportError(f"Please install optimum>=1.6.0. The version {optimum_version} was found.")
PRETRAINED_SQUAD_MODELS = [
"bert-large-uncased-whole-word-masking-finetuned-squad",
"deepset/roberta-base-squad2",
@ -62,7 +67,7 @@ def load_onnx_model(
model = ORTModelForQuestionAnswering.from_pretrained(model_id, from_transformers=True)
if onnx_path is not None:
model.latest_model_name = Path(onnx_path).name
model.model_name = Path(onnx_path).name
if provider != "CPUExecutionProvider":
model.device = torch.device("cuda:0")
@ -71,7 +76,7 @@ def load_onnx_model(
model.device = torch.device("cpu")
model.model = ORTModel.load_model(onnx_path)
else:
onnx_path = os.path.join(model.model_save_dir.as_posix(), model.latest_model_name)
onnx_path = os.path.join(model.model_save_dir.as_posix(), model.model_name)
if provider != "CPUExecutionProvider":
model.to("cuda")