PizzAI/frontend/core.py

from vad_recorder import VADRecorder
from tts_stream import TTSStream
from chatgpt_wrap import ChatGPTWrap
from faster_whisper import WhisperModel
import torch
import time
import numpy as np
import re



class Core:
    def __init__(self, whisper_model_name = "large-v3", use_chatgpt_placeholder = False):
        self.use_chatgpt_placeholder = use_chatgpt_placeholder

        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        print("\n=======================================")
        print(f"Using {self.device.capitalize()} for:")
        print(" - Faster Whisper")
        print(" - TTS")
        print("=======================================\n")

        print("Loading Faster Whisper model... ", end="")
        self.whisper_model = WhisperModel(whisper_model_name, device=self.device, compute_type="float16")
        print("Done!")

        # VADRecorder, TTSStream and ChatGPTWrap have their own console loading messages
        self.vad_rec = VADRecorder()
        self.tts = TTSStream(device=self.device)
        self.gpt_wrap = ChatGPTWrap(use_chatgpt_placeholder)


    def set_order_settings(self, phone_number, order_items, delivery_address, payment_method):
        self.phone_number = phone_number
        self.order_items = order_items
        self.delivery_address = delivery_address
        self.payment_method = payment_method


    #def set_speech_recog_settings(self, speech_recog_timeout, audio_input_device_name, audio_output_device_name, window_size_sec, vad_threshold, min_silence_duration_ms, speech_pad_ms):
    def set_speech_recog_settings(self, speech_recog_timeout, window_size_sec, vad_threshold, min_silence_duration_ms, speech_pad_ms):
        self.speech_recog_timeout = speech_recog_timeout
        #self.audio_input_device_name = audio_input_device_name
        #self.audio_output_device_name = audio_output_device_name
        self.window_size_sec = window_size_sec
        self.vad_threshold = vad_threshold
        self.min_silence_duration_ms = min_silence_duration_ms
        self.speech_pad_ms = speech_pad_ms


    def set_tts_settings(self, speaker_wav):
        #self.speaker_wav = speaker_wav
        print("Setting TTS speaker... ", end="")
        self.tts.change_speaker(speaker_wav)
        print("Done!")


    def assistant_start(self):
        print("Starting assistant...")

        #print("Setting TTS speaker... ", end="")
        #self.tts.change_speaker(self.speaker_wav)
        #print("Done!")

        print("Starting VAD recording thread... ", end="")
        self.vad_rec.start_vad_recorder(
            #target_device_name = self.audio_input_device_name,
            window_size_sec = self.window_size_sec,
            vad_threshold = self.vad_threshold,
            min_silence_duration_ms = self.min_silence_duration_ms,
            speech_pad_ms = self.speech_pad_ms
        )
        print("Done!")

        self.gpt_wrap.init_order(
            self.phone_number, 
            self.order_items, 
            self.delivery_address, 
            self.payment_method
        )



        print("LISTENING!!!")


        last_recog_time = time.perf_counter()
        speech_recog_text = ""

        self.assistant_running = True

        while True:
            if self.vad_rec.speech:
                last_recog_time = time.perf_counter()

            if len(self.vad_rec.audios_for_whisper) > 0:
                #stream_out.write(audios_for_whisper.pop(0))
                audio = np.array(self.vad_rec.audios_for_whisper.pop(0), dtype=np.float32)

                segments, _ = self.whisper_model.transcribe(audio, language="pl")
                if not self.assistant_running:
                    break

                text = "".join([segment.text for segment in segments])
                #speech_recog_text += " " if len(speech_recog_text) else "" + text

                if len(text) == 0:
                    continue

                if not text[-1] in ".,!?":
                    text += "."

                speech_recog_text += text.strip() + "\n"

                print("=========================================")
                print(text)

                last_recog_time = time.perf_counter()



            elif time.perf_counter() - last_recog_time > self.speech_recog_timeout and len(speech_recog_text) > 0:
                speech_recog_text = speech_recog_text.strip()

                print("=========================================\n\n")
                print("-----------------------------------------")
                print("!!!!!!!!!! SENDING TO CHATGPT !!!!!!!!!!!")
                print("-----------------------------------------")
                print(speech_recog_text)
                print("-----------------------------------------\n\n")

                

                gpt_response = self.gpt_wrap.get_response(speech_recog_text)

                # separate long sequences of numbers in text string (for example 123456789) into packets of 3 (123 456 789)
                gpt_response = re.sub(r"(\d{3})(?=\d)", r"\1 ", gpt_response)

                # Add space on the right side of numbers
                gpt_response = re.sub(r'(\d)([^\d\s])', r'\1 \2', gpt_response)
                # Add space on the left side of numbers
                gpt_response = re.sub(r'([^\d\s])(\d)', r'\1 \2', gpt_response)

                # replace "ul." with "ulica" (non case sensitive)
                gpt_response = re.sub(r"ul\.", "ulica", gpt_response, flags=re.IGNORECASE)


                print("-----------------------------------------")
                if self.use_chatgpt_placeholder:
                    print("!!!!! CHATGPT PLACEHOLDER RESPONSE !!!!!!")
                else:
                    print("!!!!!!!!!!! CHATGPT RESPONSE !!!!!!!!!!!!")
                print("-----------------------------------------")
                print(gpt_response)
                print("-----------------------------------------\n\n")

                if not self.assistant_running:
                    break

                speech_recog_text = ""


                # tts
                print("Speech synthesis stream started!")
                self.tts.tts_speak(gpt_response.replace(" CALLEND", ""))


                if "CALLEND" in gpt_response:
                    self.assistant_stop()


            #print(len(audios_for_whisper), time.perf_counter() - last_recog_time, len(speech_recog_text))


            time.sleep(0.01)

            if not self.assistant_running:
                break

        # set assistant_running back to True to indicate that the loop has exited

    def assistant_stop(self):
        print("Stopping assistant... ", end="")

        self.assistant_running = False

        self.vad_rec.stop_vad_recorder()
        
        print("Done!")