From c97134a45f7c36fd2057432a0a8637cccfd66862 Mon Sep 17 00:00:00 2001
From: Yuta Hayashibe <yuta@hayashibe.jp>
Date: Fri, 23 Sep 2022 20:45:18 +0900
Subject: [PATCH] Add --num_block option

---
 whisper_streaming/cli.py | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/whisper_streaming/cli.py b/whisper_streaming/cli.py
index 5d0607c..754da19 100644
--- a/whisper_streaming/cli.py
+++ b/whisper_streaming/cli.py
@@ -21,6 +21,7 @@ def transcribe_from_mic(
     *,
     config: WhisperConfig,
     sd_device: Optional[Union[int, str]],
+    num_block: int,
 ) -> None:
     wsp = WhisperStreamingTranscriber(config=config)
     q = queue.Queue()
@@ -33,7 +34,7 @@ def transcribe_from_mic(
     logger.info("Ready to transcribe")
     with sd.InputStream(
         samplerate=SAMPLE_RATE,
-        blocksize=N_FRAMES * 10,  # FIXME
+        blocksize=N_FRAMES * num_block,
         device=sd_device,
         dtype="float32",
         channels=1,
@@ -72,6 +73,13 @@ def get_opts() -> argparse.Namespace:
         type=int,
         default=5,
     )
+    parser.add_argument(
+        "--num_block",
+        "-b",
+        type=int,
+        default=20,
+        help="Number of operation unit. Larger values can improve accuracy but consume more memory.",
+    )
     parser.add_argument(
         "--mic",
     )
@@ -98,6 +106,7 @@ def main() -> None:
     transcribe_from_mic(
         config=config,
         sd_device=opts.mic,
+        num_block=opts.num_block,
     )