diff --git a/README.md b/README.md index f6db3c1..e2f6844 100644 --- a/README.md +++ b/README.md @@ -20,7 +20,7 @@ Now a pip package! ## Available models and languages -There are five model sizes, four with English-only versions, offering speed and accuracy tradeoffs. Below are the names of the available models and their approximate memory requirements and relative speed. +There are five model sizes, four with English-only versions, offering speed and accuracy tradeoffs. Below are the names of the available models and their approximate memory requirements and relative speed. | Size | Parameters | English-only model | Multilingual model | Required VRAM | Relative speed | @@ -63,6 +63,11 @@ If you are having issues, try the following: ``` sudo apt install portaudio19-dev python3-pyaudio ``` +You'll also need the python3-dev for your specific python version... +example for python 3.10: +``` +sudo apt-get install python3.10-dev +``` ## Contributing diff --git a/requirements.txt b/requirements.txt index 9deb1e9..73ac430 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,4 @@ +faster-whisper numpy tqdm more-itertools diff --git a/whisper_mic/cli.py b/whisper_mic/cli.py index bba7bd4..af50506 100755 --- a/whisper_mic/cli.py +++ b/whisper_mic/cli.py @@ -1,4 +1,5 @@ #!/usr/bin/env python3 +# flake8:noqa: E304 import click import torch @@ -7,26 +8,135 @@ from whisper_mic import WhisperMic + @click.command() -@click.option("--model", default="base", help="Model to use", type=click.Choice(["tiny","base", "small","medium","large","large-v2","large-v3"])) -@click.option("--device", default=("cuda" if torch.cuda.is_available() else "cpu"), help="Device to use", type=click.Choice(["cpu","cuda","mps"])) -@click.option("--english", default=False, help="Whether to use English model",is_flag=True, type=bool) -@click.option("--verbose", default=False, help="Whether to print verbose output", is_flag=True,type=bool) -@click.option("--energy", default=300, help="Energy level for mic to detect", type=int) -@click.option("--dynamic_energy", default=False,is_flag=True, help="Flag to enable dynamic energy", type=bool) -@click.option("--pause", default=0.8, help="Pause time before entry ends", type=float) -@click.option("--save_file",default=False, help="Flag to save file", is_flag=True,type=bool) -@click.option("--loop", default=False, help="Flag to loop", is_flag=True,type=bool) -@click.option("--dictate", default=False, help="Flag to dictate (implies loop)", is_flag=True,type=bool) -@click.option("--mic_index", default=None, help="Mic index to use", type=int) -@click.option("--list_devices",default=False, help="Flag to list devices", is_flag=True,type=bool) -@click.option("--faster",default=False, help="Use faster_whisper implementation", is_flag=True,type=bool) -@click.option("--hallucinate_threshold",default=400, help="Raise this to reduce hallucinations. Lower this to activate more often.", is_flag=False,type=int) -def main(model: str, english: bool, verbose: bool, energy: int, pause: float, dynamic_energy: bool, save_file: bool, device: str, loop: bool, dictate: bool,mic_index:Optional[int],list_devices: bool,faster: bool,hallucinate_threshold:int) -> None: +@click.option( + "--model", + default="base", + help="Model to use", + type=click.Choice(["tiny", "base", "small", "medium", "large", "large-v2", "large-v3"]),) + +@click.option( + "--device", + default=("cuda" if torch.cuda.is_available() else "cpu"), + help="Device to use", + type=click.Choice(["cpu", "cuda", "mps"]),) + +@click.option( + "--english", + default=False, + help="Whether to use English model", + is_flag=True, + type=bool,) + +@click.option( + "--verbose", + default=False, + help="Whether to print verbose output", + is_flag=True, + type=bool,) + +@click.option( + "--energy", + default=300, + help="Energy level for mic to detect", + type=int,) + +@click.option( + "--dynamic_energy", + default=False, + is_flag=True, + help="Flag to enable dynamic energy", + type=bool,) + +@click.option( + "--pause", + default=0.8, + help="Pause time before entry ends", + type=float,) + +@click.option( + "--save_file", + default=False, + help="Flag to save file", + is_flag=True, + type=bool,) + +@click.option( + "--loop", + default=False, + help="Flag to loop", + is_flag=True, + type=bool,) + +@click.option( + "--dictate", + default=False, + help="Flag to dictate (implies loop)", + is_flag=True, + type=bool,) + +@click.option( + "--mic_index", + default=None, + help="Mic index to use", + type=int,) + +@click.option( + "--list_devices", + default=False, + help="Flag to list devices", + is_flag=True, + type=bool,) + +@click.option( + "--faster", + default=False, + help="Use faster_whisper implementation", + is_flag=True, + type=bool,) + +@click.option( + "--hallucinate_threshold", + default=400, + help="Raise this to reduce hallucinations. Lower this to activate more often.", + is_flag=False, + type=int,) + + +def main( + model: str, + english: bool, + verbose: bool, + energy: int, + pause: float, + dynamic_energy: bool, + save_file: bool, + device: str, + loop: bool, + dictate: bool, + mic_index: Optional[int], + list_devices: bool, + faster: bool, + hallucinate_threshold: int +) -> None: + if list_devices: - print("Possible devices: ",sr.Microphone.list_microphone_names()) + print("Possible devices: ", sr.Microphone.list_microphone_names()) return - mic = WhisperMic(model=model, english=english, verbose=verbose, energy=energy, pause=pause, dynamic_energy=dynamic_energy, save_file=save_file, device=device,mic_index=mic_index,implementation=("faster_whisper" if faster else "whisper"),hallucinate_threshold=hallucinate_threshold) + + mic = WhisperMic( + model=model, + english=english, + verbose=verbose, + energy=energy, + pause=pause, + dynamic_energy=dynamic_energy, + save_file=save_file, + device=device, + mic_index=mic_index, + implementation=("faster_whisper" if faster else "whisper"), + hallucinate_threshold=hallucinate_threshold,) if not loop: try: @@ -39,12 +149,13 @@ def main(model: str, english: bool, verbose: bool, energy: int, pause: float, d mic.file.close() else: try: - mic.listen_loop(dictate=dictate,phrase_time_limit=2) + mic.listen_loop(dictate=dictate, phrase_time_limit=2) except KeyboardInterrupt: print("Operation interrupted successfully") finally: if save_file: mic.file.close() + if __name__ == "__main__": main() diff --git a/whisper_mic/whisper_mic.py b/whisper_mic/whisper_mic.py index 1ae74f8..b7372e0 100755 --- a/whisper_mic/whisper_mic.py +++ b/whisper_mic/whisper_mic.py @@ -1,19 +1,20 @@ -import torch -import queue -import speech_recognition as sr -import threading import numpy as np import os -import time -import tempfile import platform import pynput.keyboard +import queue +import speech_recognition as sr +import tempfile +import threading +import time +import torch +from typing import AsyncGenerator from typing import Optional # from ctypes import * from whisper_mic.utils import get_logger -#TODO: This is a linux only fix and needs to be testd. Have one for mac and windows too. +# TODO: This is a linux only fix and needs to be testd. Have one for mac and windows too. # Define a null error handler for libasound to silence the error message spam # def py_error_handler(filename, line, function, err, fmt): # None @@ -23,10 +24,26 @@ # asound = cdll.LoadLibrary('libasound.so') # asound.snd_lib_error_set_handler(c_error_handler) -class WhisperMic: - def __init__(self,model="base",device=("cuda" if torch.cuda.is_available() else "cpu"),english=False,verbose=False,energy=300,pause=2,dynamic_energy=False,save_file=False, model_root="~/.cache/whisper",mic_index=None,implementation="whisper",hallucinate_threshold=300): - self.logger = get_logger("whisper_mic", "info") + +class WhisperMic: + def __init__( + self, + model='base', + device='cuda' if torch.cuda.is_available() else 'cpu', + english=False, + verbose=False, + energy=300, + pause=2, + dynamic_energy=False, + save_file=False, + model_root='~/.cache/whisper', + mic_index=None, + implementation='whisper', + hallucinate_threshold=300 + ): + self.device = device + self.logger = get_logger('whisper_mic', 'info') self.energy = energy self.hallucinate_threshold = hallucinate_threshold self.pause = pause @@ -37,55 +54,64 @@ def __init__(self,model="base",device=("cuda" if torch.cuda.is_available() else self.keyboard = pynput.keyboard.Controller() self.platform = platform.system().lower() - if self.platform == "darwin": - if device == "cuda" or device == "mps": - self.logger.warning("CUDA is not supported on MacOS and mps does not work. Using CPU instead.") - device = "cpu" - else: - device = "cuda" if torch.cuda.is_available() else "cpu" + if self.platform == 'darwin': + if self.device == 'cuda' or self.device == 'mps': + self.logger.warning('CUDA is not supported on MacOS and mps does not work. Using CPU instead.') + self.device = 'cpu' - if (model != "large" and model != "large-v2" and model!= "large-v3") and self.english: - model = model + ".en" + if (model != 'large' and model != 'large-v2' and model != 'large-v3') and self.english: + model = model + '.en' model_root = os.path.expanduser(model_root) self.faster = False - if (implementation == "faster_whisper"): + if (implementation == 'faster_whisper'): try: from faster_whisper import WhisperModel - self.audio_model = WhisperModel(model, download_root=model_root, device="auto", compute_type="int8") - self.faster = True # Only set the flag if we succesfully imported the library and opened the model. + self.audio_model = WhisperModel( + model, + download_root=model_root, + device=self.device, + compute_type='int8',) + self.faster = True # Only set the flag if we succesfully imported the library and opened the model. except ImportError: - self.logger.error("faster_whisper not installed, falling back to whisper") - self.logger.info("To install faster_whisper, run 'pip install faster_whisper'") - import whisper - self.audio_model = whisper.load_model(model, download_root=model_root).to(device) + self.logger.error('faster_whisper not installed, falling back to whisper') + self.logger.info('To install faster_whisper, run "pip install faster_whisper"') + import whisper + self.audio_model = whisper.load_model( + model, + download_root=model_root, + device=self.device, + ).to(self.device) else: import whisper - self.audio_model = whisper.load_model(model, download_root=model_root).to(device) - + self.audio_model = whisper.load_model( + model, + download_root=model_root, + device=self.device, + ).to(self.device) + self.temp_dir = tempfile.mkdtemp() if save_file else None self.audio_queue = queue.Queue() - self.result_queue: "queue.Queue[str]" = queue.Queue() - + self.result_queue: 'queue.Queue[str]' = queue.Queue() + self.break_threads = False self.mic_active = False - self.banned_results = [""," ","\n",None] + self.banned_results = ['', ' ', '\n', None] if save_file: - self.file = open("transcribed_text.txt", "w+", encoding="utf-8") + self.file = open('transcribed_text.txt', 'w+', encoding='utf-8') self.__setup_mic(mic_index) - def __setup_mic(self, mic_index): if mic_index is None: - self.logger.info("No mic index provided, using default") - self.source = sr.Microphone(sample_rate=16000, device_index=mic_index) + self.logger.info('No mic index provided, using default') + self.source = sr.Microphone(sample_rate=16000, device_index=mic_index) self.recorder = sr.Recognizer() self.recorder.energy_threshold = self.energy self.recorder.pause_threshold = self.pause @@ -94,37 +120,44 @@ def __setup_mic(self, mic_index): with self.source: self.recorder.adjust_for_ambient_noise(self.source) - self.logger.info("Mic setup complete") + self.logger.info('Mic setup complete') - # Whisper takes a Tensor while faster_whisper only wants an NDArray + """ Whisper takes a Tensor while faster_whisper only wants an NDArray """ def __preprocess(self, data): is_audio_loud_enough = self.is_audio_loud_enough(data) + if self.faster: - return np.frombuffer(data, np.int16).flatten().astype(np.float32) / 32768.0,is_audio_loud_enough + return np.frombuffer( + data, np.int16, + ).flatten().astype(np.float32) / 32768.0, is_audio_loud_enough else: - return torch.from_numpy(np.frombuffer(data, np.int16).flatten().astype(np.float32) / 32768.0),is_audio_loud_enough - + return torch.from_numpy(np.frombuffer( + data, np.int16, + ).flatten().astype(np.float32) / 32768.0), is_audio_loud_enough + def is_audio_loud_enough(self, frame) -> bool: audio_frame = np.frombuffer(frame, dtype=np.int16) amplitude = np.mean(np.abs(audio_frame)) return amplitude > self.hallucinate_threshold - def __get_all_audio(self, min_time: float = -1.): audio = bytes() got_audio = False time_start = time.time() + while not got_audio or time.time() - time_start < min_time: while not self.audio_queue.empty(): audio += self.audio_queue.get() got_audio = True - data = sr.AudioData(audio,16000,2) + data = sr.AudioData(audio, 16000, 2) data = data.get_raw_data() return data - - # Handles the task of getting the audio input via microphone. This method has been used for listen() method + """ + Handles the task of getting the audio input via microphone. + This method has been used for listen() method + """ def __listen_handler(self, timeout, phrase_time_limit): try: with self.source as microphone: @@ -133,54 +166,56 @@ def __listen_handler(self, timeout, phrase_time_limit): audio_data = self.__get_all_audio() self.__transcribe(data=audio_data) except sr.WaitTimeoutError: - self.result_queue.put_nowait("Timeout: No speech detected within the specified time.") + self.result_queue.put_nowait('Timeout: No speech detected within the specified time.') except sr.UnknownValueError: - self.result_queue.put_nowait("Speech recognition could not understand audio.") - + self.result_queue.put_nowait('Speech recognition could not understand audio.') - # This method is similar to the __listen_handler() method but it has the added ability for recording the audio for a specified duration of time + """ + This method is similar to the __listen_handler() method + but it has the added ability for recording the audio for a specified duration of time + """ def __record_handler(self, duration=2, offset=None): with self.source as microphone: audio = self.recorder.record(source=microphone, duration=duration, offset=offset) - + self.__record_load(0, audio) audio_data = self.__get_all_audio() self.__transcribe(data=audio_data) - - # This method takes the recorded audio data, converts it into raw format and stores it in a queue. - def __record_load(self,_, audio: sr.AudioData) -> None: + """ + This method takes the recorded audio data, + converts it into raw format and stores it in a queue. + """ + def __record_load(self, _, audio: sr.AudioData) -> None: data = audio.get_raw_data() self.audio_queue.put_nowait(data) - def __transcribe_forever(self) -> None: while True: if self.break_threads: break self.__transcribe() - - def __transcribe(self,data=None, realtime: bool = False) -> None: + def __transcribe(self, data=None, realtime: bool = False) -> None: if data is None: audio_data = self.__get_all_audio() else: audio_data = data - audio_data,is_audio_loud_enough = self.__preprocess(audio_data) + audio_data, is_audio_loud_enough = self.__preprocess(audio_data) if is_audio_loud_enough: predicted_text = '' - # faster_whisper returns an iterable object rather than a string - if self.faster: + + if self.faster: # faster_whisper returns an iterable object rather than a string segments, info = self.audio_model.transcribe(audio_data) for segment in segments: predicted_text += segment.text else: if self.english: - result = self.audio_model.transcribe(audio_data,language='english',suppress_tokens="") + result = self.audio_model.transcribe(audio_data, language='english', suppress_tokens='') else: - result = self.audio_model.transcribe(audio_data,suppress_tokens="") - predicted_text = result["text"] + result = self.audio_model.transcribe(audio_data, suppress_tokens='') + predicted_text = result['text'] if not self.verbose: if predicted_text not in self.banned_results: @@ -189,22 +224,28 @@ def __transcribe(self,data=None, realtime: bool = False) -> None: if predicted_text not in self.banned_results: self.result_queue.put_nowait(result) - if self.save_file: # os.remove(audio_data) self.file.write(predicted_text) else: - # If the audio is not loud enough, we put None in the queue to indicate that we need to listen again or return None + """ + If the audio is not loud enough, + we put None in the queue to indicate that we need to listen again or return None + """ self.result_queue.put_nowait(None) - async def listen_loop_async(self, dictate: bool = False, phrase_time_limit=None) -> Optional[str]: + async def listen_loop_async( + self, + dictate: bool = False, + phrase_time_limit=None, + ) -> Optional[AsyncGenerator[str, None]]: + for result in self.listen_continuously(phrase_time_limit=phrase_time_limit): if dictate: self.keyboard.type(result) else: yield result - def listen_loop(self, dictate: bool = False, phrase_time_limit=None) -> None: for result in self.listen_continuously(phrase_time_limit=phrase_time_limit): if result is not None: @@ -213,53 +254,57 @@ def listen_loop(self, dictate: bool = False, phrase_time_limit=None) -> None: else: print(result) - def listen_continuously(self, phrase_time_limit=None): self.recorder.listen_in_background(self.source, self.__record_load, phrase_time_limit=phrase_time_limit) - self.logger.info("Listening...") + self.logger.info('Listening...') threading.Thread(target=self.__transcribe_forever, daemon=True).start() while True: yield self.result_queue.get() - - def listen(self, timeout = None, phrase_time_limit=None,try_again=True): - self.logger.info("Listening...") + def listen(self, timeout=None, phrase_time_limit=None, try_again=True): + self.logger.info('Listening...') self.__listen_handler(timeout, phrase_time_limit) + while True: if not self.result_queue.empty(): result = self.result_queue.get() + if result is None and try_again: - self.logger.info("Too quiet, listening again...") - result = self.listen(timeout=timeout, phrase_time_limit=phrase_time_limit,try_again=True) + self.logger.info('Too quiet, listening again...') + result = self.listen(timeout=timeout, phrase_time_limit=phrase_time_limit, try_again=True) return result else: return result + """ + This method is similar to the listen() method, + but it has the ability to listen for a specified duration, + mentioned in the "duration" parameter. + """ + def record(self, duration=2, offset=None, try_again=True): + self.logger.info('Listening...') - # This method is similar to the listen() method, but it has the ability to listen for a specified duration, mentioned in the "duration" parameter. - def record(self, duration=2, offset=None,try_again=True): - self.logger.info("Listening...") if duration is None: - self.logger.warning("Duration not provided, may hang indefinitely.") + self.logger.warning('Duration not provided, may hang indefinitely.') self.__record_handler(duration, offset) + while True: if not self.result_queue.empty(): result = self.result_queue.get() if result is None and try_again: - self.logger.info("Too quiet, listening again...") - result = self.record(duration=duration, offset=offset,try_again=True) + self.logger.info('Too quiet, listening again...') + result = self.record(duration=duration, offset=offset, try_again=True) return result else: return result - def toggle_microphone(self) -> None: - #TO DO: make this work + # TODO: make this work self.mic_active = not self.mic_active if self.mic_active: - print("Mic on") + print('Mic on') else: - print("turning off mic") + print('turning off mic') self.mic_thread.join() - print("Mic off") + print('Mic off')