use enum to validate backend validity in server.run

This commit is contained in:
berkaybilik
2024-06-26 23:59:07 +01:00
parent 815441e8bb
commit 5e3906fc7b
+24
View File
@@ -4,6 +4,9 @@ import threading
import json import json
import functools import functools
import logging import logging
from enum import Enum
from typing import List
import torch import torch
import numpy as np import numpy as np
from websockets.sync.server import serve from websockets.sync.server import serve
@@ -121,6 +124,25 @@ class ClientManager:
return False return False
class BackendType(Enum):
FASTER_WHISPER = "faster_whisper"
TENSORRT = "tensorrt"
@staticmethod
def valid_types() -> List[str]:
return [backend_type.value for backend_type in BackendType]
@staticmethod
def check_validity_of(backend: str) -> bool:
return backend in BackendType.valid_types()
def is_faster_whisper(self) -> bool:
return self == BackendType.FASTER_WHISPER
def is_tensorrt(self) -> bool:
return self == BackendType.TENSORRT
class TranscriptionServer: class TranscriptionServer:
RATE = 16000 RATE = 16000
@@ -311,6 +333,8 @@ class TranscriptionServer:
# TODO: load model initially # TODO: load model initially
else: else:
logging.info("Single model mode currently only works with custom models.") logging.info("Single model mode currently only works with custom models.")
if BackendType.check_validity_of(backend):
raise ValueError(f"{backend} is not a valid backend type. Choose backend from {BackendType.valid_types()}")
with serve( with serve(
functools.partial( functools.partial(
self.recv_audio, self.recv_audio,