diff --git a/.gitignore b/.gitignore index 5d36044..7ca39b3 100644 --- a/.gitignore +++ b/.gitignore @@ -11,4 +11,7 @@ wheels/ main.todo sample.mp3 -*.spec \ No newline at end of file +*.spec + +# лютый завоз +ui/assets/question_mark.png \ No newline at end of file diff --git a/README.md b/README.md index d71ed6d..3ed4fc9 100644 --- a/README.md +++ b/README.md @@ -11,5 +11,5 @@ conda env create -f environment.yml or for CUDA/Nvidia: ``` conda create -n notecast -c pytorch -c nvidia pytorch torchvision torchaudio transformers python=3.12 -conda install ffmpeg customtkinter -c conda-forge -c bioconda +conda install ffmpeg customtkinter openai -c conda-forge -c bioconda ``` \ No newline at end of file diff --git a/env_cuda.yml b/env_cuda.yml index bcd6f71..5879e5b 100644 --- a/env_cuda.yml +++ b/env_cuda.yml @@ -6,9 +6,10 @@ channels: - bioconda dependencies: - pytorch + - ffmpeg - torchvision - torchaudio - transformers - python=3.12 - - ffmpeg - - customtkinter \ No newline at end of file + - customtkinter + - openai \ No newline at end of file diff --git a/make_venv.ps1 b/make_venv.ps1 deleted file mode 100644 index b6b327a..0000000 --- a/make_venv.ps1 +++ /dev/null @@ -1,13 +0,0 @@ -$envName = "notecast" - -Write-Output >>> Creating environment $envName from environment.yml" -conda env create -f environment.yml -n $envName -if ($LASTEXITCODE -ne 0) { - Write-Output ">>> Environment already exists, updating..." - conda env update -f environment.yml --prune -} - -Write-Output ">>> Activating environment" -conda activate $envName - -Write-Output ">>> Download completed!" \ No newline at end of file diff --git a/make_venv.sh b/make_venv.sh deleted file mode 100644 index 7bd82e8..0000000 --- a/make_venv.sh +++ /dev/null @@ -1,13 +0,0 @@ -#!/bin/bash -set -e - -ENV_NAME="notecast" - -echo ">>> Creating environment $ENV_NAME from environment.yml" -conda env create -f environment.yml || conda env update -f environment.yml --prune - -echo ">>> Activating environment" -eval "$(conda shell.bash hook)" -conda activate "$ENV_NAME" - -echo ">>> Download completed!" \ No newline at end of file diff --git a/transcription/audio_transcription.py b/transcription/audio_transcription.py index 7bc0238..8415b49 100644 --- a/transcription/audio_transcription.py +++ b/transcription/audio_transcription.py @@ -1,6 +1,8 @@ import logging import math import time +import sys +import gc import torch import torchaudio @@ -27,8 +29,8 @@ class AudioTranscription: custom_batch_length: int device = "cuda" - processor: WhisperProcessor - model: WhisperForConditionalGeneration + processor: WhisperProcessor | None + model: WhisperForConditionalGeneration | None logger: logging.Logger torch_dtype: torch.dtype @@ -41,13 +43,13 @@ class AudioTranscription: filepath: str, device_configuration: DeviceConfiguration, logger: logging.Logger, - language="ru", + language: str = "ru", ) -> None: # TODO: add pretty docs here self.filepath = filepath self.language = language self.logger = logger - + # setting file extension self.file_format = filepath.split(".")[-1] @@ -61,19 +63,33 @@ class AudioTranscription: self.chunks: list = [] self.batches: list = [] self.all_transcription: list = [] - + def _load_model(self) -> None: self.logger.info("Loading model WhisperProcessor...") - try: + 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 + self.model_name, torch_dtype=self.torch_dtype ).to(self.device) - - self.logger.info("Model loaded.") + + 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: @@ -81,7 +97,7 @@ class AudioTranscription: self.filepath, format=self.file_format, backend="ffmpeg" ) self.logger.info(f"Successfully loaded file {self.filepath}.") - except Exception as e: + except Exception as e: self.logger.error(f"Unable to load file {self.filepath}: {e}") def _resample(self) -> None: @@ -118,11 +134,14 @@ class AudioTranscription: 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") + self.logger.info(f"Total: {len(self.batches)} batches, weight = {sys.getsizeof(self.batches)}") 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))): @@ -149,6 +168,15 @@ class AudioTranscription: 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" @@ -165,4 +193,5 @@ class AudioTranscription: self._split_to_chunks() self._resplit_chunks_to_batches() self._process_all_batches() + self._unload_model() return " ".join(self.all_transcription) diff --git a/ui/assets/logo.png b/ui/assets/logo.png new file mode 100644 index 0000000..acfa38e Binary files /dev/null and b/ui/assets/logo.png differ diff --git a/ui/tooltip.py b/ui/tooltip.py index 8af3341..46eb0ba 100644 --- a/ui/tooltip.py +++ b/ui/tooltip.py @@ -1,5 +1,6 @@ import customtkinter as ctk + class ToolTip(ctk.CTkToplevel): def __init__(self, widget, text): super().__init__() @@ -7,7 +8,9 @@ class ToolTip(ctk.CTkToplevel): self.overrideredirect(True) self.attributes("-topmost", True) - self.label = ctk.CTkLabel(self, text=text, fg_color="gray20", corner_radius=6, padx=10, pady=5) + self.label = ctk.CTkLabel( + self, text=text, fg_color="gray20", corner_radius=6, padx=10, pady=5 + ) self.label.pack() self.widget = widget @@ -21,4 +24,4 @@ class ToolTip(ctk.CTkToplevel): self.deiconify() def hide_tooltip(self, event=None): - self.withdraw() \ No newline at end of file + self.withdraw() diff --git a/ui/ui.py b/ui/ui.py index 69a3ece..d20808e 100644 --- a/ui/ui.py +++ b/ui/ui.py @@ -1,8 +1,11 @@ import os import queue +import sys +import json import threading import tkinter as tk from tkinter import filedialog, messagebox, scrolledtext +from PIL import Image, ImageTk import customtkinter as ctk import torch @@ -10,11 +13,12 @@ import torch from transcription.audio_transcription import AudioTranscription from transcription.device_configuration import DeviceConfiguration from transcription.torch_checker import check_torch -from ui.ui_log_handler import setup_ui_logger from ui.tooltip import ToolTip +from ui.ui_log_handler import setup_ui_logger +from utils.requests_to_api import LLMrequest -WINDOW_WIDTH = 900 -WINDOW_HEIGHT = 650 +WINDOW_WIDTH = 1000 +WINDOW_HEIGHT = 725 class TranscriberApp(ctk.CTk): @@ -25,16 +29,60 @@ class TranscriberApp(ctk.CTk): ctk.set_appearance_mode("System") ctk.set_default_color_theme("blue") + # TODO: fix this stuff + base_path = getattr(sys, "_MEIPASS", os.path.dirname(os.path.abspath(__file__))) + icon_path = os.path.join(base_path, "assets", "logo.png") + + applied_method = None + + try: + self.iconbitmap(icon_path) + applied_method = "iconbitmap" + except Exception as e: + print("iconbitmap did not work:", e) + + if applied_method is None: + try: + img = tk.PhotoImage(file=icon_path) + self.iconphoto(False, img) + applied_method = "iconphoto" + except Exception as e: + print("iconphoto did not work:", e) + + if applied_method is None: + try: + pil = Image.open(icon_path) + photo = ImageTk.PhotoImage(pil) + self.tk.call('wm', 'iconphoto', self._w, photo) + applied_method = "tk.call" + except Exception as e: + print("wm iconphoto via tk.call did not work:", e) + + if applied_method: + print(f"Applied method: {applied_method}") + else: + print("Can't set an icon.") + # states self.progress_queue = queue.Queue() self.transcribe_thread = None self.stop_flag = threading.Event() - # user variables - self.model_var = tk.StringVar(value="openai/whisper-large-v2") + # USER VARIABLES + # transcription model settings + self.model_var = tk.StringVar(value="openai/whisper-large-v3-turbo") self.batch_var = tk.StringVar(value="32") - self.dtype_var = tk.StringVar(value="torch.float16") self.chunk_var = tk.StringVar(value="30") + self.dtype_var = tk.StringVar(value="torch.float16") + self.transcription_lang_var = tk.StringVar(value="ru") + # llm settings + self.conspect_transcription_lang_var = tk.StringVar(value="Russian") + self.api_key_var = tk.StringVar(value="") + self.base_url_var = tk.StringVar(value="") + self.api_model_var = tk.StringVar(value="") + # checkboxes + self.create_conspect = tk.BooleanVar(value=False) + self.remove_transcription = tk.BooleanVar(value=False) # settings device options device_opts = [] @@ -62,7 +110,7 @@ class TranscriberApp(ctk.CTk): # logger self.ui_logger = setup_ui_logger(self.log_box) - # main transcription tab + ### TRANSCRIPTION TAB def _build_transcription_tab(self): # file selectors file_frame = ctk.CTkFrame(self.transcript_tab, corner_radius=10) @@ -93,16 +141,34 @@ class TranscriberApp(ctk.CTk): ctk.CTkButton(ctrl_frame, text="Check Torch", command=self._check_torch).pack( side="left", padx=10, pady=5 ) - self.start_btn = ctk.CTkButton( - ctrl_frame, text="Start transcription", command=self._start_transcription + self.start_button = ctk.CTkButton( + ctrl_frame, text="Start", command=self._start_transcription ) - self.start_btn.pack(side="left", padx=10, pady=5) + self.start_button.pack(side="left", padx=10, pady=5) - self.stop_btn = ctk.CTkButton( + self.stop_button = ctk.CTkButton( ctrl_frame, text="Stop", command=self._stop_transcription, state="disabled" ) - self.stop_btn.pack(side="left", padx=10, pady=5) + self.stop_button.pack(side="left", padx=10, pady=5) + + self.create_conspect_checkbox = ctk.CTkCheckBox( + ctrl_frame, + text="Create conspect", + variable=self.create_conspect, + onvalue=True, + offvalue=False, + ) + self.create_conspect_checkbox.pack(side="left", padx=10, pady=5) + self.remove_transcription_checkbox = ctk.CTkCheckBox( + ctrl_frame, + text="Remove transcription file after", + variable=self.remove_transcription, + onvalue=True, + offvalue=False, + ) + self.remove_transcription_checkbox.pack(side="left", padx=10, pady=5) + # TODO: add unload model button here # log box @@ -114,65 +180,181 @@ class TranscriberApp(ctk.CTk): ) self.log_box.pack(padx=20, pady=10, expand=True, fill="both") + ### SETTINGS TAB def _build_settings_tab(self): - # TODO: add tooltips here pad = 20 - ### Model - ctk.CTkLabel(self.settings_tab, text="Model:").pack( - anchor="w", padx=pad, pady=(pad, 5) - ) - ctk.CTkOptionMenu( - self.settings_tab, + def add_setting(parent, row, col, text, tooltip, variable, values: list | None): + frame = ctk.CTkFrame(parent) + frame.grid(row=row, column=col, padx=pad, pady=(pad, 5), sticky="nsew") + + label = ctk.CTkLabel(frame, text=text) + label.grid(row=0, column=0, sticky="w") + help_icon = ctk.CTkLabel(frame, text="?", width=20, cursor="question_arrow") + help_icon.grid(row=0, column=1, sticky="w", padx=(5, 0)) + ToolTip(help_icon, tooltip) + + if values: + ctk.CTkOptionMenu(frame, variable=variable, values=values).grid( + row=1, column=0, columnspan=2, sticky="ew", pady=(5, 0) + ) + else: + ctk.CTkEntry(frame, textvariable=variable).grid( + row=1, column=0, columnspan=2, sticky="ew", pady=(5, 0) + ) + + frame.grid_columnconfigure(0, weight=1) + + grid = ctk.CTkFrame(self.settings_tab) + grid.pack(fill="both", expand=True) + + # first row + ### Model setting + add_setting( + parent=grid, + row=0, + col=0, + text="Model:", + tooltip="Choose model for speed recognition", variable=self.model_var, values=[ + # maybe delete this option? + "openai/whisper-large-v3-turbo", "openai/whisper-large-v2", "openai/whisper-large", "openai/whisper-medium", "openai/whisper-small", "openai/whisper-tiny", ], - ).pack(fill="x", padx=pad, pady=5) - - ### Batch size - ctk.CTkLabel(self.settings_tab, text="Batch size:").pack( - anchor="w", padx=pad, pady=(pad, 5) ) - ctk.CTkOptionMenu( - self.settings_tab, + ### Batch size setting + add_setting( + parent=grid, + row=0, + col=1, + text="Batch size:", + tooltip="Chunks count for one iteration", variable=self.batch_var, values=["32", "16", "8", "4", "2"], - ).pack(fill="x", padx=pad, pady=5) - - ### Data type - ctk.CTkLabel(self.settings_tab, text="Data type:").pack( - anchor="w", padx=pad, pady=(pad, 5) ) - ctk.CTkOptionMenu( - self.settings_tab, + + # second row + ### Data type setting + add_setting( + parent=grid, + row=1, + col=0, + text="Data type:", + tooltip="Data type for calculations", variable=self.dtype_var, values=["torch.float16", "torch.float32", "torch.bfloat16"], - ).pack(fill="x", padx=pad, pady=5) - - ### Chunk length - ctk.CTkLabel(self.settings_tab, text="Chunk length (s):").pack( - anchor="w", padx=pad, pady=(pad, 5) ) - ctk.CTkOptionMenu( - self.settings_tab, + ### Chunk Length setting + add_setting( + parent=grid, + row=1, + col=1, + text="Chunk length (s):", + tooltip="Maximum length of processing audio fragment", variable=self.chunk_var, values=["30", "25", "20", "15", "10", "5"], - ).pack(fill="x", padx=pad, pady=5) - - ### Device - ctk.CTkLabel(self.settings_tab, text="Device:").pack( - anchor="w", padx=pad, pady=(pad, 5) ) - ctk.CTkOptionMenu( - self.settings_tab, + + # third row + ### Device setting + add_setting( + parent=grid, + row=2, + col=0, + text="Device:", + tooltip="Choose device\n- CUDA for CUDA & ROCm\n- MPS for Apple Silicon \n- CPU for CPU-only mode", variable=self.device_var, values=self.device_opts, - ).pack(fill="x", padx=pad, pady=5) + ) + ### Transcription language setting + add_setting( + parent=grid, + row=2, + col=1, + text="Transcription language:", + tooltip="Choose the transcription language", + variable=self.transcription_lang_var, + values=["ru", "en"], + ) + + # fourth row + ### OpenAI API key setting + add_setting( + parent=grid, + row=3, + col=0, + text="Insert OpenAI API key here:", + tooltip="Give this programm access to LLM that would create a fully prepared conspect with AI overviews", + variable=self.api_key_var, + values=None, + ) + ### Model name setting + add_setting( + parent=grid, + row=3, + col=1, + text="Model name:", + tooltip="Name of the model that you are going to use", + variable=self.api_model_var, + values=None, + ) + + # fifth row + ### Base URL setting + add_setting( + parent=grid, + row=4, + col=0, + text="Base URL:", + tooltip="OpenAI base URL. Blank for None.", + variable=self.base_url_var, + values=None, + ) + ### Output (conspect) language setting + add_setting( + parent=grid, + row=4, + col=1, + text="Conspect language:", + tooltip="Conspect language. Blank for English (default)", + variable=self.conspect_transcription_lang_var, + values=None, + ) + + ### Custom Prompt setting + customPromptFrame = ctk.CTkFrame(grid) + customPromptFrame.grid( + row=5, column=0, columnspan=2, padx=20, pady=(20, 5), sticky="nsew" + ) + + label = ctk.CTkLabel(customPromptFrame, text="Custom Prompt:") + label.grid(row=0, column=0, sticky="w") + + help_icon = ctk.CTkLabel( + customPromptFrame, text="?", width=20, cursor="question_arrow" + ) + help_icon.grid(row=0, column=1, sticky="w", padx=(5, 0)) + ToolTip( + help_icon, + "Enter your custom prompt for model.", + ) + + self.custom_prompt_textbox = ctk.CTkTextbox( + customPromptFrame, width=400, height=150 + ) + self.custom_prompt_textbox.grid( + row=1, column=0, columnspan=2, sticky="nsew", pady=(5, 0) + ) + + customPromptFrame.grid_columnconfigure(0, weight=1) + customPromptFrame.grid_rowconfigure(1, weight=1) + + grid.grid_columnconfigure((0, 1), weight=1) # action buttons def _browse_input(self): @@ -183,7 +365,7 @@ class TranscriberApp(ctk.CTk): ("Audio files", "*.wav *.mp3 *.m4a *.flac *.ogg"), ("Video files", "*.mp4 *.mkv *.avi"), ("All files", "*.*"), - ] + ], ) if path: self.input_file_var.set(path) @@ -204,6 +386,26 @@ class TranscriberApp(ctk.CTk): def _check_torch(self): check_torch(self.ui_logger) + + self.ui_logger.info(f"==== Transcription ====") + self.ui_logger.info(f"Transcription model: {self.model_var.get()}") + self.ui_logger.info(f"Batch size (in chunks): {self.batch_var.get()}") + self.ui_logger.info(f"Chunk size (in seconds): {self.chunk_var.get()}") + self.ui_logger.info(f"Data type: {self.dtype_var.get()}") + self.ui_logger.info(f"Transcription language: {self.transcription_lang_var.get()}") + self.ui_logger.info(f"=======================") + + debug_api_key_var: str = self.api_key_var.get() if self.api_key_var.get() else "Not set" + debug_api_model_var: str = self.api_model_var.get() if self.api_model_var.get() else "Not set" + debug_base_url_var: str = self.base_url_var.get() if self.base_url_var.get() else "Not set" + debug_transcription_lang_var = self.transcription_lang_var.get() if self.transcription_lang_var.get() else "Not set" + + self.ui_logger.info(f"===== LLM Setting =====") + self.ui_logger.info(f"API key: {debug_api_key_var}") + self.ui_logger.info(f"Model name setting: {debug_api_model_var}") + self.ui_logger.info(f"Base URL setting: {debug_base_url_var}") + self.ui_logger.info(f"=======================") + def _start_transcription(self): infile = self.input_file_var.get().strip() @@ -211,8 +413,8 @@ class TranscriberApp(ctk.CTk): messagebox.showerror("Error", "Please select a valid input file.") return - self.start_btn.configure(state="disabled") - self.stop_btn.configure(state="normal") + self.start_button.configure(state="disabled") + self.stop_button.configure(state="normal") self.ui_logger.info("Starting transcription...") self.stop_flag.clear() @@ -221,10 +423,12 @@ class TranscriberApp(ctk.CTk): ) self.transcribe_thread.start() - def _stop_transcription(self): + def _stop_transcription(self, Audio: AudioTranscription): self.stop_flag.set() self.ui_logger.info("Stopping transcription...") - self.stop_btn.configure(state="disabled") + self.ui_logger.info("Unloading model...") + Audio._unload_model() + self.stop_button.configure(state="disabled") def _transcribe_worker(self, infile: str): try: @@ -239,20 +443,44 @@ class TranscriberApp(ctk.CTk): filepath=infile, device_configuration=config, logger=self.ui_logger, + language=self.transcription_lang_var.get(), ) transcription = Audio.transcribe_audio() - + outfile = self.output_file_var.get().strip() if not outfile: - outfile = infile + ".txt" + outfile = infile + + if not self.remove_transcription.get(): + outfile += ".txt" - with open(outfile, "w", encoding="utf-8") as f: - f.write(transcription) + with open(outfile, "w", encoding="utf-8") as f: + f.write(transcription) + self.ui_logger.info(f"Transcription saved to {outfile}.") + + if self.create_conspect.get(): + # TODO: add custom prompt ability here + # TODO: add logging here + self.ui_logger.info(f"Starting creating conspect via {self.api_model_var.get()}...") + with open("utils/default_prompt.json", "r", encoding="utf-8") as f: + default_prompt = json.load(f)["prompt"] + + prompt = transcription + default_prompt # if custom prompt is empty do something + request = LLMrequest( + api_key=self.api_key_var.get(), + model_name=self.api_model_var.get(), + base_url=self.base_url_var.get(), + ) + response = request.get_response(prompt=prompt) + outfile += ".md" + with open(outfile, "w", encoding="utf-8") as f: + f.write(response) + + self.ui_logger.info(f"Conspect saved to {outfile}.") - self.ui_logger.info(f"Transcription saved to {outfile}") except Exception as e: self.ui_logger.error(f"Error: {e}") messagebox.showerror("Error", str(e)) finally: - self.start_btn.configure(state="normal") - self.stop_btn.configure(state="disabled") + self.start_button.configure(state="normal") + self.stop_button.configure(state="disabled") diff --git a/utils/default_prompt.json b/utils/default_prompt.json new file mode 100644 index 0000000..057ca36 --- /dev/null +++ b/utils/default_prompt.json @@ -0,0 +1,3 @@ +{ + "prompt": "Используй ТОЛЬКО текст из транскрибации. Любая информация, которой нет в транскрибации, запрещена. Если какой-то факт или определение раскрыт не до конца, или лектор сказал что-то неполно, или тему нужно дополнить для лучшего понимания — вставляй дополнение сразу в нужное место, но только под отдельной подписью: AI overview: текст дополнения (confidence: low|medium|high) Ответ должен быть строго в формате Markdown (.md). Никаких 'привет', 'вот результат' и т. п. — только готовый текст. Все цитаты из транскрибации — буквально, максимум 25 слов подряд." +} \ No newline at end of file diff --git a/utils/requests_to_api.py b/utils/requests_to_api.py new file mode 100644 index 0000000..a3d9ab3 --- /dev/null +++ b/utils/requests_to_api.py @@ -0,0 +1,17 @@ +import openai + + +class LLMrequest: + def __init__(self, api_key: str, model_name: str, base_url: str = None): + if base_url: + self.client = openai.OpenAI(api_key=api_key, base_url=base_url) + else: + self.client = openai.OpenAI(api_key=api_key) + self.model = model_name + + def get_response(self, prompt: str) -> str: + response = self.client.chat.completions.create( + model=self.model, + messages=[{"role": "user", "content": prompt}], + ) + return response.choices[0].message.content diff --git a/utils/save_to_file.py b/utils/save_to_file.py new file mode 100644 index 0000000..6776752 --- /dev/null +++ b/utils/save_to_file.py @@ -0,0 +1,9 @@ +def save_to_file( + thing: str | dict, + path: str, +): + if isinstance(thing, dict): + pass + elif isinstance(thing, str): + with open(path, "w") as outputfile: + outputfile.write(thing)