Views
No views yet
1from transformers import Wav2Vec2Processor, AutoConfig
2import onnxruntime as rt
3import torch
4import torch.nn.functional as F
5import numpy as np
6import os
7import torchaudio
8
9
10class EndOfSpeechDetection:
11 processor: Wav2Vec2Processor
12 config: AutoConfig
13 session: rt.InferenceSession
14
15 def load_model(self, path, use_gpu=False):
16 processor = Wav2Vec2Processor.from_pretrained(path)
17 config = AutoConfig.from_pretrained(path)
18
19 sess_options = rt.SessionOptions()
20 sess_options.graph_optimization_level = rt.GraphOptimizationLevel.ORT_ENABLE_ALL
21
22 providers = ["ROCMExecutionProvider"] if use_gpu else ["CPUExecutionProvider"]
23 session = rt.InferenceSession(
24 os.path.join(path, "model.onnx"), sess_options, providers=providers
25 )
26 return processor, config, session
27
28 def predict(self, segment, file_type="pcm"):
29 if file_type == "pcm":
30 # pcm files
31 speech_array = np.memmap(segment, dtype="float32", mode="r").astype(
32 np.float32
33 )
34 else:
35 # wave files
36 speech_array, _ = torchaudio.load(segment)
37 speech_array = speech_array[0].numpy()
38
39 features = self.processor(
40 speech_array, sampling_rate=16000, return_tensors="pt", padding=True
41 )
42 input_values = features.input_values
43 outputs = self.session.run(
44 [self.session.get_outputs()[-1].name],
45 {self.session.get_inputs()[-1].name: input_values.detach().cpu().numpy()},
46 )[0]
47 softmax_output = F.softmax(torch.tensor(outputs), dim=1)
48
49 both_classes_with_prob = {
50 self.config.id2label[i]: softmax_output[0][i].item()
51 for i in range(len(softmax_output[0]))
52 }
53
54 return both_classes_with_prob
55
56
57if __name__ == "__main__":
58 eos = EndOfSpeechDetection()
59 eos.processor, eos.config, eos.session = eos.load_model("eos-model-onnx")
60 print(eos.predict("some.pcm", file_type="pcm"))
61| classes | precision | recall | f1-score | support |
|---|---|---|---|---|
| eos | 0.94 | 0.95 | 0.95 | 4060 |
| not_eos | 0.95 | 0.94 | 0.95 | 4060 |