stupidog04
commited on
Commit
•
d6c86e9
1
Parent(s):
62716ae
Delete pipeline.py
Browse files- pipeline.py +0 -46
pipeline.py
DELETED
@@ -1,46 +0,0 @@
|
|
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 |
-
}
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|