import numpy as np
import pandas as pd
import re
import os
import cloudpickle
from transformers import (DebertaTokenizerFast,
TFAutoModelForTokenClassification,
BartTokenizerFast,
TFAutoModelForSeq2SeqLM)
import tensorflow as tf
import spacy
import streamlit as st
from scraper import scrape_text
os.environ['TF_USE_LEGACY_KERAS'] = "1"
class NERLabelEncoder:
'''
Label Encoder to encode and decode the entity labels
'''
def __init__(self):
self.label_mapping = {'O': 0,
'B-geo': 1,
'I-geo': 2,
'B-gpe': 3,
'I-gpe': 4,
'B-per': 5,
'I-per': 6,
'B-org': 7,
'I-org': 8,
'B-tim': 9,
'I-tim': 10,
'B-art': 11,
'I-art': 12,
'B-nat': 13,
'I-nat': 14,
'B-eve': 15,
'I-eve': 16,
'[CLS]': -100,
'[SEP]': -100}
self.inverse_label_mapping = {}
def fit(self):
self.inverse_label_mapping = {value: key for key, value in self.label_mapping.items()}
return self
def transform(self, x: pd.Series):
x = x.map(self.label_mapping)
return x
def inverse_transform(self, x: pd.Series):
x = x.map(self.inverse_label_mapping)
return x
############ NER MODEL & VARS INITIALIZATION START ####################
NER_CHECKPOINT = "microsoft/deberta-base"
NER_N_TOKENS = 50
NER_N_LABELS = 18
NER_COLOR_MAP = {'GEO': '#DFFF00', 'GPE': '#FFBF00', 'PER': '#9FE2BF',
'ORG': '#40E0D0', 'TIM': '#CCCCFF', 'ART': '#FFC0CB', 'NAT': '#FFE4B5', 'EVE': '#DCDCDC'}
@st.cache_resource
def load_ner_models():
ner_model = TFAutoModelForTokenClassification.from_pretrained(NER_CHECKPOINT, num_labels=NER_N_LABELS, attention_probs_dropout_prob=0.4, hidden_dropout_prob=0.4)
ner_model.load_weights(os.path.join("models", "general_ner_deberta_weights.h5"), by_name=True)
ner_label_encoder = NERLabelEncoder()
ner_label_encoder.fit()
ner_tokenizer = DebertaTokenizerFast.from_pretrained(NER_CHECKPOINT, add_prefix_space=True)
nlp = spacy.load(os.path.join('.', 'en_core_web_sm-3.6.0'))
print('Loaded NER models')
return ner_model, ner_label_encoder, ner_tokenizer, nlp
ner_model, ner_label_encoder, ner_tokenizer, nlp = load_ner_models()
############ NER MODEL & VARS INITIALIZATION END ####################
############ NER LOGIC START ####################
def softmax(x):
return tf.exp(x) / tf.math.reduce_sum(tf.exp(x))
def ner_process_output(res):
'''
Function to concatenate sub-word tokens, labels and
compute mean prediction probability of tokens
'''
d = {}
result = []
pred_prob = []
res.append(['-', 'B-b', 0])
for n, i in enumerate(res):
try:
split = i[1].split('-')
token = i[0]
token_prob = i[2]
prefix, suffix = split
if prefix == 'B':
if len(d) != 0:
result.append([(re.sub(r"[^\x00-\x7F]+", '', token.replace("Ġ", " ").strip()), label, np.mean(pred_prob))
for label, token in d.items()][0])
d = {}
pred_prob = []
pred_prob.append(token_prob)
d[suffix] = token
else:
d[suffix] = d[suffix] + token
pred_prob.append(token_prob)
except:
continue
return result
def ner_inference(txt):
'''
Function that returns model prediction and prediction probabitliy
'''
test_data = [txt]
# tokenizer = DebertaTokenizerFast.from_pretrained(NER_CHECKPOINT, add_prefix_space=True)
tokens = ner_tokenizer.tokenize(txt)
tokenized_data = ner_tokenizer(test_data, is_split_into_words=True, max_length=NER_N_TOKENS,
truncation=True, padding="max_length")
token_idx_to_consider = tokenized_data.word_ids()
token_idx_to_consider = [i for i in range(len(token_idx_to_consider)) if token_idx_to_consider[i] is not None]
input_ = [tokenized_data['input_ids'], tokenized_data['attention_mask']]
pred_logits = ner_model.predict(input_, verbose=0).logits[0]
pred_prob = tf.map_fn(softmax, pred_logits)
pred_idx = tf.argmax(pred_prob, axis=-1).numpy()
pred_idx = pred_idx[token_idx_to_consider]
pred_prob = tf.math.reduce_max(pred_prob, axis=-1).numpy()
pred_prob = np.round(pred_prob[token_idx_to_consider], 3)
pred_labels = ner_label_encoder.inverse_transform(pd.Series(pred_idx))
result = [[token, label, prob] for token, label,
prob in zip(tokens, pred_labels, pred_prob) if label.find('-') >= 0]
output = ner_process_output(result)
return output
def ner_inference_long_text(txt):
entities = []
doc = nlp(txt)
n_sents = len([_ for _ in doc.sents])
n = 0
progress_bar = st.progress(0, text=f'Processed 0 / {n_sents} sentences')
for sent in doc.sents:
entities.extend(ner_inference(sent.text))
n += 1
progress_bar.progress(n / n_sents, text=f'Processed {n} / {n_sents} sentences')
# progress_bar.empty()
return entities
def get_ner_text(article_txt, ner_result):
res_txt = ''
start = 0
prev_start = 0
for i in ner_result:
try:
span = next(re.finditer(fr'{i[0]}', article_txt)).span()
start = span[0]
end = span[1]
res_txt += article_txt[prev_start:start]
repl_str = f'''{article_txt[start:end].strip()}
{i[1]} ({str(np.round(i[2], 3))})'''
res_txt += article_txt[start:end].replace(article_txt[start:end], repl_str)
prev_start = 0
article_txt = article_txt[end:]
except:
continue
res_txt += article_txt
return res_txt
############ NER LOGIC END ####################
############ SUMMARIZATION MODEL & VARS INITIALIZATION START ####################
SUMM_CHECKPOINT = "facebook/bart-base"
SUMM_INPUT_N_TOKENS = 400
SUMM_TARGET_N_TOKENS = 300
@st.cache_resource
def load_summarizer_models():
summ_tokenizer = BartTokenizerFast.from_pretrained(SUMM_CHECKPOINT)
summ_model = TFAutoModelForSeq2SeqLM.from_pretrained(SUMM_CHECKPOINT)
summ_model.load_weights(os.path.join("models", "bart_en_summarizer.h5"), by_name=True)
print('Loaded summarizer models')
return summ_tokenizer, summ_model
summ_tokenizer, summ_model = load_summarizer_models()
def summ_preprocess(txt):
txt = re.sub(r'^By \. [\w\s]+ \. ', ' ', txt) # By . Ellie Zolfagharifard .
txt = re.sub(r'\d{1,2}\:\d\d [a-zA-Z]{3}', ' ', txt) # 10:30 EST
txt = re.sub(r'\d{1,2} [a-zA-Z]+ \d{4}', ' ', txt) # 10 November 1990
txt = txt.replace('PUBLISHED:', ' ')
txt = txt.replace('UPDATED', ' ')
txt = re.sub(r' [\,\.\:\'\;\|] ', ' ', txt) # remove puncts with spaces before and after
txt = txt.replace(' : ', ' ')
txt = txt.replace('(CNN)', ' ')
txt = txt.replace('--', ' ')
txt = re.sub(r'^\s*[\,\.\:\'\;\|]', ' ', txt) # remove puncts at beginning of sent
txt = re.sub(r' [\,\.\:\'\;\|] ', ' ', txt) # remove puncts with spaces before and after
txt = re.sub(r'\n+',' ', txt)
txt = " ".join(txt.split())
return txt
def summ_inference_tokenize(input_: list, n_tokens: int):
tokenized_data = summ_tokenizer(text=input_, max_length=SUMM_TARGET_N_TOKENS, truncation=True, padding="max_length", return_tensors="tf")
return summ_tokenizer, tokenized_data
def summ_inference(txt: str):
txt = summ_preprocess(txt)
test_data = [txt]
inference_tokenizer, tokenized_data = summ_inference_tokenize(input_=test_data, n_tokens=SUMM_INPUT_N_TOKENS)
pred = summ_model.generate(**tokenized_data, max_new_tokens=SUMM_TARGET_N_TOKENS)
result = inference_tokenizer.decode(pred[0])
result = re.sub("<.*?>", "", result).strip()
return result
############ SUMMARIZATION MODEL & VARS INITIALIZATION END ####################
############## ENTRY POINT START #######################
def main():
st.markdown('''News Summarizer and NER
README
''', unsafe_allow_html=True)
input_type = st.radio('Select an option:', ['Paste news URL', 'Paste news text'],
horizontal=True)
scrape_error = None
summary_error = None
ner_error = None
summ_result = None
ner_result = None
ner_df = None
article_txt = None
if input_type == 'Paste news URL':
article_url = st.text_input("Paste the URL of a news article", "")
if (st.button("Submit")) or (article_url):
with st.status("Processing...", expanded=True) as status:
status.empty()
# Scraping data Start
try:
st.info("Scraping data from the URL.", icon="ℹ️")
article_txt = scrape_text(article_url)
st.success("Successfully scraped the data.", icon="✅")
except Exception as e:
article_txt = None
scrape_error = str(e)
# Scraping data End
if article_txt is not None:
article_txt = re.sub(r'\n+',' ', article_txt)
# Generating summary start
try:
st.info("Generating the summary.", icon="ℹ️")
summ_result = summ_inference(article_txt)
except Exception as e:
summ_result = None
summary_error = str(e)
if summ_result is not None:
st.success("Successfully generated the summary.", icon="✅")
else:
st.error("Encountered an error while generating the summary.", icon="🚨")
# Generating summary end
# NER start
try:
st.info("Recognizing the entites.", icon="ℹ️")
ner_result = [[ent, label.upper(), np.round(prob, 3)]
for ent, label, prob in ner_inference_long_text(article_txt)]
ner_df = pd.DataFrame(ner_result, columns=['entity', 'label', 'confidence'])
ner_result = get_ner_text(article_txt, ner_result).replace('$', '\$')
except Exception as e:
ner_result = None
ner_error = str(e)
if ner_result is not None:
st.success("Successfully recognized the entites.", icon="✅")
else:
st.error("Encountered an error while recognizing the entites.", icon="🚨")
# NER end
else:
st.error("Encountered an error while scraping the data.", icon="🚨")
if (scrape_error is None) and (summary_error is None) and (ner_error is None):
status.update(label="Done", state="complete", expanded=False)
else:
status.update(label="Error", state="error", expanded=False)
if scrape_error is not None:
st.error(f"Scrape Error: \n{scrape_error}", icon="🚨")
else:
if summary_error is not None:
st.error(f"Summary Error: \n{summary_error}", icon="🚨")
else:
st.markdown(f"SUMMARY:
{summ_result}", unsafe_allow_html=True)
if ner_error is not None:
st.error(f"NER Error \n{ner_error}", icon="🚨")
else:
st.markdown(f"ENTITIES:
{ner_result}", unsafe_allow_html=True)
# st.dataframe(ner_df, use_container_width=True)
st.markdown(f"SCRAPED TEXT:
{article_txt}", unsafe_allow_html=True)
else:
article_txt = st.text_area("Paste the text of a news article", "", height=150)
if (st.button("Submit")) or (article_txt):
with st.status("Processing...", expanded=True) as status:
article_txt = re.sub(r'\n+',' ', article_txt)
# Generating summary start
try:
st.info("Generating the summary.", icon="ℹ️")
summ_result = summ_inference(article_txt)
except Exception as e:
summ_result = None
summary_error = str(e)
if summ_result is not None:
st.success("Successfully generated the summary.", icon="✅")
else:
st.error("Encountered an error while generating the summary.", icon="🚨")
# Generating summary end
# NER start
try:
st.info("Recognizing the entites.", icon="ℹ️")
ner_result = [[ent, label.upper(), np.round(prob, 3)]
for ent, label, prob in ner_inference_long_text(article_txt)]
ner_df = pd.DataFrame(ner_result, columns=['entity', 'label', 'confidence'])
ner_result = get_ner_text(article_txt, ner_result).replace('$', '\$')
except Exception as e:
ner_result = None
ner_error = str(e)
if ner_result is not None:
st.success("Successfully recognized the entites.", icon="✅")
else:
st.error("Encountered an error while recognizing the entites.", icon="🚨")
# NER end
if (summary_error is None) and (ner_error is None):
status.update(label="Done", state="complete", expanded=False)
else:
status.update(label="Error", state="error", expanded=False)
if summary_error is not None:
st.error(f"Summary Error: \n{summary_error}", icon="🚨")
else:
st.markdown(f"SUMMARY:
{summ_result}", unsafe_allow_html=True)
if ner_error is not None:
st.error(f"NER Error \n{ner_error}", icon="🚨")
else:
st.markdown(f"ENTITIES:
{ner_result}", unsafe_allow_html=True)
# st.dataframe(ner_df, use_container_width=True)
############## ENTRY POINT END #######################
if __name__ == "__main__":
main()