nickmuchi's picture
Update functions.py
ebc87ab verified
import whisper
import os
import random
import openai
from openai import OpenAI
import yt_dlp
from pytube import YouTube, extract
import pandas as pd
import plotly_express as px
import nltk
import plotly.graph_objects as go
from optimum.onnxruntime import ORTModelForSequenceClassification
from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification, AutoModelForSeq2SeqLM, AutoModelForTokenClassification
from sentence_transformers import SentenceTransformer, CrossEncoder, util
import streamlit as st
import en_core_web_lg
import validators
import re
import itertools
import numpy as np
from bs4 import BeautifulSoup
import base64, time
from annotated_text import annotated_text
import pickle, math
import torch
from pydub import AudioSegment
from langchain.docstore.document import Document
from langchain.embeddings import HuggingFaceEmbeddings, HuggingFaceBgeEmbeddings, HuggingFaceInstructEmbeddings
from langchain.vectorstores import FAISS
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.chat_models import ChatOpenAI
from langchain.chains import QAGenerationChain
from langchain.callbacks import StreamlitCallbackHandler
from langchain.agents import OpenAIFunctionsAgent, AgentExecutor
from langchain.agents.agent_toolkits import create_retriever_tool
from langchain.agents.openai_functions_agent.agent_token_buffer_memory import (
AgentTokenBufferMemory,
)
from langchain.prompts import MessagesPlaceholder
from langchain.prompts.chat import (
ChatPromptTemplate,
SystemMessagePromptTemplate,
AIMessagePromptTemplate,
HumanMessagePromptTemplate,
)
from langchain.schema import (
AIMessage,
HumanMessage,
SystemMessage
)
from langchain.prompts import PromptTemplate
from langsmith import Client
client = Client()
openai_audio = OpenAI()
nltk.download('punkt')
from nltk import sent_tokenize
OPEN_AI_KEY = os.environ.get('OPEN_AI_KEY')
time_str = time.strftime("%d%m%Y-%H%M%S")
HTML_WRAPPER = """<div style="overflow-x: auto; border: 1px solid #e6e9ef; border-radius: 0.25rem; padding: 1rem;
margin-bottom: 2.5rem">{}</div> """
###################### Functions #######################################################################################
#load all required models and cache
@st.cache_resource
def load_models():
'''Load and cache all the models to be used'''
q_model = ORTModelForSequenceClassification.from_pretrained("nickmuchi/quantized-optimum-finbert-tone")
ner_model = AutoModelForTokenClassification.from_pretrained("xlm-roberta-large-finetuned-conll03-english")
q_tokenizer = AutoTokenizer.from_pretrained("nickmuchi/quantized-optimum-finbert-tone")
ner_tokenizer = AutoTokenizer.from_pretrained("xlm-roberta-large-finetuned-conll03-english")
sent_pipe = pipeline("text-classification",model=q_model, tokenizer=q_tokenizer)
sum_pipe = pipeline("summarization",model="philschmid/flan-t5-base-samsum",clean_up_tokenization_spaces=True)
ner_pipe = pipeline("ner", model=ner_model, tokenizer=ner_tokenizer, grouped_entities=True)
cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-12-v2') #cross-encoder/ms-marco-MiniLM-L-12-v2
sbert = SentenceTransformer('all-MiniLM-L6-v2')
return sent_pipe, sum_pipe, ner_pipe, cross_encoder, sbert
@st.cache_data
def load_asr_model(model_name):
'''Load the open source whisper model in cases where the API is not working'''
model = whisper.load_model(model_name)
return model
@st.cache_resource
def get_spacy():
nlp = en_core_web_lg.load()
return nlp
nlp = get_spacy()
sent_pipe, sum_pipe, ner_pipe, cross_encoder, sbert = load_models()
@st.cache_data
def get_yt_audio(url):
'''Get YT video from given URL link'''
yt = YouTube(url)
title = yt.title
# Get the first available audio stream and download it
audio_stream = yt.streams.filter(only_audio=True).first().download()
return audio_stream, title
@st.cache_data
def get_yt_audio_dl(url):
'''Back up for when pytube is down'''
temp_audio_file = os.path.join('output', 'audio')
ydl_opts = {
'format': 'bestaudio/best',
'postprocessors': [{
'key': 'FFmpegExtractAudio',
'preferredcodec': 'mp3',
'preferredquality': '192',
}],
'outtmpl': temp_audio_file,
'quiet': True,
}
with yt_dlp.YoutubeDL(ydl_opts) as ydl:
info = ydl.extract_info(url, download=False)
title = info.get('title', None)
ydl.download([url])
#with open(temp_audio_file+'.mp3', 'rb') as file:
audio_file = os.path.join('output', 'audio.mp3')
return audio_file, title
@st.cache_data
def load_whisper_api(audio):
'''Transcribe YT audio to text using Open AI API'''
file = open(audio, "rb")
transcript = openai_audio.audio.transcriptions.create(model="whisper-1", file=file,response_format="text")
return transcript
@st.cache_data
def transcribe_yt_video(link, py_tube=True):
'''Transcribe YouTube video'''
if py_tube:
audio_file, title = get_yt_audio(link)
print(f'audio_file:{audio_file}')
st.session_state['audio'] = audio_file
print(f"audio_file_session_state:{st.session_state['audio'] }")
#Get size of audio file
audio_size = round(os.path.getsize(st.session_state['audio'])/(1024*1024),1)
#Check if file is > 24mb, if not then use Whisper API
if audio_size <= 25:
st.info("`Transcribing YT audio...`")
#Use whisper API
results = load_whisper_api(st.session_state['audio'])
else:
st.warning('File size larger than 24mb, applying chunking and transcription',icon="⚠️")
song = AudioSegment.from_file(st.session_state['audio'], format='mp3')
# PyDub handles time in milliseconds
twenty_minutes = 20 * 60 * 1000
chunks = song[::twenty_minutes]
transcriptions = []
video_id = extract.video_id(link)
print(video_id)
for i, chunk in enumerate(chunks):
chunk.export(f'output/chunk_{i}_{video_id}.mp4', format='mp3')
transcriptions.append(load_whisper_api(f'output/chunk_{i}_{video_id}.mp3'))
results = ','.join(transcriptions)
print(results)
else:
audio_file, title = get_yt_audio_dl(link)
print(f'audio_file:{audio_file}')
st.session_state['audio'] = audio_file
print(f"audio_file_session_state:{st.session_state['audio'] }")
#Get size of audio file
audio_size = round(os.path.getsize(st.session_state['audio'])/(1024*1024),1)
#Check if file is > 24mb, if not then use Whisper API
if audio_size <= 25:
st.info("`Transcribing YT audio...`")
#Use whisper API
results = load_whisper_api(st.session_state['audio'])
else:
st.warning('File size larger than 24mb, applying chunking and transcription',icon="⚠️")
song = AudioSegment.from_file(st.session_state['audio'], format='mp3')
# PyDub handles time in milliseconds
twenty_minutes = 20 * 60 * 1000
chunks = song[::twenty_minutes]
transcriptions = []
video_id = extract.video_id(link)
for i, chunk in enumerate(chunks):
chunk.export(f'output/chunk_{i}_{video_id}.mp3', format='mp3')
transcriptions.append(load_whisper_api(f'output/chunk_{i}_{video_id}.mp3'))
results = ','.join(transcriptions)
st.info("`YT Video transcription process complete...`")
return results, title
@st.cache_data
def inference(link, upload):
'''Convert Youtube video or Audio upload to text'''
try:
if validators.url(link):
st.info("`Downloading YT audio...`")
results, title = transcribe_yt_video(link)
return results, title
elif _upload:
#Get size of audio file
audio_size = round(os.path.getsize(_upload)/(1024*1024),1)
#Check if file is > 24mb, if not then use Whisper API
if audio_size <= 25:
st.info("`Transcribing uploaded audio...`")
#Use whisper API
results = load_whisper_api(_upload)['text']
else:
st.write('File size larger than 24mb, applying chunking and transcription')
song = AudioSegment.from_file(_upload)
# PyDub handles time in milliseconds
twenty_minutes = 20 * 60 * 1000
chunks = song[::twenty_minutes]
transcriptions = []
st.info("`Transcribing uploaded audio...`")
for i, chunk in enumerate(chunks):
chunk.export(f'output/chunk_{i}.mp4', format='mp4')
transcriptions.append(load_whisper_api(f'output/chunk_{i}.mp4')['text'])
results = ','.join(transcriptions)
st.info("`Uploaded audio transcription process complete...`")
return results, "Transcribed Earnings Audio"
except Exception as e:
st.error(f'''PyTube Error: {e},
Using yt_dlp module, might take longer than expected''',icon="🚨")
results, title = transcribe_yt_video(link, py_tube=False)
# results = _asr_model.transcribe(st.session_state['audio'], task='transcribe', language='en')
return results, title
@st.cache_resource
def send_feedback(run_id, score):
client.create_feedback(run_id, "user_score", score=score)
@st.cache_data
def clean_text(text):
'''Clean all text after inference'''
text = text.encode("ascii", "ignore").decode() # unicode
text = re.sub(r"https*\S+", " ", text) # url
text = re.sub(r"@\S+", " ", text) # mentions
text = re.sub(r"#\S+", " ", text) # hastags
text = re.sub(r"\s{2,}", " ", text) # over spaces
return text
@st.cache_data
def chunk_long_text(text,threshold,window_size=3,stride=2):
'''Preprocess text and chunk for sentiment analysis'''
#Convert cleaned text into sentences
sentences = sent_tokenize(text)
out = []
#Limit the length of each sentence to a threshold
for chunk in sentences:
if len(chunk.split()) < threshold:
out.append(chunk)
else:
words = chunk.split()
num = int(len(words)/threshold)
for i in range(0,num*threshold+1,threshold):
out.append(' '.join(words[i:threshold+i]))
passages = []
#Combine sentences into a window of size window_size
for paragraph in [out]:
for start_idx in range(0, len(paragraph), stride):
end_idx = min(start_idx+window_size, len(paragraph))
passages.append(" ".join(paragraph[start_idx:end_idx]))
return passages
@st.cache_data
def sentiment_pipe(earnings_text):
'''Determine the sentiment of the text'''
earnings_sentences = chunk_long_text(earnings_text,150,1,1)
earnings_sentiment = sent_pipe(earnings_sentences)
return earnings_sentiment, earnings_sentences
@st.cache_data
def chunk_and_preprocess_text(text, model_name= 'philschmid/flan-t5-base-samsum'):
'''Chunk and preprocess text for summarization'''
tokenizer = AutoTokenizer.from_pretrained(model_name)
sentences = sent_tokenize(text)
# initialize
length = 0
chunk = ""
chunks = []
count = -1
for sentence in sentences:
count += 1
combined_length = len(tokenizer.tokenize(sentence)) + length # add the no. of sentence tokens to the length counter
if combined_length <= tokenizer.max_len_single_sentence: # if it doesn't exceed
chunk += sentence + " " # add the sentence to the chunk
length = combined_length # update the length counter
# if it is the last sentence
if count == len(sentences) - 1:
chunks.append(chunk) # save the chunk
else:
chunks.append(chunk) # save the chunk
# reset
length = 0
chunk = ""
# take care of the overflow sentence
chunk += sentence + " "
length = len(tokenizer.tokenize(sentence))
return chunks
@st.cache_data
def summarize_text(text_to_summarize,max_len,min_len):
'''Summarize text with HF model'''
summarized_text = sum_pipe(text_to_summarize,
max_length=max_len,
min_length=min_len,
do_sample=False,
early_stopping=True,
num_beams=4)
summarized_text = ' '.join([summ['summary_text'] for summ in summarized_text])
return summarized_text
@st.cache_data
def get_all_entities_per_sentence(text):
doc = nlp(''.join(text))
sentences = list(doc.sents)
entities_all_sentences = []
for sentence in sentences:
entities_this_sentence = []
# SPACY ENTITIES
for entity in sentence.ents:
entities_this_sentence.append(str(entity))
# XLM ENTITIES
entities_xlm = [entity["word"] for entity in ner_pipe(str(sentence))]
for entity in entities_xlm:
entities_this_sentence.append(str(entity))
entities_all_sentences.append(entities_this_sentence)
return entities_all_sentences
@st.cache_data
def get_all_entities(text):
all_entities_per_sentence = get_all_entities_per_sentence(text)
return list(itertools.chain.from_iterable(all_entities_per_sentence))
@st.cache_data
def get_and_compare_entities(article_content,summary_output):
all_entities_per_sentence = get_all_entities_per_sentence(article_content)
entities_article = list(itertools.chain.from_iterable(all_entities_per_sentence))
all_entities_per_sentence = get_all_entities_per_sentence(summary_output)
entities_summary = list(itertools.chain.from_iterable(all_entities_per_sentence))
matched_entities = []
unmatched_entities = []
for entity in entities_summary:
if any(entity.lower() in substring_entity.lower() for substring_entity in entities_article):
matched_entities.append(entity)
elif any(
np.inner(sbert.encode(entity, show_progress_bar=False),
sbert.encode(art_entity, show_progress_bar=False)) > 0.9 for
art_entity in entities_article):
matched_entities.append(entity)
else:
unmatched_entities.append(entity)
matched_entities = list(dict.fromkeys(matched_entities))
unmatched_entities = list(dict.fromkeys(unmatched_entities))
matched_entities_to_remove = []
unmatched_entities_to_remove = []
for entity in matched_entities:
for substring_entity in matched_entities:
if entity != substring_entity and entity.lower() in substring_entity.lower():
matched_entities_to_remove.append(entity)
for entity in unmatched_entities:
for substring_entity in unmatched_entities:
if entity != substring_entity and entity.lower() in substring_entity.lower():
unmatched_entities_to_remove.append(entity)
matched_entities_to_remove = list(dict.fromkeys(matched_entities_to_remove))
unmatched_entities_to_remove = list(dict.fromkeys(unmatched_entities_to_remove))
for entity in matched_entities_to_remove:
matched_entities.remove(entity)
for entity in unmatched_entities_to_remove:
unmatched_entities.remove(entity)
return matched_entities, unmatched_entities
@st.cache_data
def highlight_entities(article_content,summary_output):
markdown_start_red = "<mark class=\"entity\" style=\"background: rgb(238, 135, 135);\">"
markdown_start_green = "<mark class=\"entity\" style=\"background: rgb(121, 236, 121);\">"
markdown_end = "</mark>"
matched_entities, unmatched_entities = get_and_compare_entities(article_content,summary_output)
for entity in matched_entities:
summary_output = re.sub(f'({entity})(?![^rgb\(]*\))',markdown_start_green + entity + markdown_end,summary_output)
for entity in unmatched_entities:
summary_output = re.sub(f'({entity})(?![^rgb\(]*\))',markdown_start_red + entity + markdown_end,summary_output)
print("")
print("")
soup = BeautifulSoup(summary_output, features="html.parser")
return HTML_WRAPPER.format(soup)
def summary_downloader(raw_text):
'''Download the summary generated'''
b64 = base64.b64encode(raw_text.encode()).decode()
new_filename = "new_text_file_{}_.txt".format(time_str)
st.markdown("#### Download Summary as a File ###")
href = f'<a href="data:file/txt;base64,{b64}" download="{new_filename}">Click to Download!!</a>'
st.markdown(href,unsafe_allow_html=True)
@st.cache_data
def generate_eval(raw_text, N, chunk):
# Generate N questions from context of chunk chars
# IN: text, N questions, chunk size to draw question from in the doc
# OUT: eval set as JSON list
# raw_text = ','.join(raw_text)
update = st.empty()
ques_update = st.empty()
update.info("`Generating sample questions ...`")
n = len(raw_text)
starting_indices = [random.randint(0, n-chunk) for _ in range(N)]
sub_sequences = [raw_text[i:i+chunk] for i in starting_indices]
chain = QAGenerationChain.from_llm(ChatOpenAI(temperature=0))
eval_set = []
for i, b in enumerate(sub_sequences):
try:
qa = chain.run(b)
eval_set.append(qa)
ques_update.info(f"Creating Question: {i+1}")
except Exception as e:
print(e)
st.warning(f'Error in generating Question: {i+1}...', icon="⚠️")
continue
eval_set_full = list(itertools.chain.from_iterable(eval_set))
update.empty()
ques_update.empty()
return eval_set_full
@st.cache_resource
def create_prompt_and_llm():
'''Create prompt'''
llm = ChatOpenAI(temperature=0, streaming=True, model="gpt-4-turbo-preview")
message = SystemMessage(
content=(
"You are a helpful chatbot who is tasked with answering questions acuurately about earnings call transcript provided. "
"Unless otherwise explicitly stated, it is probably fair to assume that questions are about the earnings call transcript. "
"If there is any ambiguity, you probably assume they are about that."
"Do not use any information not provided in the earnings context and remember you are a to speak like a finance expert."
"If you don't know the answer, just say 'There is no relevant answer in the given earnings call transcript'"
"don't try to make up an answer"
)
)
prompt = OpenAIFunctionsAgent.create_prompt(
system_message=message,
extra_prompt_messages=[MessagesPlaceholder(variable_name="history")],
)
return prompt, llm
@st.cache_resource
def gen_embeddings(embedding_model):
'''Generate embeddings for given model'''
if 'hkunlp' in embedding_model:
embeddings = HuggingFaceInstructEmbeddings(model_name=embedding_model,
query_instruction='Represent the Financial question for retrieving supporting paragraphs: ',
embed_instruction='Represent the Financial paragraph for retrieval: ')
elif 'mpnet' in embedding_model:
embeddings = HuggingFaceEmbeddings(model_name=embedding_model)
elif 'FlagEmbedding' in embedding_model:
encode_kwargs = {'normalize_embeddings': True}
embeddings = HuggingFaceBgeEmbeddings(model_name=embedding_model,
encode_kwargs = encode_kwargs
)
return embeddings
@st.cache_data
def create_vectorstore(corpus, title, embedding_model, chunk_size=1000, overlap=50):
'''Process text for Semantic Search'''
text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size,chunk_overlap=overlap)
texts = text_splitter.split_text(corpus)
embeddings = gen_embeddings(embedding_model)
vectorstore = FAISS.from_texts(texts, embeddings, metadatas=[{"source": i} for i in range(len(texts))])
return vectorstore
@st.cache_resource
def create_memory_and_agent(_docsearch):
'''Embed text and generate semantic search scores'''
#create vectorstore
vectorstore = _docsearch.as_retriever(search_kwargs={"k": 4})
#create retriever tool
tool = create_retriever_tool(
vectorstore,
"earnings_call_search",
"Searches and returns documents using the earnings context provided as a source, relevant to the user input question.",
)
tools = [tool]
prompt,llm = create_prompt_and_llm()
agent = OpenAIFunctionsAgent(llm=llm, tools=tools, prompt=prompt)
agent_executor = AgentExecutor(
agent=agent,
tools=tools,
verbose=True,
return_intermediate_steps=True,
)
memory = AgentTokenBufferMemory(llm=llm)
return memory, agent_executor
@st.cache_data
def gen_sentiment(text):
'''Generate sentiment of given text'''
return sent_pipe(text)[0]['label']
@st.cache_data
def gen_annotated_text(df):
'''Generate annotated text'''
tag_list=[]
for row in df.itertuples():
label = row[2]
text = row[1]
if label == 'Positive':
tag_list.append((text,label,'#8fce00'))
elif label == 'Negative':
tag_list.append((text,label,'#f44336'))
else:
tag_list.append((text,label,'#000000'))
return tag_list
def display_df_as_table(model,top_k,score='score'):
'''Display the df with text and scores as a table'''
df = pd.DataFrame([(hit[score],passages[hit['corpus_id']]) for hit in model[0:top_k]],columns=['Score','Text'])
df['Score'] = round(df['Score'],2)
return df
def make_spans(text,results):
results_list = []
for i in range(len(results)):
results_list.append(results[i]['label'])
facts_spans = []
facts_spans = list(zip(sent_tokenizer(text),results_list))
return facts_spans
##Fiscal Sentiment by Sentence
def fin_ext(text):
results = remote_clx(sent_tokenizer(text))
return make_spans(text,results)
## Knowledge Graphs code
def get_article(url):
article = Article(url)
article.download()
article.parse()
return article