refactor: use State to keep multithreading within controller

capgen/transcriber/protocol.py

@@ -0,0 +1,32 @@
+from typing import BinaryIO, Protocol
+
+
+class TranscriberProtocol(Protocol):
+    """
+    Summary
+    -------
+    a protocol for all transcriber(s)
+
+    Methods
+    -------
+    transcribe(file: str | BinaryIO, caption_format: str) -> str | None:
+        converts transcription segments into a specific caption format
+    """
+
+    __slots__ = ('model',)
+
+    def transcribe(self, file: str | BinaryIO, caption_format: str) -> str | None:
+        """
+        Summary
+        -------
+        transcribes a compatible audio/video into a chosen caption format
+
+        Parameters
+        ----------
+        file (str | BinaryIO) : the file to transcribe
+        caption_format (str) : the chosen caption format
+
+        Returns
+        -------
+        transcription (str | None) : the transcribed text in the chosen caption format
+        """

capgen/transcriber/transcriber.py

@@ -3,6 +3,7 @@
 from faster_whisper import WhisperModel
 
 from capgen.transcriber.converter import Converter
+from capgen.transcriber.protocol import TranscriberProtocol
 
 
 class WhisperParameters(TypedDict):
@@ -19,7 +20,7 @@ class WhisperParameters(TypedDict):
     num_workers: int
 
 
-class Transcriber:
+class Transcriber(TranscriberProtocol):
     """
     Summary
     -------

server/api/v1/transcribe.py

@@ -1,3 +1,5 @@
+from asyncio import wrap_future
+from concurrent.futures import ThreadPoolExecutor
 from io import BytesIO
 from typing import Annotated, Literal
 
@@ -8,8 +10,8 @@
 from litestar.params import Body
 from litestar.status_codes import HTTP_200_OK
 
-from server.features import Transcriber
 from server.schemas.v1 import Transcribed
+from server.state import AppState
 
 
 class TranscriberController(Controller):
@@ -20,10 +22,12 @@ class TranscriberController(Controller):
     """
 
     path = '/transcribe'
+    thread_pool = ThreadPoolExecutor()
 
     @post(status_code=HTTP_200_OK)
     async def transcribe(
         self,
+        state: AppState,
         data: Annotated[UploadFile, Body(media_type=RequestEncodingType.MULTI_PART)],
         caption_format: Literal['srt', 'vtt'] = 'srt',
     ) -> Transcribed:
@@ -32,7 +36,10 @@ async def transcribe(
         -------
         the POST variant of the `/transcribe` route
         """
-        if not (result := await Transcriber.transcribe(BytesIO(await data.read()), caption_format)):
+        audio = BytesIO(await data.read())
+        transcription = await wrap_future(self.thread_pool.submit(state.transcriber.transcribe, audio, caption_format))
+
+        if not transcription:
             raise ClientException(detail=f'Invalid format: {caption_format}!')
 
-        return Transcribed(result=result)
+        return Transcribed(result=transcription)

server/lifespans/load_model.py

@@ -1,15 +1,22 @@
 from contextlib import asynccontextmanager
 from typing import AsyncIterator
 
-from server.features import Transcriber
+from litestar import Litestar
+
+from capgen.transcriber import Transcriber
+from server.config import Config
 
 
 @asynccontextmanager
-async def load_model(_) -> AsyncIterator[None]:
+async def load_model(app: Litestar) -> AsyncIterator[None]:
     """
     Summary
     -------
     download and load the model
     """
-    Transcriber.load()
-    yield
+    app.state.transcriber = Transcriber('cpu', number_of_workers=Config.worker_count)
+
+    try:
+        yield
+    finally:
+        pass

server/state.py

@@ -0,0 +1,13 @@
+from litestar.datastructures import State
+
+from capgen.transcriber.protocol import TranscriberProtocol
+
+
+class AppState(State):
+    """
+    Summary
+    -------
+    the Litestar application state that will be injected into the routers
+    """
+
+    transcriber: TranscriberProtocol