app.py

import time
import pandas as pd

import torch
from keybert import KeyBERT
from sentence_transformers import SentenceTransformer
from keyphrase_vectorizers import KeyphraseCountVectorizer
from transformers import T5ForConditionalGeneration,T5Tokenizer
#from fastT5 import export_and_get_onnx_model, set_auth_token

import nltk
from nltk.tokenize import sent_tokenize

from huggingface_hub import snapshot_download, HfFolder
import streamlit as st

import traceback
import logging

logger = logging.getLogger(__name__)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

HfFolder.save_token(st.secrets["hf-auth-token"])

@st.cache(allow_output_mutation=True)
def load_base_model():
    try:
        nltk.download('stopwords')
        nltk.download('punkt')
        # Load KeyBert Model
        tmp_model = SentenceTransformer('valurank/MiniLM-L6-Keyword-Extraction', use_auth_token=True)
        kw_extractor = KeyBERT(tmp_model)

        # Load T5 for Paraphrasing
        t5_tokenizer = T5Tokenizer.from_pretrained('t5-base')
        t5_model = T5ForConditionalGeneration.from_pretrained('valurank/t5-paraphraser', use_auth_token=True)
        t5_model = t5_model.to(device)
        return kw_extractor, t5_model, t5_tokenizer
    except Exception:
        st.error('Error Loading Models. Please contact admin')
        logger.error(traceback.format_exc())
    
    

def get_keybert_results_with_vectorizer(text, number_of_results=20):
    try:
        keywords = kw_extractor.extract_keywords(text, vectorizer=KeyphraseCountVectorizer(), stop_words=None, top_n=number_of_results)
        keywords = [i for i in keywords if i[1] >= 0.25]
        return keywords
    except Exception:
        st.error('Error running Keybert. Please contact admin')
        logger.error(traceback.format_exc())



def t5_paraphraser(text, number_of_results=5):
    try:
        text =  "paraphrase: " + text
        encoding = t5_tokenizer.encode_plus(text, pad_to_max_length=True, return_tensors="pt")
        input_ids, attention_masks = encoding["input_ids"].to(device), encoding["attention_mask"].to(device)

        beam_outputs = t5_model.generate(
            input_ids=input_ids, attention_mask=attention_masks,
            do_sample=True,
            max_length=1024,
            top_k=50,
            top_p=0.95,
            early_stopping=True,
            num_return_sequences=number_of_results
        )

        final_outputs =[]
        for beam_output in beam_outputs:
            sent = t5_tokenizer.decode(beam_output, skip_special_tokens=True, clean_up_tokenization_spaces=True)
            final_outputs.append(sent)
        return final_outputs
    except Exception:
        st.error('Error running T5 Paraphrasing. Please contact admin')
        logger.error(traceback.format_exc())
    
    
    
def run_long_extraction(article, number_of_paraphrases):
    try:
        start1 = time.time()
        with st.spinner('Extraction Keywords from Original Document...'):
            original_keywords = get_keybert_results_with_vectorizer(article, number_of_results=30)
            article_sentences = sent_tokenize(article)
            target_sentences = [sent for sent in article_sentences if any(kw[0] in sent for kw in original_keywords)]

        st.success('Keyword Extraction from Original Document finished in {}'.format(time.time() - start1))
        st.info(f'Total Sentences in Article : {len(article_sentences)}')
        st.info(f'Total Target Sentences Selected : {len(target_sentences)}')

        start2 = time.time()
        with st.spinner('Extracting Keywords from Paraphrased Target Sentences...'):
            t5_paraphrasing_keywords = []
            for sent in target_sentences:
                ### T5
                t5_paraphrased = t5_paraphraser(sent, number_of_results = number_of_paraphrases)
                t5_keywords = [get_keybert_results_with_vectorizer(i) for i in t5_paraphrased]
                t5_keywords = [(word[0], word[1]) for s in t5_keywords for word in s]
                t5_paraphrasing_keywords.extend(t5_keywords)
        st.success('Keyword Extraction from Paraphrased Target Sentences finished in {}'.format(time.time() - start2))

        original_keywords_df = pd.DataFrame(original_keywords, columns=['Keyword', 'Score'])

        t5_keywords_df = pd.DataFrame(t5_paraphrasing_keywords, columns=['Keyword', 'Score']).sort_values(by='Score', ascending=False).drop_duplicates(subset=['Keyword'], keep='first').reset_index(drop=True)

        unique_keywords_df = pd.DataFrame([i for i in t5_paraphrasing_keywords if not original_keywords_df['Keyword'].str.contains(i[0], regex=False, case=False).any()], columns=['Keyword', 'Score']).sort_values(by='Score', ascending=False).drop_duplicates(subset=['Keyword'], keep='first').reset_index(drop=True)

        total_end = time.time()-start1

        return t5_keywords_df, original_keywords_df, unique_keywords_df, total_end
    except Exception:
        st.error('Error running Extraction Pipeline. Please contact admin')
        logger.error(traceback.format_exc())     
          
          
          
