Update app.py

app.py

@@ -13,7 +13,7 @@ from transformers import pipeline, Wav2Vec2ProcessorWithLM
 from pyannote.audio import Pipeline
 import whisperx
 
-from utils import split_into_sentences, sentiment, color_map
+from utils import split_into_sentences, create_fig, color_map
 from utils import speech_to_text as stt
 
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
@@ -47,6 +47,59 @@ speech_to_text = partial(
     whisper_device=whisper_device
     )
 
+def sentiment(diarized, emotion_pipeline):
+    """
+    diarized: a list of tuples. Each tuple has a string to be displayed and a label for highlighting.
+        The start/end times are not highlighted [(speaker text, speaker id), (start time/end time, None)]
+
+    This function gets the customer's sentiment and returns a list for highlighted text as well
+    as a plot of sentiment over time.
+    """
+
+    customer_sentiments = []
+    plot_sentences = []
+    to_plot = []
+
+    # used to set the x range of ticks on the plot
+    x_min = 100
+    x_max = 0
+
+    for i in range(0, len(diarized), 2):
+        speaker_speech, speaker_id = diarized[i]
+        times, _ = diarized[i + 1]
+
+        sentences = split_into_sentences(speaker_speech)
+        start_time, end_time = times[5:].split("-")
+        start_time, end_time = float(start_time), float(end_time)
+        interval_size = (end_time - start_time) / len(sentences)
+
+        if "Customer" in speaker_id:
+
+            outputs = emotion_pipeline(sentences)
+
+            for idx, (o, t) in enumerate(zip(outputs, sentences)):
+                sent = "neutral"
+                if o["score"] > thresholds[o["label"]]:
+                    customer_sentiments.append(
+                        (t + f"({round(idx*interval_size+start_time,1)} s)", o["label"])
+                    )
+                    if o["label"] in {"joy", "love", "surprise"}:
+                        sent = "positive"
+                    elif o["label"] in {"sadness", "anger", "fear"}:
+                        sent = "negative"
+                if sent != "neutral":
+                    to_plot.append((start_time + idx * interval_size, sent))
+                    plot_sentences.append(t)
+
+            if start_time < x_min:
+                x_min = start_time
+            if end_time > x_max:
+                x_max = end_time
+
+    fig = create_fig(x_min, x_max, to_plot, plot_sentences)
+
+    return customer_sentiments, fig
+
 with gr.Blocks() as demo:
 
     with gr.Row():
@@ -79,5 +132,4 @@ with gr.Blocks() as demo:
     # when sentiment button clicked, display highlighted text and plot
     sentiment_btn.click(fn=partial(sentiment, emotion_pipeline=emotion_pipeline), inputs=diarized, outputs=[analyzed, plot])
 
-
 demo.launch(debug=1)