Spaces:
Sleeping
Sleeping
Commit
•
c82bebf
1
Parent(s):
0e1a57d
Create app.py
Browse files
app.py
ADDED
@@ -0,0 +1,225 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import numpy as np
|
2 |
+
import pandas as pd
|
3 |
+
import re
|
4 |
+
import os
|
5 |
+
import cloudpickle
|
6 |
+
from transformers import (DebertaTokenizerFast,
|
7 |
+
TFAutoModelForTokenClassification,
|
8 |
+
BartTokenizerFast,
|
9 |
+
TFAutoModelForSeq2SeqLM)
|
10 |
+
import tensorflow as tf
|
11 |
+
import spacy
|
12 |
+
import streamlit as st
|
13 |
+
|
14 |
+
|
15 |
+
class NERLabelEncoder:
|
16 |
+
'''
|
17 |
+
Label Encoder to encode and decode the entity labels
|
18 |
+
'''
|
19 |
+
def __init__(self):
|
20 |
+
self.label_mapping = {'O': 0,
|
21 |
+
'B-geo': 1,
|
22 |
+
'I-geo': 2,
|
23 |
+
'B-gpe': 3,
|
24 |
+
'I-gpe': 4,
|
25 |
+
'B-per': 5,
|
26 |
+
'I-per': 6,
|
27 |
+
'B-org': 7,
|
28 |
+
'I-org': 8,
|
29 |
+
'B-tim': 9,
|
30 |
+
'I-tim': 10,
|
31 |
+
'B-art': 11,
|
32 |
+
'I-art': 12,
|
33 |
+
'B-nat': 13,
|
34 |
+
'I-nat': 14,
|
35 |
+
'B-eve': 15,
|
36 |
+
'I-eve': 16,
|
37 |
+
'[CLS]': -100,
|
38 |
+
'[SEP]': -100}
|
39 |
+
|
40 |
+
self.inverse_label_mapping = {}
|
41 |
+
|
42 |
+
def fit(self):
|
43 |
+
self.inverse_label_mapping = {value: key for key, value in self.label_mapping.items()}
|
44 |
+
return self
|
45 |
+
|
46 |
+
def transform(self, x: pd.Series):
|
47 |
+
x = x.map(self.label_mapping)
|
48 |
+
return x
|
49 |
+
|
50 |
+
def inverse_transform(self, x: pd.Series):
|
51 |
+
x = x.map(self.inverse_label_mapping)
|
52 |
+
return x
|
53 |
+
|
54 |
+
|
55 |
+
############ NER MODEL & VARS INITIALIZATION START ####################
|
56 |
+
NER_CHECKPOINT = "microsoft/deberta-base"
|
57 |
+
NER_N_TOKENS = 50
|
58 |
+
NER_N_LABELS = 18
|
59 |
+
ner_model = TFAutoModelForTokenClassification.from_pretrained(NER_CHECKPOINT, num_labels=NER_N_LABELS, attention_probs_dropout_prob=0.4, hidden_dropout_prob=0.4)
|
60 |
+
ner_model.load_weights(os.path.join("models", "general_ner_deberta_weights.h5"), by_name=True)
|
61 |
+
ner_label_encoder = NERLabelEncoder()
|
62 |
+
ner_label_encoder.fit()
|
63 |
+
nlp = spacy.load(os.path.join('.', 'en_core_web_sm-3.6.0'))
|
64 |
+
NER_COLOR_MAP = {'GEO': '#DFFF00', 'GPE': '#FFBF00', 'PER': '#9FE2BF',
|
65 |
+
'ORG': '#40E0D0', 'TIM': '#CCCCFF', 'ART': '#FFC0CB', 'NAT': '#FFE4B5', 'EVE': '#DCDCDC'}
|
66 |
+
############ NER MODEL & VARS INITIALIZATION END ####################
|
67 |
+
|
68 |
+
############ NER LOGIC START ####################
|
69 |
+
def softmax(x):
|
70 |
+
return tf.exp(x) / tf.math.reduce_sum(tf.exp(x))
|
71 |
+
|
72 |
+
def ner_process_output(res):
|
73 |
+
'''
|
74 |
+
Function to concatenate sub-word tokens, labels and
|
75 |
+
compute mean prediction probability of tokens
|
76 |
+
'''
|
77 |
+
d = {}
|
78 |
+
result = []
|
79 |
+
pred_prob = []
|
80 |
+
res.append(['-', 'B-b', 0])
|
81 |
+
for n, i in enumerate(res):
|
82 |
+
try:
|
83 |
+
split = i[1].split('-')
|
84 |
+
token = i[0]
|
85 |
+
token_prob = i[2]
|
86 |
+
prefix, suffix = split
|
87 |
+
if prefix == 'B':
|
88 |
+
if len(d) != 0:
|
89 |
+
result.append([(re.sub(r"[^\x00-\x7F]+", '', token.replace("Ġ", " ").strip()), label, np.mean(pred_prob))
|
90 |
+
for label, token in d.items()][0])
|
91 |
+
d = {}
|
92 |
+
pred_prob = []
|
93 |
+
pred_prob.append(token_prob)
|
94 |
+
d[suffix] = token
|
95 |
+
|
96 |
+
else:
|
97 |
+
d[suffix] = d[suffix] + token
|
98 |
+
pred_prob.append(token_prob)
|
99 |
+
except:
|
100 |
+
continue
|
101 |
+
|
102 |
+
return result
|
103 |
+
|
104 |
+
|
105 |
+
def ner_inference(txt):
|
106 |
+
'''
|
107 |
+
Function that returns model prediction and prediction probabitliy
|
108 |
+
'''
|
109 |
+
test_data = [txt]
|
110 |
+
tokenizer = DebertaTokenizerFast.from_pretrained(NER_CHECKPOINT, add_prefix_space=True)
|
111 |
+
tokens = tokenizer.tokenize(txt)
|
112 |
+
tokenized_data = tokenizer(test_data, is_split_into_words=True, max_length=NER_N_TOKENS,
|
113 |
+
truncation=True, padding="max_length")
|
114 |
+
|
115 |
+
token_idx_to_consider = tokenized_data.word_ids()
|
116 |
+
token_idx_to_consider = [i for i in range(len(token_idx_to_consider)) if token_idx_to_consider[i] is not None]
|
117 |
+
|
118 |
+
input_ = [tokenized_data['input_ids'], tokenized_data['attention_mask']]
|
119 |
+
pred_logits = ner_model.predict(input_, verbose=0).logits[0]
|
120 |
+
|
121 |
+
pred_prob = tf.map_fn(softmax, pred_logits)
|
122 |
+
|
123 |
+
pred_idx = tf.argmax(pred_prob, axis=-1).numpy()
|
124 |
+
pred_idx = pred_idx[token_idx_to_consider]
|
125 |
+
|
126 |
+
pred_prob = tf.math.reduce_max(pred_prob, axis=-1).numpy()
|
127 |
+
pred_prob = np.round(pred_prob[token_idx_to_consider], 3)
|
128 |
+
pred_labels = ner_label_encoder.inverse_transform(pd.Series(pred_idx))
|
129 |
+
|
130 |
+
result = [[token, label, prob] for token, label,
|
131 |
+
prob in zip(tokens, pred_labels, pred_prob) if label.find('-') >= 0]
|
132 |
+
|
133 |
+
output = ner_process_output(result)
|
134 |
+
return output
|
135 |
+
|
136 |
+
|
137 |
+
def ner_inference_long_text(txt):
|
138 |
+
entities = []
|
139 |
+
doc = nlp(txt)
|
140 |
+
for sent in doc.sents:
|
141 |
+
entities.extend(ner_inference(sent.text))
|
142 |
+
return entities
|
143 |
+
|
144 |
+
|
145 |
+
def get_ner_text(article_txt, ner_result):
|
146 |
+
res_txt = ''
|
147 |
+
start = 0
|
148 |
+
prev_start = 0
|
149 |
+
for i in ner_result:
|
150 |
+
try:
|
151 |
+
span = next(re.finditer(fr'{i[0]}', article_txt)).span()
|
152 |
+
start = span[0]
|
153 |
+
end = span[1]
|
154 |
+
res_txt += article_txt[prev_start:start]
|
155 |
+
repl_str = f'''<span style="background-color:{NER_COLOR_MAP[i[1]]}; border-radius: 3px">{article_txt[start:end].strip()}
|
156 |
+
<span style="font-size:10px; font-weight:bold; display:inline-block; vertical-align: middle;">
|
157 |
+
{i[1]} ({str(np.round(i[2], 3))})</span></span>'''
|
158 |
+
res_txt += article_txt[start:end].replace(article_txt[start:end], repl_str)
|
159 |
+
prev_start = 0
|
160 |
+
article_txt = article_txt[end:]
|
161 |
+
except:
|
162 |
+
continue
|
163 |
+
res_txt += article_txt
|
164 |
+
return res_txt
|
165 |
+
|
166 |
+
############ NER LOGIC END ####################
|
167 |
+
|
168 |
+
############ SUMMARIZATION MODEL & VARS INITIALIZATION START ####################
|
169 |
+
SUMM_CHECKPOINT = "facebook/bart-base"
|
170 |
+
SUMM_INPUT_N_TOKENS = 400
|
171 |
+
SUMM_TARGET_N_TOKENS = 100
|
172 |
+
summ_model = TFAutoModelForSeq2SeqLM.from_pretrained(SUMM_CHECKPOINT)
|
173 |
+
summ_model.load_weights(os.path.join("models", "bart_en_summarizer.h5"), by_name=True)
|
174 |
+
|
175 |
+
def summ_preprocess(txt):
|
176 |
+
txt = re.sub(r'^By \. [\w\s]+ \. ', ' ', txt) # By . Ellie Zolfagharifard .
|
177 |
+
txt = re.sub(r'\d{1,2}\:\d\d [a-zA-Z]{3}', ' ', txt) # 10:30 EST
|
178 |
+
txt = re.sub(r'\d{1,2} [a-zA-Z]+ \d{4}', ' ', txt) # 10 November 1990
|
179 |
+
txt = txt.replace('PUBLISHED:', ' ')
|
180 |
+
txt = txt.replace('UPDATED', ' ')
|
181 |
+
txt = re.sub(r' [\,\.\:\'\;\|] ', ' ', txt) # remove puncts with spaces before and after
|
182 |
+
txt = txt.replace(' : ', ' ')
|
183 |
+
txt = txt.replace('(CNN)', ' ')
|
184 |
+
txt = txt.replace('--', ' ')
|
185 |
+
txt = re.sub(r'^\s*[\,\.\:\'\;\|]', ' ', txt) # remove puncts at beginning of sent
|
186 |
+
txt = re.sub(r' [\,\.\:\'\;\|] ', ' ', txt) # remove puncts with spaces before and after
|
187 |
+
txt = " ".join(txt.split())
|
188 |
+
return txt
|
189 |
+
|
190 |
+
def summ_inference_tokenize(input_: list, n_tokens: int):
|
191 |
+
tokenizer = BartTokenizerFast.from_pretrained(SUMM_CHECKPOINT)
|
192 |
+
tokenized_data = tokenizer(text=input_, max_length=SUMM_TARGET_N_TOKENS, truncation=True, padding="max_length", return_tensors="tf")
|
193 |
+
return tokenizer, tokenized_data
|
194 |
+
|
195 |
+
def summ_inference(txt: str):
|
196 |
+
txt = summ_preprocess(txt)
|
197 |
+
test_data = [txt]
|
198 |
+
inference_tokenizer, tokenized_data = summ_inference_tokenize(input_=test_data, n_tokens=SUMM_INPUT_N_TOKENS)
|
199 |
+
pred = summ_model.generate(**tokenized_data, max_new_tokens=SUMM_TARGET_N_TOKENS)
|
200 |
+
result = inference_tokenizer.decode(pred[0])
|
201 |
+
result = re.sub("<.*?>", "", result).strip()
|
202 |
+
return result
|
203 |
+
############ SUMMARIZATION MODEL & VARS INITIALIZATION END ####################
|
204 |
+
|
205 |
+
############## ENTRY POINT START #######################
|
206 |
+
def main():
|
207 |
+
st.title("News Summarization & NER")
|
208 |
+
article_txt = st.text_area("Paste the text of a news article:", "", height=200)
|
209 |
+
if st.button("Submit"):
|
210 |
+
ner_result = [[ent, label.upper(), np.round(prob, 3)]
|
211 |
+
for ent, label, prob in ner_inference_long_text(article_txt)]
|
212 |
+
|
213 |
+
ner_df = pd.DataFrame(ner_result, columns=['entity', 'label', 'confidence'])
|
214 |
+
summ_result = summ_inference(article_txt)
|
215 |
+
|
216 |
+
ner_txt = get_ner_text(article_txt, ner_result)
|
217 |
+
|
218 |
+
st.markdown(f"<h4>SUMMARY:</h4>{summ_result}<h4>ENTITIES:</h4>", unsafe_allow_html=True)
|
219 |
+
st.markdown(f"{ner_txt}", unsafe_allow_html=True)
|
220 |
+
st.dataframe(ner_df, use_container_width=True)
|
221 |
+
|
222 |
+
############## ENTRY POINT END #######################
|
223 |
+
|
224 |
+
if __name__ == "__main__":
|
225 |
+
main()
|