diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 0000000..4b5a294 --- /dev/null +++ b/.vscode/settings.json @@ -0,0 +1,4 @@ +{ + "python-envs.defaultEnvManager": "ms-python.python:conda", + "python-envs.defaultPackageManager": "ms-python.python:conda" +} \ No newline at end of file diff --git a/env.yml b/env.yml index c4f1f12..bf116d9 100644 --- a/env.yml +++ b/env.yml @@ -13,3 +13,6 @@ dependencies: - python=3.12 - customtkinter - openai + - pip: + - vosk + diff --git a/transcription/audio_transcription.py b/transcription/audio_transcription.py index 8fb5f13..525185c 100644 --- a/transcription/audio_transcription.py +++ b/transcription/audio_transcription.py @@ -1,7 +1,7 @@ from transcription.audio import Audio from transcription.preprocessing.audio_preprocessor import AudioPreprocessor from transcription.preprocessing.splitter import Splitter -from transcription.engines.whisper import WhisperEngine +from transcription.engines.whisper_engine import WhisperEngine from transcription.configuration import Configuration # maybe inherit from AudioTranscription and rename to something like WhisperTranscription? @@ -19,11 +19,13 @@ class AudioTranscription: # self.logger = logger self.audio = Audio() - self.preprocessor = AudioPreprocessor() + self.preprocessor = AudioPreprocessor(16000) + self.splitter = Splitter( chunkSize=config.chunkSize, batchSize=config.batchSize, ) + self.engine = WhisperEngine( modelName=config.modelName, language=self.language, diff --git a/transcription/configuration.py b/transcription/configuration.py index bfdac85..9ec0cbd 100644 --- a/transcription/configuration.py +++ b/transcription/configuration.py @@ -7,8 +7,8 @@ class Configuration: # add new models device: str = "cuda" modelName: str = "openai/whisper-large-v2" - chunkSize: int = 30 - batchSize: int = 16 + chunkSize: int = 30 ## in seconds + batchSize: int = 16 ## in batches dataType: str = "torch.float16" _dtype_map = { diff --git a/transcription/engines/BaseEngine.py b/transcription/engines/BaseEngine.py new file mode 100644 index 0000000..8392347 --- /dev/null +++ b/transcription/engines/BaseEngine.py @@ -0,0 +1,11 @@ +import torch +from abc import ABC, abstractmethod + +class BaseEngine(ABC): + @abstractmethod + def loadModel(self) -> None: + pass + + @abstractmethod + def unloadModel(self) -> None: + pass \ No newline at end of file diff --git a/transcription/engines/base_engine.py b/transcription/engines/BatchSTT.py similarity index 50% rename from transcription/engines/base_engine.py rename to transcription/engines/BatchSTT.py index 431ff1d..4538ef9 100644 --- a/transcription/engines/base_engine.py +++ b/transcription/engines/BatchSTT.py @@ -1,30 +1,18 @@ -import torch -from abc import ABC, abstractmethod +from transcription.engines.BaseEngine import BaseEngine -class BaseEngine(ABC): +class BatchSTTEngine(BaseEngine): def __init__( self, modelName: str, language: str, dType: torch.dtype, device: str - ): + ) -> None: 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: + def transcribeBatch(self) -> None: pass \ No newline at end of file diff --git a/transcription/engines/StreamingSTT.py b/transcription/engines/StreamingSTT.py new file mode 100644 index 0000000..e6cf7ac --- /dev/null +++ b/transcription/engines/StreamingSTT.py @@ -0,0 +1,5 @@ +from transcription.engines.BaseEngine import BaseEngine + +class StreamingSTTEngine(BaseEngine): + def __init__(self) -> None: + ... \ No newline at end of file diff --git a/transcription/engines/vosk_engine.py b/transcription/engines/vosk_engine.py new file mode 100644 index 0000000..4960364 --- /dev/null +++ b/transcription/engines/vosk_engine.py @@ -0,0 +1,53 @@ +from transcription.engines.BaseEngine import BaseEngine +import wave, json +from vosk import Model, KaldiRecognizer + +""" +import wave +import json +import sys + +from multiprocessing.dummy import Pool +from vosk import Model, KaldiRecognizer + +model = Model("en-us") + +def recognize(line): + uid, fn = line.split() + wf = wave.open(fn, "rb") + rec = KaldiRecognizer(model, wf.getframerate()) + + text = "" + while True: + data = wf.readframes(1000) + if len(data) == 0: + break + if rec.AcceptWaveform(data): + jres = json.loads(rec.Result()) + text = text + " " + jres["text"] + jres = json.loads(rec.FinalResult()) + text = text + " " + jres["text"] + return uid + text + +def main(): + p = Pool(8) + texts = p.map(recognize, open(sys.argv[1], encoding="utf-8").readlines()) + print ("\n".join(texts)) + +main() +""" + +class VoskEngine(BaseEngine): + TARGET_SAMPLING_RATE = 16000 + + def loadModel(self) -> None: + ... + + def unloadModel(self) -> None: + ... + + def transcribeBatch( + self, + batch, + ) -> str: + ... \ No newline at end of file diff --git a/transcription/engines/whisper.py b/transcription/engines/whisper_engine.py similarity index 90% rename from transcription/engines/whisper.py rename to transcription/engines/whisper_engine.py index 623b677..07a8a6b 100644 --- a/transcription/engines/whisper.py +++ b/transcription/engines/whisper_engine.py @@ -4,9 +4,11 @@ import torch import gc from transformers import WhisperForConditionalGeneration, WhisperProcessor -from transcription.engines.base_engine import BaseEngine +from transcription.engines.BatchSTT import BatchSTT -class WhisperEngine(BaseEngine): +class WhisperEngine(BatchSTT): + TARGET_SAMPLING_RATE = 16000 + def loadModel(self) -> None: self.processor = WhisperProcessor.from_pretrained(self.modelName) self.model = WhisperForConditionalGeneration.from_pretrained( @@ -31,7 +33,7 @@ class WhisperEngine(BaseEngine): inputs = self.processor( batch, - sampling_rate=16000, + sampling_rate=self.TARGET_SAMPLING_RATE, return_tensors="pt", padding=True, ) diff --git a/transcription/preprocessing/audio_preprocessor.py b/transcription/preprocessing/audio_preprocessor.py index 63ddd7d..a6f1652 100644 --- a/transcription/preprocessing/audio_preprocessor.py +++ b/transcription/preprocessing/audio_preprocessor.py @@ -7,6 +7,9 @@ class AudioPreprocessor: # for different models in future # def __init__(self, model): # pass + + def __init__(self, target_sr: int) -> None: + self.TARGET_SAMPLING_RATE = target_sr def _resample( self, diff --git a/transcription/preprocessing/splitter.py b/transcription/preprocessing/splitter.py index ece97b2..f93dc9a 100644 --- a/transcription/preprocessing/splitter.py +++ b/transcription/preprocessing/splitter.py @@ -1,11 +1,14 @@ import torch from typing import List +from logging import Logger +# TODO: add logging here class Splitter: def __init__( self, chunkSize: int, batchSize: int, + # logger: Logger ) -> None: self.chunkSize = chunkSize * 16000 # 16 kHz after resampling self.batchSize = batchSize @@ -19,7 +22,7 @@ class Splitter: chunksCount = (totalSamples + self.chunkSize - 1) // self.chunkSize chunks: List = [] - # tqdm or something here? + # tqdm for logger or something here? for chunkNum in range(chunksCount): start = chunkNum * self.chunkSize end = min((chunkNum + 1) * self.chunkSize, totalSamples) @@ -34,11 +37,9 @@ class Splitter: 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( diff --git a/utils/timings.py b/utils/timings.py new file mode 100644 index 0000000..d176ee1 --- /dev/null +++ b/utils/timings.py @@ -0,0 +1,4 @@ +import time + +class Time: + \ No newline at end of file