From 111e3a94a64be69643f62846c7c8bb02c8df4be5 Mon Sep 17 00:00:00 2001
From: Yuta Hayashibe <yuta@hayashibe.jp>
Date: Fri, 23 Sep 2022 22:01:40 +0900
Subject: [PATCH] Fix logger

---
 whisper_streaming/cli.py         | 15 +++++++++++++--
 whisper_streaming/transcriber.py |  4 ++++
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/whisper_streaming/cli.py b/whisper_streaming/cli.py
index 8b3a428..9bd8b29 100644
--- a/whisper_streaming/cli.py
+++ b/whisper_streaming/cli.py
@@ -2,7 +2,8 @@
 
 import argparse
 import queue
-from logging import INFO, getLogger
+import sys
+from logging import DEBUG, INFO, StreamHandler, getLogger
 from typing import Optional, Union
 
 import sounddevice as sd
@@ -83,13 +84,23 @@ def get_opts() -> argparse.Namespace:
     parser.add_argument(
         "--mic",
     )
+    parser.add_argument(
+        "--debug",
+        action="store_true",
+    )
 
     return parser.parse_args()
 
 
 def main() -> None:
     opts = get_opts()
-    logger.setLevel(INFO)
+
+    root_logger = getLogger()
+    root_logger.addHandler(StreamHandler(sys.stderr))
+    if opts.debug:
+        root_logger.setLevel(DEBUG)
+    else:
+        root_logger.setLevel(INFO)
     if opts.beam_size <= 0:
         opts.beam_size = None
     try:
diff --git a/whisper_streaming/transcriber.py b/whisper_streaming/transcriber.py
index 4120459..d444ab7 100644
--- a/whisper_streaming/transcriber.py
+++ b/whisper_streaming/transcriber.py
@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 
+from logging import INFO, getLogger
 from typing import Iterator, List, Optional, Union
 
 import numpy as np
@@ -19,6 +20,8 @@ from whisper.utils import exact_div
 
 from whisper_streaming.schema import ParsedChunk, WhisperConfig
 
+logger = getLogger(__name__)
+
 
 class WhisperStreamingTranscriber:
     def __init__(self, *, config: WhisperConfig):
@@ -214,6 +217,7 @@ class WhisperStreamingTranscriber:
         seek: int = 0
         rest_start: Optional[int] = None
         while seek < mel.shape[-1]:
+            logger.debug(seek)
             segment = (
                 pad_or_trim(mel[:, :, seek:], N_FRAMES)
                 .to(self.model.device)  # type: ignore