Update app.py

app.py

@@ -34,6 +34,11 @@ emotion_pipeline = pipeline(
     model="bhadresh-savani/distilbert-base-uncased-emotion",
     device=device,
 )
+summarization_pipeline = pipeline(
+    "summarization",
+    model="knkarthick/MEETING_SUMMARY",
+    device=device
+)
 
 EXAMPLES = [["Customer_Support_Call.wav"]]
 
@@ -47,6 +52,30 @@ speech_to_text = partial(
     whisper_device=whisper_device
     )
 
+def summarize(diarized, check, summarization_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)]
+    check is a list of speaker ids whose speech will get summarized
+    """
+
+    if len(check) == 0:
+        return ""
+
+    text = ""
+    for d in diarized:
+        if len(check) == 2 and d[1] is not None:
+            text += f"\n{d[1]}: {d[0]}"
+        elif d[1] in check:
+            text += f"\n{d[0]}"
+
+    # inner function cached because outer function cannot be cached
+    @functools.lru_cache(maxsize=128)
+    def call_summarize_api(text):
+        return summarization_pipeline(text)[0]["summary_text"]
+
+    return call_summarize_api(text)
+
 def sentiment(diarized, emotion_pipeline):
     """
     diarized: a list of tuples. Each tuple has a string to be displayed and a label for highlighting.
@@ -57,22 +86,12 @@ def sentiment(diarized, emotion_pipeline):
     """
 
     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)
@@ -87,18 +106,8 @@ def sentiment(diarized, emotion_pipeline):
                         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
+    return customer_sentiments
 
 with gr.Blocks() as demo:
 
@@ -109,6 +118,11 @@ with gr.Blocks() as demo:
 
             gr.Markdown("**Call Transcript:**")
             diarized = gr.HighlightedText(label="Call Transcript")
+            gr.Markdown("Choose speaker to summarize:")
+            check = gr.CheckboxGroup(
+                choices=["Customer", "Support"], show_label=False, type="value"
+            )
+            summary = gr.Textbox(lines=4)
             sentiment_btn = gr.Button("Get Customer Sentiment")
             analyzed = gr.HighlightedText(color_map=color_map)
             plot = gr.Plot(label="Sentiment over time", type="plotly")
@@ -128,6 +142,8 @@ with gr.Blocks() as demo:
         inputs=audio,
         outputs=diarized,
     )
+    # when summarize checkboxes are changed, create summary
+    check.change(fn=partial(summarize, summarization_pipeline=summarization_pipeline), inputs=[diarized, check], outputs=summary)
 
     # when sentiment button clicked, display highlighted text and plot
     sentiment_btn.click(fn=partial(sentiment, emotion_pipeline=emotion_pipeline), inputs=diarized, outputs=[analyzed, plot])