stash commit - wip

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(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,
         )