stash commit - wip

.vscode/settings.json

@@ -0,0 +1,4 @@
+{
+    "python-envs.defaultEnvManager": "ms-python.python:conda",
+    "python-envs.defaultPackageManager": "ms-python.python:conda"
+}

env.yml

@@ -13,3 +13,6 @@ dependencies:
   - python=3.12
   - customtkinter
   - openai
+  - pip: 
+    - vosk
+

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,

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 = {

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

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

transcription/engines/StreamingSTT.py

@@ -0,0 +1,5 @@
+from transcription.engines.BaseEngine import BaseEngine
+
+class StreamingSTTEngine(BaseEngine):
+    def __init__(self) -> None:
+        ...

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:
+        ...

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(BatchSTT):
+    TARGET_SAMPLING_RATE = 16000
     
-class WhisperEngine(BaseEngine):
     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,
         )

transcription/preprocessing/audio_preprocessor.py

@@ -8,6 +8,9 @@ class AudioPreprocessor:
     # def __init__(self, model):
     #     pass
     
+    def __init__(self, target_sr: int) -> None:
+        self.TARGET_SAMPLING_RATE = target_sr
+
     def _resample(
         self,
         audio: Audio

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(

utils/timings.py

@@ -0,0 +1,4 @@
+import time
+
+class Time:
+