Spaces:
Paused
Paused
Commit
·
75b7975
1
Parent(s):
d03d236
fix: app.py
Browse files
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=
|
27 |
-
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=
|
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 |
# ---
|