svlqd/notecast
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

isort & black

Browse Source 
- trying to fix environment issues
This commit is contained in:
German Mikheev
2025-09-14 03:07:06 +04:00
parent 7cf4c5259b
commit 7291b148e5
11 changed files with 177 additions and 99 deletions
Download Patch File Download Diff File Expand all files Collapse all files
transcription/audio_transcription.py
+60 -47
View File
@@ -1,43 +1,46 @@
from transformers import WhisperProcessor, WhisperForConditionalGeneration
import logging
import math
import time
import torch
import torchaudio
from ui.ui_log_handler import UILogHandler
from transcription.device_configuration import DeviceConfiguration
import logging
import time
import math
from tqdm import tqdm
from transformers import WhisperForConditionalGeneration, WhisperProcessor
from transcription.device_configuration import DeviceConfiguration
from ui.ui_log_handler import UILogHandler
# TODO: implement transcription with shift
class AudioTranscription:
model_name = "openai/whisper-large-v2"
model_name = "openai/whisper-large-v2"
filepath: str
waveform: torch.Tensor
sampling_rate: int
chunks: list = []
batches: list = []
chunk_size: int
chunk_size: int
custom_chunk_length: int
custom_batch_length: int
device = "cuda"
processor: WhisperProcessor
model: WhisperForConditionalGeneration
logger: logging.Logger
torch_dtype: torch.dtype
language = "ru"
all_transcription: list = []
def __init__(
self,
self,
filepath: str,
device_configuration: DeviceConfiguration,
logger: logging.Logger,
language = "ru"
language="ru",
) -> None:
# TODO: add pretty docs here
self.filepath = filepath
@@ -50,44 +53,49 @@ class AudioTranscription:
self.custom_chunk_length = device_configuration.chunk_length_s
self.custom_batch_length = device_configuration.batch_size
self.torch_dtype = device_configuration.torch_dtype
self.chunks: list = []
self.batches: list = []
self.all_transcription: list = []
try:
logger.info("Loading model WhisperProcessor...")
self.processor = WhisperProcessor.from_pretrained(self.model_name)
self.model = WhisperForConditionalGeneration.from_pretrained(
self.model_name,
torch_dtype=self.torch_dtype
self.model_name, torch_dtype=self.torch_dtype
).to(self.device)
logger.info("Model loaded.")
self.waveform, self.sampling_rate = torchaudio.load(filepath, format="mp3", backend="ffmpeg")
self.waveform, self.sampling_rate = torchaudio.load(
filepath, format="mp3", backend="ffmpeg"
)
logger.info(f"Successfully loaded file {filepath}.")
except Exception as e:
logger.error(f"Unable to load file {self.filepath}: {e}")
raise
def _resample(self) -> None:
self.waveform = torchaudio.functional.resample(self.waveform, self.sampling_rate, 16000)
self.waveform = torchaudio.functional.resample(
self.waveform, self.sampling_rate, 16000
)
def _to_mono(self):
if self.waveform.shape[0] > 1:
self.waveform = self.waveform.mean(dim=0, keepdim=True)
self.waveform = self.waveform.squeeze(0)
def _split_to_chunks(self, chunk_length_s: int = 30) -> None:
self.logger.info(f"Splitting audio on chunks...")
self.chunk_size = chunk_length_s * 16000 # 16kHz after resampling
total_samples = self.waveform.shape[0]
chunks_count = (total_samples + self.chunk_size - 1) // self.chunk_size
self.logger.info(f"File length - {total_samples / 16000:.1f} seconds, splitting on {chunks_count} chunks by {chunk_length_s} seconds per chunk.")
self.logger.info(
f"File length - {total_samples / 16000:.1f} seconds, splitting on {chunks_count} chunks by {chunk_length_s} seconds per chunk."
)
self.chunks = []
for idx in tqdm(range(chunks_count)):
@@ -95,46 +103,51 @@ class AudioTranscription:
end = min((idx + 1) * self.chunk_size, total_samples)
chunk = self.waveform[start:end].cpu().numpy().astype("float32")
self.chunks.append(chunk)
def _resplit_to_batches(self) -> None:
self.logger.info(f"Splitting chunks into batches...")
self.batches = []
for i in range(0, len(self.chunks), self.custom_batch_length):
batch = self.chunks[i:i + self.custom_batch_length]
batch = self.chunks[i : i + self.custom_batch_length]
self.batches.append(batch)
def _process_all_batches(self) -> None:
start_time = time.time()
try:
self.all_transcription = []
for idx in tqdm(range(len(self.batches))):
inputs = self.processor(
self.batches[idx],
self.batches[idx],
sampling_rate=16000,
return_tensors="pt",
padding=True
return_tensors="pt",
padding=True,
)
input_features = inputs.input_features.to(self.device).to(self.torch_dtype)
input_features = inputs.input_features.to(self.device).to(
self.torch_dtype
)
with torch.no_grad():
predicted_ids = self.model.generate(
input_features,
language=self.language,
task="transcribe",
temperature=0.0
temperature=0.0,
)
texts = self.processor.batch_decode(predicted_ids, skip_special_tokens=True)
texts = self.processor.batch_decode(
predicted_ids, skip_special_tokens=True
)
self.all_transcription.extend(texts)
end_time = time.time()
self.logger.info(f"Transcription completed in {end_time - start_time:.2f} seconds")
self.logger.info(
f"Transcription completed in {end_time - start_time:.2f} seconds"
)
except Exception as e:
self.logger.error(f"Errors occured while processing chunks: {e}")
def transcribe_audio(self) -> str:
# TODO: maybe something else, not str?
self._resample()
@@ -142,4 +155,4 @@ class AudioTranscription:
self._split_to_chunks()
self._resplit_to_batches()
self._process_all_batches()
return " ".join(self.all_transcription)
return " ".join(self.all_transcription)
transcription/device_configuration.py
+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]
transcription/torch_checker.py
+19 -13
View File
@@ -1,32 +1,38 @@
import torch
from dataclasses import dataclass
import logging
from dataclasses import dataclass
import torch
def check_torch(logger: logging.Logger) -> None:
logger.info("=== Checking PyTorch ===")
logger.info(f"Torch version: {torch.__version__}")
# === NVIDIA / AMD (CUDA API) ===
# NVIDIA / AMD (CUDA API)
if torch.cuda.is_available():
backend = "CUDA"
if torch.version.hip is not None:
backend = "ROCm (AMD HIP)"
logger.info(f"{backend} backend is available")
logger.info(f"Compiled with: CUDA {torch.version.cuda}, ROCm {torch.version.hip}")
logger.info(
f"Compiled with: CUDA {torch.version.cuda}, ROCm {torch.version.hip}"
)
logger.info(f"Number of devices: {torch.cuda.device_count()}")
for i in range(torch.cuda.device_count()):
logger.info(f"GPU {i}: {torch.cuda.get_device_name(i)}")
# === Apple Silicon (MPS) ===
# Apple Silicon (MPS)
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
logger.info("MPS backend is available (Apple Silicon)")
logger.info(f"MPS version: {getattr(torch.backends.mps, '__version__', 'unknown')}")
logger.info(
f"MPS version: {getattr(torch.backends.mps, '__version__', 'unknown')}"
)
logger.info("GPU: Apple Silicon (Metal)")
# === CPU only mode ===
# CPU only mode
else:
logger.info("Only CPU is available")
logger.info("=== Check completed ===")