S-Fry commited on
Commit
c5153e0
1 Parent(s): d352bb6

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +12 -10
handler.py CHANGED
@@ -8,20 +8,21 @@ SAMPLE_RATE = 16000
8
  MODEL_NAME = "openai/whisper-large" #this always needs to stay in line 8 :D sorry for the hackiness
9
  lang = "dk"
10
 
11
- device = 0 if torch.cuda.is_available() else "cpu"
12
- pipe = pipeline(
13
- task="automatic-speech-recognition",
14
- model=MODEL_NAME,
15
- chunk_length_s=30,
16
- device=device,
17
- )
18
 
19
- pipe.model.config.forced_decoder_ids = pipe.tokenizer.get_decoder_prompt_ids(language=lang, task="transcribe")
20
 
21
  class EndpointHandler():
22
  def __init__(self, path=""):
 
 
 
 
 
 
 
23
  # load the model
24
- self.model = whisper.load_model("medium")
 
25
 
26
 
27
  def __call__(self, data: Dict[str, bytes]) -> Dict[str, str]:
@@ -37,8 +38,9 @@ class EndpointHandler():
37
  audio_nparray = ffmpeg_read(inputs, SAMPLE_RATE)
38
  audio_tensor= torch.from_numpy(audio_nparray)
39
 
40
- # run inference pipeline
41
  result = self.model.transcribe(audio_nparray)
42
 
 
43
  # postprocess the prediction
44
  return {"tekst": result["text"]}
 
8
  MODEL_NAME = "openai/whisper-large" #this always needs to stay in line 8 :D sorry for the hackiness
9
  lang = "dk"
10
 
 
 
 
 
 
 
 
11
 
12
+
13
 
14
  class EndpointHandler():
15
  def __init__(self, path=""):
16
+ pipe = pipeline(
17
+ task="automatic-speech-recognition",
18
+ model=MODEL_NAME,
19
+ chunk_length_s=30,
20
+ device=device,
21
+ )
22
+
23
  # load the model
24
+ #self.model = whisper.load_model("large")
25
+ self.model.config.forced_decoder_ids = pipe.tokenizer.get_decoder_prompt_ids(language=lang, task="transcribe")
26
 
27
 
28
  def __call__(self, data: Dict[str, bytes]) -> Dict[str, str]:
 
38
  audio_nparray = ffmpeg_read(inputs, SAMPLE_RATE)
39
  audio_tensor= torch.from_numpy(audio_nparray)
40
 
41
+ # run inference pipeline
42
  result = self.model.transcribe(audio_nparray)
43
 
44
+
45
  # postprocess the prediction
46
  return {"tekst": result["text"]}