juancopi81 commited on
Commit
25733b9
1 Parent(s): 6472070

Add device_dict

Browse files
Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -16,6 +16,7 @@ TRANSCRIBER_MODEL_NAME = "juancopi81/whisper-medium-es"
16
  lang = "es"
17
 
18
  device = "cuda" if torch.cuda.is_available() else "cpu"
 
19
  dtype = torch.float16 if device == "cuda" else torch.float32
20
 
21
  # Detect if code is running in Colab
@@ -32,7 +33,7 @@ transcription_pipe = pipeline(
32
  task="automatic-speech-recognition",
33
  model=TRANSCRIBER_MODEL_NAME,
34
  chunk_length_s=30,
35
- device=device,
36
  )
37
  transcription_pipe.model.config.forced_decoder_ids = transcription_pipe.tokenizer.get_decoder_prompt_ids(language=lang, task="transcribe")
38
  audio_transcriber = Transcriber(transcription_pipe)
 
16
  lang = "es"
17
 
18
  device = "cuda" if torch.cuda.is_available() else "cpu"
19
+ device_dict = {"cuda": 0, "cpu": -1}
20
  dtype = torch.float16 if device == "cuda" else torch.float32
21
 
22
  # Detect if code is running in Colab
 
33
  task="automatic-speech-recognition",
34
  model=TRANSCRIBER_MODEL_NAME,
35
  chunk_length_s=30,
36
+ device=device_dict[device],
37
  )
38
  transcription_pipe.model.config.forced_decoder_ids = transcription_pipe.tokenizer.get_decoder_prompt_ids(language=lang, task="transcribe")
39
  audio_transcriber = Transcriber(transcription_pipe)