MikeTangoEcho commited on
Commit
75b7975
·
1 Parent(s): d03d236

fix: app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -6
app.py CHANGED
@@ -13,27 +13,27 @@ logger = logging.get_logger("transformers")
13
 
14
  # Pipelines
15
 
16
- device = 0 if torch.cuda.is_available() else "cpu"
17
-
18
  ## Automatic Speech Recognition
19
  ## https://huggingface.co/docs/transformers/task_summary#automatic-speech-recognition
20
  ## Require ffmpeg to be installed
21
-
 
22
  asr_model = "openai/whisper-tiny"
23
  asr = pipeline(
24
  "automatic-speech-recognition",
25
  model=asr_model,
26
- torch_dtype=torch.float16,
27
- device=device
28
  )
29
 
30
  ## Token Classification / Name Entity Recognition
31
  ## https://huggingface.co/docs/transformers/task_summary#token-classification
 
32
  tc_model = "dslim/distilbert-NER"
33
  tc = pipeline(
34
  "token-classification", # ner
35
  model=tc_model,
36
- device=device
37
  )
38
 
39
  # ---
 
13
 
14
  # Pipelines
15
 
 
 
16
  ## Automatic Speech Recognition
17
  ## https://huggingface.co/docs/transformers/task_summary#automatic-speech-recognition
18
  ## Require ffmpeg to be installed
19
+ asr_device = "cuda:0" if torch.cuda.is_available() else "cpu"
20
+ asr_torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
21
  asr_model = "openai/whisper-tiny"
22
  asr = pipeline(
23
  "automatic-speech-recognition",
24
  model=asr_model,
25
+ torch_dtype=asr_torch_dtype,
26
+ device=asr_device
27
  )
28
 
29
  ## Token Classification / Name Entity Recognition
30
  ## https://huggingface.co/docs/transformers/task_summary#token-classification
31
+ tc_device = 0 if torch.cuda.is_available() else "cpu"
32
  tc_model = "dslim/distilbert-NER"
33
  tc = pipeline(
34
  "token-classification", # ner
35
  model=tc_model,
36
+ device=tc_device
37
  )
38
 
39
  # ---