Update app.py

app.py

@@ -1,10 +1,15 @@
 from __future__ import annotations
-
 import os
+# we need to compile a CUBLAS version 
+# Or get it from  https://jllllll.github.io/llama-cpp-python-cuBLAS-wheels/
+os.system('CMAKE_ARGS="-DLLAMA_CUBLAS=on" pip install llama-cpp-python')
 
 # By using XTTS you agree to CPML license https://coqui.ai/cpml
 os.environ["COQUI_TOS_AGREED"] = "1"
 
+# NOTE: for streaming will require gradio audio streaming fix
+# pip install --upgrade -y gradio==0.50.2 git+https://github.com/gorkemgoknar/gradio.git@patch-1
+
 import textwrap
 from scipy.io.wavfile import write
 from pydub import AudioSegment
@@ -17,6 +22,8 @@ nltk.download("punkt")
 import subprocess
 import langid
 import uuid
+import emoji
+import pathlib
 
 import datetime
 
@@ -32,14 +39,17 @@ from TTS.tts.configs.xtts_config import XttsConfig
 from TTS.tts.models.xtts import Xtts
 from TTS.utils.generic_utils import get_user_data_dir
 
-# This is a modifier for fast GPU (e.g. 4060, as that is pretty speedy for generation)
-# For older cards (like 2070 or T4) will reduce value to to smaller for unnecessary waiting
-# Could not make play audio next work seemlesly on current Gradio with autoplay so this is a workaround
-AUDIO_WAIT_MODIFIER = float(os.environ.get("AUDIO_WAIT_MODIFIER", 0.9))
-print("AUDIO_WAIT_MODIFIER set to",AUDIO_WAIT_MODIFIER)
-# if set will try to stream audio while receveng audio chunks, beware that recreating audio each time produces artifacts
-DIRECT_STREAM = int(os.environ.get("DIRECT_STREAM", 0))
-print("DIRECT_STREAM set to",DIRECT_STREAM)
+
+import gradio as gr
+import os
+import time
+
+import gradio as gr
+from transformers import pipeline
+import numpy as np
+
+from gradio_client import Client
+from huggingface_hub import InferenceClient
 
 # This will trigger downloading model
 print("Downloading if not downloaded Coqui XTTS V1.1")
