Mirroristas/whisper
2
0
Fork 0
mirror of https://github.com/openai/whisper.git synced 2025-11-29 16:38:40 +00:00
Code Issues Packages Projects Releases Wiki Activity

allow test_transcribe to run on CPU when CUDA is not available

Browse Source
This commit is contained in:
Jong Wook Kim 2023-01-17 13:43:36 -08:00
parent 493dfffa37
commit b1d213c0c7
2 changed files with 5 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
.github/workflows/test.yml vendored
View File

@ -18,7 +18,7 @@ jobs:
pytorch-version: 1.10.2 pytorch-version: 1.10.2
steps: steps:
- uses: conda-incubator/setup-miniconda@v2 - uses: conda-incubator/setup-miniconda@v2
- run: conda install -n test python=${{ matrix.python-version }} pytorch=${{ matrix.pytorch-version }} cpuonly -c pytorch - run: conda install -n test ffmpeg python=${{ matrix.python-version }} pytorch=${{ matrix.pytorch-version }} cpuonly -c pytorch
- uses: actions/checkout@v2 - uses: actions/checkout@v2
- run: echo "$CONDA/envs/test/bin" >> $GITHUB_PATH - run: echo "$CONDA/envs/test/bin" >> $GITHUB_PATH
- run: pip install pytest - run: pip install pytest

6
tests/test_transcribe.py
View File

@ -1,13 +1,15 @@
import os import os
import pytest import pytest
import torch
import whisper import whisper
@pytest.mark.parametrize('model_name', whisper.available_models()) @pytest.mark.parametrize("model_name", whisper.available_models())
def test_transcribe(model_name: str): def test_transcribe(model_name: str):
model = whisper.load_model(model_name).cuda() device = "cuda" if torch.cuda.is_available() else "cpu"
model = whisper.load_model(model_name).to(device)
audio_path = os.path.join(os.path.dirname(__file__), "jfk.flac") audio_path = os.path.join(os.path.dirname(__file__), "jfk.flac")
language = "en" if model_name.endswith(".en") else None language = "en" if model_name.endswith(".en") else None