def run_short_extraction(article, number_of_paraphrases):
    try:
        start1 = time.time()
        original_keywords = get_keybert_results_with_vectorizer(article)
        article_sentences = sent_tokenize(article)
        st.info(f'Total Sentences in Article : {len(article_sentences)}')

        target_sentences = []
        tmp = []
        token_count = 0
        for i in article_sentences:
            enc = t5_tokenizer.encode(i)
            if token_count + len(enc) <= 96:
                tmp.append(i)
                token_count += len(enc)
            else:
                target_sentences.append(' '.join(tmp))
                token_count = len(enc)
                tmp = [i]          
     
        start2 = time.time()
        with st.spinner('Extracting Keywords from Paraphrased Sentences Groups...'):
            t5_paraphrasing_keywords = []
            for sent in target_sentences:
                ### T5
                t5_paraphrased = t5_paraphraser(sent, number_of_results = number_of_paraphrases)
                t5_keywords = [get_keybert_results_with_vectorizer(i) for i in t5_paraphrased]
                t5_keywords = [(word[0], word[1]) for s in t5_keywords for word in s]
                t5_paraphrasing_keywords.extend(t5_keywords)
        st.success('Keyword Extraction from Paraphrased Grouped Sentences finished in {}'.format(time.time() - start2))

        original_keywords_df = pd.DataFrame(original_keywords, columns=['Keyword', 'Score'])

        t5_keywords_df = pd.DataFrame(t5_paraphrasing_keywords, columns=['Keyword', 'Score']).sort_values(by='Score', ascending=False).drop_duplicates(subset=['Keyword'], keep='first').reset_index(drop=True)

        unique_keywords_df = pd.DataFrame([i for i in t5_paraphrasing_keywords if not original_keywords_df['Keyword'].str.contains(i[0], regex=False, case=False).any()], columns=['Keyword', 'Score']).sort_values(by='Score', ascending=False).drop_duplicates(subset=['Keyword'], keep='first').reset_index(drop=True)

        total_end = time.time()-start1

        return t5_keywords_df, original_keywords_df, unique_keywords_df, total_end
    except Exception:
        st.error('Error running Extraction Pipeline. Please contact admin')
        logger.error(traceback.format_exc())     
        



def check_document_length(article, number_of_paraphrases):
    total_tokens = len(t5_tokenizer.encode(article))
    st.info(f'Token Counts for Encoded Document: {total_tokens}')

    if total_tokens >= 512:
        st.info('Running Extraction for Long Document')
        t5_keywords_df, original_keywords_df, unique_keywords_df, total_end = run_long_extraction(article, number_of_paraphrases)
    else:
        st.info('Running Extraction for Short Document')
        t5_keywords_df, original_keywords_df, unique_keywords_df, total_end = run_short_extraction(article, number_of_paraphrases)
    
    return t5_keywords_df, original_keywords_df, unique_keywords_df, total_end
    
    

kw_extractor, t5_model, t5_tokenizer = load_base_model()


st.title('Exhaustive Keyword Extraction with Paraphrasing')
with st.sidebar:
    st.header('Overview')
    st.markdown('This demo allows users to input text article and generate synonym-aware keywords. The pipeline includes the use of T5 Model for paraphrasing target sentences, and Sentence-transformers based Keyword Extraction')

    st.header('Parameters')
    # number_of_keywords = st.slider('Number of Keywords to extract for each target sentence', min_value=5, max_value=50, step=5, value=20)
    number_of_paraphrases = st.slider('Number of Paraphrased versions to generate for each target sentence', min_value=1, max_value=20, step=1, value=5)

    st.header('Specifications')
    # st.markdown('To generate context aware and OOV keywords for long, we first run KeyBert for keyword extraction on the original article. The sentences which had Keywords are then passed through T5 for generating multiple paraphrased versions. These paraphrased sentences are then run through Keyword Extraction again to generate the final results')



doc = st.text_area("Enter a custom document")
if doc:
    t5_keywords_df, original_keywords_df, unique_keywords_df, total_end = check_document_length(doc, number_of_paraphrases)

    # extract_paraphrased_article(input_list[0])
    st.text(f'PIPELINE RUNTIME: {total_end}\n')

    st.subheader('\nOriginal Keywords Extracted:\n\n')
    st.dataframe(original_keywords_df)

    st.subheader('\nT5 Unique New Keywords Extracted:\n\n')
    st.dataframe(unique_keywords_df)

    st.subheader('\nT5 Keywords Extracted:\n\n')
    st.dataframe(t5_keywords_df)