isort & black

- trying to fix environment issues
This commit is contained in:
2025-09-14 03:07:06 +04:00
parent 7cf4c5259b
commit 7291b148e5
11 changed files with 177 additions and 99 deletions
+12 -9
View File
@@ -1,6 +1,8 @@
from dataclasses import dataclass
import torch
@dataclass
class DeviceConfiguration:
"""
@@ -8,36 +10,37 @@ class DeviceConfiguration:
Attributes:
device (str): Type of device. Possible options: "cuda", "cpu", "mps".
model_name (str): Whisper models. Possible options:
- "openai/whisper-tiny"
- "openai/whisper-small"
- "openai/whisper-medium"
- "openai/whisper-large"
- "openai/whisper-large-v2"
batch_size (int): Chunks in one batch. Selected for VRAM.
chunk_length_s (int): Length of one audio chunk in seconds. Smaller -> less VRAM.
data_type (str): custom data type of model. Variants:
- torch.float16 - for GPUs
- torch.float32 - for CPU / weak GPU
- torch.bfloat16 - for GPUs which has BF16 support
"""
device: str = "cuda"
model_name: str = "openai/whisper-large-v2"
batch_size: int = 16
chunk_length_s: int = 30
data_type: str = "torch.float16"
_dtype_map = {
"torch.float16": torch.float16,
"torch.float32": torch.float32,
"torch.bfloat16": torch.bfloat16
"torch.bfloat16": torch.bfloat16,
}
torch_dtype: torch.dtype = None
def __post_init__(self):
self.torch_dtype = self._dtype_map[self.data_type]
self.torch_dtype = self._dtype_map[self.data_type]