add agent

pages/agent.py

@@ -0,0 +1,3 @@
+from src.content.agent import agent_page
+
+agent_page()

requirements.txt

@@ -2,4 +2,8 @@ librosa==0.10.2.post1
 streamlit==1.40.2
 openai==1.57.1
 streamlit_mic_recorder==0.0.8
-sshtunnel
+sshtunnel
+accelerate==1.3.0
+FlagEmbedding==1.3.3
+sentence-transformers==3.4.0
+sentencepiece==0.1.99

src/content/agent.py

@@ -0,0 +1,240 @@
+import copy
+import base64
+
+import streamlit as st
+
+from src.generation import MAX_AUDIO_LENGTH
+from src.retrieval import retrieve_relevant_docs
+from src.utils import bytes_to_array, array_to_bytes
+from src.content.common import (
+    MODEL_NAMES,
+    AUDIO_SAMPLES_W_INSTRUCT, 
+    STANDARD_QUERIES,
+    DEFAULT_DIALOGUE_STATES,
+    init_state_section,
+    header_section,
+    sidebar_fragment,
+    retrive_response_with_ui
+)
+
+
+LLM_PROMPT_TEMPLATE = """User asked a question about the audio clip.
+
+## User question
+{user_question}
+
+{audio_information_prompt}Please reply this user question with an friendly, accurate, and helpful answer."""
+
+AUDIO_INFO_TEMPLATE = """Here are some information about this audio clip.
+
+## Audio Information
+{audio_information}
+
+This may or may not contain relevant information to the user question, please use with caution.
+
+"""
+
+
+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 = [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, 
+            ag_messages=[],
+            ag_model_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, ag_messages=[], ag_model_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, ag_messages=[], ag_model_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()
+
+
+def bottom_input_section():
+    bottom_cols = st.columns([0.03, 0.03, 0.94])
+    with bottom_cols[0]:
+        st.button(
+            'Clear', 
+            disabled=st.session_state.disprompt,
+            on_click=lambda: st.session_state.update(copy.deepcopy(DEFAULT_DIALOGUE_STATES))
+        )
+
+    with bottom_cols[1]:
+        if st.button("\+ Audio", disabled=st.session_state.disprompt):
+            audio_attach_dialogue()
+
+    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)
+        ):
+            st.session_state.new_prompt = chat_input
+
+
+def conversation_section():
+    chat_message_container = st.container(height=480)
+    if st.session_state.ag_audio_array.size:
+        with chat_message_container.chat_message("user"):
+            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):
+            if message.get("error"):
+                st.error(message["error"])
+            for warning_msg in message.get("warnings", []):
+                st.warning(warning_msg)
+            if process := message.get("process", []):
+                with st.status("Thought process...", expanded=True, state="complete"):
+                    for proc in process:
+                        if proc.get("error"):
+                            st.error(proc["error"])
+                        for proc_warning_msg in proc.get("warnings", []):
+                            st.warning(proc_warning_msg)
+                        if proc.get("content"):
+                            st.write(proc["content"])
+            if message.get("content"):
+                st.write(message["content"])
+    
+    with st._bottom:
+        bottom_input_section()
+
+    if one_time_prompt := st.session_state.new_prompt:
+        st.session_state.update(new_prompt="")
+
+        with chat_message_container.chat_message("user"):
+            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)
+            
+            relevant_query_indices = retrieve_relevant_docs(one_time_prompt, STANDARD_QUERIES)
+            if len(st.session_state.ag_messages) <= 2:
+                relevant_query_indices.append(0)
+
+            relevant_query_indices = list(set(relevant_query_indices).difference(st.session_state.ag_visited_query_indices))
+
+            audio_info = []
+            if relevant_query_indices:
+                with st.status("Thought process...", expanded=True) as status:
+                    for idx in relevant_query_indices:
+                        error_msg, warnings, response = retrive_response_with_ui(
+                            model_name=MODEL_NAMES["with_lora"]["vllm_name"],
+                            prompt=STANDARD_QUERIES[idx]["query_text"], 
+                            array_audio=st.session_state.ag_audio_array,
+                            base64_audio=st.session_state.ag_audio_base64, 
+                            prefix=f"**{STANDARD_QUERIES[idx]['ui_text']}** :speech_balloon: : ",
+                            stream=True
+                        )
+                        audio_info.append(STANDARD_QUERIES[idx]["response_prefix_text"] + response)
+                        
+                        assistant_message["process"].append({
+                            "error": error_msg,
+                            "warnings": warnings, 
+                            "content": response
+                        })
+                    
+                    status.update(state="complete")
+
+            audio_information_prompt = ""
+            if audio_info:
+                audio_information_prompt = AUDIO_INFO_TEMPLATE.format(
+                    audio_information="\n".join(audio_info)
+                )
+
+            prompt = LLM_PROMPT_TEMPLATE.format(
+                user_question=one_time_prompt, 
+                audio_information_prompt=audio_information_prompt
+            )
+            
+            error_msg, warnings, response = retrive_response_with_ui(
+                model_name=MODEL_NAMES["wo_lora"]["vllm_name"],
+                prompt=prompt, 
+                array_audio=st.session_state.ag_audio_array,
+                base64_audio="", 
+                stream=True,
+                history=st.session_state.ag_model_messages
+            )
+            
+            assistant_message.update({"error": error_msg, "warnings": warnings, "content": response})
+            st.session_state.ag_model_messages.extend([
+                {"role": "user", "content": prompt},
+                {"role": "assistant", "content": response}
+            ])
+
+        st.session_state.disprompt=False
+        st.rerun(scope="app")
+
+
+def agent_page():
+    init_state_section()
+    header_section(component_name="Agent System", icon="👥")
+
+    with st.sidebar:
+        sidebar_fragment()
+
+    successful_example_section()
+    conversation_section()  