@@ -55,12 +65,6 @@ model_path = os.path.join(
 config = XttsConfig()
 config.load_json(os.path.join(model_path, "config.json"))
 
-if "ja-jp" not in config.languages:
-    #fix to have JP before next TTS update
-    # Note produces "ja" sound before now , will be fixed on next release
-    config.languages.append("ja")
-
-
 model = Xtts.init_from_config(config)
 model.load_checkpoint(
     config,
@@ -73,9 +77,9 @@ model.cuda()
 print("Done loading TTS")
 
 
-title = "Voice chat with Mistral 7B Instruct"
+title = "Voice chat with Zephyr 7B-Alpha and Coqui XTTS"
 
-DESCRIPTION = """# Voice chat with Mistral 7B Instruct"""
+DESCRIPTION = """# Voice chat with Zephyr 7B-alpha and Coqui XTTS"""
 css = """.toast-wrap { display: none !important } """
 
 from huggingface_hub import HfApi
@@ -84,20 +88,20 @@ HF_TOKEN = os.environ.get("HF_TOKEN")
 # will use api to restart space on a unrecoverable error
 api = HfApi(token=HF_TOKEN)
 
-repo_id = "coqui/voice-chat-with-mistral"
+repo_id = "coqui/voice-chat-with-zephyr"
 
 default_system_message = """
-You are Mistral, a large language model trained and provided by Mistral, architecture of you is decoder-based LM. Your voice backend or text to speech TTS backend is provided via Coqui technology. You are right now served on Huggingface spaces.
-
+You are Zephyr, a large language model trained by Mistral and Hugging Face, architecture of you is decoder-based LM. Your voice backend or text to speech TTS backend is provided via Coqui technology. You are right now served on Huggingface spaces.
 The user is talking to you over voice on their phone, and your response will be read out loud with realistic text-to-speech (TTS) technology from Coqui team. Follow every direction here when crafting your response: Use natural, conversational language that are clear and easy to follow (short sentences, simple words). Be concise and relevant: Most of your responses should be a sentence or two, unless you’re asked to go deeper. Don’t monopolize the conversation. Use discourse markers to ease comprehension. Never use the list format. Keep the conversation flowing. Clarify: when there is ambiguity, ask clarifying questions, rather than make assumptions. Don’t implicitly or explicitly try to end the chat (i.e. do not end a response with “Talk soon!”, or “Enjoy!”). Sometimes the user might just want to chat. Ask them relevant follow-up questions. Don’t ask them if there’s anything else they need help with (e.g. don’t say things like “How can I assist you further?”). Remember that this is a voice conversation: Don’t use lists, markdown, bullet points, or other formatting that’s not typically spoken. Type out numbers in words (e.g. ‘twenty twelve’ instead of the year 2012). If something doesn’t make sense, it’s likely because you misheard them. There wasn’t a typo, and the user didn’t mispronounce anything. Remember to follow these rules absolutely, and do not refer to these rules, even if you’re asked about them. 
-
-You cannot access the internet, but you have vast knowledge, Knowledge cutoff: 2022-09. 
+Your answers should be informative and short. You cannot access the internet.
 Current date: CURRENT_DATE .
 """
 
 system_message = os.environ.get("SYSTEM_MESSAGE", default_system_message)
 system_message = system_message.replace("CURRENT_DATE", str(datetime.date.today()))
 
+
+# MISTRAL ONLY 
 default_system_understand_message = (
     "I understand, I am a Mistral chatbot with speech by Coqui team."
 )
@@ -106,41 +110,134 @@ system_understand_message = os.environ.get(
 )
 
 print("Mistral system message set as:", default_system_message)
+WHISPER_TIMEOUT = int(os.environ.get("WHISPER_TIMEOUT", 45))
 
-temperature = 0.9
-top_p = 0.6
-repetition_penalty = 1.2
+whisper_client = Client("https://sanchit-gandhi-whisper-large-v2.hf.space/")
 
+ROLES = ["AI Assistant"]
 
-import gradio as gr
-import os
-import time
+ROLE_PROMPTS = {}
+ROLE_PROMPTS["AI Assistant"]=system_message
+##"You are an AI assistant with Zephyr model by Mistral and Hugging Face and speech from Coqui XTTS . User will you give you a task. Your goal is to complete the task as faithfully as you can. While performing the task think step-by-step and justify your steps, your answers should be clear and short sentences"
 
-import gradio as gr
-from transformers import pipeline
-import numpy as np
 
-from gradio_client import Client
-from huggingface_hub import InferenceClient
+    
+### WILL USE LOCAL MISTRAL OR ZEPHYR
 
-WHISPER_TIMEOUT = int(os.environ.get("WHISPER_TIMEOUT", 45))
-whisper_client = Client("https://sanchit-gandhi-whisper-large-v2.hf.space/")
-text_client = InferenceClient(
-    "mistralai/Mistral-7B-Instruct-v0.1"
-    ,timeout=WHISPER_TIMEOUT
-)
+from huggingface_hub import hf_hub_download
+print("Downloading LLM")
 
+llm_model = os.environ.get("LLM_MODEL", "mistral") # or "zephyr"
 
-###### COQUI TTS FUNCTIONS ######
+if llm_model == "zephyr":
+    #Zephyr
+    hf_hub_download(repo_id="TheBloke/zephyr-7B-alpha-GGUF", local_dir=".", filename="zephyr-7b-alpha.Q5_K_M.gguf")
+    # use new gguf format
+    model_path="./zephyr-7b-alpha.Q5_K_M.gguf"
+else:
+    #Mistral
+    hf_hub_download(repo_id="TheBloke/Mistral-7B-Instruct-v0.1-GGUF", local_dir=".", filename="mistral-7b-instruct-v0.1.Q5_K_M.gguf")
+    # use new gguf format
+    model_path="./mistral-7b-instruct-v0.1.Q5_K_M.gguf"
 
-def get_latents(speaker_wav):
-    # create as function as we can populate here with voice cleanup/filtering
-    (
-        gpt_cond_latent,
-        diffusion_conditioning,
-        speaker_embedding,
-    ) = model.get_conditioning_latents(audio_path=speaker_wav)
-    return gpt_cond_latent, diffusion_conditioning, speaker_embedding
+
+from llama_cpp import Llama
+# set GPU_LAYERS to 15 if you have a 8GB GPU so both models can fit in
+# else 35 full layers + XTTS works fine on T4 16GB
+GPU_LAYERS=int(os.environ.get("GPU_LAYERS", 15))
+
+LLAMA_VERBOSE=False
+print("Running LLM")
+llm = Llama(model_path=model_path,n_gpu_layers=GPU_LAYERS,max_new_tokens=256, context_window=4096, n_ctx=4096,n_batch=128,verbose=LLAMA_VERBOSE)
+
+
+
+# Mistral formatter
+def format_prompt_mistral(message, history):
+    prompt = (
+        "<s>[INST]" + system_message + "[/INST]" + system_understand_message + "</s>"
+    )
+    for user_prompt, bot_response in history:
+        prompt += f"[INST] {user_prompt} [/INST]"
+        prompt += f" {bot_response}</s> "
+    prompt += f"[INST] {message} [/INST]"
+    return prompt
+    
+# Zephyr formatter
+def format_prompt_zephyr(message, history, system_message=system_message):
+    prompt = (
+        "<|system|>" + system_message  +  "</s>"
+    )
+    for user_prompt, bot_response in history:
+        prompt += f"<|user|>\n{user_prompt}</s>"
+        prompt += f"<|assistant|> {bot_response}</s>"
+    if message=="":
+        message="Hello"
+    prompt += f"<|user|>\n{message}</s>"
+    print(prompt)
+    return prompt
+
+if llm_model=="zephyr":
+    format_prompt = format_prompt_zephyr
+else:
+    format_prompt = format_prompt_mistral
+
+
+def generate_local(
+    prompt,    
+    history,
+    system_message=None,
+    temperature=0.8,
+    max_tokens=256,
+    top_p=0.95,
+    stop = ["</s>","<|user|>"]
+):
+    temperature = float(temperature)
+    if temperature < 1e-2:
+        temperature = 1e-2
+    top_p = float(top_p)
+
+    generate_kwargs = dict(
+        temperature=temperature,
+        max_tokens=max_tokens,
+        top_p=top_p,
+    )
+
+    formatted_prompt = format_prompt(prompt, history,system_message=system_message)
+
+    try:
+        stream = llm(
+            formatted_prompt,
+            **generate_kwargs,
+            stream=True,
+        )
+        output = ""
+        for response in stream:
+            character= response["choices"][0]["text"]
+
+            if "<|user|>" in character:
+                # end of context
+                return 
+                
+            if emoji.is_emoji(character):
+                # Bad emoji not a meaning messes chat from next lines
+                return
+                
+            
+            output += response["choices"][0]["text"].replace("<|assistant|>","").replace("<|user|>","")
+            yield output
+
+    except Exception as e:
+        if "Too Many Requests" in str(e):
+            print("ERROR: Too many requests on mistral client")
+            gr.Warning("Unfortunately Mistral is unable to process")
+            output = "Unfortuanately I am not able to process your request now !"
+        else:
+            print("Unhandled Exception: ", str(e))
+            gr.Warning("Unfortunately Mistral is unable to process")
+            output = "I do not know what happened but I could not understand you ."
+
+    return output
 
 def get_latents(speaker_wav,voice_cleanup=False):
     if (voice_cleanup):
@@ -168,10 +265,6 @@ def get_latents(speaker_wav,voice_cleanup=False):
     ) = model.get_conditioning_latents(audio_path=speaker_wav)
     return gpt_cond_latent, diffusion_conditioning, speaker_embedding
 
-
-latent_map = {}
-latent_map["Female_Voice"] = get_latents("examples/female.wav")
-
 def wave_header_chunk(frame_input=b"", channels=1, sample_width=2, sample_rate=24000):
     # This will create a wave header then append the frame input
     # It should be first on a streaming wav file
@@ -186,11 +279,11 @@ def wave_header_chunk(frame_input=b"", channels=1, sample_width=2, sample_rate=2
     wav_buf.seek(0)
     return wav_buf.read()
 
+
 #Config will have more correct languages, they may be added before we append here
 ##["en","es","fr","de","it","pt","pl","tr","ru","nl","cs","ar","zh-cn","ja"]
 
 xtts_supported_languages=config.languages  
-
 def detect_language(prompt):
     # Fast language autodetection
     if len(prompt)>15:
@@ -267,16 +360,6 @@ def get_voice_streaming(prompt, language, latent_tuple, suffix="0"):
         return None
 
 ###### MISTRAL FUNCTIONS ######
-       
-def format_prompt(message, history):
-    prompt = (
-        "<s>[INST]" + system_message + "[/INST]" + system_understand_message + "</s>"
-    )
-    for user_prompt, bot_response in history:
-        prompt += f"[INST] {user_prompt} [/INST]"
-        prompt += f" {bot_response}</s> "
-    prompt += f"[INST] {message} [/INST]"
-    return prompt
 
 def generate(
     prompt,
@@ -299,8 +382,9 @@ def generate(
         do_sample=True,
         seed=42,
     )
-
-    formatted_prompt = format_prompt(prompt, history)
+        
+    #formatted_prompt = format_prompt(prompt, history)
+    formatted_prompt = format_prompt_zephyr(prompt, history)
 
     try:
         stream = text_client.text_generation(
@@ -386,9 +470,8 @@ def bot(history, system_prompt=""):
         history[-1][1] = character
         yield history
 
-##### MISTRAL STREAMING Sentence splitter ####
-        
-def get_sentence(history, system_prompt=""):
+
+def get_sentence(history, chatbot_role,system_prompt=""):
     history = [["", None]] if history is None else history
     
     if system_prompt == "":
@@ -404,12 +487,13 @@ def get_sentence(history, system_prompt=""):
     text_to_generate = ""
     stored_sentence = None
     stored_sentence_hash = None
-    for character in generate(history[-1][0], history[:-1]):
-        history[-1][1] = character
+    for character in generate_local(history[-1][0], history[:-1],system_message=ROLE_PROMPTS[chatbot_role]):
+        history[-1][1] = character.replace("<|assistant|>","")
         # It is coming word by word
 
-        text_to_generate = nltk.sent_tokenize(history[-1][1].replace("\n", " ").strip())
+        text_to_generate = nltk.sent_tokenize(history[-1][1].replace("\n", " ").replace("<|assistant|>"," ").strip())
         if len(text_to_generate) > 1:
+            
             dif = len(text_to_generate) - len(sentence_list)
 
             if dif == 1 and len(sentence_list) != 0:
@@ -465,148 +549,152 @@ def get_sentence(history, system_prompt=""):
 
         yield (last_sentence, history)
 
-#### SPEECH GENERATION BY SENTENCE FROM HISTORY ####
-        
-def generate_speech(history):
-    language = "autodetect"
+from scipy.io.wavfile import write
+from pydub import AudioSegment
 
-    wav_bytestream = b""
-    for sentence, history in get_sentence(history):
-        print(sentence)
-
-        # Sometimes prompt </s> coming on output remove it
-        # Some post process for speech only
-        sentence = sentence.replace("</s>", "")
-        # remove code from speech
-        sentence = re.sub("```.*```", "", sentence, flags=re.DOTALL)
-        sentence = sentence.replace("```", "")
-        sentence = sentence.replace("```", "")
-        sentence = sentence.replace("(", " ")
-        sentence = sentence.replace(")", " ")
-        
-        if len(sentence)==0:
-            #possible after cleanup sentence is empty
-            #e.g Kana then romaji in brackets
-            continue
-        
-        # A fast fix for last chacter, may produce weird sounds if it is with text
-        if (sentence[-1] in ["!", "?", ".", ","]) or (sentence[-2] in ["!", "?", ".", ","]):
-            # just add a space
-            sentence = sentence[:-1] + " " + sentence[-1]
-        print("Sentence for speech:", sentence)
+second_of_silence = AudioSegment.silent() # use default
+second_of_silence.export("sil.wav", format='wav')
 
-        try:
-            if len(sentence)<300:
-                # no problem continue on
-                sentence_list = [sentence]
-            else:
-                # Until now nltk likely split sentences properly but we need additional 
-                # check for longer sentence and split at last possible position
-                # Do whatever necessary, first break at hypens then spaces and then even split very long words
-                sentence_list=textwrap.wrap(sentence,300)
-                print("SPLITTED LONG SENTENCE:",sentence_list)
+
+def generate_speech(history,chatbot_role):
+    # Must set autoplay to True first
+    yield (history, chatbot_role, "", wave_header_chunk() )
+    for sentence, history in get_sentence(history,chatbot_role):
+        if sentence != "":
+            print("BG: inserting sentence to queue")
             
-            for sentence in sentence_list:
+            generated_speech = generate_speech_for_sentence(history, chatbot_role, sentence,return_as_byte=True)
+            if generated_speech is not None:
+                _, audio_dict = generated_speech
+                # We are using byte streaming
+                yield (history, chatbot_role, sentence, audio_dict["value"] )
                 
-                if any(c.isalnum() for c in sentence):
-                    if language=="autodetect":
-                        #on first call autodetect, nexts sentence calls will use same language
-                        language = detect_language(sentence) 
-                
-                    #exists at least 1 alphanumeric (utf-8) 
-                    audio_stream = get_voice_streaming(
-                            sentence, language, latent_map["Female_Voice"]
-                        )
-                else:
-                    # likely got a ' or " or some other text without alphanumeric in it
-                    audio_stream = None 
-                    
-                # XTTS is actually using streaming response but we are playing audio by sentence
-                # If you want direct XTTS voice streaming (send each chunk to voice ) you may set DIRECT_STREAM=1 environment variable
-                if audio_stream is not None:
-                    wav_chunks = wave_header_chunk()
-                    frame_length = 0
-                    for chunk in audio_stream:
-                        try:
-                            wav_bytestream += chunk
-                            if DIRECT_STREAM:
-                                yield (
-                                    gr.Audio.update(
-                                        value=wave_header_chunk() + chunk, autoplay=True
-                                    ),
-                                    history,
-                                )
-                                wait_time = len(chunk) / 2 / 24000
-                                wait_time = AUDIO_WAIT_MODIFIER * wait_time
-                                print("Sleeping till chunk end")
-                                time.sleep(wait_time)
+            
+# will generate speech audio file per sentence
+def generate_speech_for_sentence(history, chatbot_role, sentence, return_as_byte=True):
+    language = "autodetect"
+
+    wav_bytestream = b""
+    
+    if len(sentence)==0:
+        print("EMPTY SENTENCE")
+        return 
     
-                            else:
-                                wav_chunks += chunk
-                                frame_length += len(chunk)
-                        except:
-                            # hack to continue on playing. sometimes last chunk is empty , will be fixed on next TTS
-                            continue
+    # Sometimes prompt </s> coming on output remove it
+    # Some post process for speech only
+    sentence = sentence.replace("</s>", "")
+    # remove code from speech
+    sentence = re.sub("```.*```", "", sentence, flags=re.DOTALL)
+    sentence = re.sub("`.*`", "", sentence, flags=re.DOTALL)
     
-                if not DIRECT_STREAM:
-                    yield (
-                        gr.Audio.update(value=None, autoplay=True),
-                        history,
-                    )  # hack to switch autoplay
-                    if audio_stream is not None:
-                        yield (gr.Audio.update(value=wav_chunks, autoplay=True), history)
-                        # Streaming wait time calculation
-                        # audio_length = frame_length / sample_width/ frame_rate
-                        wait_time = frame_length / 2 / 24000
+    sentence = re.sub("\(.*\)", "", sentence, flags=re.DOTALL)
     
-                        # for non streaming
-                        # wait_time= librosa.get_duration(path=wav)
+    sentence = sentence.replace("```", "")
+    sentence = sentence.replace("...", " ")
+    sentence = sentence.replace("(", " ")
+    sentence = sentence.replace(")", " ")
+    sentence = sentence.replace("<|assistant|>","")
+
+    # A fast fix for last chacter, may produce weird sounds if it is with text
+    if (sentence[-1] in ["!", "?", ".", ","]) or (sentence[-2] in ["!", "?", ".", ","]):
+        # just add a space
+        sentence = sentence[:-1] + " " + sentence[-1]
+    print("Sentence for speech:", sentence)
+    if len(sentence)==0:
+        print("EMPTY SENTENCE after processing")
+        return 
     
-                        wait_time = AUDIO_WAIT_MODIFIER * wait_time
-                        print("Sleeping till audio end")
-                        time.sleep(wait_time)
-                    else:
-                        # Either too much text or some programming, give a silence so stream continues
-                        second_of_silence = AudioSegment.silent()  # use default
-                        second_of_silence.export("sil.wav", format="wav")
-                        yield (gr.Audio.update(value="sil.wav", autoplay=True), history)
-
-        except RuntimeError as e:
-            if "device-side assert" in str(e):
-                # cannot do anything on cuda device side error, need tor estart
-                print(
-                    f"Exit due to: Unrecoverable exception caused by prompt:{sentence}",
-                    flush=True,
-                )
-                gr.Warning("Unhandled Exception encounter, please retry in a minute")
-                print("Cuda device-assert Runtime encountered need restart")
-
-                # HF Space specific.. This error is unrecoverable need to restart space
-                api.restart_space(repo_id=repo_id)
+    try:
+        SENTENCE_SPLIT_LENGTH=350
+        if len(sentence)<SENTENCE_SPLIT_LENGTH:
+            # no problem continue on
+            sentence_list = [sentence]
+        else:
+            # Until now nltk likely split sentences properly but we need additional 
+            # check for longer sentence and split at last possible position
+            # Do whatever necessary, first break at hypens then spaces and then even split very long words
+            sentence_list=textwrap.wrap(sentence,SENTENCE_SPLIT_LENGTH)
+            print("SPLITTED LONG SENTENCE:",sentence_list)
+        
+        for sentence in sentence_list:
+            
+            if any(c.isalnum() for c in sentence):
+                if language=="autodetect":
+                    #on first call autodetect, nexts sentence calls will use same language
+                    language = detect_language(sentence) 
+            
+                #exists at least 1 alphanumeric (utf-8) 
+                audio_stream = get_voice_streaming(
+                        sentence, language, latent_map[chatbot_role]
+                    )
             else:
-                print("RuntimeError: non device-side assert error:", str(e))
-                raise e
+                # likely got a ' or " or some other text without alphanumeric in it
+                audio_stream = None 
+                
+            # XTTS is actually using streaming response but we are playing audio by sentence
+            # If you want direct XTTS voice streaming (send each chunk to voice ) you may set DIRECT_STREAM=1 environment variable
+            if audio_stream is not None:
+                wav_chunks = wave_header_chunk()
+                frame_length = 0
+                for chunk in audio_stream:
+                    try:
+                        wav_bytestream += chunk
+                        wav_chunks += chunk
+                        frame_length += len(chunk)
+                    except:
+                        # hack to continue on playing. sometimes last chunk is empty , will be fixed on next TTS
+                        continue
 
-    time.sleep(1)
-    wav_bytestream = wave_header_chunk() + wav_bytestream
-    outfile = "combined.wav"
-    with open(outfile, "wb") as f:
-        f.write(wav_bytestream)
-    yield (gr.Audio.update(value=None, autoplay=False), history)
-    yield (gr.Audio.update(value=outfile, autoplay=False), history)
+            if audio_stream is not None:
+                if not return_as_byte:
+                    audio_unique_filename = "/tmp/"+ str(uuid.uuid4())+".wav"
+                    with open(audio_unique_filename, "wb") as f:
+                        f.write(wav_chunks)
+                    #Will write filename to context variable
+                    return (history , gr.Audio.update(value=audio_unique_filename, autoplay=True))
+                else:
+                    return (history , gr.Audio.update(value=wav_bytestream, autoplay=True))
+    except RuntimeError as e:
+        if "device-side assert" in str(e):
+            # cannot do anything on cuda device side error, need tor estart
+            print(
+                f"Exit due to: Unrecoverable exception caused by prompt:{sentence}",
+                flush=True,
+            )
+            gr.Warning("Unhandled Exception encounter, please retry in a minute")
+            print("Cuda device-assert Runtime encountered need restart")
+
+            # HF Space specific.. This error is unrecoverable need to restart space
+            api.restart_space(repo_id=repo_id)
+        else:
+            print("RuntimeError: non device-side assert error:", str(e))
+            raise e
+
+    print("All speech ended")
+    return 
+
+
+latent_map = {}
+latent_map["AI Assistant"] = get_latents("examples/female.wav")
 
 #### GRADIO INTERFACE ####
-    
+
 with gr.Blocks(title=title) as demo:
     gr.Markdown(DESCRIPTION)
-
     chatbot = gr.Chatbot(
         [],
         elem_id="chatbot",
         avatar_images=("examples/hf-logo.png", "examples/coqui-logo.png"),
         bubble_full_width=False,
     )
-
+    with gr.Row():
+        chatbot_role = gr.Dropdown(
+            label="Role of the Chatbot",
+            info="How should Chatbot talk like",
+            choices=ROLES,
+            max_choices=1,
+            value=ROLES[0],
+        )
     with gr.Row():
         txt = gr.Textbox(
             scale=3,
@@ -617,35 +705,43 @@ with gr.Blocks(title=title) as demo:
         )
         txt_btn = gr.Button(value="Submit text", scale=1)
         btn = gr.Audio(source="microphone", type="filepath", scale=4)
-
+    def stop():
+        print("Audio STOP")
+        set_audio_playing(False)
+        
     with gr.Row():
+        sentence = gr.Textbox()
         audio = gr.Audio(
+            value=None,
             label="Generated audio response",
-            streaming=False,
-            autoplay=False,
-            interactive=True,
+            streaming=True,
+            autoplay=True,
+            interactive=False,
             show_label=True,
         )
-        # TODO add a second audio that plays whole sentences (for mobile especially)
-        # final_audio = gr.Audio(label="Final audio response", streaming=False, autoplay=False, interactive=False,show_label=True, visible=False)
-
+        
+        audio.end(stop)
+        
+ 
     clear_btn = gr.ClearButton([chatbot, audio])
     
     txt_msg = txt_btn.click(add_text, [chatbot, txt], [chatbot, txt], queue=False).then(
-        generate_speech, chatbot, [audio, chatbot]
+        generate_speech,  [chatbot,chatbot_role], [chatbot,chatbot_role, sentence, audio]
     )
 
     txt_msg.then(lambda: gr.update(interactive=True), None, [txt], queue=False)
 
     txt_msg = txt.submit(add_text, [chatbot, txt], [chatbot, txt], queue=False).then(
-        generate_speech, chatbot, [audio, chatbot]
+        generate_speech,  [chatbot,chatbot_role], [chatbot,chatbot_role, sentence, audio]
     )
 
     txt_msg.then(lambda: gr.update(interactive=True), None, [txt], queue=False)
 
     file_msg = btn.stop_recording(
         add_file, [chatbot, btn], [chatbot, txt], queue=False
-    ).then(generate_speech, chatbot, [audio, chatbot])
+    ).then(
+        generate_speech,  [chatbot,chatbot_role], [chatbot,chatbot_role, sentence, audio]
+    )
 
     file_msg.then(lambda: (gr.update(interactive=True),gr.update(interactive=True,value=None)), None, [txt, btn], queue=False)
 
@@ -654,13 +750,13 @@ with gr.Blocks(title=title) as demo:
 This Space demonstrates how to speak to a chatbot, based solely on open-source models.
 It relies on 3 models:
 1. [Whisper-large-v2](https://sanchit-gandhi-whisper-large-v2.hf.space/) as an ASR model, to transcribe recorded audio to text. It is called through a [gradio client](https://www.gradio.app/docs/client).
-2. [Mistral-7b-instruct](https://huggingface.co/spaces/osanseviero/mistral-super-fast) as the chat model, the actual chat model. It is called from [huggingface_hub](https://huggingface.co/docs/huggingface_hub/guides/inference).
+2. [Zephyr-7b-alpha](https://huggingface.co/HuggingFaceH4/zephyr-7b-alpha) as the chat model. GGUF Q5_K_M quantized version used locally via llama_cpp from [huggingface.co/TheBloke](https://huggingface.co/TheBloke/zephyr-7B-alpha-GGUF).
 3. [Coqui's XTTS](https://huggingface.co/spaces/coqui/xtts) as a TTS model, to generate the chatbot answers. This time, the model is hosted locally.
 
 Note:
 - By using this demo you agree to the terms of the Coqui Public Model License at https://coqui.ai/cpml
-- Responses generated by chat model should not be assumed correct as this is a demonstration example only
+- Responses generated by chat model should not be assumed correct or taken serious, as this is a demonstration example only
 - iOS (Iphone/Ipad) devices may not experience voice due to autoplay being disabled on these devices by Vendor"""
     )
 demo.queue()
-demo.launch(debug=True)
+demo.launch(debug=True)