Add/Mod: Add new audio recorder, threshold hparam, min and max beam size single slider.

app.py

@@ -1,17 +1,25 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 
-import os.path as osp
-
-from tempfile import NamedTemporaryFile
-from typing import Any
+from tempfile import NamedTemporaryFile, _TemporaryFileWrapper
+from typing import Any, Optional
 
 import streamlit as st
 
-from audiorecorder import audiorecorder
+from st_audiorec import st_audiorec
 from streamlit.runtime.uploaded_file_manager import UploadedFile
 
 from conette import CoNeTTEModel, conette
+from conette.utils.collections import dict_list_to_list_dict
+
+
+ALLOW_REP_MODES = ("stopwords", "all", "none")
+MAX_BEAM_SIZE = 20
+MAX_PRED_SIZE = 30
+MAX_BATCH_SIZE = 32
+RECORD_AUDIO_FNAME = "record.wav"
+DEFAULT_THRESHOLD = 0.3
+THRESHOLD_PRECISION = 100
 
 
 @st.cache_resource
@@ -26,46 +34,86 @@ def format_candidate(candidate: str) -> str:
         return f"{candidate[0].title()}{candidate[1:]}."
 
 
+def format_tags(tags: Optional[list[str]]) -> str:
+    if tags is None or len(tags) == 0:
+        return "None."
+    else:
+        return ", ".join(tags)
+
+
+def get_result_hash(audio_fname: str, generate_kwds: dict[str, Any]) -> str:
+    return f"{audio_fname}-{generate_kwds}"
+
+
 def get_results(
     model: CoNeTTEModel,
-    audios: list[UploadedFile],
+    audio_files: dict[str, bytes],
     generate_kwds: dict[str, Any],
-) -> tuple[list[str], list[str]]:
-    audio_to_predict = []
-    cands = [""] * len(audios)
-    tmp_files = []
-    tmp_fpaths = []
-    audio_fnames = []
-
-    for i, audio in enumerate(audios):
-        audio_fname = audio.name
-        audio_fnames.append(audio_fname)
-        cand_key = f"{audio_fname}-{generate_kwds}"
-
-        if cand_key in st.session_state:
-            cand = st.session_state[cand_key]
-            cands[i] = cand
-        else:
-            tmp_file = NamedTemporaryFile()
-            tmp_file.write(audio.getvalue())
-            tmp_files.append(tmp_file)
-            audio_to_predict.append((i, cand_key, tmp_file))
-
-            tmp_fpath = tmp_file.name
-            tmp_fpaths.append(tmp_fpath)
-
-    if len(tmp_fpaths) > 0:
-        outputs = model(
-            tmp_fpaths,
+) -> dict[str, dict[str, Any]]:
+    # Get audio to be processed
+    audio_to_predict: dict[str, bytes] = {}
+    for audio_fname, audio in audio_files.items():
+        result_hash = get_result_hash(audio_fname, generate_kwds)
+        if result_hash not in st.session_state or audio_fname == RECORD_AUDIO_FNAME:
+            audio_to_predict[result_hash] = audio
+
+    # Save audio to be processed
+    tmp_files: dict[str, _TemporaryFileWrapper] = {}
+    for result_hash, audio in audio_to_predict.items():
+        tmp_file = NamedTemporaryFile()
+        tmp_file.write(audio)
+        tmp_files[result_hash] = tmp_file
+
+    # Generate predictions and store them in session state
+    for start in range(0, len(tmp_files), MAX_BATCH_SIZE):
+        end = min(start + MAX_BATCH_SIZE, len(tmp_files))
+        result_hashes_j = list(tmp_files.keys())[start:end]
+        tmp_files_j = list(tmp_files.values())[start:end]
+        tmp_paths_j = [tmp_file.name for tmp_file in tmp_files_j]
+        outputs_j = model(
+            tmp_paths_j,
             **generate_kwds,
         )
-        for i, (j, cand_key, tmp_file) in enumerate(audio_to_predict):
-            cand = outputs["cands"][i]
-            cands[j] = cand
-            st.session_state[cand_key] = cand
+        for tmp_file in tmp_files_j:
             tmp_file.close()
+        outputs_lst = dict_list_to_list_dict(outputs_j)  # type: ignore
+        for result_hash, output_i in zip(result_hashes_j, outputs_lst):
+            st.session_state[result_hash] = output_i
+
+    # Get outputs
+    outputs = {}
+    for audio_fname in audio_files.keys():
+        result_hash = get_result_hash(audio_fname, generate_kwds)
+        output_i = st.session_state[result_hash]
+        outputs[audio_fname] = output_i
+
+    return outputs
+
+
+def show_results(outputs: dict[str, dict[str, Any]]) -> None:
+    st.divider()
 
