DmitrySobolevsky commited on
Commit
b8500cb
1 Parent(s): 7b257e3

Adding handler

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