Fix: During ONNX inference, the tensor of the return value of the tokenizer is now obtained in Numpy's NDArray format to eliminate the dependency on PyTorch

This commit is contained in:
tsukumi
2024-09-23 07:40:17 +09:00
parent dcca37b9a8
commit b7cd512978
3 changed files with 20 additions and 20 deletions

View File

@@ -98,7 +98,7 @@ def extract_bert_feature_onnx(
""" """
tokenizer = onnx_bert_models.load_tokenizer(Languages.ZH) tokenizer = onnx_bert_models.load_tokenizer(Languages.ZH)
inputs = tokenizer(text, return_tensors="pt") inputs = tokenizer(text, return_tensors="np")
session = onnx_bert_models.load_model( session = onnx_bert_models.load_model(
language=Languages.ZH, language=Languages.ZH,
@@ -108,21 +108,21 @@ def extract_bert_feature_onnx(
res = session.run( res = session.run(
[output_name], [output_name],
{ {
"input_ids": inputs["input_ids"].detach().numpy(), # type: ignore "input_ids": inputs["input_ids"],
"token_type_ids": inputs["token_type_ids"].detach().numpy(), # type: ignore "token_type_ids": inputs["token_type_ids"],
"attention_mask": inputs["attention_mask"].detach().numpy(), # type: ignore "attention_mask": inputs["attention_mask"],
}, },
)[0] )[0]
style_res_mean = None style_res_mean = None
if assist_text: if assist_text:
style_inputs = tokenizer(assist_text, return_tensors="pt") style_inputs = tokenizer(assist_text, return_tensors="np")
style_res = session.run( style_res = session.run(
[output_name], [output_name],
{ {
"input_ids": style_inputs["input_ids"].detach().numpy(), # type: ignore "input_ids": style_inputs["input_ids"],
"token_type_ids": style_inputs["token_type_ids"].detach().numpy(), # type: ignore "token_type_ids": style_inputs["token_type_ids"],
"attention_mask": style_inputs["attention_mask"].detach().numpy(), # type: ignore "attention_mask": style_inputs["attention_mask"],
}, },
)[0] )[0]
style_res_mean = np.mean(style_res, axis=0) style_res_mean = np.mean(style_res, axis=0)

View File

@@ -98,7 +98,7 @@ def extract_bert_feature_onnx(
""" """
tokenizer = onnx_bert_models.load_tokenizer(Languages.EN) tokenizer = onnx_bert_models.load_tokenizer(Languages.EN)
inputs = tokenizer(text, return_tensors="pt") inputs = tokenizer(text, return_tensors="np")
session = onnx_bert_models.load_model( session = onnx_bert_models.load_model(
language=Languages.EN, language=Languages.EN,
@@ -108,19 +108,19 @@ def extract_bert_feature_onnx(
res = session.run( res = session.run(
[output_name], [output_name],
{ {
"input_ids": inputs["input_ids"].detach().numpy(), # type: ignore "input_ids": inputs["input_ids"],
"attention_mask": inputs["attention_mask"].detach().numpy(), # type: ignore "attention_mask": inputs["attention_mask"],
}, },
)[0] )[0]
style_res_mean = None style_res_mean = None
if assist_text: if assist_text:
style_inputs = tokenizer(assist_text, return_tensors="pt") style_inputs = tokenizer(assist_text, return_tensors="np")
style_res = session.run( style_res = session.run(
[output_name], [output_name],
{ {
"input_ids": style_inputs["input_ids"].detach().numpy(), # type: ignore "input_ids": style_inputs["input_ids"],
"attention_mask": style_inputs["attention_mask"].detach().numpy(), # type: ignore "attention_mask": style_inputs["attention_mask"],
}, },
)[0] )[0]
style_res_mean = np.mean(style_res, axis=0) style_res_mean = np.mean(style_res, axis=0)

View File

@@ -111,7 +111,7 @@ def extract_bert_feature_onnx(
assist_text = "".join(text_to_sep_kata(assist_text, raise_yomi_error=False)[0]) assist_text = "".join(text_to_sep_kata(assist_text, raise_yomi_error=False)[0])
tokenizer = onnx_bert_models.load_tokenizer(Languages.JP) tokenizer = onnx_bert_models.load_tokenizer(Languages.JP)
inputs = tokenizer(text, return_tensors="pt") inputs = tokenizer(text, return_tensors="np")
session = onnx_bert_models.load_model( session = onnx_bert_models.load_model(
language=Languages.JP, language=Languages.JP,
@@ -121,19 +121,19 @@ def extract_bert_feature_onnx(
res = session.run( res = session.run(
[output_name], [output_name],
{ {
"input_ids": inputs["input_ids"].detach().numpy(), # type: ignore "input_ids": inputs["input_ids"],
"attention_mask": inputs["attention_mask"].detach().numpy(), # type: ignore "attention_mask": inputs["attention_mask"],
}, },
)[0] )[0]
style_res_mean = None style_res_mean = None
if assist_text: if assist_text:
style_inputs = tokenizer(assist_text, return_tensors="pt") style_inputs = tokenizer(assist_text, return_tensors="np")
style_res = session.run( style_res = session.run(
[output_name], [output_name],
{ {
"input_ids": style_inputs["input_ids"].detach().numpy(), # type: ignore "input_ids": style_inputs["input_ids"],
"attention_mask": style_inputs["attention_mask"].detach().numpy(), # type: ignore "attention_mask": style_inputs["attention_mask"],
}, },
)[0] )[0]
style_res_mean = np.mean(style_res, axis=0) style_res_mean = np.mean(style_res, axis=0)