facebook

/

nllb-200-distilled-600M

like 516

Follow

AI at Meta 3,395

I don't know what the input variables are when I use onnxruntime inference after exporting the onnx model
among
inputs = tokenizer(text, return_tensors="pt", padding=True)
target_lang_id = np.array([[tokenizer.convert_tokens_to_ids('zho_Hans')]])
decoder_input_ids = target_lang_id
ort_input = {
"input_ids": inputs['input_ids'].numpy(),
"attention_mask": inputs['attention_mask'].numpy(),
'decoder_input_ids':decoder_input_ids
}
But it cannot be accurately translated
The complete code is as follows：
import onnxruntime as ort
import numpy as np
from transformers import , AutoTokenizer

ort_session = ort.InferenceSession('nllb-200-distilled-600M-ONNX/model_quantized.onnx')
text = "hello"
tokenizer = AutoTokenizer.from_pretrained('nllb-200-distilled-600M-ONNX', use_auth_token=True, src_lang='eng_Latn')
inputs = tokenizer(text, return_tensors="pt", padding=True)
target_lang_id = np.array([[tokenizer.convert_tokens_to_ids('zho_Hans')]])
decoder_input_ids = target_lang_id
print(inputs)
import onnxruntime as ort

inputs = ort_session.get_inputs()

for input in inputs:
print(f"{input}")

outputs = ort_session.run(None, {
"input_ids": inputs['input_ids'].numpy(),
"attention_mask": inputs['attention_mask'].numpy(),
'decoder_input_ids':decoder_input_ids
})
print(outputs)

logits = outputs[0]
predicted_ids = np.argmax(logits, axis=-1)
print(predicted_ids)

output_text = tokenizer.decode(predicted_ids[0], skip_special_tokens=True)
print("output: ", output_text)
What are the input variables and how can I use the onnx model to obtain accurate inference