1import numpy as np
2import librosa
3import onnxruntime as ort
4from huggingface_hub import hf_hub_download
5
6# 1. Download model
7repo_id = "Bao2311/wav2vec2-vi-pronunciation-onnx"
8onnx_path = hf_hub_download(repo_id, "phoneme_classifier.onnx")
9data_path = hf_hub_download(repo_id, "phoneme_classifier.onnx.data")
10
11# 2. Load ONNX session
12session = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
13
14# 3. Load & preprocess audio (16kHz, 500ms)
15audio, _ = librosa.load("your_audio.wav", sr=16000, mono=True)
16audio = audio[:8000] # crop to 500ms = 8000 samples
17if len(audio) < 8000:
18 audio = np.pad(audio, (0, 8000 - len(audio)))
19
20# 4. Normalize
21audio = audio.astype(np.float32)
22mean, std = audio.mean(), audio.std()
23if std > 0:
24 audio = (audio - mean) / std
25
26# 5. Inference
27input_values = audio.reshape(1, -1)
28logits = session.run(None, {"input_values": input_values})[0][0]
29
30# 6. Softmax → prediction
31exp_logits = np.exp(logits - np.max(logits))
32probs = exp_logits / exp_logits.sum()
33pred = int(np.argmax(probs))
34
35labels = {0: "correct ✅", 1: "wrong ❌"}
36print(f"Prediction: {labels[pred]} (confidence: {probs[pred]:.1%})")
