diff --git a/node-hub/dora-distil-whisper/dora_distil_whisper/main.py b/node-hub/dora-distil-whisper/dora_distil_whisper/main.py index 6ea69e2d..1b56ff6e 100644 --- a/node-hub/dora-distil-whisper/dora_distil_whisper/main.py +++ b/node-hub/dora-distil-whisper/dora_distil_whisper/main.py @@ -1,15 +1,15 @@ -from dora import Node -import pyarrow as pa import os import sys from pathlib import Path +import pyarrow as pa +import torch +from dora import Node + DEFAULT_PATH = "openai/whisper-large-v3-turbo" -TARGET_LANGUAGE = os.getenv("TARGET_LANGUAGE", "chinese") +TARGET_LANGUAGE = os.getenv("TARGET_LANGUAGE", "english") TRANSLATE = bool(os.getenv("TRANSLATE", "False") in ["True", "true"]) -import torch - def load_model(): from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline @@ -47,9 +47,7 @@ def load_model(): def load_model_mlx(): - from lightning_whisper_mlx import ( - LightningWhisperMLX, - ) # pylint: disable=import-error + from lightning_whisper_mlx import LightningWhisperMLX # noqa whisper = LightningWhisperMLX(model="distil-large-v3", batch_size=12, quant=None) return whisper diff --git a/node-hub/dora-distil-whisper/pyproject.toml b/node-hub/dora-distil-whisper/pyproject.toml index 956f4b15..f8af54c2 100644 --- a/node-hub/dora-distil-whisper/pyproject.toml +++ b/node-hub/dora-distil-whisper/pyproject.toml @@ -26,6 +26,11 @@ lightning-whisper-mlx = { version = "^0.0.10", platform = "darwin" } [tool.poetry.scripts] dora-distil-whisper = "dora_distil_whisper.main:main" +[tool.poetry.dev-dependencies] +pytest = ">= 6.3.4" +pylint = ">= 3.3.2" +black = ">= 22.10" + [build-system] requires = ["poetry-core>=1.8.0"] build-backend = "poetry.core.masonry.api"