src/content/common.py

@@ -1,10 +1,13 @@
 import copy
+import itertools
+from collections import OrderedDict
 
 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
 
 
@@ -15,6 +18,11 @@ DEFAULT_DIALOGUE_STATES = dict(
     vc_audio_base64='',
     vc_audio_array=np.array([]),
     vc_messages=[], 
+    ag_audio_base64='',
+    ag_audio_array=np.array([]),
+    ag_visited_query_indices=[],
+    ag_messages=[], 
+    ag_model_messages=[],
     disprompt = False,
     new_prompt = "",
     on_select=False, 
@@ -24,17 +32,16 @@ DEFAULT_DIALOGUE_STATES = dict(
 )
 
 
-DEFAULT_VOICE_CHAT_STATES = dict(
-    audio_base64='',
-    audio_array=np.array([]),
-    disprompt = False,
-    new_prompt = "",
-    messages=[], 
-    on_select=False, 
-    on_upload=False, 
-    on_record=False, 
-    on_select_quick_action=False
-)
+MODEL_NAMES = OrderedDict({
+    "with_lora": {
+        "vllm_name": "MERaLiON-lora", 
+        "ui_name": "MERaLiON-AudioLLM (more accurate)"
+    },
+    "wo_lora": {
+        "vllm_name": "MERaLiON_local/MERaLiON-AudioLLM-Whisper-SEA-LION-wo-lora", 
+        "ui_name": "MERaLiON-AudioLLM-instruction-following (more flexible)"
+    }
+})
 
 
 AUDIO_SAMPLES_W_INSTRUCT = {
@@ -295,6 +302,46 @@ AUDIO_SAMPLES_W_INSTRUCT = {
 }
 
 
+STANDARD_QUERIES = [
+    {
+        "query_text": "Please transcribe this speech.",
+        "doc_text": "Listen to a speech and write down exactly what is being said in text form. It's essentially converting spoken words into written words. Provide the exact transcription of the given audio. Record whatever the speaker has said into written text.",
+        "response_prefix_text": "The transcription of the speech is: ",
+        "ui_text": "speech trancription"
+    },
+    {
+        "query_text": "Please describe what happended in this audio",
+        "doc_text": "Text captions describing the sound events and environments in the audio clips, describing the events and actions happened in the audio.",
+        "response_prefix_text": "Events in this audio clip: ",
+        "ui_text": "audio caption"
+    },
+    {
+        "query_text": "May I know the gender of the speakers",
+        "doc_text": "Please identify speaker gender by analyzing pitch, formants, harmonics, and prosody features, which reflect physiological and speech pattern differences between genders.",
+        "response_prefix_text": "By analyzing pitch, formants, harmonics, and prosody features, which reflect physiological and speech pattern differences between genders: ",
+        "ui_text": "gender recognition"
+    },
+    {
+        "query_text": "May I know the nationality of the speakers",
+        "doc_text": "Discover speakers' nationality, country, or the place he is coming from. Analyze speakers' accent, pronunciation patterns, intonation, rhythm, phoneme usage, and language-specific speech features influenced by cultural and linguistic backgrounds.",
+        "response_prefix_text": "By analyzing accent, pronunciation patterns, intonation, rhythm, phoneme usage, and language-specific speech features influenced by cultural and linguistic backgrounds: ",
+        "ui_text": "accent recognition"
+    },
+    {
+        "query_text": "Can you guess which ethnic group this person is from based on their accent.",
+        "doc_text": "Discover speakers' ethnic group, home country, or the place he is coming from, from speech features like accent, tone, intonation, phoneme variations, and vocal characteristics influenced by cultural, regional, and linguistic factors.",
+        "response_prefix_text": "By analyzing speech features like accent, tone, intonation, phoneme variations, and vocal characteristics influenced by cultural, regional, and linguistic factors: ",
+        "ui_text": "accent recognition"
+    },
+    {
+        "query_text": "What do you think the speakers are feeling.",
+        "doc_text": "What do you think the speakers are feeling. Please identify speakers' emotions by analyzing vocal features like pitch, tone, volume, speech rate, rhythm, and spectral energy, which reflect emotional states such as happiness, anger, sadness, or fear.",
+        "response_prefix_text": "By analyzing vocal features like pitch, tone, volume, speech rate, rhythm, and spectral energy: ",
+        "ui_text": "emotion recognition"
+    },
+]
+
+
 def init_state_section():
     st.set_page_config(page_title='MERaLiON-AudioLLM', page_icon = "🔥", layout='wide')
 
@@ -317,7 +364,10 @@ def init_state_section():
         st.session_state.server = start_server()
 
     if "client" not in st.session_state or 'model_name' not in st.session_state:
-        st.session_state.client, _ = load_model()
+        st.session_state.client = load_model()
+
+    if "retriever" not in st.session_state:
+        st.session_state.retriever = load_retriever()
 
     for key, value in FIXED_GENERATION_CONFIG.items():
         if key not in st.session_state:
@@ -364,8 +414,8 @@ def header_section(component_name="Playground", icon="🤖"):
 def sidebar_fragment():
     with st.container(height=256, border=False):
         st.page_link("pages/playground.py", disabled=st.session_state.disprompt, label="🚀 Playground")
+        st.page_link("pages/agent.py", disabled=st.session_state.disprompt, label="👥 Multi-Agent System")
         st.page_link("pages/voice_chat.py", disabled=st.session_state.disprompt, label="🗣️ Voice Chat (experimental)")
-
     
     st.divider()
 
@@ -376,9 +426,9 @@ def sidebar_fragment():
     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):
+def retrive_response_with_ui(model_name, prompt, array_audio, base64_audio, prefix="", **kwargs):
     generation_params = dict(
-        model=st.session_state.model_name,
+        model=model_name,
         max_completion_tokens=st.session_state.max_completion_tokens,
         temperature=st.session_state.temperature,
         top_p=st.session_state.top_p,
@@ -390,26 +440,33 @@ def retrive_response_with_ui(prompt, array_audio, base64_audio, stream):
         seed=st.session_state.seed
     )
 
-    error_msg, warnings, stream = retrive_response(
+    error_msg, warnings, response_obj = retrive_response(
         prompt, 
         array_audio,
         base64_audio, 
-        params=generation_params,
-        stream=True
+        **generation_params,
+        **kwargs
     )
-    response = ""
 
     if error_msg:
         st.error(error_msg)
     for warning_msg in warnings:
         st.warning(warning_msg)
-    if stream:
-        response = st.write_stream(stream)
+
+    response = ""
+    if response_obj is not None:
+        if kwargs.get("stream", ""):
+            response_obj = itertools.chain([prefix], response_obj)
+            response = st.write_stream(response_obj)
+        else:
+            response = response_obj.choices[0].message.content
+            st.write(prefix+response)
 
     st.session_state.logger.register_query(
         session_id=st.session_state.session_id,
         base64_audio=base64_audio,
         text_input=prompt,
+        history=kwargs.get("history", []),
         params=generation_params,
         response=response,
         warnings=warnings,

src/content/playground.py

@@ -6,6 +6,7 @@ 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, 
     DEFAULT_DIALOGUE_STATES,
     init_state_section,
@@ -104,19 +105,15 @@ def audio_attach_dialogue():
 
 @st.fragment
 def select_model_variants_fradment():
-    display_mapper = {
-        'MERaLiON-lora': "MERaLiON-AudioLLM (better transcription)",
-        'MERaLiON_local/MERaLiON-AudioLLM-Whisper-SEA-LION-wo-lora': "MERaLiON-AudioLLM-instruction-following (more flexible)"
-    }
+    display_mapper = {value["vllm_name"]: value["ui_name"] for value in MODEL_NAMES.values()}
 
     st.selectbox(
         label=":fire: Explore more MERaLiON-AudioLLM variants!",
-        # label_visibility="collapsed",
-        options=['MERaLiON-lora', 'MERaLiON_local/MERaLiON-AudioLLM-Whisper-SEA-LION-wo-lora'],
+        options=[value["vllm_name"] for value in MODEL_NAMES.values()],
         index=0,
         format_func=lambda o: display_mapper[o],
-        key="model_name",
-        placeholder=":fire: Explore more model variants!",
+        key="pg_model_name",
+        placeholder=":fire: Explore more MERaLiON-AudioLLM variants!",
         disabled=st.session_state.disprompt,
     )
 
@@ -196,9 +193,10 @@ def conversation_section():
         with st.chat_message("assistant"):
             with st.spinner("Thinking..."):
                 error_msg, warnings, response = retrive_response_with_ui(
-                    one_time_prompt, 
-                    st.session_state.pg_audio_array,
-                    st.session_state.pg_audio_base64, 
+                    model_name=st.session_state.pg_model_name,
+                    prompt=one_time_prompt, 
+                    array_audio=st.session_state.pg_audio_array,
+                    base64_audio=st.session_state.pg_audio_base64, 
                     stream=True
                 )
 

src/content/voice_chat.py

@@ -7,6 +7,7 @@ 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,
     DEFAULT_DIALOGUE_STATES,
     init_state_section,
     header_section,
@@ -122,9 +123,10 @@ def conversation_section():
         with st.chat_message("assistant"):
             with st.spinner("Thinking..."):
                 error_msg, warnings, response = retrive_response_with_ui(
-                    one_time_prompt, 
-                    one_time_array,
-                    one_time_base64, 
+                    model_name=MODEL_NAMES["wo_lora"]["vllm_name"],
+                    prompt=one_time_prompt, 
+                    array_audio=one_time_array,
+                    base64_audio=one_time_base64, 
                     stream=True
                 )
 
@@ -141,7 +143,6 @@ def conversation_section():
 
 def voice_chat_page():
     init_state_section()
-    st.session_state.model_name = 'MERaLiON_local/MERaLiON-AudioLLM-Whisper-SEA-LION-wo-lora'
     header_section(component_name="Voice Chat")
 
     with st.sidebar:

src/generation.py

@@ -6,7 +6,7 @@ from typing import List
 import streamlit as st
 from openai import OpenAI, APIConnectionError
 
-from src.exceptions import NoAudioException, TunnelNotRunningException
+from src.exceptions import TunnelNotRunningException
 
 
 local_port = int(os.getenv('LOCAL_PORT'))
@@ -33,48 +33,38 @@ def load_model():
         api_key=openai_api_key,
         base_url=openai_api_base,
     )
-
-    models = client.models.list()
-    model_name = models.data[0].id
     
-    return client, model_name
+    return client
 
 def _retrive_response(text_input: str, base64_audio_input: str, **kwargs):
     """
     Send request through OpenAI client. 
     """
-    return st.session_state.client.chat.completions.create(
-        messages=[{
-            "role":
-            "user",
-            "content": [
-                {
-                    "type": "text",
-                    "text": f"Text instruction: {text_input}"
+    history = kwargs.pop("history", [])
+    if base64_audio_input:
+        content = [
+            {
+                "type": "text",
+                "text": f"Text instruction: {text_input}"
+            },
+            {
+                "type": "audio_url",
+                "audio_url": {
+                    "url": f"data:audio/ogg;base64,{base64_audio_input}"
                 },
-                {
-                    "type": "audio_url",
-                    "audio_url": {
-                        "url": f"data:audio/ogg;base64,{base64_audio_input}"
-                    },
-                },
-            ],
-        }],
+            },
+        ]
+    else:
+        content = text_input
+    return st.session_state.client.chat.completions.create(
+        messages=history + [{"role": "user", "content": content}],
         **kwargs
     )
 
 
-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.")
-
+def _retry_retrive_response_throws_exception(retry=3, **kwargs):
     try:
-        response_object = _retrive_response(
-            text_input=text_input, 
-            base64_audio_input=base64_audio_input, 
-            stream=stream,
-            **params
-        )
+        response_object = _retrive_response(**kwargs)
     except APIConnectionError as e:
         if not st.session_state.server.is_running():
             if retry == 0:
@@ -87,7 +77,7 @@ def _retry_retrive_response_throws_exception(text_input, base64_audio_input, par
             elif st.session_state.server.is_starting():
                 time.sleep(2)
                 
-            return _retry_retrive_response_throws_exception(text_input, retry-1)
+            return _retry_retrive_response_throws_exception(retry-1, **kwargs)
         raise e
 
     return response_object
@@ -104,25 +94,37 @@ def _validate_input(text_input, array_audio_input) -> List[str]:
     if re.search(r'[\u4e00-\u9fff]+', text_input):
         warnings.append("NOTE: Please try to prompt in English for the best performance.")  
 
+    if array_audio_input.shape[0] == 0:
+        warnings.append("NOTE: Please specify audio from examples or local files.")
+
     if array_audio_input.shape[0] / 16000 > 30.0:
         warnings.append((
-            "MERaLiON-AudioLLM is trained to process audio up to **30 seconds**."
+            "WARNING: MERaLiON-AudioLLM is trained to process audio up to **30 seconds**."
             f" Audio longer than **{MAX_AUDIO_LENGTH} seconds** will be truncated."
         ))
 
     return warnings
 
 
-def retrive_response(text_input, array_audio_input, base64_audio_input, params, stream=False):
+def retrive_response(
+        text_input, 
+        array_audio_input, 
+        base64_audio_input, 
+        stream=True,
+        history=[],
+        **kwargs
+    ):
     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, params, stream
+            text_input=text_input, 
+            base64_audio_input=base64_audio_input, 
+            stream=stream,
+            history=history,
+            **kwargs
         )
-    except NoAudioException:
-        error_msg = "Please specify audio first!"
     except TunnelNotRunningException:
         error_msg = "Internet connection cannot be established. Please contact the administrator."
     except Exception as e:

src/logger.py

@@ -49,25 +49,22 @@ class Logger:
                        session_id, 
                        base64_audio, 
                        text_input, 
-                       params,
                        response,
-                       warnings,
-                       error_msg
+                       **kwargs
                        ):
         new_query_id = self.query_increment
         current_time = get_current_strftime()
 
         with logger_lock:
-            self.query_data.append({
+            current_query_data = {
                 "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,
-            })
+            }
+            current_query_data.update(kwargs)
+            self.query_data.append(current_query_data)
 
             self.audio_data.append({
                 "session_id": session_id,
@@ -98,13 +95,13 @@ class Logger:
                     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')
-                )
+                # 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')
+                # )
 
                 buffer.close()
 

src/retrieval.py

@@ -0,0 +1,20 @@
+from typing import Dict, List
+
+import numpy as np
+import streamlit as st
+from FlagEmbedding import FlagReranker
+
+
+@st.cache_resource()
+def load_retriever():
+    reranker = FlagReranker('BAAI/bge-reranker-v2-m3', use_fp16=True)
+    reranker.compute_score([["test", "test"]], normalize=True)
+    return reranker
+
+
+def retrieve_relevant_docs(user_question, docs: List[Dict]) -> List[int]:
+    scores = st.session_state.retriever.compute_score([[user_question, d["doc_text"]] for d in docs], normalize=True)
+    normalized_scores = np.array(scores) / np.sum(scores)
+
+    selected_indices = np.where((np.array(scores) > 0.02) & (normalized_scores > 0.3))[0]
+    return selected_indices.tolist()

style/small_window.css

@@ -15,4 +15,10 @@
     div[data-testid="stSidebarCollapsedControl"] button[data-testid="stBaseButton-headerNoPadding"]::after {
         content: "More Use Cases"
     }
+}
+
+@media (max-width: 916px) and (max-height: 958px) {
+    div[height="480"][data-testid="stVerticalBlockBorderWrapper"] {
+        height: 380px;
+    }
 }