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

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