msawant/audio_intelligence
0
1import re2 3import requests4import time5from scipy.io.wavfile import write6import io7import plotly.express as px8 9 10upload_endpoint = "https://api.assemblyai.com/v2/upload"11transcript_endpoint = "https://api.assemblyai.com/v2/transcript"12 13# Colors for sentiment analysis highlighting14green = "background-color: #159609"15red = "background-color: #cc0c0c"16 17# Converts Gradio checkboxes to AssemlbyAI header arguments18transcription_options_headers = {19 'Automatic Language Detection': 'language_detection',20 'Speaker Labels': 'speaker_labels',21 'Filter Profanity': 'filter_profanity',22}23 24# Converts Gradio checkboxes to AssemblyAI header arguments25audio_intelligence_headers = {26 'Summarization': 'auto_chapters',27 'Auto Highlights': 'auto_highlights',28 'Topic Detection': 'iab_categories',29 'Entity Detection': 'entity_detection',30 'Sentiment Analysis': 'sentiment_analysis',31 'PII Redaction': 'redact_pii',32 'Content Moderation': 'content_safety',33}34 35# Converts selected language in Gradio to language code for AssemblyAI header argument36language_headers = {37 'Global English': 'en',38 'US English': 'en_us',39 'British English': 'en_uk',40 'Australian English': 'en_au',41 'Spanish': 'es',42 'French': 'fr',43 'German': 'de',44 'Italian': 'it',45 'Portuguese': 'pt',46 'Dutch': 'nl',47 'Hindi': 'hi',48 'Japanese': 'jp',49}50 51 52def make_header(api_key):53 return {54 'authorization': api_key,55 'content-type': 'application/json'56 }57 58 59def _read_file(filename, chunk_size=5242880):60 """Helper for `upload_file()`"""61 with open(filename, "rb") as f:62 while True:63 data = f.read(chunk_size)64 if not data:65 break66 yield data67 68 69def _read_array(audio, chunk_size=5242880):70 """Like _read_file but for array - creates temporary unsaved "file" from sample rate and audio np.array"""71 sr, aud = audio72 73 # Create temporary "file" and write data to it74 bytes_wav = bytes()75 temp_file = io.BytesIO(bytes_wav)76 write(temp_file, sr, aud)77 78 while True:79 data = temp_file.read(chunk_size)80 if not data:81 break82 yield data83 84 85def upload_file(audio_file, header, is_file=True):86 """Uploads a file to AssemblyAI for analysis"""87 upload_response = requests.post(88 upload_endpoint,89 headers=header,90 data=_read_file(audio_file) if is_file else _read_array(audio_file)91 )92 if upload_response.status_code != 200:93 upload_response.raise_for_status()94 # Returns {'upload_url': <URL>}95 return upload_response.json()96 97 98def request_transcript(upload_url, header, **kwargs):99 """Request a transcript/audio analysis from AssemblyAI"""100 101 # If input is a dict returned from `upload_file` rather than a raw upload_url string102 if type(upload_url) is dict:103 upload_url = upload_url['upload_url']104 105 # Create request106 transcript_request = {107 'audio_url': upload_url,108 **kwargs109 }110 111 # POST request112 transcript_response = requests.post(113 transcript_endpoint,114 json=transcript_request,115 headers=header116 )117 118 return transcript_response.json()119 120 121def make_polling_endpoint(transcript_id):122 """Create a polling endpoint from a transcript ID to check on the status of the transcript"""123 # If upload response is input rather than raw upload_url string124 if type(transcript_id) is dict:125 transcript_id = transcript_id['id']126 127 polling_endpoint = "https://api.assemblyai.com/v2/transcript/" + transcript_id128 return polling_endpoint129 130 131def wait_for_completion(polling_endpoint, header):132 """Given a polling endpoint, waits for the transcription/audio analysis to complete"""133 while True:134 polling_response = requests.get(polling_endpoint, headers=header)135 polling_response = polling_response.json()136 137 if polling_response['status'] == 'completed':138 break139 elif polling_response['status'] == 'error':140 raise Exception(f"Error: {polling_response['error']}")141 142 time.sleep(5)143 144 145def make_true_dict(transcription_options, audio_intelligence_selector):146 """Given transcription / audio intelligence Gradio options, create a dictionary to be used in AssemblyAI request"""147 # Convert Gradio checkbox names to AssemblyAI API keys148 aai_tran_keys = [transcription_options_headers[elt] for elt in transcription_options]149 aai_audint_keys = [audio_intelligence_headers[elt] for elt in audio_intelligence_selector]150 151 # For each checked box, set it to true in the JSON used POST request to AssemblyAI152 aai_tran_dict = {key: 'true' for key in aai_tran_keys}153 aai_audint_dict = {key: 'true' for key in aai_audint_keys}154 155 return {**aai_tran_dict, **aai_audint_dict}156 157 158def make_final_json(true_dict, language):159 """Takes in output of `make_true_dict()` and adds all required other key-value pairs"""160 # If automatic language detection selected but no language specified, default to US english161 if 'language_detection' not in true_dict:162 if language is None:163 language = "US English"164 true_dict = {**true_dict, 'language_code': language_headers[language]}165 # If PII Redaction is enabled, add default redaction policies166 if 'redact_pii' in true_dict:167 true_dict = {**true_dict, 'redact_pii_policies': ['drug', 'injury', 'person_name', 'money_amount']}168 return true_dict, language169 170 171def _split_on_capital(string):172 """Adds spaces between capitalized words of a string via regex. 'HereAreSomeWords' -> 'Here Are Some Words'"""173 return ' '.join(re.findall("[A-Z][^A-Z]*", string))174 175 176def _make_tree(c, ukey=''):177 '''178 Given a list whose elements are nested topic lists, generates a JSON-esque dictionary tree of topics and179 subtopics180 181 E.g. the input182 183 [184 185 ['Education', 'CollegeEducation', 'PostgraduateEducation'],186 187 ['Education', 'CollegeEducation', 'UndergraduateEducation']188 189 ]190 191 Would output a dictionary corresponding to a tree with two leaves, 'UndergraduateEducation' and192 'PostgraduateEducation', which fall under a node 'CollegeEducation' which in turn falls under the node 'Education'193 194 :param c: List of topics195 :param ukey: "Upper key". For recursion - name of upper level key whose value (list) is being recursed on196 :return: Dictionary that defines a tree structure197 '''198 199 # Create empty dict for current sublist200 d = dict()201 202 # If leaf, return None203 if c is None and ukey is None:204 return None205 elif c is None:206 return {None: None}207 else:208 # For each elt of the input (itself a list),209 for n, i in enumerate(c):210 # For topics with sublist e.g. if ['NewsAndPolitics' 'Politics'] and211 # ['NewsAndPolitics' 'Politics', 'Elections'] are both in list - need way to signify politics itself212 # included213 if i is None:214 d[None] = None215 # If next subtopic not in dict, add it. If the remaining list empty, make value None216 elif i[0] not in d.keys():217 topic = i.pop(0)218 d[topic] = None if i == [] else [i]219 # If subtopic already in dict220 else:221 # If the value for this subtopic is only None (i.e. subject itself is a leaf), then append sublist222 if d[i[0]] is None:223 d[i[0]] = [None, i[1:]]224 # If value for this subtopic is a list itself, then append the remaining list225 else:226 d[i[0]].append(i[1:])227 # Recurse on remaining leaves228 for key in d:229 d[key] = _make_tree(d[key], key)230 return d231 232 233def _make_html_tree(dic, level=0, HTML=''):234 """Generates an HTML tree from an output of _make_tree"""235 HTML += "<ul>"236 for key in dic:237 # Add the topic to HTML, specifying the current level and whether it is a topic238 if type(dic[key]) == dict:239 HTML += "<li>"240 if None in dic[key].keys():241 del dic[key][None]242 HTML += f'<p class="topic-L{level} istopic">{_split_on_capital(key)}</p>'243 else:244 HTML += f'<p class="topic-L{level}">{_split_on_capital(key)}</p>'245 HTML += "</li>"246 247 HTML = _make_html_tree(dic[key], level=level + 1, HTML=HTML)248 else:249 HTML += "<li>"250 HTML += f'<p class="topic-L{level} istopic">{_split_on_capital(key)}</p>'251 HTML += "</li>"252 HTML += "</ul>"253 return HTML254 255 256def _make_html_body(dic):257 """Makes an HTML body from an output of _make_tree"""258 HTML = '<body>'259 HTML += _make_html_tree(dic)260 HTML += "</body>"261 return HTML262 263 264def _make_html(dic):265 """Makes a full HTML document from an output of _make_tree using styles.css styling"""266 HTML = '<!DOCTYPE html>' \267 '<html>' \268 '<head>' \269 '<title>Another simple example</title>' \270 '<link rel="stylesheet" type="text/css" href="styles.css"/>' \271 '</head>'272 HTML += _make_html_body(dic)273 HTML += "</html>"274 return HTML275 276 277# make_html_from_topics(j['iab_categories_result']['summary'])278def make_html_from_topics(dic, threshold=0.0):279 """Given a topics dictionary from AAI Topic Detection API, generates appropriate corresponding structured HTML.280 Input is `response.json()['iab_categories_result']['summary']` from GET request on AssemblyAI `v2/transcript`281 endpoint."""282 # Potentially filter some items out283 cats = [k for k, v in dic.items() if float(v) >= threshold]284 285 # Sort remaining topics286 cats.sort()287 288 # Split items into lists289 cats = [i.split(">") for i in cats]290 291 # Make topic tree292 tree = _make_tree(cats)293 294 # Return formatted HTML295 return _make_html(tree)296 297 298def make_paras_string(transc_id, header):299 """ Makes a string by concatenating paragraphs newlines in between. Input is response.json()['paragraphs'] from300 from AssemblyAI paragraphs endpoint """301 endpoint = transcript_endpoint + "/" + transc_id + "/paragraphs"302 paras = requests.get(endpoint, headers=header).json()['paragraphs']303 paras = '\n\n'.join(i['text'] for i in paras)304 return paras305 306 307def create_highlighted_list(paragraphs_string, highlights_result, rank=0):308 """Outputs auto highlights information in appropriate format for `gr.HighlightedText()`. `highlights_result` is309 response.json()['auto_highlights_result]['results'] where response from GET request on AssemblyAI v2/transcript310 endpoint"""311 # Max and min opacities to highlight to312 MAX_HIGHLIGHT = 1 # Max allowed = 1313 MIN_HIGHLIGHT = 0.25 # Min allowed = 0314 315 # Filter list for everything above the input rank316 highlights_result = [i for i in highlights_result if i['rank'] >= rank]317 318 # Get max/min ranks and find scale/shift we'll need so ranks are mapped to [MIN_HIGHLIGHT, MAX_HIGHLIGHT]319 max_rank = max([i['rank'] for i in highlights_result])320 min_rank = min([i['rank'] for i in highlights_result])321 scale = (MAX_HIGHLIGHT - MIN_HIGHLIGHT) / (max_rank - min_rank)322 shift = (MAX_HIGHLIGHT - max_rank * scale)323 324 # Isolate only highlight text and rank325 highlights_result = [(i['text'], i['rank']) for i in highlights_result]326 327 entities = []328 for highlight, rank in highlights_result:329 # For each highlight, find all starting character instances330 starts = [c.start() for c in re.finditer(highlight, paragraphs_string)]331 # Create list of locations for this highlight with entity value (highlight opacity) scaled properly332 e = [{"entity": rank * scale + shift,333 "start": start,334 "end": start + len(highlight)}335 for start in starts]336 entities += e337 338 # Create dictionary339 highlight_dict = {"text": paragraphs_string, "entities": entities}340 341 # Sort entities by start char. A bug in Gradio requires this342 highlight_dict['entities'] = sorted(highlight_dict['entities'], key=lambda x: x['start'])343 344 return highlight_dict345 346 347def make_summary(chapters):348 """Makes HTML for "Summary" `gr.Tab()` tab. Input is `response.json()['chapters']` where response is from GET349 request to AssemblyAI's v2/transcript endpoint"""350 html = "<div>"351 for chapter in chapters:352 html += "<details>" \353 f"<summary><b>{chapter['headline']}</b></summary>" \354 f"{chapter['summary']}" \355 "</details>"356 html += "</div>"357 return html358 359 360def to_hex(num, max_opacity=128):361 """Converts a confidence value in the range [0, 1] to a hex value"""362 return hex(int(max_opacity * num))[2:]363 364 365def make_sentiment_output(sentiment_analysis_results):366 """Makes HTML output of sentiment analysis info for display with `gr.HTML()`. Input is367 `response.json()['sentiment_analysis_results']` from GET request on AssemblyAI v2/transcript."""368 p = '<p>'369 for sentiment in sentiment_analysis_results:370 if sentiment['sentiment'] == 'POSITIVE':371 p += f'<mark style="{green + to_hex(sentiment["confidence"])}">' + sentiment['text'] + '</mark> '372 elif sentiment['sentiment'] == "NEGATIVE":373 p += f'<mark style="{red + to_hex(sentiment["confidence"])}">' + sentiment['text'] + '</mark> '374 else:375 p += sentiment['text'] + ' '376 p += "</p>"377 return p378 379 380def make_entity_dict(entities, t, offset=40):381 """Creates dictionary that will be used to generate HTML for Entity Detection `gr.Tab()` tab.382 Inputs are response.json()['entities'] and response.json()['text'] for response of GET request383 on AssemblyAI v2/transcript endpoint"""384 len_text = len(t)385 386 d = {}387 for entity in entities:388 # Find entity in the text389 s = t.find(entity['text'])390 if s == -1:391 p = None392 else:393 len_entity = len(entity['text'])394 # Get entity context (colloquial sense)395 p = t[max(0, s - offset):min(s + len_entity + offset, len_text)]396 # Make sure start and end with a full word397 p = '... ' + ' '.join(p.split(' ')[1:-1]) + ' ...'398 # Add to dict399 label = ' '.join(entity['entity_type'].split('_')).title()400 if label in d:401 d[label] += [[p, entity['text']]]402 else:403 d[label] = [[p, entity['text']]]404 405 return d406 407 408def make_entity_html(d, highlight_color="#FFFF0080"):409 """Input is output of `make_entity_dict`. Creates HTML for Entity Detection info"""410 h = "<ul>"411 for i in d:412 h += f"""<li style="color: #6b2bd6; font-size: 20px;">{i}"""413 h += "<ul>"414 for sent, ent in d[i]:415 if sent is None:416 h += f"""<li style="color: black; font-size: 16px;">[REDACTED]</li>"""417 else:418 h += f"""<li style="color: black; font-size: 16px;">{sent.replace(ent, f'<mark style="background-color: {highlight_color}">{ent}</mark>')}</li>"""419 h += '</ul>'420 h += '</li>'421 h += "</ul>"422 return h423 424 425def make_content_safety_fig(cont_safety_summary):426 """Creates content safety figure from response.json()['content_safety_labels']['summary'] from GET request on427 AssemblyAI v2/transcript endpoint"""428 # Create dictionary as demanded by plotly429 d = {'label': [], 'severity': [], 'color': []}430 431 # For each sentitive topic, add the (formatted) name, severity, and plot color432 for key in cont_safety_summary:433 d['label'] += [' '.join(key.split('_')).title()]434 d['severity'] += [cont_safety_summary[key]]435 d['color'] += ['rgba(107, 43, 214, 1)']436 437 # Create the figure (n.b. repetitive color info but was running into plotly bugs)438 content_fig = px.bar(d, x='severity', y='label', color='color', color_discrete_map={439 'Crime Violence': 'rgba(107, 43, 214, 0.1)',440 'Alcohol': 'rgba(107, 43, 214, 0.1)',441 'Accidents': 'rgba(107, 43, 214, 0.1)'})442 443 # Update the content figure plot444 content_fig.update_layout({'plot_bgcolor': 'rgba(107, 43, 214, 0.1)'})445 446 # Scales axes appropriately447 content_fig.update_xaxes(range=[0, 1])448 return content_fig