-    return audio_fnames, cands
+    for audio_fname, output in outputs.items():
+        cand = output["cands"]
+        lprobs = output["lprobs"]
+        tags = output.get("tags")
+
+        cand = format_candidate(cand)
+        tags = format_tags(tags)
+        prob = lprobs.exp().tolist()
+
+        if audio_fname == RECORD_AUDIO_FNAME:
+            header = "##### Result for microphone input:"
+        else:
+            header = f'##### Result for "{audio_fname}"'
+
+        content = f"""
+        {header}
+        - **Description:** "{cand}"
+        - **Mean confidence:** {prob*100:.0f}%
+        - **Tags:** {tags}"""
+        st.markdown(content)
+        st.divider()
 
 
 def main() -> None:
@@ -73,46 +121,35 @@ def main() -> None:
 
     model = load_conette(model_kwds=dict(device="cpu"))
 
-    st.warning(
-        "Recommanded audio: lasting from **1 to 30s**, sampled at **32 kHz** minimum."
-    )
-    audios = st.file_uploader(
-        "**Upload audio files here:**",
+    # st.warning(
+    #     "Recommanded audio: lasting from **1 to 30s**, sampled at **32 kHz** minimum."
+    # )
+
+    record_data = st_audiorec()
+    audio_files: Optional[list[UploadedFile]] = st.file_uploader(
+        "**Or upload audio files here:**",
         type=["wav", "flac", "mp3", "ogg", "avi"],
         accept_multiple_files=True,
     )
-    st.write("**OR**")
-    record = audiorecorder(
-        start_prompt="Start recording", stop_prompt="Stop recording", pause_prompt=""
-    )
-
-    record_fpath = "record.wav"
-    if len(record) > 0:
-        record.export(record_fpath, format="wav")
-        st.write(
-            f"Record frame rate: {record.frame_rate}Hz, record duration: {record.duration_seconds:.2f}s"
-        )
-        st.audio(record.export().read())  # type: ignore
 
     with st.expander("Model hyperparameters"):
         task = st.selectbox("Task embedding input", model.tasks, 0)
-        allow_rep_mode = st.selectbox(
-            "Allow repetition of words", ["stopwords", "all", "none"], 0
-        )
+        allow_rep_mode = st.selectbox("Allow repetition of words", ALLOW_REP_MODES, 0)
         beam_size: int = st.select_slider(  # type: ignore
             "Beam size",
-            list(range(1, 21)),
+            list(range(1, MAX_BEAM_SIZE + 1)),
             model.config.beam_size,
         )
-        min_pred_size: int = st.select_slider(  # type: ignore
-            "Minimal number of words",
-            list(range(1, 31)),
-            model.config.min_pred_size,
+        min_pred_size, max_pred_size = st.slider(
+            "Minimal and maximal number of words",
+            1,
+            MAX_PRED_SIZE,
+            (model.config.min_pred_size, model.config.max_pred_size),
         )
-        max_pred_size: int = st.select_slider(  # type: ignore
-            "Maximal number of words",
-            list(range(1, 31)),
-            model.config.max_pred_size,
+        threshold = st.select_slider(
+            "Tags threshold",
+            [(i / THRESHOLD_PRECISION) for i in range(THRESHOLD_PRECISION + 1)],
+            DEFAULT_THRESHOLD,
         )
 
         if allow_rep_mode == "all":
@@ -122,7 +159,6 @@ def main() -> None:
         elif allow_rep_mode == "stopwords":
             forbid_rep_mode = "content_words"
         else:
-            ALLOW_REP_MODES = ("all", "none", "stopwords")
             raise ValueError(
                 f"Unknown option {allow_rep_mode=}. (expected one of {ALLOW_REP_MODES})"
             )
@@ -134,23 +170,18 @@ def main() -> None:
             min_pred_size=min_pred_size,
             max_pred_size=max_pred_size,
             forbid_rep_mode=forbid_rep_mode,
+            threshold=threshold,
         )
 
-    if audios is not None and len(audios) > 0:
-        audio_fnames, cands = get_results(model, audios, generate_kwds)
-
-        for audio_fname, cand in zip(audio_fnames, cands):
-            st.success(f"**Output for {audio_fname}:**\n- {format_candidate(cand)}")
+    audios: dict[str, bytes] = {}
+    if audio_files is not None:
+        audios |= {audio.name: audio.getvalue() for audio in audio_files}
+    if record_data is not None:
+        audios |= {RECORD_AUDIO_FNAME: record_data}
 
-    if len(record) > 0:
-        outputs = model(
-            record_fpath,
-            **generate_kwds,
-        )
-        cand = outputs["cands"][0]
-        st.success(
-            f"**Output for {osp.basename(record_fpath)}:**\n- {format_candidate(cand)}"
-        )
+    if len(audios) > 0:
+        outputs = get_results(model, audios, generate_kwds)
+        show_results(outputs)
 
 
 if __name__ == "__main__":

requirements.txt

@@ -1,2 +1,2 @@
-conette~=0.2.0
-streamlit-audiorecorder
+conette~=0.2.2
+streamlit-audiorec~=0.1.3