Views
No views yet
hubert_base.onnxfloat32 | shape: [batch_size, sequence_length]
bool | shape: [batch_size, sequence_length]
padding_mask = np.zeros(waveform.shape, dtype=np.bool_)float32 | shape: [batch_size, sequence_length, 768 ]1import numpy as np
2import onnxruntime as ort
3
4class OnnxHubert:
5 """
6 Class to load and run the ONNX model exported by Hubert.
7
8 Attributes:
9 session (ort.InferenceSession): The ONNX Runtime session.
10 input_name (str): The name of the input node.
11 output_name (str): The name of the output node.
12
13 Methods:
14 extract_features_batch (source, padding_mask): Run the ONNX model and extract features from the batch.
15 extract_features (source, padding_mask): Run the ONNX model and extract features from a single input.
16 """
17 def __init__(self, model_path: str, thread_num: int = None):
18 """
19 Initialize the OnnxHubert object.
20
21 Parameters:
22 model_path (str): The path to the ONNX model file.
23 thread_num (int, optional): The number of threads to use for inference. Defaults to None.
24
25 Attributes:
26 session (ort.InferenceSession): The ONNX Runtime session.
27 input_name (str): The name of the input node.
28 output_name (str): The name of the output node.
29 """
30 self.session = ort.InferenceSession(model_path)
31
32 self.input_name = self.session.get_inputs()[0].name
33 self.output_name = self.session.get_outputs()[0].name
34 def extract_features(
35 self,
36 source: np.ndarray,
37 padding_mask: np.ndarray
38 ) -> np.ndarray:
39 """
40 Extract features from the batch using the ONNX model.
41
42 Inputs:
43 source: ndarray of shape (batch_size, sequence_length) float32
44 padding_mask: ndarray of shape (batch_size, sequence_length) bool
45
46 Returns:
47 ndarray of shape (D, 768) with the extracted features
48 """
49 result = self.session.run(None, {
50 "source": source,
51 "padding_mask": padding_mask
52 })
53 return result[0]pip install onnxruntime numpy