Update duplex.py

duplex.py

@@ -3,8 +3,11 @@ import json
 import random
 import string
 
+import numpy as np
 import gradio as gr
 import requests
+import soundfile as sf
+
 from transformers import pipeline, set_seed
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import logging
@@ -14,11 +17,14 @@ import gradio as gr
 from transformers import pipeline, AutoModelForCTC, Wav2Vec2Processor, Wav2Vec2ProcessorWithLM
 
 DEBUG = os.environ.get("DEBUG", "false")[0] in "ty1"
-HF_AUTH_TOKEN = os.environ.get("HF_AUTH_TOKEN", None)
 MAX_LENGTH = int(os.environ.get("MAX_LENGTH", 1024))
+DEFAULT_LANG = os.environ.get("DEFAULT_LANG", "English")
 
 HEADER = """
 # Poor Man's Duplex
+
+Talk to a language model like you talk on a Walkie-Talkie! Well, with larger latencies.
+The models are [EleutherAI's GPT-J-6B](https://huggingface.co/EleutherAI/gpt-j-6B) for English, and [BERTIN GPT-J-6B](https://huggingface.co/bertin-project/bertin-gpt-j-6B) for Spanish.
 """.strip()
 
 FOOTER = """
@@ -45,7 +51,8 @@ def generate_es(text, **kwargs):
     api_uri = "https://hf.space/embed/bertin-project/bertin-gpt-j-6B/+/api/predict/"
     response = requests.post(api_uri, data=json.dumps({"data": [text, 100, 100, 50, 0.95, True, True]}))
     if response.ok:
-        print(response.json())
+        if DEBUG:
+            print(response.json())
         return response.json()["data"][0]
     else:
         return ""
@@ -65,9 +72,14 @@ speak_en = gr.Interface.load(f"huggingface/{tts_model_name}")
 transcribe_en = lambda input_file: asr_en(input_file, chunk_length_s=5, stride_length_s=1)["text"]
 generate_iface = gr.Interface.load("huggingface/EleutherAI/gpt-j-6B")
 
+empty_audio = 'empty.flac'
+sf.write(empty_audio, [], 16000)
+deuncase = gr.Interface.load("huggingface/pere/DeUnCaser")
+
 def generate_en(text, **kwargs):
     response = generate_iface(text)
-    print(response)
+    if DEBUG:
+        print(response)
     return response or ""
 
 
@@ -97,11 +109,28 @@ def select_lang_vars(lang):
     return AGENT, USER, CONTEXT
 
 
+def format_chat(history):
+    interventions = []
+    for user, bot in history:
+        interventions.append(f"""
+            <div data-testid="user" style="background-color:#16a34a" class="px-3 py-2 rounded-[22px] rounded-bl-none place-self-start text-white ml-7 text-sm">{user}</div>
+            <div data-testid="bot" style="background-color:gray" class="px-3 py-2 rounded-[22px] rounded-br-none  text-white ml-7 text-sm">{bot}</div>
+        """)
+    return f"""<details><summary>Conversation log</summary>
+    <div class="overflow-y-auto h-[40vh]">
+        <div class="flex flex-col items-end space-y-4 p-3">
+        {"".join(interventions)}
+        </div>
+    </div>
+    </summary>"""
+
 
 def chat_with_gpt(lang, agent, user, context, audio_in, history):
+    if not audio_in:
+        return history, history, empty_audio, format_chat(history)
     generate, transcribe, speak = select_lang(lang)
     AGENT, USER, _ = select_lang_vars(lang)
-    user_message = transcribe(audio_in)
+    user_message = deuncase(transcribe(audio_in))
     # agent = AGENT
     # user = USER
     generation_kwargs = {
@@ -156,17 +185,17 @@ def chat_with_gpt(lang, agent, user, context, audio_in, history):
     if not response.strip():
         response = "Lo siento, no puedo hablar ahora" if lang.lower() == "Spanish" else "Sorry, can't talk right now"
     history.append((user_message, response))
-    return history, history, speak(response)
+    return history, history, speak(response), format_chat(history)
 
 
 with gr.Blocks() as demo:
     gr.Markdown(HEADER)
-    lang = gr.Radio(label="Language", choices=["English", "Spanish"], default="English", type="value")
-    AGENT, USER, CONTEXT = select_lang_vars("English")
+    lang = gr.Radio(label="Language", choices=["English", "Spanish"], value=DEFAULT_LANG, type="value")
+    AGENT, USER, CONTEXT = select_lang_vars(DEFAULT_LANG)
     context = gr.Textbox(label="Context", lines=5, value=CONTEXT)
     with gr.Row():
         audio_in = gr.Audio(label="User", source="microphone", type="filepath")
-        audio_out = gr.Audio(label="Agent", interactive=False)
+        audio_out = gr.Audio(label="Agent", interactive=False, value=empty_audio)
         # chat_btn = gr.Button("Submit")
     with gr.Row():
         user = gr.Textbox(label="User", value=USER)
@@ -175,7 +204,8 @@ with gr.Blocks() as demo:
     history = gr.Variable(value=[])
     chatbot = gr.Variable()  # gr.Chatbot(color_map=("green", "gray"), visible=False)
     # chat_btn.click(chat_with_gpt, inputs=[lang, agent, user, context, audio_in, history], outputs=[chatbot, history, audio_out])
-    audio_in.change(chat_with_gpt, inputs=[lang, agent, user, context, audio_in, history], outputs=[chatbot, history, audio_out])
+    log = gr.HTML()
+    audio_in.change(chat_with_gpt, inputs=[lang, agent, user, context, audio_in, history], outputs=[chatbot, history, audio_out, log])
     gr.Markdown(FOOTER)
 
 demo.launch()