mirror of
https://github.com/shirayu/whispering.git
synced 2024-05-20 01:08:14 +00:00
106 lines
2.4 KiB
Python
106 lines
2.4 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import argparse
|
|
import queue
|
|
from logging import INFO, getLogger
|
|
from typing import Optional, Union
|
|
|
|
import sounddevice as sd
|
|
import torch
|
|
from whisper import available_models
|
|
from whisper.audio import N_FRAMES, SAMPLE_RATE
|
|
from whisper.tokenizer import LANGUAGES, TO_LANGUAGE_CODE
|
|
|
|
from whisper_streaming.schema import WhisperConfig
|
|
from whisper_streaming.transcriber import WhisperStreamingTranscriber
|
|
|
|
logger = getLogger(__name__)
|
|
|
|
|
|
def transcribe_from_mic(
|
|
*,
|
|
config: WhisperConfig,
|
|
sd_device: Optional[Union[int, str]],
|
|
) -> None:
|
|
wsp = WhisperStreamingTranscriber(config=config)
|
|
q = queue.Queue()
|
|
|
|
def sd_callback(indata, frames, time, status):
|
|
if status:
|
|
logger.warning(status)
|
|
q.put(indata.ravel())
|
|
|
|
logger.info("Ready to transcribe")
|
|
with sd.InputStream(
|
|
samplerate=SAMPLE_RATE,
|
|
blocksize=N_FRAMES * 10, # FIXME
|
|
device=sd_device,
|
|
dtype="float32",
|
|
channels=1,
|
|
callback=sd_callback,
|
|
):
|
|
while True:
|
|
segment = q.get()
|
|
for chunk in wsp.transcribe(segment=segment):
|
|
print(f"{chunk.start}->{chunk.end}\t{chunk.text}")
|
|
|
|
|
|
def get_opts() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"--language",
|
|
type=str,
|
|
default=None,
|
|
choices=sorted(LANGUAGES.keys())
|
|
+ sorted([k.title() for k in TO_LANGUAGE_CODE.keys()]),
|
|
required=True,
|
|
)
|
|
parser.add_argument(
|
|
"--model",
|
|
type=str,
|
|
choices=available_models(),
|
|
required=True,
|
|
)
|
|
parser.add_argument(
|
|
"--device",
|
|
default="cuda" if torch.cuda.is_available() else "cpu",
|
|
help="device to use for PyTorch inference",
|
|
)
|
|
parser.add_argument(
|
|
"--beam_size",
|
|
"-b",
|
|
type=int,
|
|
default=5,
|
|
)
|
|
parser.add_argument(
|
|
"--mic",
|
|
)
|
|
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> None:
|
|
opts = get_opts()
|
|
logger.setLevel(INFO)
|
|
if opts.beam_size <= 0:
|
|
opts.beam_size = None
|
|
try:
|
|
opts.mic = int(opts.mic)
|
|
except Exception:
|
|
pass
|
|
|
|
config = WhisperConfig(
|
|
model_name=opts.model,
|
|
language=opts.language,
|
|
device=opts.device,
|
|
beam_size=opts.beam_size,
|
|
)
|
|
transcribe_from_mic(
|
|
config=config,
|
|
sd_device=opts.mic,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|