TenzinGayche commited on
Commit
ef9dd5a
1 Parent(s): 64d3b03

Create handler.py

Browse files
Files changed (1) hide show
  1. handler.py +36 -0
handler.py ADDED
@@ -0,0 +1,36 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Dict
2
+ from transformers.pipelines.audio_utils import ffmpeg_read
3
+ import torch
4
+ import pyewts
5
+ from transformers import pipeline
6
+ converter = pyewts.pyewts()
7
+
8
+ SAMPLE_RATE = 16000
9
+
10
+
11
+
12
+ class EndpointHandler():
13
+ def __init__(self, path=""):
14
+ # load the model
15
+ self.pipe = pipeline(model="TenzinGayche/whisper-small-3",device='cuda')
16
+
17
+
18
+ def __call__(self, data: Dict[str, bytes]) -> Dict[str, str]:
19
+ """
20
+ Args:
21
+ data (:obj:):
22
+ includes the deserialized audio file as bytes
23
+ Return:
24
+ A :obj:`dict`:. base64 encoded image
25
+ """
26
+ # process input
27
+ inputs = data.pop("inputs", data)
28
+ audio_nparray = ffmpeg_read(inputs, SAMPLE_RATE)
29
+ audio_tensor= torch.from_numpy(audio_nparray)
30
+ text = pipe(audio_tensor)["text"]
31
+
32
+ # run inference pipeline
33
+ result = converter.toUnicode(text)
34
+
35
+ # postprocess the prediction
36
+ return {"text": result}