feat: add datatype in context

shirayu · Oct 17, 2022 · 77ed1c7 · 77ed1c7
1 parent 4fee6ee
commit 77ed1c7
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/whispering/schema.py b/whispering/schema.py
@@ -48,6 +48,8 @@ class Context(BaseModel, arbitrary_types_allowed=True):
     vad_threshold: float
     max_nospeech_skip: int
 
+    data_type: Optional[str] = "float32"
+
 
 class ParsedChunk(BaseModel):
     start: float

diff --git a/whispering/serve.py b/whispering/serve.py
@@ -66,7 +66,7 @@ async def serve_with_websocket_main(websocket):
             continue
 
         logger.debug(f"Message size: {len(message)}")
-        audio = np.frombuffer(message, dtype=np.float32)
+        audio = np.frombuffer(message, dtype=np.dtype(ctx.data_type)).astype(np.float32)
         if ctx is None:
             await websocket.send(
                 json.dumps(