add mic button

src/content/agent.py

@@ -1,19 +1,18 @@
-import copy
-import base64
-
+import numpy as np
 import streamlit as st
 
-from src.generation import MAX_AUDIO_LENGTH
 from src.retrieval import STANDARD_QUERIES, retrieve_relevant_docs
-from src.utils import bytes_to_array, array_to_bytes
 from src.content.common import (
     MODEL_NAMES,
     AUDIO_SAMPLES_W_INSTRUCT, 
     AGENT_DIALOGUE_STATES,
+    reset_states,
+    update_voice_instruction_state,
     init_state_section,
     header_section,
     sidebar_fragment,
-    reset_states,
+    successful_example_section,
+    audio_attach_dialogue,
     retrive_response_with_ui
 )
 
@@ -42,103 +41,29 @@ However, the audio analysis may or may not contain relevant information to the u
 AUDIO_ANALYSIS_STATUS = "MERaLiON-AudioLLM Analysis"
 
 
-def _update_audio(audio_bytes):
-    origin_audio_array = bytes_to_array(audio_bytes)
-    truncated_audio_array = origin_audio_array[: MAX_AUDIO_LENGTH*16000]
-    truncated_audio_bytes = array_to_bytes(truncated_audio_array)
-
-    st.session_state.ag_audio_array = origin_audio_array
-    st.session_state.ag_audio_base64 = base64.b64encode(truncated_audio_bytes).decode('utf-8')
-
-
-@st.fragment
-def successful_example_section():
-    audio_sample_names = [name for name in AUDIO_SAMPLES_W_INSTRUCT.keys() if "Paral" in name]
-    
-    st.markdown(":fire: **Successful Tasks and Examples**")
-    
-    sample_name = st.selectbox(
-        label="**Select Audio:**",
-        label_visibility="collapsed",
-        options=audio_sample_names,
-        format_func=lambda o: AUDIO_SAMPLES_W_INSTRUCT[o]["apperance"],
-        index=None,
-        placeholder="Select an audio sample:",
-        on_change=lambda: st.session_state.update(
-            on_select=True, 
-            ag_messages=[],
-            ag_model_messages=[],
-            ag_visited_query_indices=[],
-            disprompt=True
-        ),
-        key='select')
-    
-    if sample_name and st.session_state.on_select:
-        audio_bytes = open(f"audio_samples/{sample_name}.wav", "rb").read()
-        st.session_state.update(
-            on_select=False,
-            new_prompt=AUDIO_SAMPLES_W_INSTRUCT[sample_name]["instructions"][0]
-        )
-        _update_audio(audio_bytes)
-        st.rerun(scope="app")
-
-
-@st.dialog("Specify Audio")
-def audio_attach_dialogue():
-    st.markdown("**Upload**")
-
-    uploaded_file = st.file_uploader(
-        label="**Upload Audio:**", 
-        label_visibility="collapsed",
-        type=['wav', 'mp3'],
-        on_change=lambda: st.session_state.update(
-            on_upload=True, 
-            ag_messages=[], 
-            ag_model_messages=[],
-            ag_visited_query_indices=[]
-            ),
-        key='upload'
-    )
-    
-    if uploaded_file and st.session_state.on_upload:
-        audio_bytes = uploaded_file.read()
-        _update_audio(audio_bytes)
-        st.session_state.on_upload = False
-        st.rerun()
-
-    st.markdown("**Record**")
-    
-    uploaded_file = st.audio_input(
-        label="**Record Audio:**",
-        label_visibility="collapsed",
-        on_change=lambda: st.session_state.update(
-            on_record=True, 
-            ag_messages=[], 
-            ag_model_messages=[],
-            ag_visited_query_indices=[]
-            ),
-        key='record'
-    )
-    
-    if uploaded_file and st.session_state.on_record:
-        audio_bytes = uploaded_file.read()
-        _update_audio(audio_bytes)
-        st.session_state.on_record = False
-        st.rerun()
+AG_CONVERSATION_STATES = dict(
+    ag_messages=[],
+    ag_model_messages=[],
+    ag_visited_query_indices=[],
+)
 
 
 def bottom_input_section():
