Leslie Jones-Dove commited on
Commit
f0386c6
1 Parent(s): e08502d

add custom handler

Browse files
Files changed (3) hide show
  1. __pycache__/handler.cpython-39.pyc +0 -0
  2. handler.py +26 -0
  3. test.py +10 -0
__pycache__/handler.cpython-39.pyc ADDED
Binary file (1.35 kB). View file
 
handler.py ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Dict, List, Any
2
+ from transformers import pipeline, WhisperProcessor, WhisperForConditionalGeneration
3
+ import torch
4
+
5
+ class EndpointHandler():
6
+ def __init__(self, path=""):
7
+ device = 0 if torch.cuda.is_available() else "cpu"
8
+ self.pipe = pipeline(
9
+ task="automatic-speech-recognition",
10
+ model=path,
11
+ chunk_length_s=30,
12
+ device=device,
13
+ )
14
+
15
+ def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
16
+ """
17
+ data args:
18
+ inputs (:obj: `str` | `PIL.Image` | `np.array`)
19
+ kwargs
20
+ Return:
21
+ A :obj:`list` | `dict`: will be serialized and returned
22
+ """
23
+ inputs = data.pop("inputs",data)
24
+ print("inputs", inputs)
25
+ prediction = self.pipe(inputs, return_timestamps=True)
26
+ return prediction
test.py ADDED
@@ -0,0 +1,10 @@
 
 
 
 
 
 
 
 
 
 
 
1
+ from handler import EndpointHandler
2
+ from datasets import load_dataset
3
+
4
+ # init handler
5
+ my_handler = EndpointHandler(path=".")
6
+
7
+ # print(type(sample["data"]))
8
+ prediction = my_handler({"inputs": "/Users/lesliejones-dove/Desktop/projects/Whisper/whisperx/whisperCustomHandler/Saffored_short.wav"})
9
+ # /Users/lesliejones-dove/Desktop/projects/Whisper/whisperx/whisperCustomHandler
10
+ print(prediction)