use enum to validate backend validity in server.run
This commit is contained in:
@@ -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,
|
||||||
|
|||||||
Reference in New Issue
Block a user