-    bottom_cols = st.columns([0.03, 0.03, 0.94])
+    bottom_cols = st.columns([0.03, 0.03, 0.91, 0.03])
     with bottom_cols[0]:
         st.button(
-            'Clear', 
+            ':material/delete:', 
             disabled=st.session_state.disprompt,
             on_click=lambda: reset_states(AGENT_DIALOGUE_STATES)
         )
 
     with bottom_cols[1]:
-        if st.button("\+ Audio", disabled=st.session_state.disprompt):
-            audio_attach_dialogue()
+        if st.button(":material/add:", disabled=st.session_state.disprompt):
+            audio_attach_dialogue(
+                audio_array_state="ag_audio_array",
+                audio_base64_state="ag_audio_base64",
+                restore_state=AG_CONVERSATION_STATES
+            )
 
     with bottom_cols[2]:
         if chat_input := st.chat_input(
@@ -148,6 +73,23 @@ def bottom_input_section():
         ):
             st.session_state.new_prompt = chat_input
 
+    with bottom_cols[3]:
+        uploaded_voice = st.audio_input(
+            label="voice_instruction",
+            label_visibility="collapsed", 
+            disabled=st.session_state.disprompt, 
+            on_change=lambda: st.session_state.update(
+                disprompt=True,
+                on_record_voice_instruction=True
+                ),
+            key='voice_instruction'  
+        )
+
+        if uploaded_voice and st.session_state.on_record_voice_instruction:
+            voice_bytes = uploaded_voice.read()
+            update_voice_instruction_state(voice_bytes)
+            st.session_state.on_record_voice_instruction = False
+
 
 def _prepare_final_prompt_with_ui(one_time_prompt):
     if st.session_state.ag_audio_array.shape[0] == 0:
@@ -216,9 +158,7 @@ def conversation_section():
             st.audio(st.session_state.ag_audio_array, format="audio/wav", sample_rate=16000)
 
     for message in st.session_state.ag_messages:
-        message_name = "assistant" if "assistant" in message["role"] else message["role"]
-
-        with chat_message_container.chat_message(name=message_name):
+        with chat_message_container.chat_message(name=message["role"]):
             if message.get("error"):
                 st.error(message["error"])
             for warning_msg in message.get("warnings", []):
@@ -238,38 +178,73 @@ def conversation_section():
     with st._bottom:
         bottom_input_section()
 
-    if one_time_prompt := st.session_state.new_prompt:
-        st.session_state.update(new_prompt="")
+    if (not st.session_state.new_prompt) and (not st.session_state.new_vi_base64):
+        return
+    
+    one_time_prompt = st.session_state.new_prompt
+    one_time_vi_array = st.session_state.new_vi_array
+    one_time_vi_base64 = st.session_state.new_vi_base64
+
+    st.session_state.update(
+        new_prompt="", 
+        new_vi_array=np.array([]),
+        new_vi_base64="",
+    )
 
-        with chat_message_container.chat_message("user"):
+    with chat_message_container.chat_message("user"):
+        if one_time_vi_base64:
+            with st.spinner("Transcribing..."):
+                error_msg, warnings, one_time_prompt = retrive_response_with_ui(
+                    model_name=MODEL_NAMES["audiollm"]["vllm_name"],
+                    text_input="Write out the dialogue as text.", 
+                    array_audio_input=one_time_vi_array,
+                    base64_audio_input=one_time_vi_base64,
+                    stream=False,
+                    normalise_response=True
+                )
+        else:
+            error_msg, warnings = "", []
             st.write(one_time_prompt)
-        st.session_state.ag_messages.append({"role": "user", "content": one_time_prompt})
-
-        with chat_message_container.chat_message("assistant"):
-            assistant_message = {"role": "assistant", "process": []}
-            st.session_state.ag_messages.append(assistant_message)
 
