stupidog04
commited on
Commit
•
02cf913
1
Parent(s):
aa1dbe1
Update pipeline.py
Browse files- pipeline.py +46 -32
pipeline.py
CHANGED
@@ -1,32 +1,46 @@
|
|
1 |
-
|
2 |
-
from
|
3 |
-
|
4 |
-
|
5 |
-
|
6 |
-
|
7 |
-
|
8 |
-
|
9 |
-
|
10 |
-
|
11 |
-
|
12 |
-
|
13 |
-
|
14 |
-
|
15 |
-
|
16 |
-
#
|
17 |
-
|
18 |
-
|
19 |
-
|
20 |
-
|
21 |
-
|
22 |
-
|
23 |
-
|
24 |
-
|
25 |
-
|
26 |
-
|
27 |
-
|
28 |
-
|
29 |
-
|
30 |
-
|
31 |
-
|
32 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import numpy as np
|
2 |
+
from typing import Dict
|
3 |
+
|
4 |
+
from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC
|
5 |
+
from pyctcdecode import Alphabet, BeamSearchDecoderCTC
|
6 |
+
|
7 |
+
class PreTrainedPipeline():
|
8 |
+
def __init__(self, path):
|
9 |
+
"""
|
10 |
+
Initialize model
|
11 |
+
"""
|
12 |
+
self.processor = Wav2Vec2Processor.from_pretrained(path)
|
13 |
+
self.model = Wav2Vec2ForCTC.from_pretrained(path)
|
14 |
+
vocab_list = list(self.processor.tokenizer.get_vocab().keys())
|
15 |
+
|
16 |
+
# convert ctc blank character representation
|
17 |
+
vocab_list[0] = ""
|
18 |
+
|
19 |
+
# replace special characters
|
20 |
+
vocab_list[1] = "⁇"
|
21 |
+
vocab_list[2] = "⁇"
|
22 |
+
vocab_list[3] = "⁇"
|
23 |
+
|
24 |
+
# convert space character representation
|
25 |
+
vocab_list[4] = " "
|
26 |
+
|
27 |
+
alphabet = Alphabet.build_alphabet(vocab_list, ctc_token_idx=0)
|
28 |
+
|
29 |
+
self.decoder = BeamSearchDecoderCTC(alphabet)
|
30 |
+
self.sampling_rate = 16000
|
31 |
+
|
32 |
+
|
33 |
+
def __call__(self, inputs)-> Dict[str, str]:
|
34 |
+
"""
|
35 |
+
Args:
|
36 |
+
inputs (:obj:`np.array`):
|
37 |
+
The raw waveform of audio received. By default at 16KHz.
|
38 |
+
Return:
|
39 |
+
A :obj:`dict`:. The object return should be liked {"text": "XXX"} containing
|
40 |
+
the detected text from the input audio.
|
41 |
+
"""
|
42 |
+
input_values = self.processor(inputs, return_tensors="pt", sampling_rate=self.sampling_rate).input_values # Batch size 1
|
43 |
+
logits = self.model(input_values).logits.cpu().detach().numpy()[0]
|
44 |
+
return {
|
45 |
+
"text": self.decoder.decode(logits)
|
46 |
+
}
|