1from transformers import AutoTokenizer
2import onnxruntime
3import numpy as np
4
5
6def infer_onnx(text: str, lang: str, onnx_model_path: str = "byt5_g2p_model.onnx"):
7 """
8 Exports the ByT5 model to ONNX format and then performs inference using ONNX Runtime.
9
10 Args:
11 text (str): The input text to convert to phonemes.
12 lang (str): The language tag (e.g., "en").
13 onnx_model_path (str): The path to save/load the ONNX model.
14 """
15 model_name = 'fdemelo/g2p-multilingual-byt5-tiny-8l-ipa-childes'
16 tokenizer = AutoTokenizer.from_pretrained(model_name)
17
18 # --- Step 2: Perform Inference with ONNX Runtime ---
19 print("\n--- Performing inference with ONNX Runtime ---")
20
21 # Create an ONNX Runtime session
22 try:
23 session = onnxruntime.InferenceSession(onnx_model_path, providers=['CPUExecutionProvider'])
24 except Exception as e:
25 print(f"Error loading ONNX model: {e}")
26 return
27
28 # Get input and output names from the ONNX model
29 onnx_input_names = [inp.name for inp in session.get_inputs()]
30 onnx_output_names = [out.name for out in session.get_outputs()]
31
32 # Prepare actual input for ONNX inference
33 input_text_for_onnx = f"<{lang}>: {text}"
34 inputs_for_onnx = tokenizer([input_text_for_onnx], return_tensors="pt", add_special_tokens=False)
35
36 input_ids_np = inputs_for_onnx["input_ids"].cpu().numpy()
37 attention_mask_np = inputs_for_onnx["attention_mask"].cpu().numpy()
38
39 # Manual greedy decoding loop for ONNX Runtime
40 # This simulates the 'generate' method's greedy decoding.
41 generated_ids = []
42 # T5 models typically use pad_token_id as the initial token for generation
43 # or a specific decoder_start_token_id.
44 # For T5, the decoder_start_token_id is usually the pad_token_id.
45 current_decoder_input_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0
46
47 # Ensure it's a batch of 1
48 decoder_input_ids_np = np.array([[current_decoder_input_id]])
49
50 max_length = 512 # Same as in the original predict_byt5
51
52 # Store encoder outputs if needed for cross-attention in decoder (T5 does this)
53 # When exporting the full T5 model's forward pass, the encoder_hidden_states
54 # are implicitly handled within the graph. We just need to feed the decoder_input_ids.
55
56 for _ in range(max_length):
57 # Prepare inputs for the current step
58 onnx_inputs = {
59 "input_ids": input_ids_np,
60 "attention_mask": attention_mask_np,
61 "decoder_input_ids": decoder_input_ids_np
62 }
63
64 # Run inference
65 outputs = session.run(onnx_output_names, onnx_inputs)
66 logits = outputs[0] # Get the logits
67
68 # Get the logits for the last token in the sequence
69 next_token_logits = logits[0, -1, :] # Batch 0, last token, all vocab logits
70
71 # Greedy decoding: pick the token with the highest logit
72 next_token_id = np.argmax(next_token_logits)
73 generated_ids.append(next_token_id)
74
75 # Check for end-of-sequence token
76 if next_token_id == tokenizer.eos_token_id:
77 break
78
79 # Update decoder input for the next step
80 # Append the new token to the decoder input sequence
81 decoder_input_ids_np = np.concatenate((decoder_input_ids_np, np.array([[next_token_id]])), axis=1)
82
83 # Decode the generated ONNX phoneme IDs
84 onnx_phones = tokenizer.batch_decode([generated_ids], skip_special_tokens=True)
85 print(f"ONNX Runtime Inference: {onnx_phones}")
86 return onnx_phones
87``
