import pyttsx3 import numpy as np import whisper import pyaudio import sys import torch import requests import json import yaml from yaml import Loader import pygame, sys import pygame.locals BACK_COLOR = (0,0,0) REC_COLOR = (255,0,0) TEXT_COLOR = (255,255,255) REC_SIZE = 80 FONT_SIZE = 24 WIDTH = 320 HEIGHT = 240 INPUT_DEFAULT_DURATION_SECONDS = 5 INPUT_FORMAT = pyaudio.paInt16 INPUT_CHANNELS = 1 INPUT_RATE = 16000 INPUT_CHUNK = 1024 OLLAMA_REST_HEADERS = {'Content-Type': 'application/json',} INPUT_CONFIG_PATH ="assistant.yaml" class Assistant: def __init__(self): self.config = self.initConfig() programIcon = pygame.image.load('assistant.png') self.clock = pygame.time.Clock() pygame.display.set_icon(programIcon) pygame.display.set_caption("Assistant") self.windowSurface = pygame.display.set_mode((WIDTH, HEIGHT), 0, 32) self.font = pygame.font.SysFont(None, FONT_SIZE) self.audio = pyaudio.PyAudio() try: self.audio.open(format=INPUT_FORMAT, channels=INPUT_CHANNELS, rate=INPUT_RATE, input=True, frames_per_buffer=INPUT_CHUNK).close() except : self.wait_exit() self.display_message(self.config.messages.loadingModel) self.model = whisper.load_model(self.config.whisperRecognition.modelPath) self.tts = pyttsx3.init() #self.conversation_history = [self.config.conversation.context, # self.config.conversation.greeting] self.context = [] self.display_ready() self.text_to_speech(self.config.conversation.greeting) def wait_exit(self): while True: self.display_message(self.config.messages.noAudioInput) self.clock.tick(60) for event in pygame.event.get(): if event.type == pygame.locals.QUIT: self.shutdown() def shutdown(self): self.audio.terminate() pygame.quit() sys.exit() def initConfig(self): class Inst: pass config=Inst(); config.messages = Inst() config.messages.pressSpace = "Pressez sur espace pour parler puis relachez." config.messages.loadingModel = "Loading model..." config.messages.noAudioInput = "Erreur: Pas d'entrée son" config.whisperRecognition = Inst() config.whisperRecognition.modelPath = "whisper/large-v3.pt" config.whisperRecognition.lang = "fr" config.ollama = Inst() config.ollama.url = "http://localhost:11434/api/generate" config.ollama.model = 'mistral' config.conversation = Inst() config.conversation.context = "This is a discussion in french.\n" config.conversation.greeting = "Je vous écoute." config.conversation.recognitionWaitMsg = "J'interprète votre demande." config.conversation.llmWaitMsg = "Laissez moi réfléchir." stream = open(INPUT_CONFIG_PATH, 'r', encoding="utf-8") dic = yaml.load(stream, Loader=Loader) #dic depth 2: map values to attributes def dic2Object(dic, object): for key in dic: if hasattr(object, key): setattr(object, key, dic[key]) else: print("Ignoring unknow setting ", key) #dic depth 1: fill depth 2 attributes for key in dic: if hasattr(config, key): dic2Object(dic[key], getattr(config, key)) else: print("Ignoring unknow setting ", key) return config def display_rec_start(self): self.windowSurface.fill(BACK_COLOR) pygame.draw.circle(self.windowSurface, REC_COLOR, (WIDTH/2, HEIGHT/2), REC_SIZE) pygame.display.flip() def display_message(self, text): self.windowSurface.fill(BACK_COLOR) label = self.font.render(text, 1, TEXT_COLOR) size = label.get_rect()[2:4] self.windowSurface.blit(label, (WIDTH/2 - size[0]/2, HEIGHT/2 - size[1]/2)) pygame.display.flip() def display_ready(self): self.display_message(self.config.messages.pressSpace) def waveform_from_mic(self, key = pygame.K_SPACE) -> np.ndarray: self.display_rec_start() stream = self.audio.open(format=INPUT_FORMAT, channels=INPUT_CHANNELS, rate=INPUT_RATE, input=True, frames_per_buffer=INPUT_CHUNK) frames = [] while True: pygame.event.pump() # process event queue pressed = pygame.key.get_pressed() if pressed[key]: data = stream.read(INPUT_CHUNK) frames.append(data) else: break stream.stop_stream() stream.close() self.display_ready() return np.frombuffer(b''.join(frames), np.int16).astype(np.float32) * (1 / 32768.0) def speech_to_text(self, waveform): self.text_to_speech(self.config.conversation.recognitionWaitMsg) transcript = self.model.transcribe(waveform, language = self.config.whisperRecognition.lang, fp16=torch.cuda.is_available()) text = transcript["text"] self.text_to_speech(text) return text def ask_ollama(self, prompt, responseCallback): #self.conversation_history.append(prompt) #full_prompt = "\n".join(self.conversation_history) full_prompt = prompt if hasattr(self, "contextSent") else (self.config.conversation.context+"\n"+prompt) self.contextSent = True jsonParam= {"model": self.config.ollama.model, "stream":True, "context":self.context, "prompt":full_prompt} response = requests.post(self.config.ollama.url, json=jsonParam, headers=OLLAMA_REST_HEADERS, stream=True) response.raise_for_status() print(jsonParam) self.text_to_speech(self.config.conversation.llmWaitMsg) tokens = [] for line in response.iter_lines(): print(line) body = json.loads(line) token = body.get('response', '') tokens.append(token) # the response streams one token at a time, process only at end of sentences if token == "." or token == ":" or token == "!" or token == "?": current_response = "".join(tokens) #self.conversation_history.append(current_response) responseCallback(current_response) tokens = [] if 'error' in body: responseCallback("Erreur: " + body['error']) if body.get('done', False): self.context = body['context'] def text_to_speech(self, text): print(text) self.tts.say(text) self.tts.runAndWait() def main(): if sys.version_info[0:3] != (3, 9, 13): print('Warning, it was only tested with python 3.9.13, it may fail') pygame.init() ass = Assistant() push_to_talk_key = pygame.K_SPACE; while True: ass.clock.tick(60) for event in pygame.event.get(): if event.type == pygame.KEYDOWN and event.key == push_to_talk_key: print('Talk to me!') speech = ass.waveform_from_mic(push_to_talk_key) transcription = ass.speech_to_text(waveform=speech) ass.ask_ollama(transcription, ass.text_to_speech) print('Done') if event.type == pygame.locals.QUIT: ass.shutdown() if __name__ == "__main__": main()