-            final_prompt = _prepare_final_prompt_with_ui(one_time_prompt)
-            
-            error_msg, warnings, response = retrive_response_with_ui(
-                model_name=MODEL_NAMES["llm"]["vllm_name"],
-                text_input=final_prompt, 
-                array_audio_input=st.session_state.ag_audio_array,
-                base64_audio_input="", 
-                prefix=f"**{MODEL_NAMES['llm']['ui_name']}**: ",
-                stream=True,
-                history=st.session_state.ag_model_messages,
-                show_warning=False
-            )
-            
-            assistant_message.update({"error": error_msg, "warnings": warnings, "content": response})
-            st.session_state.ag_model_messages.extend([
-                {"role": "user", "content": final_prompt},
-                {"role": "assistant", "content": response}
-            ])
+    st.session_state.ag_messages.append({
+        "role": "user", 
+        "error": error_msg,
+        "warnings": warnings, 
+        "content": one_time_prompt
+    })
+
+    with chat_message_container.chat_message("assistant"):
+        assistant_message = {"role": "assistant", "process": []}
+        st.session_state.ag_messages.append(assistant_message)
+
+        final_prompt = _prepare_final_prompt_with_ui(one_time_prompt)
+
+        llm_response_prefix = f"**{MODEL_NAMES['llm']['ui_name']}**: "
+        error_msg, warnings, response = retrive_response_with_ui(
+            model_name=MODEL_NAMES["llm"]["vllm_name"],
+            text_input=final_prompt, 
+            array_audio_input=st.session_state.ag_audio_array,
+            base64_audio_input="", 
+            prefix=llm_response_prefix,
+            stream=True,
+            history=st.session_state.ag_model_messages,
+            show_warning=False
+        )
+        
+        assistant_message.update({
+            "error": error_msg, 
+            "warnings": warnings, 
+            "content": response
+        })
+        
+        pure_response = response.replace(llm_response_prefix, "")
+        st.session_state.ag_model_messages.extend([
+            {"role": "user", "content": final_prompt},
+            {"role": "assistant", "content": pure_response}
+        ])
 
-        st.session_state.disprompt=False
-        st.rerun(scope="app")
+    st.session_state.disprompt=False
+    st.rerun(scope="app")
 
 
 def agent_page():
@@ -286,5 +261,12 @@ def agent_page():
     with st.sidebar:
         sidebar_fragment()
 
-    successful_example_section()
+    audio_sample_names = [name for name in AUDIO_SAMPLES_W_INSTRUCT.keys() if "Paral" in name]
+
+    successful_example_section(
+        audio_sample_names, 
+        audio_array_state="ag_audio_array",
+        audio_base64_state="ag_audio_base64",
+        restore_state=AG_CONVERSATION_STATES
+    )
     conversation_section()  

src/content/common.py

@@ -1,5 +1,6 @@
 import os
 import copy
+import base64
 import itertools
 from collections import OrderedDict
 from typing import List, Optional
@@ -8,9 +9,15 @@ import numpy as np
 import streamlit as st
 
 from src.tunnel import start_server
-from src.generation import FIXED_GENERATION_CONFIG, load_model, retrive_response
 from src.retrieval import load_retriever
 from src.logger import load_logger
