Update pages.py

pages.py

@@ -2,11 +2,12 @@ import os
 import base64
 
 import numpy as np
+from openai import APIConnectionError
 import streamlit as st
 import streamlit.components.v1 as components
 from streamlit_mic_recorder import mic_recorder
 
-from utils import load_model, generate_response, bytes_to_array, start_server
+from utils import load_model, generate_response, bytes_to_array, start_server, NoAudioException
 
 
 def audio_llm():    
@@ -56,36 +57,6 @@ def audio_llm():
     )
 
     col1, col2, col3 = st.columns([4, 4, 1.2])
-    
-    with col3:
-        st.markdown("or **Record Audio:**")
-        
-        recording = mic_recorder(
-            format="wav", 
-            use_container_width=True, 
-            callback=lambda: st.session_state.update(on_record=True, messages=[]),
-            key='record')
-        
-        if recording and st.session_state.on_record:
-            audio_bytes = recording["bytes"]
-            st.session_state.audio_base64 = base64.b64encode(audio_bytes).decode('utf-8')
-            st.session_state.audio_array = bytes_to_array(audio_bytes)
-    
-    with col2:
-        st.markdown("or **Upload Audio:**")
-
-        uploaded_file = st.file_uploader(
-            label="**Upload Audio:**", 
-            label_visibility="collapsed",
-            type=['wav', 'mp3'],
-            on_change=lambda: st.session_state.update(on_upload=True, messages=[]),
-            key='upload'
-        )
-        
-        if uploaded_file and st.session_state.on_upload:
-            audio_bytes = uploaded_file.read()
-            st.session_state.audio_base64 = base64.b64encode(audio_bytes).decode('utf-8')
-            st.session_state.audio_array = bytes_to_array(audio_bytes)
 
     with col1:
         audio_samples_w_instruct = {
@@ -165,8 +136,43 @@ def audio_llm():
             st.session_state.default_instruction = audio_samples_w_instruct[sample_name]
             st.session_state.audio_base64 = base64.b64encode(audio_bytes).decode('utf-8')
             st.session_state.audio_array = bytes_to_array(audio_bytes)
+
+
+    with col2:
+        st.markdown("or **Upload Audio:**")
+
+        uploaded_file = st.file_uploader(
+            label="**Upload Audio:**", 
+            label_visibility="collapsed",
+            type=['wav', 'mp3'],
+            on_change=lambda: st.session_state.update(on_upload=True, messages=[]),
+            key='upload'
+        )
+        
+        if uploaded_file and st.session_state.on_upload:
+            audio_bytes = uploaded_file.read()
+            st.session_state.audio_base64 = base64.b64encode(audio_bytes).decode('utf-8')
+            st.session_state.audio_array = bytes_to_array(audio_bytes)
+
+
+    with col3:
+        st.markdown("or **Record Audio:**")
+        
+        recording = mic_recorder(
+            format="wav", 
+            use_container_width=True, 
+            callback=lambda: st.session_state.update(on_record=True, messages=[]),
+            key='record')
+        
+        if recording and st.session_state.on_record:
+            audio_bytes = recording["bytes"]
+            st.session_state.audio_base64 = base64.b64encode(audio_bytes).decode('utf-8')
+            st.session_state.audio_array = bytes_to_array(audio_bytes)
     
+
     st.audio(st.session_state.audio_array, format="audio/wav", sample_rate=16000)
+    if st.session_state.audio_array.shape[0] / 16000 > 30.0:
+        st.warning("MERaLiON-AudioLLM can only process audio for up to 30 seconds. Audio longer than that will be truncated.")
     st.session_state.update(on_upload=False, on_record=False, on_select=False)
 
     if st.session_state.default_instruction:
@@ -194,7 +200,12 @@ def audio_llm():
 
     for message in st.session_state.messages[-2:]:
         with st.chat_message(message["role"]):
-            st.write(message["content"])
+            if message.get("error"):
+                st.error(message["error"])
+            for warning_msg in message.get("warnings", []):
+                st.warning(warning_msg)
+            if message.get("content"):
+                st.write(message["content"])
 
     if prompt := st.chat_input(
         placeholder="Type Your Instruction Here", 
@@ -206,19 +217,25 @@ def audio_llm():
         st.session_state.messages.append({"role": "user", "content": prompt})
     
         with st.chat_message("assistant"):
-            if not st.session_state.audio_base64:
-                response = "Please specify audio first!"
-                st.write(response)
-            else:
-                with st.spinner("Thinking..."):
-                    try:
-                        stream = generate_response(prompt)
-                        response = st.write_stream(stream)
-                    except Exception as e:
-                        response = f"Caught Exception: {repr(e)}. Please contact the administrator to restart this space."
-                        st.write(response)
-                        raise(e)
-            st.session_state.messages.append({"role": "assistant", "content": response})
+            response, error_msg, warnings = "", "", []
+            with st.spinner("Thinking..."):
+                try:
+                    stream, warnings = generate_response(prompt)
+                    for warning_msg in warnings:
+                        st.warning(warning_msg)
+                    response = st.write_stream(stream)
+                except NoAudioException:
+                    error_msg = "Please specify audio first!"
+                except APIConnectionError:
+                    error_msg = "Internet connection seems to be down. Please contact the administrator to restart the space."
+                except Exception as e:
+                    error_msg = f"Caught Exception: {repr(e)}. Please contact the administrator."
+            st.session_state.messages.append({
+                "role": "assistant", 
+                "error": error_msg,
+                "warnings": warnings, 
+                "content": response
+            })
 
         st.session_state.disprompt = False
-        st.rerun()
+        st.rerun()