new structure for transcription engine
- found problems with mps - integrated jinja2 templates for rendering (.tex?) - raw ui & bad structure
This commit is contained in:
@@ -0,0 +1,17 @@
|
||||
from dataclasses import dataclass
|
||||
import torchaudio
|
||||
import torch
|
||||
|
||||
class Audio:
|
||||
waveform: torch.Tensor
|
||||
sr: int
|
||||
|
||||
def load(self, filepath):
|
||||
"""
|
||||
Loads audio from file's path
|
||||
"""
|
||||
self.waveform, self.sr = torchaudio.load(
|
||||
filepath,
|
||||
backend="ffmpeg",
|
||||
)
|
||||
return self
|
||||
@@ -1,246 +1,50 @@
|
||||
import logging
|
||||
import math
|
||||
import time
|
||||
import sys
|
||||
import gc
|
||||
from transcription.audio import Audio
|
||||
from transcription.preprocessing.audio_preprocessor import AudioPreprocessor
|
||||
from transcription.preprocessing.splitter import Splitter
|
||||
from transcription.engines.whisper import WhisperEngine
|
||||
from transcription.configuration import Configuration
|
||||
|
||||
import torch
|
||||
import torchaudio
|
||||
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
|
||||
# maybe inherit from AudioTranscription and rename to something like WhisperTranscription?
|
||||
class AudioTranscription:
|
||||
"""
|
||||
Class for automatical audio transcription using Whisper.
|
||||
|
||||
Provides audio file loading, resampling, conversion to mono,
|
||||
splitting into chunks and batches, running the model, and compiling the final text.
|
||||
|
||||
Attributes:
|
||||
model_name (str): Whisper model name in HuggingFace.
|
||||
filepath (str): Input filepath.
|
||||
waveform (torch.Tensor): Audiosignal in tensor form.
|
||||
sampling_rate (int): Input file's sampling frequency.
|
||||
chunks (list): Audio's chunks list.
|
||||
batches (list): Batches combined from chunks.
|
||||
chunk_size (int): Chunk size in samples.
|
||||
custom_chunk_length (int): Custom chunk length (in seconds).
|
||||
custom_batch_length (int): Custom batch length (in chunks).
|
||||
device (str): Inference device ("cuda", "cpu", "mps").
|
||||
processor (WhisperProcessor | None): Tokenizator/preprocessor.
|
||||
model (WhisperForConditionalGeneration | None): Whisper model.
|
||||
logger (logging.Logger): Logger.
|
||||
torch_dtype (torch.dtype): Data type for calculations (fp16/fp32).
|
||||
language (str): Transcription language (default "ru").
|
||||
all_transcription (list): List of strings with transcription.
|
||||
|
||||
Args:
|
||||
filepath (str): Input filepath.
|
||||
device_configuration (DeviceConfiguration): Device configuration
|
||||
(GPU/CPU/MPS, model, chunk length, batch etc.).
|
||||
logger (logging.Logger): Logger.
|
||||
language (str, optional): Transcription language. Default "ru".
|
||||
|
||||
Example:
|
||||
>>> from transcription.device_configuration import DeviceConfiguration
|
||||
>>> import logging
|
||||
>>> config = DeviceConfiguration(device="cuda", model_name="openai/whisper-large-v3-turbo") # recheck this, not true i think
|
||||
>>> logger = logging.getLogger("transcription")
|
||||
>>> transcriber = AudioTranscription("audio.wav", config, logger, language="en")
|
||||
>>> text = transcriber.transcribe_audio()
|
||||
>>> print(text)
|
||||
"This is a test audio transcription."
|
||||
|
||||
Methods:
|
||||
transcribe_audio() -> str:
|
||||
Starts full transcription pipeline: model loading,
|
||||
file loading and preprocessing, splitting into chunks/batches,
|
||||
inference and model unloading. Returns the final transcription.
|
||||
"""
|
||||
model_name = "openai/whisper-large-v3-turbo"
|
||||
|
||||
filepath: str
|
||||
waveform: torch.Tensor
|
||||
sampling_rate: int
|
||||
file_format: str
|
||||
|
||||
chunks: list = []
|
||||
batches: list = []
|
||||
chunk_size: int
|
||||
custom_chunk_length: int
|
||||
custom_batch_length: int
|
||||
|
||||
device = "cuda"
|
||||
processor: WhisperProcessor | None
|
||||
model: WhisperForConditionalGeneration | None
|
||||
logger: logging.Logger
|
||||
torch_dtype: torch.dtype
|
||||
|
||||
language = "ru"
|
||||
|
||||
all_transcription: list = []
|
||||
|
||||
# add multimodel ability
|
||||
def __init__(
|
||||
self,
|
||||
filepath: str,
|
||||
device_configuration: DeviceConfiguration,
|
||||
logger: logging.Logger,
|
||||
language: str = "ru",
|
||||
config: Configuration,
|
||||
language,
|
||||
# logger
|
||||
) -> None:
|
||||
# TODO: add pretty docs here
|
||||
self.filepath = filepath
|
||||
self.language = language
|
||||
self.logger = logger
|
||||
# self.logger = logger
|
||||
|
||||
# setting file extension
|
||||
self.file_format = filepath.split(".")[-1]
|
||||
|
||||
# extracting configuration
|
||||
self.device = device_configuration.device
|
||||
self.model_name = device_configuration.model_name
|
||||
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 = []
|
||||
|
||||
def _load_model(self) -> None:
|
||||
self.logger.info("Loading model WhisperProcessor...")
|
||||
try:
|
||||
start_time = time.time()
|
||||
|
||||
self.processor = WhisperProcessor.from_pretrained(self.model_name)
|
||||
self.model = WhisperForConditionalGeneration.from_pretrained(
|
||||
self.model_name, torch_dtype=self.torch_dtype
|
||||
).to(self.device)
|
||||
|
||||
end_time = time.time()
|
||||
self.logger.info(
|
||||
f"Model loaded successfully in {end_time - start_time:.2f} seconds."
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error while loading model: {e}")
|
||||
|
||||
def _unload_model(self) -> None:
|
||||
self.logger.info("Unloading model...")
|
||||
self.model = None
|
||||
self.processor = None
|
||||
if self.device == "cuda":
|
||||
torch.cuda.empty_cache()
|
||||
# TODO: maybe do something here for MPS
|
||||
self.logger.info("Model unloaded successfully.")
|
||||
|
||||
def _load_file(self) -> None:
|
||||
self.logger.info(f"Loading file {self.filepath}")
|
||||
try:
|
||||
# self.waveform, self.sampling_rate = torchaudio.load(
|
||||
# self.filepath, format=self.file_format, backend="ffmpeg"
|
||||
# )
|
||||
self.waveform, self.sampling_rate = torchaudio.load(self.filepath)
|
||||
self.logger.info(f"Successfully loaded file {self.filepath}.")
|
||||
except Exception as e:
|
||||
self.logger.error(f"Unable to load file {self.filepath}: {e}")
|
||||
# raise RuntimeError(f"Failed to load file {self.filepath}: {e}")
|
||||
|
||||
def _resample(self) -> None:
|
||||
self.waveform = torchaudio.functional.resample(
|
||||
self.waveform, self.sampling_rate, 16000
|
||||
self.audio = Audio()
|
||||
self.preprocessor = AudioPreprocessor()
|
||||
self.splitter = Splitter(
|
||||
chunkSize=config.chunkSize,
|
||||
batchSize=config.batchSize,
|
||||
)
|
||||
|
||||
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, shift: bool = False) -> None:
|
||||
self.logger.info(f"Splitting audio on chunks...")
|
||||
|
||||
self.chunk_size = self.custom_chunk_length * 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 {self.custom_chunk_length} seconds per chunk."
|
||||
self.engine = WhisperEngine(
|
||||
modelName=config.modelName,
|
||||
language=self.language,
|
||||
dType=config.dType,
|
||||
device=config.device,
|
||||
)
|
||||
|
||||
# maybe add something like temperature here?
|
||||
def transcribeAudio(self) -> str:
|
||||
transcription: list = []
|
||||
|
||||
self.engine.loadModel()
|
||||
|
||||
self.preprocessor.prepare(self.audio.load(self.filepath))
|
||||
|
||||
self.chunks = []
|
||||
for idx in tqdm(range(chunks_count)):
|
||||
start = idx * self.chunk_size
|
||||
end = min((idx + 1) * self.chunk_size, total_samples)
|
||||
chunk = self.waveform[start:end].cpu().numpy().astype("float32")
|
||||
self.chunks.append(chunk)
|
||||
batches = self.splitter.split(self.audio.waveform)
|
||||
|
||||
def _resplit_chunks_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]
|
||||
self.batches.append(batch)
|
||||
self.logger.info(f"Total: {len(self.batches)} batches, weight = {sys.getsizeof(self.batches)}")
|
||||
for batch in batches:
|
||||
batchText: str = self.engine.transcribeBatch(batch)
|
||||
transcription.append(batchText)
|
||||
|
||||
def _process_all_batches(self) -> None:
|
||||
start_time = time.time()
|
||||
try:
|
||||
assert self.processor is not None
|
||||
assert self.model is not None
|
||||
|
||||
self.all_transcription = []
|
||||
|
||||
for idx in tqdm(range(len(self.batches))):
|
||||
inputs = self.processor(
|
||||
self.batches[idx],
|
||||
sampling_rate=16000,
|
||||
return_tensors="pt",
|
||||
padding=True,
|
||||
)
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
texts = self.processor.batch_decode(
|
||||
predicted_ids, skip_special_tokens=True
|
||||
)
|
||||
self.all_transcription.extend(texts)
|
||||
|
||||
inputs = None
|
||||
input_features = None
|
||||
predicted_ids = None
|
||||
gc.collect()
|
||||
if self.device.startswith("cuda"):
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.ipc_collect()
|
||||
|
||||
end_time = time.time()
|
||||
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:
|
||||
self._load_model()
|
||||
self._load_file()
|
||||
self._resample()
|
||||
self._to_mono()
|
||||
self._split_to_chunks()
|
||||
self._resplit_chunks_to_batches()
|
||||
self._process_all_batches()
|
||||
self._unload_model()
|
||||
return " ".join(self.all_transcription)
|
||||
self.engine.unloadModel()
|
||||
|
||||
return str(" ".join(transcription))
|
||||
|
||||
@@ -0,0 +1,23 @@
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
|
||||
@dataclass
|
||||
class Configuration:
|
||||
# add new models
|
||||
device: str = "cuda"
|
||||
modelName: str = "openai/whisper-large-v2"
|
||||
chunkSize: int = 30
|
||||
batchSize: int = 16
|
||||
dataType: str = "torch.float16"
|
||||
|
||||
_dtype_map = {
|
||||
"torch.float16": torch.float16,
|
||||
"torch.float32": torch.float32,
|
||||
"torch.bfloat16": torch.bfloat16,
|
||||
}
|
||||
|
||||
dType: torch.dtype = None
|
||||
|
||||
def __post_init__(self):
|
||||
self.dType = self._dtype_map[self.dataType]
|
||||
@@ -1,47 +0,0 @@
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
@dataclass
|
||||
class DeviceConfiguration:
|
||||
"""
|
||||
Configurations for Whisper model on different devices.
|
||||
|
||||
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"
|
||||
- "openai/whisper-large-v3-turbo"
|
||||
|
||||
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
|
||||
- torch.bfloat16 - for GPUs which has BF16 support (usually RTX 40XX+)
|
||||
"""
|
||||
|
||||
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_dtype: torch.dtype = None
|
||||
|
||||
def __post_init__(self):
|
||||
self.torch_dtype = self._dtype_map[self.data_type]
|
||||
@@ -0,0 +1,30 @@
|
||||
import torch
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
class BaseEngine(ABC):
|
||||
def __init__(
|
||||
self,
|
||||
modelName: str,
|
||||
language: str,
|
||||
dType: torch.dtype,
|
||||
device: str
|
||||
):
|
||||
self.modelName = modelName
|
||||
self.device = device
|
||||
self.language = language
|
||||
self.dType = dType
|
||||
|
||||
@abstractmethod
|
||||
def loadModel(self) -> None:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def unloadModel(self) -> None:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def transcribeBatch(
|
||||
self,
|
||||
batch
|
||||
) -> str:
|
||||
pass
|
||||
@@ -0,0 +1,65 @@
|
||||
# from logging import Logger
|
||||
import time
|
||||
import torch
|
||||
import gc
|
||||
from transformers import WhisperForConditionalGeneration, WhisperProcessor
|
||||
|
||||
from transcription.engines.base_engine import BaseEngine
|
||||
|
||||
class WhisperEngine(BaseEngine):
|
||||
def loadModel(self) -> None:
|
||||
self.processor = WhisperProcessor.from_pretrained(self.modelName)
|
||||
self.model = WhisperForConditionalGeneration.from_pretrained(
|
||||
self.modelName,
|
||||
torch_dtype = self.dType # check twice
|
||||
).to(self.device) # ??? recheck
|
||||
|
||||
def unloadModel(self) -> None:
|
||||
self.model = None
|
||||
self.processor = None
|
||||
|
||||
# TODO: MPS?
|
||||
if self.device == "cuda":
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
def transcribeBatch(
|
||||
self,
|
||||
batch,
|
||||
) -> str:
|
||||
assert self.processor is not None
|
||||
assert self.model is not None
|
||||
|
||||
inputs = self.processor(
|
||||
batch,
|
||||
sampling_rate=16000,
|
||||
return_tensors="pt",
|
||||
padding=True,
|
||||
)
|
||||
|
||||
input_features = inputs.input_features.to(self.device).to(self.dType)
|
||||
|
||||
with torch.no_grad():
|
||||
predicted_ids = self.model.generate(
|
||||
input_features,
|
||||
language=self.language,
|
||||
task="transcribe",
|
||||
temperature=0.0,
|
||||
)
|
||||
|
||||
batchText = self.processor.batch_decode(
|
||||
predicted_ids,
|
||||
skip_special_tokens=True,
|
||||
)
|
||||
|
||||
inputs = None
|
||||
input_features = None
|
||||
predicted_ids = None
|
||||
gc.collect()
|
||||
|
||||
# maybe do here something with MPS?
|
||||
if self.device.startswith("cuda"):
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.ipc_collect()
|
||||
|
||||
return " ".join(batchText)
|
||||
|
||||
@@ -0,0 +1,35 @@
|
||||
from transcription.audio import Audio
|
||||
import torchaudio
|
||||
|
||||
class AudioPreprocessor:
|
||||
TARGET_SAMPLING_RATE: int = 16000
|
||||
|
||||
# for different models in future
|
||||
# def __init__(self, model):
|
||||
# pass
|
||||
|
||||
def _resample(
|
||||
self,
|
||||
audio: Audio
|
||||
) -> None:
|
||||
if audio.sr != self.TARGET_SAMPLING_RATE:
|
||||
audio.waveform = torchaudio.functional.resample(
|
||||
audio.waveform,
|
||||
audio.sr,
|
||||
self.TARGET_SAMPLING_RATE
|
||||
)
|
||||
|
||||
def _to_mono(
|
||||
self,
|
||||
audio: Audio
|
||||
) -> None:
|
||||
if audio.waveform.shape[0] > 1:
|
||||
audio.waveform = audio.waveform.mean(dim=0, keepdim=True)
|
||||
audio.waveform = audio.waveform.squeeze(0)
|
||||
|
||||
def prepare(
|
||||
self,
|
||||
audio: Audio
|
||||
):
|
||||
self._resample(audio)
|
||||
self._to_mono(audio)
|
||||
@@ -0,0 +1,50 @@
|
||||
import torch
|
||||
from typing import List
|
||||
|
||||
class Splitter:
|
||||
def __init__(
|
||||
self,
|
||||
chunkSize: int,
|
||||
batchSize: int,
|
||||
) -> None:
|
||||
self.chunkSize = chunkSize * 16000 # 16 kHz after resampling
|
||||
self.batchSize = batchSize
|
||||
|
||||
# maybe raise some exceptions here?
|
||||
def _split_to_chunks(
|
||||
self,
|
||||
waveform: torch.Tensor,
|
||||
) -> List:
|
||||
totalSamples = waveform.shape[0]
|
||||
chunksCount = (totalSamples + self.chunkSize - 1) // self.chunkSize
|
||||
|
||||
chunks: List = []
|
||||
# tqdm or something here?
|
||||
for chunkNum in range(chunksCount):
|
||||
start = chunkNum * self.chunkSize
|
||||
end = min((chunkNum + 1) * self.chunkSize, totalSamples)
|
||||
|
||||
chunk = waveform[start : end].cpu().numpy().astype("float32")
|
||||
chunks.append(chunk)
|
||||
|
||||
return chunks
|
||||
|
||||
def _split_to_batches(
|
||||
self,
|
||||
chunks: List,
|
||||
) -> List:
|
||||
batches: List = []
|
||||
|
||||
for i in range(0, len(chunks), self.batchSize):
|
||||
batch = chunks[i : i + self.batchSize]
|
||||
batches.append(batch)
|
||||
|
||||
return batches
|
||||
|
||||
def split(
|
||||
self,
|
||||
waveform: torch.Tensor
|
||||
) -> List:
|
||||
chunks = self._split_to_chunks(waveform)
|
||||
batches = self._split_to_batches(chunks)
|
||||
return batches
|
||||
@@ -1,10 +1,8 @@
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def check_torch(logger: logging.Logger) -> None:
|
||||
def checkTorch(logger: logging.Logger) -> None:
|
||||
logger.info("=== Checking PyTorch ===")
|
||||
logger.info(f"Torch version: {torch.__version__}")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user