defining commited on
Commit
9fa76a8
1 Parent(s): e94f32c

Create handler.py

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