RealChar / client /cli.py
pycui's picture
Add RealChar deployment for HuggingFace (V0)
babeaf6
raw
history blame
10.1 kB
import os
import queue
import asyncio
import concurrent.futures
import functools
import io
import sys
import random
from threading import Thread
import time
from dotenv import load_dotenv
import pyaudio
import speech_recognition as sr
import websockets
from aioconsole import ainput # for async input
from pydub import AudioSegment
from simpleaudio import WaveObject
load_dotenv()
executor = concurrent.futures.ThreadPoolExecutor(max_workers=3)
web2_initial_message = True
CHUNK = 1024
FORMAT = pyaudio.paInt16
CHANNELS = 1
RATE = 44100
class AudioPlayer:
def __init__(self):
self.play_thread = None
self.stop_flag = False
self.queue = queue.Queue()
def play_audio(self):
while not self.stop_flag or not self.queue.empty():
try:
wav_data = self.queue.get_nowait()
except queue.Empty:
continue
wave_obj = WaveObject.from_wave_file(wav_data)
play_obj = wave_obj.play()
while play_obj.is_playing() and not self.stop_flag:
time.sleep(0.1)
if self.stop_flag:
play_obj.stop()
def start_playing(self, wav_data):
self.stop_flag = False
self.queue.put(wav_data)
if self.play_thread is None or not self.play_thread.is_alive():
self.play_thread = Thread(target=self.play_audio)
self.play_thread.start()
def stop_playing(self):
if self.play_thread and self.play_thread.is_alive():
self.stop_flag = True
self.play_thread.join()
self.play_thread = None
def add_to_queue(self, wav_data):
self.queue.put(wav_data)
audio_player = AudioPlayer()
def get_input_device_id():
p = pyaudio.PyAudio()
devices = [(i, p.get_device_info_by_index(i)['name'])
for i in range(p.get_device_count())
if p.get_device_info_by_index(i).get('maxInputChannels')]
print('Available devices:')
for id, name in devices:
print(f"Device id {id} - {name}")
return int(input('Please select device id: '))
async def handle_audio(websocket, device_id):
with sr.Microphone(device_index=device_id, sample_rate=RATE) as source:
recognizer = sr.Recognizer()
print('Source sample rate: ', source.SAMPLE_RATE)
print('Source width: ', source.SAMPLE_WIDTH)
print('Adjusting for ambient noise...Wait for 2 seconds')
recognizer.energy_threshold = 5000
recognizer.dynamic_energy_ratio = 6
recognizer.dynamic_energy_adjustment_damping = 0.85
recognizer.non_speaking_duration = 0.5
recognizer.pause_threshold = 0.8
recognizer.phrase_threshold = 0.5
recognizer.adjust_for_ambient_noise(source, duration=2)
listen_func = functools.partial(
recognizer.listen, source, phrase_time_limit=30)
print('Okay, start talking!')
while True:
print('[*]', end="") # indicate that we are listening
audio = await asyncio.get_event_loop().run_in_executor(executor, listen_func)
await websocket.send(audio.frame_data)
print('[-]', end="") # indicate that we are done listening
await asyncio.sleep(2)
async def handle_text(websocket):
print('You: ', end="", flush=False)
while True:
message = await ainput()
await websocket.send(message)
initial_message = True
async def receive_message(websocket, websocket2):
web1_init_message = await websocket.recv()
print('web1_init_message: ', web1_init_message)
web2_init_message = await websocket2.recv()
print('web1_init_message: ', web2_init_message)
message_to_websocket1 = "Suppose I'm Steve Jobs now. What question do you have for me?"
await websocket.send(message_to_websocket1)
web1_message = ''
while True:
try:
message = await websocket.recv()
print('here')
except websockets.exceptions.ConnectionClosedError as e:
print("Connection closed unexpectedly: ", e)
break
except Exception as e:
print("An error occurred: ", e)
break
if isinstance(message, str):
if message == '[end]\n':
if not web1_message:
continue
# remove everything before '> ' in the message
message_to_websocket2 = web1_message[web1_message.find('> ') + 2:]
# print('message_to_websocket2: ', message_to_websocket2)
await websocket2.send(message_to_websocket2)
web2_message = ''
j = 0
while True:
j += 1
try:
message = await websocket2.recv()
except websockets.exceptions.ConnectionClosedError as e:
print("Connection closed unexpectedly: ", e)
break
except Exception as e:
print("An error occurred: ", e)
break
if isinstance(message, str):
if message == '[end]\n':
# print('\nWebsocket2: ', end="", flush=False)
if not web2_message:
# print('skip')
continue
# remove everything before '> ' in the message
print(web2_message)
message_from_websocket2 = web2_message[web2_message.find('> ') + 2:]
await websocket.send(message_from_websocket2)
break
elif message.startswith('[+]'):
# stop playing audio
audio_player.stop_playing()
# indicate the transcription is done
# print(f"\nnWebsocket2: {message}", end="\n", flush=False)
elif message.startswith('[=]'):
# indicate the response is done
# print(f"nWebsocket2: {web2_message}", end="\n", flush=False)
pass
else:
# print('\nmessage++\n')
web2_message += message
elif isinstance(message, bytes):
global web2_initial_message
if web2_initial_message:
web2_initial_message = False
continue
audio_data = io.BytesIO(message)
audio = AudioSegment.from_mp3(audio_data)
wav_data = io.BytesIO()
audio.export(wav_data, format="wav")
# Start playing audio
audio_player.start_playing(wav_data)
elif message.startswith('[+]'):
# stop playing audio
audio_player.stop_playing()
# indicate the transcription is done
print(f"\n{message}", end="\n", flush=False)
elif message.startswith('[=]'):
# indicate the response is done
print(f"{message}", end="\n", flush=False)
else:
web1_message += message
print(f"{message}", end="", flush=False)
elif isinstance(message, bytes):
audio_data = io.BytesIO(message)
audio = AudioSegment.from_mp3(audio_data)
wav_data = io.BytesIO()
audio.export(wav_data, format="wav")
# Start playing audio
audio_player.start_playing(wav_data)
else:
print("Unexpected message")
break
def select_model():
llm_model_selection = input(
'1: gpt-3.5-turbo-16k \n'
'2: gpt-4 \n'
'3: claude-2 \n'
'Select llm model:')
if llm_model_selection == '1':
llm_model = 'gpt-3.5-turbo-16k'
elif llm_model_selection == '2':
llm_model = 'gpt-4'
elif llm_model_selection == '3':
llm_model = 'claude-2'
return llm_model
async def start_client(client_id, url):
api_key = os.getenv('AUTH_API_KEY')
llm_model = select_model()
uri = f"ws://{url}/ws/{client_id}?api_key={api_key}&llm_model={llm_model}"
async with websockets.connect(uri) as websocket:
uri2 = f"ws://{url}/ws/9999999?api_key={api_key}&llm_model={llm_model}"
# send client platform info
async with websockets.connect(uri2) as websocket2:
await websocket.send('terminal')
await websocket2.send('terminal')
print(f"Client #{client_id} connected to websocket1")
print(f"Client 9999999 connected to websocket2")
welcome_message = await websocket.recv()
welcome_message2 = await websocket2.recv()
print(f"{welcome_message}")
character = input('Select character: ')
await websocket.send(character)
await websocket2.send('6')
mode = input('Select mode (1: audio, 2: text): ')
if mode.lower() == '1':
device_id = get_input_device_id()
send_task = asyncio.create_task(handle_audio(websocket, device_id))
else:
send_task = asyncio.create_task(handle_text(websocket))
receive_task = asyncio.create_task(receive_message(websocket, websocket2))
await asyncio.gather(receive_task, send_task)
async def main(url):
client_id = random.randint(0, 1000000)
task = asyncio.create_task(start_client(client_id, url))
try:
await task
except KeyboardInterrupt:
task.cancel()
await asyncio.wait_for(task, timeout=None)
print("Client stopped by user")
if __name__ == "__main__":
url = sys.argv[1] if len(sys.argv) > 1 else 'localhost:8000'
asyncio.run(main(url))