Create app.py

app.py

@@ -0,0 +1,225 @@
+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
+
+
+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_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()
+nlp = spacy.load(os.path.join('.', 'en_core_web_sm-3.6.0'))
+NER_COLOR_MAP = {'GEO': '#DFFF00', 'GPE': '#FFBF00', 'PER': '#9FE2BF', 
+                 'ORG': '#40E0D0', 'TIM': '#CCCCFF', 'ART': '#FFC0CB', 'NAT': '#FFE4B5', 'EVE': '#DCDCDC'}
+############ 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 = tokenizer.tokenize(txt)
+    tokenized_data = 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)
+    for sent in doc.sents:
+        entities.extend(ner_inference(sent.text))
+    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
+
+############ NER LOGIC END ####################
+
+############ SUMMARIZATION MODEL & VARS INITIALIZATION START ####################
+SUMM_CHECKPOINT = "facebook/bart-base"
+SUMM_INPUT_N_TOKENS = 400
+SUMM_TARGET_N_TOKENS = 100
+summ_model = TFAutoModelForSeq2SeqLM.from_pretrained(SUMM_CHECKPOINT)
+summ_model.load_weights(os.path.join("models", "bart_en_summarizer.h5"), by_name=True)
+
+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 = " ".join(txt.split())
+    return txt
+
+def summ_inference_tokenize(input_: list, n_tokens: int):
+    tokenizer = BartTokenizerFast.from_pretrained(SUMM_CHECKPOINT)
+    tokenized_data = tokenizer(text=input_, max_length=SUMM_TARGET_N_TOKENS, truncation=True, padding="max_length", return_tensors="tf")
+    return 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.title("News Summarization & NER")
+    article_txt = st.text_area("Paste the text of a news article:", "", height=200)
+    if st.button("Submit"):
+        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'])
+        summ_result = summ_inference(article_txt)        
+        
+        ner_txt = get_ner_text(article_txt, ner_result)
+
+        st.markdown(f"<h4>SUMMARY:</h4>{summ_result}<h4>ENTITIES:</h4>", unsafe_allow_html=True)
+        st.markdown(f"{ner_txt}", unsafe_allow_html=True)
+        st.dataframe(ner_df, use_container_width=True)
+
+############## ENTRY POINT END #######################
+
+if __name__ == "__main__":
+    main()