import os import re import functools from functools import partial import requests import pandas as pd import plotly.express as px import torch import gradio as gr from transformers import pipeline, Wav2Vec2ProcessorWithLM from pyannote.audio import Pipeline import whisperx from utils import split_into_sentences, summarize, sentiment, color_map from utils import speech_to_text as stt os.environ["TOKENIZERS_PARALLELISM"] = "false" device = 0 if torch.cuda.is_available() else -1 # Audio components whisper_device = "cuda" if torch.cuda.is_available() else "cpu" whisper = whisperx.load_model("tiny.en", whisper_device) alignment_model, metadata = whisperx.load_align_model(language_code="en", device=whisper_device) speaker_segmentation = Pipeline.from_pretrained("pyannote/speaker-diarization@2.1", use_auth_token=os.environ['ENO_TOKEN']) # Text components emotion_pipeline = pipeline( "text-classification", model="bhadresh-savani/distilbert-base-uncased-emotion", device=device, ) summarization_pipeline = pipeline( "summarization", model="knkarthick/MEETING_SUMMARY", device=device ) EXAMPLES = [["Customer_Support_Call.wav"]] speech_to_text = partial( stt, speaker_segmentation=speaker_segmentation, whisper=whisper, alignment_model=alignment_model, metadata=metadata, whisper_device=whisper_device ) with gr.Blocks() as demo: with gr.Row(): with gr.Column(): audio = gr.Audio(label="Audio file", type="filepath") btn = gr.Button("Transcribe and Diarize") 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") with gr.Column(): gr.Markdown("## Example Files") gr.Examples( examples=EXAMPLES, inputs=[audio], outputs=[diarized], fn=speech_to_text, cache_examples=True ) # when example button is clicked, convert audio file to text and diarize btn.click( fn=speech_to_text, 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]) demo.launch(debug=1)