Spaces:
Running
Running
| import itertools | |
| import json | |
| import re | |
| from functools import partial | |
| from pathlib import Path | |
| import pandas as pd | |
| import requests | |
| import streamlit as st | |
| import webvtt | |
| from transformers import AutoTokenizer | |
| from generate_text_api import TextGenerator | |
| from model_inferences.utils.chunking import Truncater | |
| from model_inferences.utils.files import get_captions_from_vtt, get_transcript | |
| USE_PARAGRAPHING_MODEL = True | |
| def get_sublist_by_flattened_index(A, i): | |
| current_index = 0 | |
| for sublist in A: | |
| sublist_length = len(sublist) | |
| if current_index <= i < current_index + sublist_length: | |
| return sublist, A.index(sublist) | |
| current_index += sublist_length | |
| return None, None | |
| import requests | |
| def get_talk_metadata(video_id): | |
| url = "https://www.ted.com/graphql" | |
| headers = { | |
| "Content-Type": "application/json", | |
| "Accept": "application/json", | |
| "x-operation-name": "Transcript", # Replace with the actual operation name | |
| } | |
| data = { | |
| "query": """ | |
| query GetTalk($videoId: ID!) { | |
| video(id: $videoId) { | |
| title, | |
| presenterDisplayName, | |
| nativeDownloads {medium} | |
| } | |
| } | |
| """, | |
| "variables": { | |
| "videoId": video_id, # Corrected key to "videoId" | |
| }, | |
| } | |
| response = requests.post(url, json=data, headers=headers) | |
| if response.status_code == 200: | |
| result = response.json() | |
| return result | |
| else: | |
| print(f"Error: {response.status_code}, {response.text}") | |
| class OfflineTextSegmenterClient: | |
| def __init__(self, host_url): | |
| self.host_url = host_url.rstrip("/") + "/segment" | |
| def segment(self, text, captions=None, generate_titles=False, threshold=0.4): | |
| payload = { | |
| 'text': text, | |
| 'captions': captions, | |
| 'generate_titles': generate_titles, | |
| "prefix_titles": True, | |
| "threshold": threshold, | |
| } | |
| headers = { | |
| 'Content-Type': 'application/json' | |
| } | |
| response = requests.post(self.host_url, data=json.dumps(payload), headers=headers).json() | |
| #segments = response["annotated_segments"] if "annotated_segments" in response else response["segments"] | |
| return {'segments':response["segments"], 'titles': response["titles"], 'sentences': response["sentences"]} | |
| class Toc: | |
| def __init__(self): | |
| self._items = [] | |
| self._placeholder = None | |
| def title(self, text): | |
| self._markdown(text, "h1") | |
| def header(self, text): | |
| self._markdown(text, "h2", " " * 2) | |
| def subheader(self, text): | |
| self._markdown(text, "h3", " " * 4) | |
| def placeholder(self, sidebar=False): | |
| self._placeholder = st.sidebar.empty() if sidebar else st.empty() | |
| def generate(self): | |
| if self._placeholder: | |
| self._placeholder.markdown("\n".join(self._items), unsafe_allow_html=True) | |
| def _markdown(self, text, level, space=""): | |
| key = re.sub(r'[^\w-]', '', text.replace(" ", "-").replace("'", "-").lower()) | |
| st.markdown(f"<{level} id='{key}'>{text}</{level}>", unsafe_allow_html=True) | |
| self._items.append(f"{space}* <a href='#{key}'>{text}</a>") | |
| # custom_css = "<style type='text/css'>" + Path('style.css').read_text() + "</style>" | |
| # st.write(custom_css, unsafe_allow_html=True) | |
| def concat_prompt(prompt_text, text, model_name): | |
| if 'flan' in model_name: | |
| input_ = prompt_text + "\n\n" + text | |
| elif 'galactica' in model_name: | |
| input_ = text + "\n\n" + prompt_text | |
| return input_ | |
| endpoint = "http://hiaisc.isl.iar.kit.edu/summarize" | |
| ENDPOINTS = {"http://hiaisc.isl.iar.kit.edu/summarize": "meta-llama/Llama-2-13b-chat-hf",} | |
| client = OfflineTextSegmenterClient("http://hiaisc.isl.iar.kit.edu/chapterize") | |
| if USE_PARAGRAPHING_MODEL: | |
| paragrapher = OfflineTextSegmenterClient("http://hiaisc.isl.iar.kit.edu/paragraph") | |
| summarizer = TextGenerator(endpoint) | |
| tokenizer = AutoTokenizer.from_pretrained(ENDPOINTS[endpoint], use_fast=False) | |
| # TLDR PROMPT | |
| SYSTEM_PROMPT = "You are an assistant who replies with a summary to every message." | |
| TLDR_PROMPT_TEMPLATE = """<s>[INST] <<SYS>> | |
| {system_prompt} | |
| <</SYS>> | |
| {user_message} [/INST] Sure! Here is a summary of the research presentation in a single, short sentence:""" | |
| TLDR_USER_PROMPT = "Summarize the following research presentation in a single, short sentence:\n\n{input}" | |
| TLDR_PROMPT = TLDR_PROMPT_TEMPLATE.format(system_prompt=SYSTEM_PROMPT, user_message=TLDR_USER_PROMPT) | |
| TLDR_PROMPT_LENGTH = tokenizer(TLDR_PROMPT, return_tensors="pt")["input_ids"].size(1) | |
| BP_PROMPT_TEMPLATE = """<s>[INST] <<SYS>> | |
| {system_prompt} | |
| <</SYS>> | |
| {user_message} [/INST] Sure! Here is a summary of the research presentation using three bullet points:\n\n\u2022""" | |
| BP_USER_PROMPT = "Summarize the following research presentation using three bullet points:\n\n{input}" | |
| BP_PROMPT = BP_PROMPT_TEMPLATE.format(system_prompt=SYSTEM_PROMPT, user_message=TLDR_USER_PROMPT) | |
| BP_PROMPT_LENGTH = tokenizer(BP_PROMPT, return_tensors="pt")["input_ids"].size(1) | |
| CONTEXT_LENGTH = 3072 | |
| MAX_SUMMARY_LENGTH = 1024 | |
| TLDR_MAX_INPUT_LENGTH = CONTEXT_LENGTH - MAX_SUMMARY_LENGTH - TLDR_PROMPT_LENGTH - 1 | |
| BP_MAX_INPUT_LENGTH = CONTEXT_LENGTH - MAX_SUMMARY_LENGTH - BP_PROMPT_LENGTH - 1 | |
| text_generator = TextGenerator(endpoint) | |
| temperature = 0.7 | |
| import re | |
| def replace_newlines(text): | |
| updated_text = re.sub(r'\n+', r'\n\n', text) | |
| return updated_text | |
| def generate_summary(summarizer, generated_text_box, input_, prompt, max_input_length, prefix=""): | |
| all_generated_text = prefix | |
| truncater = Truncater(tokenizer, max_length=max_input_length) | |
| input_ = truncater(input_) | |
| input_ = prompt.format(input=input_) | |
| for generated_text in summarizer.generate_text_stream(input_, max_new_tokens=MAX_SUMMARY_LENGTH, do_sample=True, temperature=temperature): | |
| all_generated_text += replace_newlines(generated_text) | |
| generated_text_box.info(all_generated_text) | |
| print(all_generated_text) | |
| return all_generated_text.strip() | |
| st.header("Demo: Intelligent Recap") | |
| if not hasattr(st, 'global_state'): | |
| st.global_state = {'NIPS 2021 Talks': None, 'TED Talks': None} | |
| # NIPS 2021 Talks | |
| transcript_files = itertools.islice(Path("demo_data/nips-2021/").rglob("transcript_whisper_large-v2.vtt"), 15) | |
| # get titles from metadata.json | |
| transcripts_map = {} | |
| for transcript_file in transcript_files: | |
| base_path = transcript_file.parent | |
| metadata = base_path / "metadata.json" | |
| txt_file = base_path / "transcript_whisper_large-v2.txt" | |
| with open(metadata) as f: | |
| metadata = json.load(f) | |
| title = metadata["title"] | |
| transcript = get_transcript(txt_file) | |
| captions = get_captions_from_vtt(transcript_file) | |
| transcripts_map[title] = {"transcript": transcript, "captions": captions, "video": base_path / "video.mp4"} | |
| st.global_state['NIPS 2021 Talks'] = transcripts_map | |
| data = pd.read_json("demo_data/ted_talks.json") | |
| video_ids = data.talk_id.tolist() | |
| transcripts = data.text.apply(lambda x: " ".join(x)).tolist() | |
| transcripts_map = {} | |
| for video_id, transcript in zip(video_ids, transcripts): | |
| metadata = get_talk_metadata(video_id) | |
| title = metadata["data"]["video"]["title"] | |
| presenter = metadata["data"]["video"]["presenterDisplayName"] | |
| print(metadata["data"]) | |
| if metadata["data"]["video"]["nativeDownloads"] is None: | |
| continue | |
| video_url = metadata["data"]["video"]["nativeDownloads"]["medium"] | |
| transcripts_map[title] = {"transcript": transcript, "video": video_url, "presenter": presenter} | |
| st.global_state['TED Talks'] = transcripts_map | |
| def get_lecture_id(path): | |
| return int(path.parts[-2].split('-')[1]) | |
| transcript_files = Path("demo_data/lectures/").rglob("English.vtt") | |
| sorted_path_list = sorted(transcript_files, key=get_lecture_id) | |
| transcripts_map = {} | |
| for transcript_file in sorted_path_list: | |
| base_path = transcript_file.parent | |
| lecture_id = base_path.parts[-1] | |
| transcript = " ".join([c["text"].strip() for c in get_captions_from_vtt(transcript_file)]).replace("\n", " ") | |
| video_path = Path(base_path, "video.mp4") | |
| transcripts_map["Machine Translation: " + lecture_id] = {"transcript": transcript, "video": video_path} | |
| st.global_state['KIT Lectures'] = transcripts_map | |
| type_of_document = st.selectbox('What kind of document do you want to test it on?', list(st.global_state.keys())) | |
| transcripts_map = st.global_state[type_of_document] | |
| selected_talk = st.selectbox("Choose a document...", list(transcripts_map.keys())) | |
| st.video(str(transcripts_map[selected_talk]['video']), format="video/mp4", start_time=0) | |
| input_text = st.text_area("Transcript", value=transcripts_map[selected_talk]['transcript'], height=300) | |
| toc = Toc() | |
| summarization_todos = [] | |
| with st.expander("Adjust Thresholds"): | |
| threshold = st.slider('Chapter Segmentation Threshold', 0.00, 1.00, value=0.4, step=0.05) | |
| paragraphing_threshold = st.slider('Paragraphing Threshold', 0.00, 1.00, value=0.5, step=0.05) | |
| if st.button("Process Transcript"): | |
| with st.sidebar: | |
| st.header("Table of Contents") | |
| toc.placeholder() | |
| st.header(selected_talk, divider='rainbow') | |
| # if 'presenter' in transcripts_map[selected_talk]: | |
| # st.markdown(f"### *by **{transcripts_map[selected_talk]['presenter']}***") | |
| captions = transcripts_map[selected_talk]['captions'] if 'captions' in transcripts_map[selected_talk] else None | |
| result = client.segment(input_text, captions, generate_titles=True, threshold=threshold) | |
| if USE_PARAGRAPHING_MODEL: | |
| presult = paragrapher.segment(input_text, captions, generate_titles=False, threshold=paragraphing_threshold) | |
| paragraphs = presult['segments'] | |
| segments, titles, sentences = result['segments'], result['titles'], result['sentences'] | |
| if USE_PARAGRAPHING_MODEL: | |
| prev_chapter_idx = 0 | |
| prev_paragraph_idx = 0 | |
| segment = [] | |
| for i, sentence in enumerate(sentences): | |
| chapter, chapter_idx = get_sublist_by_flattened_index(segments, i) | |
| paragraph, paragraph_idx = get_sublist_by_flattened_index(paragraphs, i) | |
| if (chapter_idx != prev_chapter_idx and paragraph_idx == prev_paragraph_idx) or (paragraph_idx != prev_paragraph_idx and chapter_idx != prev_chapter_idx): | |
| print("Chapter / Chapter & Paragraph") | |
| segment_text = " ".join(segment) | |
| toc.subheader(titles[prev_chapter_idx]) | |
| if len(segment_text) > 1200: | |
| generated_text_box = st.info("") | |
| summarization_todos.append(partial(generate_summary, summarizer, generated_text_box, segment_text, BP_PROMPT, BP_MAX_INPUT_LENGTH, prefix="\u2022")) | |
| elif len(segment_text) > 450: | |
| generated_text_box = st.info("") | |
| summarization_todos.append(partial(generate_summary, summarizer, generated_text_box, segment_text, TLDR_PROMPT, TLDR_MAX_INPUT_LENGTH)) | |
| st.write(segment_text) | |
| segment = [] | |
| elif paragraph_idx != prev_paragraph_idx and chapter_idx == prev_chapter_idx: | |
| print("Paragraph") | |
| segment.append("\n\n") | |
| segment.append(sentence) | |
| prev_chapter_idx = chapter_idx | |
| prev_paragraph_idx = paragraph_idx | |
| segment_text = " ".join(segment) | |
| toc.subheader(titles[prev_chapter_idx]) | |
| if len(segment_text) > 1200: | |
| generated_text_box = st.info("") | |
| summarization_todos.append(partial(generate_summary, summarizer, generated_text_box, segment_text, BP_PROMPT, BP_MAX_INPUT_LENGTH, prefix="\u2022")) | |
| elif len(segment_text) > 450: | |
| generated_text_box = st.info("") | |
| summarization_todos.append(partial(generate_summary, summarizer, generated_text_box, segment_text, TLDR_PROMPT, TLDR_MAX_INPUT_LENGTH)) | |
| st.write(segment_text) | |
| else: | |
| segments = [" ".join([sentence for sentence in segment]) for segment in segments] | |
| for title, segment in zip(titles, segments): | |
| toc.subheader(title) | |
| if len(segment) > 1200: | |
| generated_text_box = st.info("") | |
| summarization_todos.append(partial(generate_summary, summarizer, generated_text_box, segment, BP_PROMPT, BP_MAX_INPUT_LENGTH, prefix="\u2022")) | |
| elif len(segment) > 450: | |
| generated_text_box = st.info("") | |
| summarization_todos.append(partial(generate_summary, summarizer, generated_text_box, segment, TLDR_PROMPT, TLDR_MAX_INPUT_LENGTH)) | |
| st.write(segment) | |
| toc.generate() | |
| for summarization_todo in summarization_todos: | |
| summarization_todo() | |