diff --git a/.gitignore b/.gitignore index d256603..5d36044 100644 --- a/.gitignore +++ b/.gitignore @@ -10,4 +10,5 @@ wheels/ .venv main.todo -sample.mp3 \ No newline at end of file +sample.mp3 +*.spec \ No newline at end of file diff --git a/transcription/__init__.py b/transcription/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/transcription/audio_transcription.py b/transcription/audio_transcription.py index 6f450c9..27f967e 100644 --- a/transcription/audio_transcription.py +++ b/transcription/audio_transcription.py @@ -1,13 +1,14 @@ from transformers import WhisperProcessor, WhisperForConditionalGeneration import torch import torchaudio -from utils.logger import setup_logger +from ui.ui_log_handler import UILogHandler +from transcription.device_configuration import DeviceConfiguration +import logging import time import math from tqdm import tqdm -logger = setup_logger("AudioTranscribe module") - +# TODO: rename naming class AudioTranscription: model_name = "openai/whisper-large-v2" @@ -16,10 +17,16 @@ class AudioTranscription: sampling_rate: int chunks: list = [] + batches: list = [] 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" @@ -28,23 +35,31 @@ class AudioTranscription: def __init__( self, filepath: str, - language = "ru", - device = "cuda", - model_name = "openai/whisper-large-v2" + device_configuration: DeviceConfiguration, + logger: logging.Logger, + language = "ru" ) -> None: + # TODO: add pretty docs here self.filepath = filepath self.language = language - self.device = device - self.model_name = model_name + self.logger = logger + + # 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 = [] try: logger.info("Loading model WhisperProcessor...") self.processor = WhisperProcessor.from_pretrained(self.model_name) self.model = WhisperForConditionalGeneration.from_pretrained( self.model_name, - torch_dtype=torch.float16, - device_map="auto" + torch_dtype=self.torch_dtype ).to(self.device) logger.info("Model loaded.") @@ -65,58 +80,44 @@ class AudioTranscription: self.waveform = self.waveform.squeeze(0) def split_to_chunks(self, chunk_length_s: int = 30) -> None: - logger.info(f"Splitting audio on chunks...") + 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 - 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 range(chunks_count): + 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) - def process_chunk( - self, - chunk - ) -> str: - inputs = self.processor(chunk, sampling_rate=16000, return_tensors="pt") - input_features = inputs.input_features.to(self.device).to(torch.float16) - - with torch.no_grad(): - predicted_ids = self.model.generate( - input_features, - language=self.language, - task="transcribe", - temperature=0.0 - ) + 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] + self.batches.append(batch) - text = self.processor.batch_decode(predicted_ids, skip_special_tokens=True)[0] - - return text - def process_all_chunks(self, batch_size: int = 16) -> None: + def process_all_batches(self) -> None: start_time = time.time() try: self.all_transcription = [] - - for i in tqdm(range(math.ceil(len(self.chunks) / batch_size))): - # TODO: rewrite batching as a separate function - batch = self.chunks[i*batch_size:(i+1)*batch_size] - + + for idx in tqdm(range(len(self.batches))): inputs = self.processor( - batch, - sampling_rate=16000, + self.batches[idx], + sampling_rate=16000, return_tensors="pt", padding=True ) - - input_features = inputs.input_features.to(self.device).to(torch.float16) - + + input_features = inputs.input_features.to(self.device).to(self.torch_dtype) + with torch.no_grad(): predicted_ids = self.model.generate( input_features, @@ -124,20 +125,20 @@ class AudioTranscription: task="transcribe", temperature=0.0 ) - + texts = self.processor.batch_decode(predicted_ids, skip_special_tokens=True) self.all_transcription.extend(texts) - end_time = time.time() - 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: - logger.error(f"Errors occured while processing chunks: {e}") + self.logger.error(f"Errors occured while processing chunks: {e}") def transcribe_audio(self) -> str: self.resample() self.to_mono() self.split_to_chunks() - self.process_all_chunks() + self.resplit_to_batches() + self.process_all_batches() return " ".join(self.all_transcription) \ No newline at end of file diff --git a/transcription/device_configuration.py b/transcription/device_configuration.py new file mode 100644 index 0000000..2f405c2 --- /dev/null +++ b/transcription/device_configuration.py @@ -0,0 +1,43 @@ +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" + + 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_dtype: torch.dtype = None + + def __post_init__(self): + self.torch_dtype = self._dtype_map[self.data_type] \ No newline at end of file diff --git a/transcription/torch_checker.py b/transcription/torch_checker.py index 184d7e5..16e2919 100644 --- a/transcription/torch_checker.py +++ b/transcription/torch_checker.py @@ -1,11 +1,32 @@ import torch +from dataclasses import dataclass +import logging -def check_torch() -> None: - print("=== Checking PyTorch ===") - print(f"Torch version: {torch.version}") - print(f"CUDA is available: {torch.cuda.is_available()}") +def check_torch(logger: logging.Logger) -> None: + logger.info("=== Checking PyTorch ===") + logger.info(f"Torch version: {torch.__version__}") + + # === NVIDIA / AMD (CUDA API) === if torch.cuda.is_available(): - print(f"CUDA version: {torch.version.cuda}") - print(f"Number of GPU: {torch.cuda.device_count()}") - print(f"Name of GPU: {torch.cuda.get_device_name(0)}") - print("=== Check completed ===") \ No newline at end of file + 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"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) === + 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("GPU: Apple Silicon (Metal)") + + # === CPU only mode === + else: + logger.info("Only CPU is available") + + logger.info("=== Check completed ===") diff --git a/ui.py b/ui.py deleted file mode 100644 index 3bd2f69..0000000 --- a/ui.py +++ /dev/null @@ -1,13 +0,0 @@ -import tkinter as tk - -root = tk.Tk() -root.title("Audio Transcriptor") -root.geometry("800x600") - -path_entry = tk.Entry(root, width=40) -path_entry.pack(padx=10, pady=(5, 10)) - -transcribe_button = tk.Button(text="Transcribe") -transcribe_button.pack(anchor="e", rely=25, relx=15) - -root.mainloop() \ No newline at end of file diff --git a/ui/__init__.py b/ui/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/ui/ui.py b/ui/ui.py new file mode 100644 index 0000000..782e659 --- /dev/null +++ b/ui/ui.py @@ -0,0 +1,123 @@ +import tkinter as tk +from tkinter.scrolledtext import ScrolledText +from ui.ui_log_handler import UILogHandler, setup_ui_logger +from transcription.torch_checker import check_torch +from transcription.device_configuration import DeviceConfiguration +from transcription.audio_transcription import AudioTranscription + +def main(): + root = tk.Tk() + root.title("Audio Transcriptor") + root.geometry("800x600") + + for col in range(4): + root.grid_columnconfigure(col, weight=1) + root.grid_rowconfigure(6, weight=1) + + ### Buttons selector + check_torch_baton = tk.Button(root, text="Check Torch") + check_torch_baton.grid(row=0, column=0, padx=5, pady=5, sticky="ew") + + # TODO: implement saving function + save_configuration_baton = tk.Button(root, text="Save configuration") + save_configuration_baton.grid(row=0, column=1, padx=5, pady=5, sticky="ew") + + # TODO: implement deleting function + delete_configuration_baton = tk.Button(root, text="Delete configuration") + delete_configuration_baton.grid(row=0, column=2, padx=5, pady=5, sticky="ew") + + start_transcription_baton = tk.Button(root, text="Transcript") + start_transcription_baton.grid(row=0, column=3, padx=5, pady=5, sticky="ew") + + ### Model options selector + model_options = [ + "openai/whisper-large-v2", + "openai/whisper-large", + "openai/whisper-medium", + "openai/whisper-small", + "openai/whisper-tiny" + ] + selected_model = tk.StringVar(value=model_options[0]) + label_model = tk.Label(root, text="Model name:") + label_model.grid(row=1, column=0, sticky="w", pady=5, padx=5) + dropdown_model_selection = tk.OptionMenu(root, selected_model, *model_options) + dropdown_model_selection.grid(row=1, column=1, sticky="ew", pady=5, padx=5) + + ### Batch size selector + batch_sizes = ["32", "16", "8", "4", "2"] + selected_batch_size = tk.StringVar(value=batch_sizes[0]) + label_batch_size = tk.Label(root, text="Batch size:") + label_batch_size.grid(row=1, column=2, sticky="w", pady=5, padx=5) + dropdown_batch_size_selection = tk.OptionMenu(root, selected_batch_size, *batch_sizes) + dropdown_batch_size_selection.grid(row=1, column=3, sticky="ew", pady=5, padx=5) + + ### Data type selector + data_types = ["torch.float16", "torch.float32", "torch.bfloat16"] + selected_data_type = tk.StringVar(value=data_types[0]) + label_data_type = tk.Label(root, text="Data type:") + label_data_type.grid(row=2, column=0, sticky="w", pady=5, padx=5) + dropdown_data_type_selection = tk.OptionMenu(root, selected_data_type, *data_types) + dropdown_data_type_selection.grid(row=2, column=1, sticky="ew", pady=5, padx=5) + + ### Chunk length selector + chunk_lengths = ["30", "25", "20", "15", "10", "5"] + selected_chunk_length = tk.StringVar(value=chunk_lengths[0]) + label_chunk_length = tk.Label(root, text="Chunk length:") + label_chunk_length.grid(row=2, column=2, sticky="w", pady=5, padx=5) + dropdown_chunk_length_selection = tk.OptionMenu(root, selected_chunk_length, *chunk_lengths) + dropdown_chunk_length_selection.grid(row=2, column=3, sticky="ew", pady=5, padx=5) + + # TODO: add device selector (cuda/mps/cpu) + + ### Filepath (input) + # TODO: add checker if path is valid/invalid (i think in utils or something) + label_file_path = tk.Label(root, text="Input filepath:") + label_file_path.grid(row=3, column=0, sticky="w", pady=5, padx=5) + file_path = tk.Text(root, height=1) + file_path.grid(row=3, column=1, columnspan=3, sticky="ew", pady=5, padx=5) + + ### Filepath (output) + # TODO: add question mark here with tip while mouse is on it + label_output_file_path = tk.Label(root, text="Output filepath:") + label_output_file_path.grid(row=4, column=0, sticky="w", pady=5, padx=5) + output_file_path = tk.Text(root, height=1) + output_file_path.grid(row=4, column=1, columnspan=3, sticky="ew", pady=5, padx=5) + + def show_selections(): + ui_logger.info(f"Selected model: {selected_model.get()}") + ui_logger.info(f"Selected batch size: {selected_batch_size.get()} chunks") + ui_logger.info(f"Selected data type: {selected_data_type.get()}") + + show_selections_baton = tk.Button(root, text="Show Selections", command=show_selections) + show_selections_baton.grid(row=5, column=0, columnspan=4, pady=5, sticky="ew") + + log_box = ScrolledText(root, wrap="word") + log_box.grid(row=6, column=0, columnspan=4, sticky="nsew", padx=10, pady=5) + ui_logger = setup_ui_logger(log_box) + + def transcribe(): + current_device_config = DeviceConfiguration( + device="cuda", + model_name=selected_model.get(), + batch_size=int(selected_batch_size.get()), + chunk_length_s=30, + data_type=selected_data_type.get() + ) + Audio = AudioTranscription( + filepath=file_path.get("1.0", "end-1c"), + device_configuration=current_device_config, + logger=ui_logger + ) + transcription = Audio.transcribe_audio() + with open(f"{file_path.get('1.0', 'end-1c')}.txt", "w") as output_file: + output_file.write(transcription) + + check_torch_baton.config(command=lambda: check_torch(ui_logger)) + start_transcription_baton.config(command=transcribe) + + root.mainloop() + + + +if __name__ == "__main__": + main() diff --git a/ui/ui_log_handler.py b/ui/ui_log_handler.py new file mode 100644 index 0000000..60e03d5 --- /dev/null +++ b/ui/ui_log_handler.py @@ -0,0 +1,24 @@ +import tkinter as tk +import logging + +class UILogHandler(logging.Handler): + def __init__(self, text_widget: tk.Text): + super().__init__() + self.text_widget = text_widget + + def emit(self, record): + log_entry = self.format(record) + self.text_widget.insert(tk.END, log_entry + "\n") + self.text_widget.see(tk.END) + +# TODO: maybe some tqdm here, not in console? +def setup_ui_logger(text_widget: tk.Text, level=logging.INFO): + logger = logging.getLogger("UI_LOGGER") + logger.setLevel(level) + logger.handlers.clear() + + handler = UILogHandler(text_widget) + handler.setFormatter(logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")) + logger.addHandler(handler) + + return logger \ No newline at end of file diff --git a/utils/logger.py b/utils/logger.py deleted file mode 100644 index 2bf3c0e..0000000 --- a/utils/logger.py +++ /dev/null @@ -1,36 +0,0 @@ -import logging -import sys - -def setup_logger( - name: str = __name__, - level: int = logging.INFO, - log_to_file: bool = False, - filename: str = "app.log" -) -> logging.Logger: - """ - Logger configuration with output to command line and (optional) to file - - :param name: logger name - :param level: logging level (DEBUG, INFO, WARNING, ERROR, CRITICAL) - :param log_to_file: logging to filename - :param filename: filename for logs - :return: logger instance - """ - logger = logging.getLogger(name) - logger.setLevel(level) - - formatter = logging.Formatter( - fmt="%(asctime)s | %(levelname)-8s | %(name)s | %(message)s", - datefmt="%Y-%m-%d %H:%M:%S" - ) - - console_handler = logging.StreamHandler(sys.stdout) - console_handler.setFormatter(formatter) - logger.addHandler(console_handler) - - if log_to_file: - file_handler = logging.FileHandler(filename, encoding="utf-8") - file_handler.setFormatter(formatter) - logger.addHandler(file_handler) - - return logger \ No newline at end of file