|
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_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() |
|
|
|
|
|
|
|
|
|
|
|
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] |
|
|
|
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') |
|
|
|
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'''<span style="background-color:{NER_COLOR_MAP[i[1]]}; border-radius: 3px">{article_txt[start:end].strip()} |
|
<span style="font-size:10px; font-weight:bold; display:inline-block; vertical-align: middle;"> |
|
{i[1]} ({str(np.round(i[2], 3))})</span></span>''' |
|
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 |
|
|
|
|
|
|
|
|
|
|
|
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) |
|
txt = re.sub(r'\d{1,2}\:\d\d [a-zA-Z]{3}', ' ', txt) |
|
txt = re.sub(r'\d{1,2} [a-zA-Z]+ \d{4}', ' ', txt) |
|
txt = txt.replace('PUBLISHED:', ' ') |
|
txt = txt.replace('UPDATED', ' ') |
|
txt = re.sub(r' [\,\.\:\'\;\|] ', ' ', txt) |
|
txt = txt.replace(' : ', ' ') |
|
txt = txt.replace('(CNN)', ' ') |
|
txt = txt.replace('--', ' ') |
|
txt = re.sub(r'^\s*[\,\.\:\'\;\|]', ' ', txt) |
|
txt = re.sub(r' [\,\.\:\'\;\|] ', ' ', txt) |
|
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 clean_summary(summary: str): |
|
summary = summary.strip() |
|
if summary[-1] != '.': |
|
sents = summary.split(". ") |
|
summary = ". ".join(sents[:-1]) |
|
summary += "." |
|
summary = re.sub(r'^-', "", summary) |
|
summary = summary.strip() |
|
if len(summary) <= 5: |
|
summary = "" |
|
return summary |
|
|
|
def summ_inference(txt: str): |
|
txt = summ_preprocess(txt) |
|
inference_tokenizer, tokenized_data = summ_inference_tokenize(input_=[txt], n_tokens=SUMM_INPUT_N_TOKENS) |
|
pred = summ_model.generate(**tokenized_data, max_new_tokens=SUMM_TARGET_N_TOKENS) |
|
result = "" if txt=="" else clean_summary(inference_tokenizer.decode(pred[0], skip_special_tokens=True)) |
|
return result |
|
|
|
|
|
|
|
def main(): |
|
st.markdown('''<h3>News Summarizer and NER</h3> |
|
<p><a href="https://huggingface.co/spaces/ksvmuralidhar/news_summarizer_ner/blob/main/README.md#new-summarization-and-ner" target="_blank">README</a> |
|
<br> |
|
The app works best in summarizing <a href="https://edition.cnn.com/" target="_blank">CNN</a> and |
|
<a href="https://www.dailymail.co.uk/home/index.html" target="_blank">Daily Mail</a> news articles, |
|
as the BART model is fine-tuned on them. |
|
</p> |
|
|
|
''', 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() |
|
|
|
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) |
|
|
|
|
|
|
|
if article_txt is not None: |
|
article_txt = re.sub(r'\n+',' ', article_txt) |
|
|
|
|
|
|
|
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="π¨") |
|
|
|
|
|
|
|
|
|
|
|
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="π¨") |
|
|
|
|
|
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"<h4>SUMMARY:</h4>{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"<h4>ENTITIES:</h4>{ner_result}", unsafe_allow_html=True) |
|
|
|
|
|
st.markdown(f"<h4>SCRAPED TEXT:</h4>{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) |
|
|
|
|
|
|
|
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="π¨") |
|
|
|
|
|
|
|
|
|
|
|
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="π¨") |
|
|
|
|
|
|
|
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"<h4>SUMMARY:</h4>{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"<h4>ENTITIES:</h4>{ner_result}", unsafe_allow_html=True) |
|
|
|
|
|
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
main() |