+from src.utils import array_to_bytes, bytes_to_array, postprocess_voice_transcription
+from src.generation import (
+    FIXED_GENERATION_CONFIG, 
+    MAX_AUDIO_LENGTH, 
+    load_model, 
+    retrive_response
+)
 
 
 PLAYGROUND_DIALOGUE_STATES = dict(
@@ -40,10 +47,13 @@ AGENT_DIALOGUE_STATES = dict(
 COMMON_DIALOGUE_STATES = dict(
     disprompt=False,
     new_prompt="",
+    new_vi_array=np.array([]),
+    new_vi_base64="",
     on_select=False, 
     on_upload=False, 
     on_record=False, 
-    on_select_quick_action=False
+    on_select_quick_action=False,
+    on_record_voice_instruction=False
 )
 
 
@@ -319,6 +329,26 @@ AUDIO_SAMPLES_W_INSTRUCT = {
 exec(os.getenv('APP_CONFIGS'))
 
 
+def reset_states(*state_dicts):
+    for states in state_dicts:
+        st.session_state.update(copy.deepcopy(states))
+    st.session_state.update(copy.deepcopy(COMMON_DIALOGUE_STATES))
+
+
+def process_audio_bytes(audio_bytes):
+    origin_audio_array = bytes_to_array(audio_bytes)
+    truncated_audio_array = origin_audio_array[: MAX_AUDIO_LENGTH*16000]
+    truncated_audio_bytes = array_to_bytes(truncated_audio_array)
+    audio_base64 = base64.b64encode(truncated_audio_bytes).decode('utf-8')
+
+    return origin_audio_array, audio_base64
+
+
+def update_voice_instruction_state(voice_bytes):
+    st.session_state.new_vi_array, st.session_state.new_vi_base64 = \
+        process_audio_bytes(voice_bytes)
+
+
 def init_state_section():
     st.set_page_config(page_title='MERaLiON-AudioLLM', page_icon = "🔥", layout='wide')
 
@@ -397,10 +427,75 @@ def sidebar_fragment():
     st.slider(label="Repetition Penalty", min_value=1.0, max_value=1.2, value=1.1, key="repetition_penalty")
 
 
-def reset_states(*state_dicts):
-    for states in state_dicts:
-        st.session_state.update(copy.deepcopy(states))
-    st.session_state.update(copy.deepcopy(COMMON_DIALOGUE_STATES))
+@st.fragment
+def successful_example_section(audio_sample_names, audio_array_state, audio_base64_state, restore_state={}):    
+    st.markdown(":fire: **Successful Tasks and Examples**")
+    
+    sample_name = st.selectbox(
+        label="**Select Audio:**",
+        label_visibility="collapsed",
+        options=audio_sample_names,
+        format_func=lambda o: AUDIO_SAMPLES_W_INSTRUCT[o]["apperance"],
+        index=None,
+        placeholder="Select an audio sample:",
+        on_change=lambda: st.session_state.update(
+            on_select=True, 
+            disprompt=True,
+            **copy.deepcopy(restore_state)
+        ),
+        key='select')
+    
+    if sample_name and st.session_state.on_select:
+        audio_bytes = open(f"audio_samples/{sample_name}.wav", "rb").read()
+        st.session_state.update(
+            on_select=False,
+            new_prompt=AUDIO_SAMPLES_W_INSTRUCT[sample_name]["instructions"][0]
+        )
+        st.session_state[audio_array_state], st.session_state[audio_base64_state] = \
+            process_audio_bytes(audio_bytes)
+        st.rerun(scope="app")
+
+
+@st.dialog("Specify audio context for analysis")
+def audio_attach_dialogue(audio_array_state, audio_base64_state, restore_state={}):
+    st.markdown("**Upload**")
+
+    uploaded_file = st.file_uploader(
+        label="**Upload Audio:**", 
+        label_visibility="collapsed",
+        type=['wav', 'mp3'],
+        on_change=lambda: st.session_state.update(
+            on_upload=True, 
+            **copy.deepcopy(restore_state)
+            ),
+        key='upload'
+    )
+    
+    if uploaded_file and st.session_state.on_upload:
+        audio_bytes = uploaded_file.read()
+        st.session_state[audio_array_state], st.session_state[audio_base64_state] = \
+            process_audio_bytes(audio_bytes)
+        st.session_state.on_upload = False
+        st.rerun()
+
+    st.markdown("**Record**")
+    
+    uploaded_file = st.audio_input(
+        label="**Record Audio:**",
+        label_visibility="collapsed",
+        on_change=lambda: st.session_state.update(
+            on_record=True,
+            **copy.deepcopy(restore_state)
+            ),
+        key='record'
+    )
+    
+    if uploaded_file and st.session_state.on_record:
+        audio_bytes = uploaded_file.read()
+        st.session_state[audio_array_state], st.session_state[audio_base64_state] = \
+            process_audio_bytes(audio_bytes)
+        st.session_state.on_record = False
+        st.rerun()
 
 
 def retrive_response_with_ui(
@@ -410,6 +505,7 @@ def retrive_response_with_ui(
         base64_audio_input: str, 
         prefix: str = "", 
         stream: bool = True, 
+        normalise_response: bool = False,
         history: Optional[List] = None, 
         show_warning: bool = True,
         **kwargs
@@ -455,7 +551,10 @@ def retrive_response_with_ui(
             response = st.write_stream(response_obj)
         else:
             response = response_obj.choices[0].message.content
-            st.write(prefix+response)
+            if normalise_response:
+                response = postprocess_voice_transcription(response)
+            response = prefix + response
+            st.write(response)
 
     st.session_state.logger.register_query(
         session_id=st.session_state.session_id,

src/content/playground.py

@@ -1,18 +1,17 @@
-import copy
-import base64
-
+import numpy as np
 import streamlit as st
 
-from src.generation import MAX_AUDIO_LENGTH
-from src.utils import bytes_to_array, array_to_bytes
 from src.content.common import (
     MODEL_NAMES,
     AUDIO_SAMPLES_W_INSTRUCT, 
     PLAYGROUND_DIALOGUE_STATES,
+    reset_states,
+    update_voice_instruction_state,
     init_state_section,
     header_section,
     sidebar_fragment,
-    reset_states,
+    successful_example_section,
+    audio_attach_dialogue,
     retrive_response_with_ui
 )
 
@@ -31,86 +30,22 @@ QUICK_ACTIONS = [
 ]
 
 
-def _update_audio(audio_bytes):
-    origin_audio_array = bytes_to_array(audio_bytes)
-    truncated_audio_array = origin_audio_array[: MAX_AUDIO_LENGTH*16000]
-    truncated_audio_bytes = array_to_bytes(truncated_audio_array)
-
-    st.session_state.pg_audio_array = origin_audio_array
-    st.session_state.pg_audio_base64 = base64.b64encode(truncated_audio_bytes).decode('utf-8')
-
-
-@st.fragment
-def successful_example_section():
-    audio_sample_names = [audio_sample_name for audio_sample_name in AUDIO_SAMPLES_W_INSTRUCT.keys()]
-    
-    st.markdown(":fire: **Successful Tasks and Examples**")
-    
-    sample_name = st.selectbox(
-        label="**Select Audio:**",
-        label_visibility="collapsed",
-        options=audio_sample_names,
-        format_func=lambda o: AUDIO_SAMPLES_W_INSTRUCT[o]["apperance"],
-        index=None,
-        placeholder="Select an audio sample:",
-        on_change=lambda: st.session_state.update(
-            on_select=True, 
-            pg_messages=[],
-            disprompt=True
-        ),
-        key='select')
-    
-    if sample_name and st.session_state.on_select:
-        audio_bytes = open(f"audio_samples/{sample_name}.wav", "rb").read()
-        st.session_state.update(
-            on_select=False,
-            new_prompt=AUDIO_SAMPLES_W_INSTRUCT[sample_name]["instructions"][0]
-        )
-        _update_audio(audio_bytes)
-        st.rerun(scope="app")
-
-
-@st.dialog("Specify Audio")
-def audio_attach_dialogue():
-    st.markdown("**Upload**")
-
-    uploaded_file = st.file_uploader(
-        label="**Upload Audio:**", 
-        label_visibility="collapsed",
-        type=['wav', 'mp3'],
-        on_change=lambda: st.session_state.update(on_upload=True, pg_messages=[]),
-        key='upload'
-    )
-    
-    if uploaded_file and st.session_state.on_upload:
-        audio_bytes = uploaded_file.read()
-        _update_audio(audio_bytes)
-        st.session_state.on_upload = False
-        st.rerun()
-
-    st.markdown("**Record**")
-    
-    uploaded_file = st.audio_input(
-        label="**Record Audio:**",
-        label_visibility="collapsed",
-        on_change=lambda: st.session_state.update(on_record=True, pg_messages=[]),
-        key='record'
-    )
-    
-    if uploaded_file and st.session_state.on_record:
-        audio_bytes = uploaded_file.read()
-        _update_audio(audio_bytes)
-        st.session_state.on_record = False
-        st.rerun()
+PG_CONVERSATION_STATES = dict(
+    pg_messages=[],
+)
 
 
 @st.fragment
 def select_model_variants_fradment():
-    display_mapper = {value["vllm_name"]: value["ui_name"] for value in MODEL_NAMES.values()}
+    display_mapper = {
+        value["vllm_name"]: value["ui_name"] 
+        for key, value in MODEL_NAMES.items()
+        if "audiollm" in key
+        }
 
     st.selectbox(
         label=":fire: Explore more MERaLiON-AudioLLM variants!",
-        options=[value["vllm_name"] for value in MODEL_NAMES.values()],
+        options=list(display_mapper.keys()),
         index=0,
         format_func=lambda o: display_mapper[o],
         key="pg_model_name",
@@ -122,27 +57,52 @@ def select_model_variants_fradment():
 def bottom_input_section():
     select_model_variants_fradment()
 
-    bottom_cols = st.columns([0.03, 0.03, 0.94])
+    bottom_cols = st.columns([0.03, 0.03, 0.91, 0.03])
     with bottom_cols[0]:
         st.button(
-            'Clear', 
+            ':material/delete:', 
             disabled=st.session_state.disprompt,
             on_click=lambda: reset_states(PLAYGROUND_DIALOGUE_STATES)
         )
 
     with bottom_cols[1]:
-        if st.button("\+ Audio", disabled=st.session_state.disprompt):
-            audio_attach_dialogue()
+        if st.button(":material/add:", disabled=st.session_state.disprompt):
+            audio_attach_dialogue(
+                audio_array_state="pg_audio_array",
+                audio_base64_state="pg_audio_base64",
+                restore_state=PG_CONVERSATION_STATES
+            )
 
     with bottom_cols[2]:
         if chat_input := st.chat_input(
             placeholder="Instruction...", 
             disabled=st.session_state.disprompt, 
-            on_submit=lambda: st.session_state.update(disprompt=True, pg_messages=[])
+            on_submit=lambda: st.session_state.update(
+                disprompt=True, 
+                **PG_CONVERSATION_STATES
+            )
         ):
             st.session_state.new_prompt = chat_input
 
+    with bottom_cols[3]:
+        uploaded_voice = st.audio_input(
+            label="voice_instruction",
+            label_visibility="collapsed", 
+            disabled=st.session_state.disprompt, 
+            on_change=lambda: st.session_state.update(
+                disprompt=True,
+                on_record_voice_instruction=True,
+                **PG_CONVERSATION_STATES
+                ),
+            key='voice_instruction'  
+        )
 
+        if uploaded_voice and st.session_state.on_record_voice_instruction:
+            voice_bytes = uploaded_voice.read()
+            update_voice_instruction_state(voice_bytes)
+            st.session_state.on_record_voice_instruction = False
+
+        
 @st.fragment
 def quick_actions_fragment():
     action_cols_spec = [_["width"] for _ in QUICK_ACTIONS]
@@ -184,32 +144,61 @@ def conversation_section():
     with st._bottom:
         bottom_input_section()
 
-    if one_time_prompt := st.session_state.new_prompt:
-        st.session_state.update(new_prompt="", pg_messages=[])
-
-        with st.chat_message("user"):
-            st.write(one_time_prompt)
-        st.session_state.pg_messages.append({"role": "user", "content": one_time_prompt})
+    if (not st.session_state.new_prompt) and (not st.session_state.new_vi_base64):
+        return
     
-        with st.chat_message("assistant"):
-            with st.spinner("Thinking..."):
-                error_msg, warnings, response = retrive_response_with_ui(
-                    model_name=st.session_state.pg_model_name,
-                    text_input=one_time_prompt, 
-                    array_audio_input=st.session_state.pg_audio_array,
-                    base64_audio_input=st.session_state.pg_audio_base64, 
-                    stream=True
+    one_time_prompt = st.session_state.new_prompt
+    one_time_vi_array = st.session_state.new_vi_array
+    one_time_vi_base64 = st.session_state.new_vi_base64
+
+    st.session_state.update(
+        new_prompt="", 
+        new_vi_array=np.array([]),
+        new_vi_base64="",
+        pg_messages=[]
+    )
+
+    with st.chat_message("user"):
+        if one_time_vi_base64:
+            with st.spinner("Transcribing..."):
+                error_msg, warnings, one_time_prompt = retrive_response_with_ui(
+                    model_name=MODEL_NAMES["audiollm"]["vllm_name"],
+                    text_input="Write out the dialogue as text.", 
+                    array_audio_input=one_time_vi_array,
+                    base64_audio_input=one_time_vi_base64,
+                    stream=False,
+                    normalise_response=True
                 )
+        else:
+            error_msg, warnings = "", []
+            st.write(one_time_prompt)
 
-        st.session_state.pg_messages.append({
-            "role": "assistant", 
-            "error": error_msg,
-            "warnings": warnings, 
-            "content": response
-        })
+    st.session_state.pg_messages.append({
+        "role": "user", 
+        "error": error_msg,
+        "warnings": warnings, 
+        "content": one_time_prompt
+    })
+
+    with st.chat_message("assistant"):
+        with st.spinner("Thinking..."):
+            error_msg, warnings, response = retrive_response_with_ui(
+                model_name=st.session_state.pg_model_name,
+                text_input=one_time_prompt, 
+                array_audio_input=st.session_state.pg_audio_array,
+                base64_audio_input=st.session_state.pg_audio_base64, 
+                stream=True
+            )
 
-        st.session_state.disprompt=False
-        st.rerun(scope="app")
+    st.session_state.pg_messages.append({
+        "role": "assistant", 
+        "error": error_msg,
+        "warnings": warnings, 
+        "content": response
+    })
+
+    st.session_state.disprompt=False
+    st.rerun(scope="app")
 
 
 def playground_page():
@@ -223,14 +212,18 @@ def playground_page():
         <strong>Spoken Question Answering</strong>,
         <strong>Spoken Dialogue Summarization</strong>, 
         <strong>Speech Instruction</strong>, and 
-        <strong>Paralinguistics</strong> tasks.
-        This playground currently only support <strong>single-round</strong> conversation.
-        """,
-        concise_description=" This playground currently only support <strong>single-round</strong> conversation."
+        <strong>Paralinguistics</strong> tasks.""",
+        concise_description=""
         )
 
     with st.sidebar:
         sidebar_fragment()
 
-    successful_example_section()
+    audio_sample_names = [name for name in AUDIO_SAMPLES_W_INSTRUCT.keys()]
+    successful_example_section(
+        audio_sample_names, 
+        audio_array_state="pg_audio_array",
+        audio_base64_state="pg_audio_base64",
+        restore_state=PG_CONVERSATION_STATES
+    )
     conversation_section()  

src/content/voice_chat.py

@@ -1,24 +1,20 @@
-import copy
-import base64
-
 import numpy as np
 import streamlit as st
 
 from src.generation import (
-    MAX_AUDIO_LENGTH, 
     prepare_multimodal_content, 
     change_multimodal_content
 )
 from src.content.common import (
     MODEL_NAMES,
     VOICE_CHAT_DIALOGUE_STATES,
+    reset_states,
+    process_audio_bytes,
     init_state_section,
     header_section,
     sidebar_fragment,
-    reset_states,
     retrive_response_with_ui
 )
-from src.utils import bytes_to_array, array_to_bytes
 
 
 # TODO: change this.
@@ -26,20 +22,11 @@ DEFAULT_PROMPT = "Based on the information in this user’s voice, please reply
 MAX_VC_ROUNDS = 5
 
 
-def _update_audio(audio_bytes):
-    origin_audio_array = bytes_to_array(audio_bytes)
-    truncated_audio_array = origin_audio_array[: MAX_AUDIO_LENGTH*16000]
-    truncated_audio_bytes = array_to_bytes(truncated_audio_array)
-
-    st.session_state.vc_audio_array = origin_audio_array
-    st.session_state.vc_audio_base64 = base64.b64encode(truncated_audio_bytes).decode('utf-8')
-
-
 def bottom_input_section():
     bottom_cols = st.columns([0.03, 0.97])
     with bottom_cols[0]:
         st.button(
-            'Clear', 
+            ':material/delete:', 
             disabled=st.session_state.disprompt,
             on_click=lambda: reset_states(VOICE_CHAT_DIALOGUE_STATES)
         )
@@ -48,6 +35,7 @@ def bottom_input_section():
         uploaded_file = st.audio_input(
             label="record audio",
             label_visibility="collapsed",
+            disabled=st.session_state.disprompt,
             on_change=lambda: st.session_state.update(
                 on_record=True, 
                 disprompt=True
@@ -57,7 +45,8 @@ def bottom_input_section():
 
         if uploaded_file and st.session_state.on_record:
             audio_bytes = uploaded_file.read()
-            _update_audio(audio_bytes)
+            st.session_state.vc_audio_array, st.session_state.vc_audio_base64 = \
+                process_audio_bytes(audio_bytes)
             st.session_state.update(
                 on_record=False,
             )
@@ -69,6 +58,7 @@ def system_prompt_fragment():
         st.text_area(
             label="Insert system instructions or background knowledge here.",
             label_visibility="collapsed",
+            disabled=st.session_state.disprompt,
             max_chars=5000,
             key="system_prompt",
             value=DEFAULT_PROMPT,
@@ -151,9 +141,9 @@ def voice_chat_page():
     init_state_section()
     header_section(
         component_name="Voice Chat",
-        description=""" It currently only support up to <strong>5 rounds</strong> of conversations.
+        description=""" Currently support up to <strong>5 rounds</strong> of conversations.
         Feel free to talk about anything.""",
-        concise_description=" It currently only support up to <strong>5 rounds</strong> of conversations.",
+        concise_description=" Currently support up to <strong>5 rounds</strong> of conversations.",
         icon="🗣️"
         )
 

src/utils.py

@@ -1,4 +1,5 @@
 import io
+import re
 from datetime import datetime
 from scipy.io.wavfile import write
 
@@ -21,4 +22,10 @@ def array_to_bytes(audio_array):
     bytes_wav = bytes()
     byte_io = io.BytesIO(bytes_wav)
     write(byte_io, 16000, audio_array)
-    return byte_io.read()
+    return byte_io.read()
+
+
+def postprocess_voice_transcription(text):
+    text = re.sub("<.*>:?|\(.*\)|\[.*\]", "", text)
+    text = re.sub("\s+", " ", text).strip()
+    return text

style/app_style.css

@@ -88,15 +88,15 @@ div[data-testid="stBottomBlockContainer"] div[data-testid="stHorizontalBlock"]:h
 }
 
 div[data-testid="stBottomBlockContainer"] div[data-testid="stColumn"]:has( div[data-testid="stButton"]):first-of-type {
-    width: 61px;
-    min-width: 61px;
-    flex: 0 0 61px;
+    width: 42px;
+    min-width: 42px;
+    flex: 0 0 42px;
 }
 
 div[data-testid="stBottomBlockContainer"] div[data-testid="stColumn"]:has( div[data-testid="stButton"]):nth-of-type(2) {
-    width: 76px;
-    min-width: 76px;
-    flex: 0 0 76px;
+    width: 42px;
+    min-width: 42px;
+    flex: 0 0 42px;
 }
 
 div[data-testid="stBottomBlockContainer"] div[data-testid="stColumn"]:has( div[data-testid="stChatInput"]) {
@@ -113,4 +113,38 @@ div[data-testid="stBottomBlockContainer"] div[data-testid="stColumn"]:has( div[d
 
 div[data-testid="stBottomBlockContainer"] div[data-testid="stAudioInput"]>div {
     max-height: 40px;
+}
+
+/* Mic Button */
+
+div[data-testid="stBottomBlockContainer"]:has( div[data-testid="stChatInput"]) div[data-testid="stAudioInput"]>div {
+    display: block;
+    padding: 0;
+    margin: auto;
+}
+
+div[data-testid="stBottomBlockContainer"]:has( div[data-testid="stChatInput"]) div[data-testid="stAudioInput"]>div>div:last-of-type {
+    display:none;
+}
+
+div[data-testid="stBottomBlockContainer"]:has( div[data-testid="stChatInput"]) div[data-testid="stAudioInput"]>div>div:nth-of-type(2) {
+    margin:auto;
+}
+
+div[data-testid="stBottomBlockContainer"]:has( div[data-testid="stChatInput"]) div[data-testid="stAudioInput"]>div>div:nth-of-type(2)>span:last-of-type {
+    display:none;
+}
+
+div[data-testid="stBottomBlockContainer"]:has( div[data-testid="stChatInput"]) div[data-testid="stAudioInput"]>div>div:nth-of-type(2)>span:only-of-type {
+    display:block;
+}
+
+div[data-testid="stBottomBlockContainer"]:has( div[data-testid="stChatInput"]) div[data-testid="stAudioInput"]>div>span {
+    display:none;
+}
+
+div[data-testid="stBottomBlockContainer"]:has( div[data-testid="stChatInput"]) div[data-testid="stColumn"]:has( div[data-testid="stAudioInput"]) {
+    width: 24px;
+    min-width: 24px;
+    flex: 0 0 24px;
 }