slplab commited on
Commit
e769edb
1 Parent(s): cfed61f

Update good_handler.py

Browse files
Files changed (1) hide show
  1. good_handler.py +4 -2
good_handler.py CHANGED
@@ -4,13 +4,15 @@ import torch
4
 
5
  #### USE of PIPELINE
6
 
 
 
7
  class EndpointHandler:
8
  def __init__(self, path=""):
9
- self.pipe = pipeline(task='automatic-speech-recognition', model=path)
10
 
11
 
12
  def __call__(self, data: Any) -> List[Dict[str, str]]:
13
- print('==========NEW PROCESS=========')
14
  transcribe = self.pipe
15
  transcribe.model.config.forced_decoder_ids = transcribe.tokenizer.get_decoder_prompt_ids(language="ko", task="transcribe")
16
  result = transcribe(data['inputs'])
 
4
 
5
  #### USE of PIPELINE
6
 
7
+ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
8
+
9
  class EndpointHandler:
10
  def __init__(self, path=""):
11
+ self.pipe = pipeline(task='automatic-speech-recognition', model=path, device=device)
12
 
13
 
14
  def __call__(self, data: Any) -> List[Dict[str, str]]:
15
+ inputs = data.pop("inputs", data)
16
  transcribe = self.pipe
17
  transcribe.model.config.forced_decoder_ids = transcribe.tokenizer.get_decoder_prompt_ids(language="ko", task="transcribe")
18
  result = transcribe(data['inputs'])