Highlight

app.py

@@ -1,6 +1,10 @@
 import gradio as gr
 from transformers import pipeline
 import numpy as np
+import pandas as pd
+import re
+from collections import Counter
+from functools import reduce
 
 transcriber = pipeline(
     "automatic-speech-recognition",
@@ -10,10 +14,8 @@ transcriber = pipeline(
 
 
 def transcribe_live(state, words_list, new_chunk):
-    print(f"state: {state}")
-
     try:
-        words_to_check_for = [word.strip() for word in words_list.split(",")]
+        words_to_check_for = [word.strip().lower() for word in words_list.split(",")]
     except:
         gr.Warning("Please enter a valid list of words to check for")
         words_to_check_for = []
@@ -50,24 +52,49 @@ def transcribe_live(state, words_list, new_chunk):
         print(f"Transcription failed. Error: {e}")
         return state, previous_counts_of_words, previous_transcription
 
-    print(f"new transcription: {new_transcription}")
     full_transcription_text = new_transcription["text"]
 
     full_transcription_text_lower = full_transcription_text.lower()
 
-    new_counts_of_words = {
-        word: full_transcription_text_lower.count(word) for word in words_to_check_for
+    # Use re to find all the words in the transcription, and their start and end indices
+    matches: list[re.Match] = list(
+        re.finditer(
+            r"\b(" + "|".join(words_to_check_for) + r")\b",
+            full_transcription_text_lower,
+        )
+    )
+
+    counter = Counter(
+        match.group(0) for match in matches if match.group(0) in words_to_check_for
+    )
+
+    new_counts_of_words = {word: counter.get(word, 0) for word in words_to_check_for}
+
+    new_highlighted_transcription = {
+        "text": full_transcription_text,
+        "entities": [
+            {
+                "entity": "FILLER",
+                "start": match.start(),
+                "end": match.end(),
+            }
+            for match in matches
+        ],
     }
 
     new_state = {
         "stream": stream,
         "full_transcription": full_transcription_text,
         "counts_of_words": new_counts_of_words,
+        "highlighted_transcription": new_highlighted_transcription,
     }
 
-    print(f"new state: {new_state}")
-
-    return new_state, new_counts_of_words, full_transcription_text
+    return (
+        new_state,
+        new_counts_of_words,
+        full_transcription_text,
+        new_highlighted_transcription,
+    )
 
 
 with gr.Blocks() as demo:
@@ -82,14 +109,24 @@ with gr.Blocks() as demo:
     recording = gr.Audio(streaming=True, label="Recording")
 
     word_counts = gr.JSON(label="Filler words count", value={})
-    transcription = gr.Textbox(label="Transcription", value="")
+    # word_counts = gr.BarPlot(label="Filler words count", value={})
+    transcription = gr.Textbox(label="Transcription", value="", visible=False)
+
+    highlighted_transcription = gr.HighlightedText(
+        label="Transcription",
+        value={
+            "text": "",
+            "entities": [],
+        },
+        color_map={"FILLER": "red"},
+    )
 
     recording.stream(
         transcribe_live,
         inputs=[state, filler_words, recording],
-        outputs=[state, word_counts, transcription],
+        outputs=[state, word_counts, transcription, highlighted_transcription],
         stream_every=5,
-        time_limit=60,
+        time_limit=-1,
     )
 
 demo.launch(show_error=True)