svlqd
8f41105e4b
- added logging in UI - added switcher for models, batch size etc. - added device configuration dataclass - minor improvements in audio transcription
124 lines
5.3 KiB
Python
124 lines
5.3 KiB
Python
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()
|