Initial commit

transcription-api/backends/backend.py

@@ -0,0 +1,68 @@
+
+from typing import Any, Mapping, TypedDict
+import numpy as np
+from faster_whisper.audio import decode_audio  # type: ignore
+
+SUPPORTED_MODELS = ["tiny", "tiny.en", "small", "small.en", "base", "base.en", "medium", "medium.en", "large-v2"]
+
+WordData = TypedDict(
+    "WordData", {"word": str, "start": float | str, "end": float | str, "score": float}
+)
+
+Segment = TypedDict(
+    "Segment",
+    {
+        "id": str,
+        "text": str,
+        "start": float | str,
+        "end": float | str,
+        "score": float,
+        "words": list[WordData],
+    },
+)
+
+Transcription = TypedDict(
+    "Transcription",
+    {   
+        "text": str,
+        "language": str,
+        "duration": float,
+        "segments": list[Segment],
+    },
+)
+
+class Backend:
+    name: str = "faster-whisper"
+    model_size: str
+
+    def __post_init__(self):
+        if self.model_size not in self.supported_model_sizes():
+            raise ValueError(f"model must be one of {self.supported_model_sizes()}")
+
+    def supported_backends(self):
+        """
+        This is of not much use as of the moment, If we ever support multiple
+        backends this can be utilized.
+        """
+        return ["faster-whisper"]
+
+    def model_path(self) -> str:
+        """
+        Returns the local path to the model, error-out if unavailable
+        """
+        raise NotImplementedError()
+
+    def supported_model_sizes(self) -> list[str]:
+        return SUPPORTED_MODELS
+    
+    def download_model(self):
+        raise NotImplementedError()
+
+    def load(self):
+        raise NotImplementedError()
+
+    def transcribe(self, input: np.ndarray) -> list[Segment]:
+        """
+        This should return word level transcription data.
+        """
+        raise NotImplementedError()

transcription-api/backends/fasterwhisper.py

@@ -0,0 +1,97 @@
+import numpy as np
+from .backend import Backend, Transcription, Segment
+import os, math
+from tqdm import tqdm  # type: ignore
+import uuid
+from faster_whisper import WhisperModel, download_model, decode_audio
+
+class FasterWhisperBackend(Backend):
+    device: str = "cpu"  # cpu, cuda
+    quantization: str = "int8"  # int8, float16
+    model: WhisperModel | None = None
+
+    def __init__(self, model_size, device: str = "cpu"):
+        self.model_size = model_size
+        self.device = device
+        self.__post_init__()
+
+    def model_path(self) -> str:
+        local_model_path = os.path.join(
+            os.environ["WHISPER_MODELS_DIR"], f"faster-whisper-{self.model_size}"
+        )
+
+        if os.path.exists(local_model_path):
+            return local_model_path
+        else:
+            raise RuntimeError(f"model not found in {local_model_path}")
+        
+    def load(self) -> None:
+        self.model = WhisperModel(
+            self.model_path(), device=self.device, compute_type=self.quantization
+        )
+
+    def get_model(self) -> None:
+        print(f"Downloading model {self.model_size}...")
+        local_model_path = os.path.join(os.environ["WHISPER_MODELS_DIR"], f"faster-whisper-{self.model_size}")
+        local_model_cache = os.path.join(os.environ["WHISPER_MODELS_DIR"], f"faster-whisper-{self.model_size}", "cache")
+        # Check if directory exists
+        if not os.path.exists(local_model_path):
+            os.makedirs(local_model_path)
+        try:
+            download_model(self.model_size, output_dir=local_model_path, local_files_only=True, cache_dir=local_model_cache)
+            print("Model already cached...")
+        except:
+            download_model(self.model_size, output_dir=local_model_path, local_files_only=False, cache_dir=local_model_cache)
+
+    def transcribe(
+        self, input: np.ndarray, silent: bool = False, language: str = None
+    ) -> Transcription:
+        """
+        Return word level transcription data.
+        World level probabities are calculated by ctranslate2.models.Whisper.align
+        """
+        result: list[Segment] = []
+        assert self.model is not None
+        segments, info = self.model.transcribe(
+            input,
+            beam_size=5,
+            word_timestamps=True,
+            language=language,
+        )
+        # ps = playback seconds
+        with tqdm(
+            total=info.duration, unit_scale=True, unit="ps", disable=silent
+        ) as pbar:
+            for segment in segments:
+                if segment.words is None:
+                    continue
+                id = uuid.uuid4().hex
+                segment_extract: Segment = {
+                    "id": id,
+                    "text": segment.text,
+                    "start": segment.start,
+                    "end": segment.end,
+                    "score": round(math.exp(segment.avg_logprob), 2),
+                    "words": [
+                        {
+                            "start": w.start,
+                            "end": w.end,
+                            "word": w.word,
+                            "score": round(w.probability, 2),
+                        }
+                        for w in segment.words
+                    ],
+                }
+                result.append(segment_extract)
+                if not silent:
+                    pbar.update(segment.end - pbar.last_print_n)
+        
+        text = " ".join([segment["text"] for segment in result])
+        text = ' '.join(text.strip().split())
+        transcription: Transcription = {
+            "text": text,
+            "language": info.language,
+            "duration": info.duration,
+            "segments": result,
+        }
+        return transcription