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
|
from transcription.audio import Audio
|
||||||
import math
|
from transcription.preprocessing.audio_preprocessor import AudioPreprocessor
|
||||||
import time
|
from transcription.preprocessing.splitter import Splitter
|
||||||
import sys
|
from transcription.engines.whisper import WhisperEngine
|
||||||
import gc
|
from transcription.configuration import Configuration
|
||||||
|
|
||||||
import torch
|
# maybe inherit from AudioTranscription and rename to something like WhisperTranscription?
|
||||||
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
|
|
||||||
class AudioTranscription:
|
class AudioTranscription:
|
||||||
"""
|
# add multimodel ability
|
||||||
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 = []
|
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
filepath: str,
|
filepath: str,
|
||||||
device_configuration: DeviceConfiguration,
|
config: Configuration,
|
||||||
logger: logging.Logger,
|
language,
|
||||||
language: str = "ru",
|
# logger
|
||||||
) -> None:
|
) -> None:
|
||||||
# TODO: add pretty docs here
|
|
||||||
self.filepath = filepath
|
self.filepath = filepath
|
||||||
self.language = language
|
self.language = language
|
||||||
self.logger = logger
|
# self.logger = logger
|
||||||
|
|
||||||
# setting file extension
|
self.audio = Audio()
|
||||||
self.file_format = filepath.split(".")[-1]
|
self.preprocessor = AudioPreprocessor()
|
||||||
|
self.splitter = Splitter(
|
||||||
# extracting configuration
|
chunkSize=config.chunkSize,
|
||||||
self.device = device_configuration.device
|
batchSize=config.batchSize,
|
||||||
self.model_name = device_configuration.model_name
|
)
|
||||||
self.custom_chunk_length = device_configuration.chunk_length_s
|
self.engine = WhisperEngine(
|
||||||
self.custom_batch_length = device_configuration.batch_size
|
modelName=config.modelName,
|
||||||
self.torch_dtype = device_configuration.torch_dtype
|
language=self.language,
|
||||||
|
dType=config.dType,
|
||||||
self.chunks: list = []
|
device=config.device,
|
||||||
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
|
|
||||||
)
|
)
|
||||||
|
|
||||||
def _to_mono(self):
|
# maybe add something like temperature here?
|
||||||
if self.waveform.shape[0] > 1:
|
def transcribeAudio(self) -> str:
|
||||||
self.waveform = self.waveform.mean(dim=0, keepdim=True)
|
transcription: list = []
|
||||||
self.waveform = self.waveform.squeeze(0)
|
|
||||||
|
|
||||||
def _split_to_chunks(self, shift: bool = False) -> None:
|
self.engine.loadModel()
|
||||||
self.logger.info(f"Splitting audio on chunks...")
|
|
||||||
|
|
||||||
self.chunk_size = self.custom_chunk_length * 16000 # 16kHz after resampling
|
self.preprocessor.prepare(self.audio.load(self.filepath))
|
||||||
total_samples = self.waveform.shape[0]
|
|
||||||
chunks_count = (total_samples + self.chunk_size - 1) // self.chunk_size
|
|
||||||
|
|
||||||
self.logger.info(
|
batches = self.splitter.split(self.audio.waveform)
|
||||||
f"File length - {total_samples / 16000:.1f} seconds, splitting on {chunks_count} chunks by {self.custom_chunk_length} seconds per chunk."
|
|
||||||
)
|
|
||||||
|
|
||||||
self.chunks = []
|
for batch in batches:
|
||||||
for idx in tqdm(range(chunks_count)):
|
batchText: str = self.engine.transcribeBatch(batch)
|
||||||
start = idx * self.chunk_size
|
transcription.append(batchText)
|
||||||
end = min((idx + 1) * self.chunk_size, total_samples)
|
|
||||||
chunk = self.waveform[start:end].cpu().numpy().astype("float32")
|
|
||||||
self.chunks.append(chunk)
|
|
||||||
|
|
||||||
def _resplit_chunks_to_batches(self) -> None:
|
self.engine.unloadModel()
|
||||||
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)}")
|
|
||||||
|
|
||||||
def _process_all_batches(self) -> None:
|
return str(" ".join(transcription))
|
||||||
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)
|
|
||||||
|
|||||||
@@ -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
|
import logging
|
||||||
from dataclasses import dataclass
|
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
|
|
||||||
def check_torch(logger: logging.Logger) -> None:
|
def checkTorch(logger: logging.Logger) -> None:
|
||||||
logger.info("=== Checking PyTorch ===")
|
logger.info("=== Checking PyTorch ===")
|
||||||
logger.info(f"Torch version: {torch.__version__}")
|
logger.info(f"Torch version: {torch.__version__}")
|
||||||
|
|
||||||
|
|||||||
@@ -1,4 +1,4 @@
|
|||||||
from transcription.device_configuration import DeviceConfiguration
|
from transcription.configuration import Configuration
|
||||||
|
|
||||||
|
|
||||||
# TODO: implement saving & removing configuration
|
# TODO: implement saving & removing configuration
|
||||||
|
|||||||
+36
@@ -0,0 +1,36 @@
|
|||||||
|
import tkinter as tk
|
||||||
|
import torch
|
||||||
|
|
||||||
|
class State:
|
||||||
|
def __init__(self, root):
|
||||||
|
# transcription
|
||||||
|
self.model = tk.StringVar(root, "openai/whisper-large-v3-turbo")
|
||||||
|
self.batch = tk.StringVar(root, "32")
|
||||||
|
self.chunk = tk.StringVar(root, "30")
|
||||||
|
self.dtype = tk.StringVar(root, "torch.float16")
|
||||||
|
self.language = tk.StringVar(root, "ru")
|
||||||
|
|
||||||
|
# llm
|
||||||
|
self.api_key = tk.StringVar(root, "")
|
||||||
|
self.api_model = tk.StringVar(root, "")
|
||||||
|
self.base_url = tk.StringVar(root, "")
|
||||||
|
self.conspect_lang = tk.StringVar(root, "Russian")
|
||||||
|
|
||||||
|
# flags
|
||||||
|
self.create_conspect = tk.BooleanVar(root, False)
|
||||||
|
self.remove_transcription = tk.BooleanVar(root, False)
|
||||||
|
|
||||||
|
# files
|
||||||
|
self.input_file = tk.StringVar(root)
|
||||||
|
self.output_file = tk.StringVar(root)
|
||||||
|
|
||||||
|
# device
|
||||||
|
devices = []
|
||||||
|
if torch.cuda.is_available():
|
||||||
|
devices.append("cuda")
|
||||||
|
if torch.backends.mps.is_available():
|
||||||
|
devices.append("mps")
|
||||||
|
devices.append("cpu")
|
||||||
|
|
||||||
|
self.device_opts = devices
|
||||||
|
self.device = tk.StringVar(root, devices[0])
|
||||||
@@ -0,0 +1 @@
|
|||||||
|
# new structure coming soon
|
||||||
@@ -0,0 +1 @@
|
|||||||
|
# new structure coming soon
|
||||||
@@ -0,0 +1 @@
|
|||||||
|
# new structure coming soon
|
||||||
@@ -11,8 +11,8 @@ import customtkinter as ctk
|
|||||||
import torch
|
import torch
|
||||||
|
|
||||||
from transcription.audio_transcription import AudioTranscription
|
from transcription.audio_transcription import AudioTranscription
|
||||||
from transcription.device_configuration import DeviceConfiguration
|
from transcription.configuration import Configuration
|
||||||
from transcription.torch_checker import check_torch
|
from transcription.torch_checker import checkTorch
|
||||||
from ui.tooltip import ToolTip
|
from ui.tooltip import ToolTip
|
||||||
from ui.ui_log_handler import setup_ui_logger
|
from ui.ui_log_handler import setup_ui_logger
|
||||||
from utils.requests_to_api import LLMrequest
|
from utils.requests_to_api import LLMrequest
|
||||||
@@ -20,7 +20,6 @@ from utils.requests_to_api import LLMrequest
|
|||||||
WINDOW_WIDTH = 1000
|
WINDOW_WIDTH = 1000
|
||||||
WINDOW_HEIGHT = 725
|
WINDOW_HEIGHT = 725
|
||||||
|
|
||||||
|
|
||||||
class TranscriberApp(ctk.CTk):
|
class TranscriberApp(ctk.CTk):
|
||||||
def __init__(self):
|
def __init__(self):
|
||||||
super().__init__()
|
super().__init__()
|
||||||
@@ -54,6 +53,7 @@ class TranscriberApp(ctk.CTk):
|
|||||||
# checkboxes
|
# checkboxes
|
||||||
self.create_conspect = tk.BooleanVar(value=False)
|
self.create_conspect = tk.BooleanVar(value=False)
|
||||||
self.remove_transcription = tk.BooleanVar(value=False)
|
self.remove_transcription = tk.BooleanVar(value=False)
|
||||||
|
self.latex_compiler = tk.StringVar(value="")
|
||||||
|
|
||||||
# settings device options
|
# settings device options
|
||||||
device_opts = []
|
device_opts = []
|
||||||
@@ -187,7 +187,7 @@ class TranscriberApp(ctk.CTk):
|
|||||||
row=0,
|
row=0,
|
||||||
col=0,
|
col=0,
|
||||||
text="Model:",
|
text="Model:",
|
||||||
tooltip="Choose model for speed recognition",
|
tooltip="Choose model for speech recognition",
|
||||||
variable=self.model_var,
|
variable=self.model_var,
|
||||||
values=[
|
values=[
|
||||||
"openai/whisper-large-v3-turbo",
|
"openai/whisper-large-v3-turbo",
|
||||||
@@ -296,6 +296,20 @@ class TranscriberApp(ctk.CTk):
|
|||||||
variable=self.conspect_transcription_lang_var,
|
variable=self.conspect_transcription_lang_var,
|
||||||
values=None,
|
values=None,
|
||||||
)
|
)
|
||||||
|
add_setting(
|
||||||
|
parent=grid,
|
||||||
|
row=4,
|
||||||
|
col=0,
|
||||||
|
text="LaTeX compiler:",
|
||||||
|
tooltip="Choose LaTeX compiler",
|
||||||
|
variable=self.latex_compiler,
|
||||||
|
values=[
|
||||||
|
"latexmk",
|
||||||
|
"lualatex",
|
||||||
|
"xelatex",
|
||||||
|
"bibtex"
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
### Custom Prompt setting
|
### Custom Prompt setting
|
||||||
customPromptFrame = ctk.CTkFrame(grid)
|
customPromptFrame = ctk.CTkFrame(grid)
|
||||||
@@ -350,13 +364,13 @@ class TranscriberApp(ctk.CTk):
|
|||||||
input_name = os.path.basename(self.input_file_var.get())
|
input_name = os.path.basename(self.input_file_var.get())
|
||||||
name, _ = os.path.splitext(input_name)
|
name, _ = os.path.splitext(input_name)
|
||||||
# TODO: redo output_name logic maybe?
|
# TODO: redo output_name logic maybe?
|
||||||
output_name = f"{"".join(name.split(".").pop())}.txt"
|
output_name = f"{"".join(name.split(".")[:-1:])}.txt"
|
||||||
path = os.path.join(directory, output_name)
|
path = os.path.join(directory, output_name)
|
||||||
|
|
||||||
self.output_file_var.set(path)
|
self.output_file_var.set(path)
|
||||||
|
|
||||||
def _check_torch(self):
|
def _check_torch(self):
|
||||||
check_torch(self.ui_logger)
|
checkTorch(self.ui_logger)
|
||||||
|
|
||||||
self.ui_logger.info(f"==== Transcription ====")
|
self.ui_logger.info(f"==== Transcription ====")
|
||||||
self.ui_logger.info(f"Transcription model: {self.model_var.get()}")
|
self.ui_logger.info(f"Transcription model: {self.model_var.get()}")
|
||||||
@@ -405,21 +419,21 @@ class TranscriberApp(ctk.CTk):
|
|||||||
|
|
||||||
def _transcribe_worker(self, infile: str):
|
def _transcribe_worker(self, infile: str):
|
||||||
try:
|
try:
|
||||||
config = DeviceConfiguration(
|
config = Configuration(
|
||||||
device=self.device_var.get(),
|
device=self.device_var.get(),
|
||||||
model_name=self.model_var.get(),
|
modelName=self.model_var.get(),
|
||||||
batch_size=int(self.batch_var.get()),
|
batchSize=int(self.batch_var.get()),
|
||||||
chunk_length_s=int(self.chunk_var.get()),
|
chunkSize=int(self.chunk_var.get()),
|
||||||
data_type=self.dtype_var.get(),
|
data_type=self.dtype_var.get(),
|
||||||
)
|
)
|
||||||
Audio = AudioTranscription(
|
Audio = AudioTranscription(
|
||||||
filepath=infile,
|
filepath=infile,
|
||||||
device_configuration=config,
|
config=config,
|
||||||
logger=self.ui_logger,
|
|
||||||
language=self.transcription_lang_var.get(),
|
language=self.transcription_lang_var.get(),
|
||||||
)
|
)
|
||||||
transcription = Audio.transcribe_audio()
|
transcription = Audio.transcribeAudio()
|
||||||
|
|
||||||
|
# remove from here
|
||||||
outfile = self.output_file_var.get().strip()
|
outfile = self.output_file_var.get().strip()
|
||||||
if not outfile:
|
if not outfile:
|
||||||
outfile = infile
|
outfile = infile
|
||||||
@@ -431,32 +445,32 @@ class TranscriberApp(ctk.CTk):
|
|||||||
f.write(transcription)
|
f.write(transcription)
|
||||||
self.ui_logger.info(f"Transcription saved to {outfile}.")
|
self.ui_logger.info(f"Transcription saved to {outfile}.")
|
||||||
|
|
||||||
if self.create_conspect.get():
|
# if self.create_conspect.get():
|
||||||
# TODO: add custom prompt ability here
|
# # TODO: add custom prompt ability here
|
||||||
# TODO: add logging here
|
# # TODO: add logging here
|
||||||
# TODO: add progressbar instead of tqdm
|
# # TODO: add progressbar instead of tqdm
|
||||||
self.ui_logger.info(f"Starting creating conspect via {self.api_model_var.get()}...")
|
# self.ui_logger.info(f"Starting creating conspect via {self.api_model_var.get()}...")
|
||||||
with open("utils/default_prompt.txt", "r", encoding="utf-8") as f:
|
# with open("utils/prompts/default_prompt.txt", "r", encoding="utf-8") as f:
|
||||||
default_prompt = "\n".join(f.readlines())
|
# default_prompt = "\n".join(f.readlines())
|
||||||
|
|
||||||
# if self.custom_prompt_textbox.get(): # <-- issue here
|
# # if self.custom_prompt_textbox.get(): # <-- issue here
|
||||||
# prompt = transcription + "\n" + self.custom_prompt_textbox.get()
|
# # prompt = transcription + "\n" + self.custom_prompt_textbox.get()
|
||||||
# else:
|
# # else:
|
||||||
# prompt = transcription + "\n" + default_prompt
|
# # prompt = transcription + "\n" + default_prompt
|
||||||
|
|
||||||
prompt = transcription + "\n" + default_prompt
|
# prompt = transcription + "\n" + default_prompt
|
||||||
|
|
||||||
request = LLMrequest(
|
# request = LLMrequest(
|
||||||
api_key=self.api_key_var.get(),
|
# api_key=self.api_key_var.get(),
|
||||||
model_name=self.api_model_var.get(),
|
# model_name=self.api_model_var.get(),
|
||||||
base_url=self.base_url_var.get(),
|
# base_url=self.base_url_var.get(),
|
||||||
)
|
# )
|
||||||
response = request.get_response(prompt=prompt)
|
# response = request.get_response(prompt=prompt)
|
||||||
outfile += ".md"
|
# outfile += ".md"
|
||||||
with open(outfile, "w", encoding="utf-8") as f:
|
# with open(outfile, "w", encoding="utf-8") as f:
|
||||||
f.write(response)
|
# f.write(response)
|
||||||
|
|
||||||
self.ui_logger.info(f"Conspect saved to {outfile}.")
|
# self.ui_logger.info(f"Conspect saved to {outfile}.")
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
self.ui_logger.error(f"Error: {e}")
|
self.ui_logger.error(f"Error: {e}")
|
||||||
|
|||||||
@@ -24,3 +24,16 @@ def setup_ui_logger(text_widget: tk.Text, level=logging.INFO):
|
|||||||
logger.addHandler(handler)
|
logger.addHandler(handler)
|
||||||
|
|
||||||
return logger
|
return logger
|
||||||
|
|
||||||
|
# def setup_ui_logger(
|
||||||
|
# text_widget: tk.Text,
|
||||||
|
# base_logger: logging.Logger,
|
||||||
|
# level=logging.INFO,
|
||||||
|
# ):
|
||||||
|
# handler = UILogHandler(text_widget)
|
||||||
|
# handler.setLevel(level)
|
||||||
|
# handler.setFormatter(
|
||||||
|
# logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
|
||||||
|
# )
|
||||||
|
|
||||||
|
# base_logger.addHandler(handler)
|
||||||
|
|||||||
@@ -1,84 +0,0 @@
|
|||||||
Используй **исключительно текст из транскрибации** как первоисточник.
|
|
||||||
Разрешается расширять тему и объяснять глубже, но **дополнительная информация должна быть строго отделена и явно помечена**.
|
|
||||||
|
|
||||||
### Цель
|
|
||||||
|
|
||||||
Создать структурированную, полный образовательный конспект лекции **формата Obsidian Markdown**, который:
|
|
||||||
|
|
||||||
* воспроизводит информацию из транскрибата максимально точно
|
|
||||||
* дополняет и разъясняет сложные моменты
|
|
||||||
* превращает речь в ясный учебный материал
|
|
||||||
* сохраняет смысл, порядок, и логическую структуру лектора
|
|
||||||
* улучшает академичность и связность
|
|
||||||
|
|
||||||
### Правила работы с транскриптом
|
|
||||||
|
|
||||||
* Используй только факты из транскрибации как основу.
|
|
||||||
* **Не выдумывай факты**, которых не было.
|
|
||||||
* **Дополнительные пояснения допустимы**, но строго в спец-блоках (см. ниже).
|
|
||||||
* Цитаты из оригинального текста — **дослівно**, максимум **25 слов подряд**.
|
|
||||||
* Нельзя писать «лектор сказал», «в тексте было» — пишем как учебный материал.
|
|
||||||
|
|
||||||
### Блоки для расширений
|
|
||||||
|
|
||||||
Если нужно объяснить термин, дополнить концепцию, исправить недосказанность — используй:
|
|
||||||
|
|
||||||
```
|
|
||||||
> AI clarification:
|
|
||||||
> Ваше дополнение, объяснение, аналогия или расширение темы.
|
|
||||||
```
|
|
||||||
|
|
||||||
Для коротких определений:
|
|
||||||
|
|
||||||
```
|
|
||||||
**AI note:** краткое пояснение
|
|
||||||
```
|
|
||||||
|
|
||||||
### Обязательная структура результата
|
|
||||||
|
|
||||||
```
|
|
||||||
# Название темы
|
|
||||||
|
|
||||||
## Ключевые идеи
|
|
||||||
- маркеры
|
|
||||||
- …
|
|
||||||
|
|
||||||
## Основные понятия и определения
|
|
||||||
- **Термин** — определение
|
|
||||||
- …
|
|
||||||
|
|
||||||
## Подробный конспект
|
|
||||||
(структура, логика, абзацы, тезисы)
|
|
||||||
|
|
||||||
> Цитаты (≤25 слов)
|
|
||||||
|
|
||||||
## Примеры и объяснения
|
|
||||||
- собственные пересказанные примеры из транскрибации
|
|
||||||
- при необходимости:
|
|
||||||
> AI clarification
|
|
||||||
|
|
||||||
## Выводы и ключевые инсайты
|
|
||||||
- …
|
|
||||||
|
|
||||||
## Вопросы для самопроверки
|
|
||||||
- …
|
|
||||||
```
|
|
||||||
|
|
||||||
### Обязательный формат Obsidian Markdown
|
|
||||||
|
|
||||||
Используй следующие возможности Obsidian:
|
|
||||||
|
|
||||||
| Элемент | Формат |
|
|
||||||
| ----------------- | ----------------------------------------- |
|
|
||||||
| Заголовки | `#`, `##`, `###` |
|
|
||||||
| Цитаты | `>` |
|
|
||||||
| Списки | `-` и `1.` |
|
|
||||||
| Врезка/внимание | `> [!note]` `> [!tip]` |
|
|
||||||
| Код/формулы | `\`…`` и блоки ``` |
|
|
||||||
| Выделение | `**жирный**`, `_курсив_`, `==выделение==` |
|
|
||||||
| Внутренние ссылки | `[[Термин]]` если термин важный |
|
|
||||||
| Таблицы | стандартные markdown таблицы |
|
|
||||||
| Чек-лист | `- [ ]` |
|
|
||||||
| Формулы | Для однострочных формул используй $ formula $. Для мультистрочных формул используй $$ formula $$ |
|
|
||||||
|
|
||||||
**Не добавляй** приветствий, заключений, обращений, пояснений формата. Только готовый md-текст.
|
|
||||||
@@ -0,0 +1,13 @@
|
|||||||
|
# not sure, but why not?
|
||||||
|
|
||||||
|
class AudioPreparationError(Exception):
|
||||||
|
"""Base exception for audio preparation"""
|
||||||
|
pass
|
||||||
|
|
||||||
|
class ResamplingError(AudioPreparationError):
|
||||||
|
"""Raises when resampling fails"""
|
||||||
|
pass
|
||||||
|
|
||||||
|
class DownmixingError(AudioPreparationError):
|
||||||
|
"""Raises when downmixing to mono fails"""
|
||||||
|
pass
|
||||||
@@ -0,0 +1,58 @@
|
|||||||
|
Create a complete, structured **Obsidian Markdown** educational summary based on the transcript.
|
||||||
|
|
||||||
|
The final document must be written in **{{ language }}**.
|
||||||
|
|
||||||
|
Use the transcript strictly as the factual basis:
|
||||||
|
- Do NOT invent facts.
|
||||||
|
- You may add clarifications and explanations, but they MUST be placed in special blocks.
|
||||||
|
- Preserve the logical order and meaning of the original speaker.
|
||||||
|
- Transform speech into clear academic exposition.
|
||||||
|
- Quotes must be accurate and **<= 25 words**.
|
||||||
|
- Do not reference the transcript or the speaker. Write as a stand-alone educational text.
|
||||||
|
|
||||||
|
You must follow the template below exactly.
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
# {{ title }}
|
||||||
|
|
||||||
|
## Key Ideas
|
||||||
|
- bullet points summarizing the main concepts
|
||||||
|
|
||||||
|
## Core Concepts and Definitions
|
||||||
|
- **Term** — definition
|
||||||
|
- **Another Term** — definition
|
||||||
|
|
||||||
|
## Detailed Summary
|
||||||
|
Write a logically structured, academically clear exposition.
|
||||||
|
Break it into paragraphs, subsections, lists, and quotes.
|
||||||
|
|
||||||
|
To include quotes from the transcript, use:
|
||||||
|
|
||||||
|
> "{{ quote }}"
|
||||||
|
(Ensure that every quote is ≤ 25 words.)
|
||||||
|
|
||||||
|
## Examples and Explanations
|
||||||
|
Use examples that appear in the transcript.
|
||||||
|
Additional clarifications must use the following block:
|
||||||
|
|
||||||
|
> AI clarification:
|
||||||
|
> Your extended explanation, analogy, or contextual detail.
|
||||||
|
|
||||||
|
Short notes should use:
|
||||||
|
|
||||||
|
**AI note:** short clarification
|
||||||
|
|
||||||
|
## Conclusions and Key Insights
|
||||||
|
- summary points
|
||||||
|
- what the reader should remember
|
||||||
|
|
||||||
|
## Self-Assessment Questions
|
||||||
|
- question 1
|
||||||
|
- question 2
|
||||||
|
- question 3
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
### Transcript (for your reference only):
|
||||||
|
{{ transcription }}
|
||||||
@@ -0,0 +1,75 @@
|
|||||||
|
Create a clean, structured **LaTeX educational summary** based strictly on the transcript.
|
||||||
|
|
||||||
|
The final document must be written in **{{ language }}**.
|
||||||
|
|
||||||
|
Guidelines:
|
||||||
|
- Use ONLY facts from the transcript as the base.
|
||||||
|
- You MAY add explanations, but they MUST appear in designated blocks.
|
||||||
|
- Do NOT reference the transcript or the speaker.
|
||||||
|
- Maintain the speaker`s logical order.
|
||||||
|
- Produce academically clear, structured LaTeX text.
|
||||||
|
- Quotes must be accurate and <= 25 words.
|
||||||
|
|
||||||
|
Use the structure below.
|
||||||
|
Assume that stylistic environments are already defined in the `.sty` file.
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
\section*{ {{ title }} }
|
||||||
|
|
||||||
|
\subsection*{Key Ideas}
|
||||||
|
\begin{itemize}
|
||||||
|
\item main idea
|
||||||
|
\item …
|
||||||
|
\end{itemize}
|
||||||
|
|
||||||
|
\subsection*{Core Concepts and Definitions}
|
||||||
|
\begin{description}
|
||||||
|
\item[\textbf{Term}] definition
|
||||||
|
\item[\textbf{Another Term}] definition
|
||||||
|
\end{description}
|
||||||
|
|
||||||
|
\subsection*{Detailed Summary}
|
||||||
|
Write a clear, logically structured academic explanation.
|
||||||
|
Break content into paragraphs, subsections, lists.
|
||||||
|
|
||||||
|
Quotes from the transcript must be included as:
|
||||||
|
|
||||||
|
\begin{quote}
|
||||||
|
"{{ quote }}"
|
||||||
|
\end{quote}
|
||||||
|
|
||||||
|
(Ensure each quote is ≤ 25 words.)
|
||||||
|
|
||||||
|
\subsection*{Examples and Explanations}
|
||||||
|
|
||||||
|
Examples based on the transcript.
|
||||||
|
|
||||||
|
Additional clarifications must use the following custom environment (assumed to exist in the .sty):
|
||||||
|
|
||||||
|
\begin{aiclarification}
|
||||||
|
Your additional explanation, analogy, or contextual expansion.
|
||||||
|
\end{aiclarification}
|
||||||
|
|
||||||
|
Short notes should use:
|
||||||
|
|
||||||
|
\ainote{short clarification}
|
||||||
|
|
||||||
|
\subsection*{Conclusions and Key Insights}
|
||||||
|
\begin{itemize}
|
||||||
|
\item conclusion
|
||||||
|
\item conclusion
|
||||||
|
\end{itemize}
|
||||||
|
|
||||||
|
\subsection*{Self-Assessment Questions}
|
||||||
|
\begin{enumerate}
|
||||||
|
\item question
|
||||||
|
\item question
|
||||||
|
\item question
|
||||||
|
\end{enumerate}
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
% Transcript included only for reference (model must not copy or analyze or mention this directly)
|
||||||
|
% ======================================================
|
||||||
|
% {{ transcription | replace("%", "\\%") }}
|
||||||
@@ -0,0 +1,78 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import os
|
||||||
|
from jinja2 import Environment, FileSystemLoader, StrictUndefined
|
||||||
|
|
||||||
|
|
||||||
|
class PromptManager:
|
||||||
|
"""
|
||||||
|
Self-written class for Jinja2 templates especially for summarization.
|
||||||
|
|
||||||
|
Supports:
|
||||||
|
- md
|
||||||
|
- latex
|
||||||
|
|
||||||
|
Usage examples:
|
||||||
|
```
|
||||||
|
from prompts.prompt_manager import PromptManager
|
||||||
|
|
||||||
|
pm = PromptManager("notecast/prompts")
|
||||||
|
|
||||||
|
markdown_text = pm.render(
|
||||||
|
"markdown_default_prompt.md.j2",
|
||||||
|
language="Russian",
|
||||||
|
title="Calculus - Lecture 1",
|
||||||
|
transcription=raw_transcription_text
|
||||||
|
)
|
||||||
|
|
||||||
|
latex_text = pm.render(
|
||||||
|
"latex_default_prompt.j2",
|
||||||
|
language="Russian",
|
||||||
|
title="Calculus - Lecture 1",
|
||||||
|
transcription=raw_transcription_text
|
||||||
|
)
|
||||||
|
```
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, templates_dir: str) -> None:
|
||||||
|
if not os.path.isdir(templates_dir):
|
||||||
|
raise NotADirectoryError(f"Templates directory not found: {templates_dir}")
|
||||||
|
|
||||||
|
self.templates_dir = templates_dir
|
||||||
|
|
||||||
|
self.env = Environment(
|
||||||
|
loader=FileSystemLoader(self.templates_dir),
|
||||||
|
autoescape=False,
|
||||||
|
trim_blocks=True,
|
||||||
|
lstrip_blocks=True,
|
||||||
|
undefined=StrictUndefined
|
||||||
|
)
|
||||||
|
|
||||||
|
def list_templates(self) -> list[str]:
|
||||||
|
templates = []
|
||||||
|
for file in os.listdir(self.templates_dir):
|
||||||
|
if file.endswith(".j2"):
|
||||||
|
templates.append(file)
|
||||||
|
return templates
|
||||||
|
|
||||||
|
def load(self, template_name: str):
|
||||||
|
try:
|
||||||
|
return self.env.get_template(template_name)
|
||||||
|
except Exception as e:
|
||||||
|
raise RuntimeError(f"Failed to load template '{template_name}': {e}")
|
||||||
|
|
||||||
|
def render(
|
||||||
|
self,
|
||||||
|
template_name: str,
|
||||||
|
*,
|
||||||
|
language: str,
|
||||||
|
title: str,
|
||||||
|
transcription: str
|
||||||
|
) -> str:
|
||||||
|
template = self.load(template_name)
|
||||||
|
|
||||||
|
return template.render(
|
||||||
|
language=language,
|
||||||
|
title=title,
|
||||||
|
transcription=transcription
|
||||||
|
)
|
||||||
@@ -1,6 +1,6 @@
|
|||||||
import openai
|
import openai
|
||||||
|
|
||||||
|
# maybe make it asynchronous?
|
||||||
class LLMrequest:
|
class LLMrequest:
|
||||||
def __init__(self, api_key: str, model_name: str, base_url: str = None):
|
def __init__(self, api_key: str, model_name: str, base_url: str = None):
|
||||||
if base_url:
|
if base_url:
|
||||||
|
|||||||
Reference in New Issue
Block a user