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

WIP - first working version with UI, hooray!

Browse Source 
- added logging in UI
- added switcher for models, batch size etc.
- added device configuration dataclass
- minor improvements in audio transcription
This commit is contained in:
German Mikheev
2025-09-09 00:49:21 +03:00
parent ecbbbfd5d1
commit 8f41105e4b
10 changed files with 268 additions and 104 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+1
View File
@@ -11,3 +11,4 @@ wheels/
main.todo
sample.mp3
*.spec
transcription/__init__.py
View File
transcription/audio_transcription.py
+42 -41
View File
@@ -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,57 +80,43 @@ 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)
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)
with torch.no_grad():
predicted_ids = self.model.generate(
input_features,
language=self.language,
task="transcribe",
temperature=0.0
)
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,
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(
@@ -127,17 +128,17 @@ class AudioTranscription:
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)
transcription/device_configuration.py
+43
View File
@@ -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]
transcription/torch_checker.py
+29 -8
View File
@@ -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 ===")
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 ===")
ui.py
-13
View File
@@ -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()
ui/__init__.py
View File
ui/ui.py
+123
View File
@@ -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()
ui/ui_log_handler.py
+24
View File
@@ -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
utils/logger.py
-36
View File
@@ -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