add logging function

src/content/common.py

@@ -4,7 +4,8 @@ import numpy as np
 import streamlit as st
 
 from src.tunnel import start_server
-from src.generation import FIXED_GENERATION_CONFIG, load_model
+from src.generation import FIXED_GENERATION_CONFIG, load_model, retrive_response
+from src.logger import load_logger
 
 
 DEFAULT_DIALOGUE_STATES = dict(
@@ -308,6 +309,10 @@ def init_state_section():
         unsafe_allow_html=True
     )
 
+    if "logger" not in st.session_state:
+        st.session_state.logger = load_logger()
+        st.session_state.session_id = st.session_state.logger.register_session()
+
     if "server" not in st.session_state:
         st.session_state.server = start_server()
 
@@ -328,7 +333,7 @@ def header_section(component_name="Playground", icon="🤖"):
         f"<h1 style='text-align: center;'>MERaLiON-AudioLLM {component_name} {icon}</h1>", 
         unsafe_allow_html=True
         )
-
+    
     st.markdown(
         f"""<div class="main-intro-normal-window">
         <p>This {component_name.lower()} is based on 
@@ -336,7 +341,8 @@ def header_section(component_name="Playground", icon="🤖"):
         target="_blank" rel="noopener noreferrer"> MERaLiON-AudioLLM</a>, 
         developed by I2R, A*STAR, in collaboration with AISG, Singapore. 
         It is tailored for Singapore’s multilingual and multicultural landscape.
-        MERaLiON-AudioLLM supports <strong>Automatic Speech Recognation</strong>, 
+        MERaLiON-AudioLLM supports 
+        <strong>Automatic Speech Recognation</strong>, 
         <strong>Speech Translation</strong>, 
         <strong>Spoken Question Answering</strong>,
         <strong>Spoken Dialogue Summarization</strong>, 
@@ -356,9 +362,9 @@ def header_section(component_name="Playground", icon="🤖"):
 
 @st.fragment
 def sidebar_fragment():
-    with st.container(height=300, border=False):
-        st.page_link("pages/playground.py", disabled=st.session_state.disprompt, label="Playground")
-        st.page_link("pages/voice_chat.py", disabled=st.session_state.disprompt, label="Voice Chat (experimental)")
+    with st.container(height=256, border=False):
+        st.page_link("pages/playground.py", disabled=st.session_state.disprompt, label="🚀 Playground")
+        st.page_link("pages/voice_chat.py", disabled=st.session_state.disprompt, label="🗣️ Voice Chat (experimental)")
 
     
     st.divider()
@@ -367,4 +373,47 @@ def sidebar_fragment():
 
     st.slider(label='Top P', min_value=0.0, max_value=1.0, value=0.9, key='top_p')
 
-    st.slider(label="Repetition Penalty", min_value=1.0, max_value=1.2, value=1.1, key="repetition_penalty")
+    st.slider(label="Repetition Penalty", min_value=1.0, max_value=1.2, value=1.1, key="repetition_penalty")
+
+
+def retrive_response_with_ui(prompt, array_audio, base64_audio, stream):
+    generation_params = dict(
+        model=st.session_state.model_name,
+        max_completion_tokens=st.session_state.max_completion_tokens,
+        temperature=st.session_state.temperature,
+        top_p=st.session_state.top_p,
+        extra_body={
+            "repetition_penalty": st.session_state.repetition_penalty,
+            "top_k": st.session_state.top_k,
+            "length_penalty": st.session_state.length_penalty
+        },
+        seed=st.session_state.seed
+    )
+
+    error_msg, warnings, stream = retrive_response(
+        prompt, 
+        array_audio,
+        base64_audio, 
+        params=generation_params,
+        stream=True
+    )
+    response = ""
+
+    if error_msg:
+        st.error(error_msg)
+    for warning_msg in warnings:
+        st.warning(warning_msg)
+    if stream:
+        response = st.write_stream(stream)
+
+    st.session_state.logger.register_query(
+        session_id=st.session_state.session_id,
+        base64_audio=base64_audio,
+        text_input=prompt,
+        params=generation_params,
+        response=response,
+        warnings=warnings,
+        error_msg=error_msg
+    )
+
+    return error_msg, warnings, response

src/content/playground.py

@@ -3,14 +3,15 @@ import base64
 
 import streamlit as st
 
-from src.generation import MAX_AUDIO_LENGTH, retrive_response
+from src.generation import MAX_AUDIO_LENGTH
 from src.utils import bytes_to_array, array_to_bytes
 from src.content.common import (
     AUDIO_SAMPLES_W_INSTRUCT, 
     DEFAULT_DIALOGUE_STATES,
     init_state_section,
     header_section,
-    sidebar_fragment
+    sidebar_fragment,
+    retrive_response_with_ui
 )
 
 
@@ -173,20 +174,12 @@ def conversation_section():
     
         with st.chat_message("assistant"):
             with st.spinner("Thinking..."):
-                error_msg, warnings, stream = retrive_response(
+                error_msg, warnings, response = retrive_response_with_ui(
                     one_time_prompt, 
                     st.session_state.pg_audio_array,
                     st.session_state.pg_audio_base64, 
                     stream=True
                 )
-                response = ""
-
-                if error_msg:
-                    st.error(error_msg)
-                for warning_msg in warnings:
-                    st.warning(warning_msg)
-                if stream:
-                    response = st.write_stream(stream)
 
         st.session_state.pg_messages.append({
             "role": "assistant", 

src/content/voice_chat.py

@@ -4,13 +4,14 @@ import base64
 import numpy as np
 import streamlit as st
 
-from src.generation import retrive_response
+from src.generation import MAX_AUDIO_LENGTH
 from src.utils import bytes_to_array, array_to_bytes
 from src.content.common import (
     DEFAULT_DIALOGUE_STATES,
     init_state_section,
     header_section,
-    sidebar_fragment
+    sidebar_fragment,
+    retrive_response_with_ui
 )
 
 
@@ -18,9 +19,6 @@ from src.content.common import (
 DEFAULT_PROMPT = "Please follow the instruction in the speech."
 
 
-MAX_AUDIO_LENGTH = 120
-
-
 def _update_audio(audio_bytes):
     origin_audio_array = bytes_to_array(audio_bytes)
     truncated_audio_array = origin_audio_array[: MAX_AUDIO_LENGTH*16000]
@@ -123,20 +121,12 @@ def conversation_section():
     
         with st.chat_message("assistant"):
             with st.spinner("Thinking..."):
-                error_msg, warnings, stream = retrive_response(
+                error_msg, warnings, response = retrive_response_with_ui(
                     one_time_prompt, 
                     one_time_array,
                     one_time_base64, 
                     stream=True
                 )
-                response = ""
-
-                if error_msg:
-                    st.error(error_msg)
-                for warning_msg in warnings:
-                    st.warning(warning_msg)
-                if stream:
-                    response = st.write_stream(stream)
 
         st.session_state.vc_messages.append({
             "role": "assistant", 

src/generation.py

@@ -64,7 +64,7 @@ def _retrive_response(text_input: str, base64_audio_input: str, **kwargs):
     )
 
 
-def _retry_retrive_response_throws_exception(text_input, base64_audio_input, stream=False, retry=3):
+def _retry_retrive_response_throws_exception(text_input, base64_audio_input, params, stream=False, retry=3):
     if not base64_audio_input:
         raise NoAudioException("audio is empty.")
 
@@ -72,17 +72,8 @@ def _retry_retrive_response_throws_exception(text_input, base64_audio_input, str
         response_object = _retrive_response(
             text_input=text_input, 
             base64_audio_input=base64_audio_input, 
-            model=st.session_state.model_name,
-            max_completion_tokens=st.session_state.max_completion_tokens,
-            temperature=st.session_state.temperature,
-            top_p=st.session_state.top_p,
-            extra_body={
-                "repetition_penalty": st.session_state.repetition_penalty,
-                "top_k": st.session_state.top_k,
-                "length_penalty": st.session_state.length_penalty
-            },
-            seed=st.session_state.seed,
-            stream=stream
+            stream=stream,
+            **params
         )
     except APIConnectionError as e:
         if not st.session_state.server.is_running():
@@ -122,13 +113,13 @@ def _validate_input(text_input, array_audio_input) -> List[str]:
     return warnings
 
 
-def retrive_response(text_input, array_audio_input, base64_audio_input, stream=False):
+def retrive_response(text_input, array_audio_input, base64_audio_input, params, stream=False):
     warnings = _validate_input(text_input, array_audio_input)
 
     response_object, error_msg = None, ""
     try:
         response_object = _retry_retrive_response_throws_exception(
-            text_input, base64_audio_input, stream
+            text_input, base64_audio_input, params, stream
         )
     except NoAudioException:
         error_msg = "Please specify audio first!"

src/logger.py

@@ -0,0 +1,112 @@
+import io
+import os
+import time
+import json
+from threading import Thread, Lock
+
+import streamlit as st
+from huggingface_hub import HfApi
+
+from src.utils import get_current_strftime
+
+
+logger_lock = Lock()
+
+
+def threaded(fn):
+    def wrapper(*args, **kwargs):
+        thread = Thread(target=fn, args=args, kwargs=kwargs)
+        thread.start()
+        return thread
+    return wrapper
+
+
+class Logger:
+    def __init__(self):
+        self.app_id = get_current_strftime()
+        self.session_increment = 0
+        self.query_increment = 0
+        self.sync_interval = 180
+
+        self.session_data = []
+        self.query_data = []
+        self.audio_data = []
+
+        self.sync_data()
+
+    def register_session(self) -> str:
+        new_session_id = f"{self.app_id}+{self.session_increment}"
+        with logger_lock:
+            self.session_data.append({
+                "session_id": new_session_id,
+                "creation_time": get_current_strftime()
+            })
+
+            self.session_increment += 1
+        return new_session_id
+
+    def register_query(self, 
+                       session_id, 
+                       base64_audio, 
+                       text_input, 
+                       params,
+                       response,
+                       warnings,
+                       error_msg
+                       ):
+        new_query_id = self.query_increment
+        current_time = get_current_strftime()
+
+        with logger_lock:
+            self.query_data.append({
+                "session_id": session_id,
+                "query_id": new_query_id,
+                "creation_time": current_time,
+                "text": text_input,
+                "params": params,
+                "response": response,
+                "warnings": warnings,
+                "error": error_msg,
+            })
+
+            self.audio_data.append({
+                "session_id": session_id,
+                "query_id": new_query_id,
+                "creation_time": current_time,
+                "audio": base64_audio,
+            })
+            self.query_increment += 1
+
+
+    @threaded
+    def sync_data(self):
+        api = HfApi()
+
+        while True:
+            time.sleep(self.sync_interval)
+
+            for data_name in ["session_data", "query_data", "audio_data"]:
+                with logger_lock:
+                    last_data = getattr(self, data_name, [])
+                    setattr(self, data_name, [])
+
+                if not last_data:
+                    continue
+
+                buffer = io.BytesIO()
+                for row in last_data:
+                    row_str = json.dumps(row, ensure_ascii=False)+"\n"
+                    buffer.write(row_str.encode("utf-8"))
+
+                api.upload_file(
+                    path_or_fileobj=buffer,
+                    path_in_repo=f"{data_name}/{get_current_strftime()}.json",
+                    repo_id=os.getenv("LOGGING_REPO_NAME"),
+                    repo_type="dataset",
+                    token=os.getenv('HF_TOKEN')
+                )
+
+
+@st.cache_resource()
+def load_logger():
+    return Logger()

src/utils.py

@@ -1,9 +1,14 @@
 import io
+from datetime import datetime
 from scipy.io.wavfile import write
 
 import librosa
 
 
+def get_current_strftime():
+    return datetime.now().strftime(r'%d-%m-%y-%H-%M-%S')
+
+
 def bytes_to_array(audio_bytes):
     audio_array, _ = librosa.load(
         io.BytesIO(audio_bytes), 

style/app_style.css

@@ -1,7 +1,16 @@
+div[data-testid="stMainBlockContainer"] {
+    padding-top: 2rem;
+    padding-bottom: 1rem;
+}
+
 div[data-testid="stMainBlockContainer"] div[data-testid="stAudioInput"]>div {
     max-height: 3rem;
 }
 
+div[data-testid="stMainBlockContainer"] h1 {
+    padding-top: 0.25rem;
+}
+
 div[class="sidebar-intro"] p {
     margin-bottom: 0.75rem;
 }
@@ -16,8 +25,15 @@ div[data-testid="stChatMessage"]:has(> div[data-testid="stChatMessageAvatarUser"
     text-align: right;
 }
 
+/* audio quick actions */
+
+div[data-testid="stChatMessage"] div[data-testid="stVerticalBlock"]:has( audio[data-testid="stAudio"]) {
+    gap: 2px;
+}
+
 div[data-testid="stChatMessage"] div[data-testid="stHorizontalBlock"]:has(> div[data-testid="stColumn"]) {
     flex-direction: row-reverse;
+    gap: 4px;
 }
 
 div[data-testid="stChatMessage"] div[data-testid="stHorizontalBlock"]>div[data-testid="stColumn"]:has( div[data-testid="stButton"]) {

style/normal_window.css

@@ -1,6 +1,7 @@
 @media(min-width: 576px) {
-    .stMainBlockContainer {
-        padding: 2rem 5rem 1rem;
+    div[data-testid="stMainBlockContainer"] {
+        padding-left: 5rem;
+        padding-bottom: 5rem;
     }
 
     div[data-testid="stBottomBlockContainer"] {

style/small_window.css

@@ -1,4 +1,9 @@
 @media(max-width: 576px) {
+    div[data-testid="stMainBlockContainer"] {
+        padding-left: 1rem;
+        padding-bottom: 1rem;
+    }
+
     div[data-testid="stMainBlockContainer"] div[data-testid="stVerticalBlock"]>div[data-testid="stElementContainer"]:has( div[data-testid="stHeadingWithActionElements"]) {
         display: none;
     }
@@ -6,4 +11,8 @@
     div[class="main-intro-normal-window"] {
         display: none;
     }
+    
+    div[data-testid="stSidebarCollapsedControl"] button[data-testid="stBaseButton-headerNoPadding"]::after {
+        content: "More Use Cases"
+    }
 }