|
"""Demo for NER4OPT.""" |
|
import re |
|
import warnings |
|
|
|
warnings.filterwarnings("ignore", category=DeprecationWarning) |
|
|
|
import spacy |
|
from spacy import displacy |
|
from spacy.training import iob_to_biluo, biluo_tags_to_offsets |
|
from spacy.tokenizer import Tokenizer |
|
|
|
import streamlit as st |
|
|
|
from simpletransformers.ner import NERModel, NERArgs |
|
|
|
HTML_WRAPPER = """<div style="overflow-x: auto; border: 1px solid #e6e9ef; border-radius: 0.25rem; padding: 1rem; margin-bottom: 2.5rem">{}</div>""" |
|
|
|
|
|
@st.cache_resource |
|
def load_models(): |
|
"""Load custom built NER4OPT model.""" |
|
custom_labels = [ |
|
'O', |
|
'B-CONST_DIR', |
|
'I-CONST_DIR', |
|
'B-LIMIT', |
|
'I-LIMIT', |
|
'B-VAR', |
|
'I-VAR', |
|
'B-OBJ_DIR', |
|
'B-OBJ_NAME', |
|
'I-OBJ_NAME', |
|
'B-PARAM', |
|
'I-PARAM', |
|
] |
|
|
|
model_args = NERArgs() |
|
model_args.use_early_stopping = True |
|
model_args.early_stopping_delta = 0.01 |
|
model_args.early_stopping_metric = "eval_loss" |
|
model_args.early_stopping_metric_minimize = False |
|
model_args.early_stopping_patience = 5 |
|
model_args.evaluate_during_training_steps = 2000 |
|
model_args.overwrite_output_dir = True |
|
model_args.reprocess_input_data = True |
|
model_args.num_train_epochs = 11 |
|
model_args.adafactor_beta1 = 0.9 |
|
model_args.weight_decay = 0.01 |
|
model_args.max_seq_length = 512 |
|
model_args.learning_rate = 4e-5 |
|
model_args.train_batch_size = 1 |
|
model_args.eval_batch_size = 1 |
|
model_args.manual_seed = 123456789 |
|
model_args.output_dir = "trained_transformer_model" |
|
model_args.use_cuda = True |
|
model_args.use_multiprocessing = False |
|
model = NERModel("roberta", |
|
"skadio/ner4opt-roberta", |
|
labels=custom_labels, |
|
use_cuda=True, |
|
args=model_args) |
|
spacy_model = spacy.load("en_core_web_sm") |
|
spacy_model.tokenizer = Tokenizer(spacy_model.vocab, |
|
token_match=re.compile(r'\S+').match) |
|
spacy_blank_model = spacy.blank('en') |
|
spacy_blank_model.tokenizer = Tokenizer( |
|
spacy_blank_model.vocab, token_match=re.compile(r'\S+').match) |
|
return model, spacy_model, spacy_blank_model |
|
|
|
|
|
def main(): |
|
|
|
st.sidebar.title(""" |
|
NER4OPT Demo: \nFull code will be available at https://github.com/skadio/Ner4Opt |
|
""") |
|
|
|
text = st.text_area( |
|
"Text", |
|
"Cautious Asset Investment has a total of $ 150,000 to manage and decides to invest it in money market fund , which yields a 2 % return as well as in foreign bonds , which gives and average rate of return of 10.2 % . Internal policies require PAI to diversify the asset allocation so that the minimum investment in money market fund is 40 % of the total investment . Due to the risk of default of foreign countries , no more than 40 % of the total investment should be allocated to foreign bonds . How much should the Cautious Asset Investment allocate in each asset so as to maximize its average return ?" |
|
) |
|
if text == "": |
|
st.write("Please write a valid sentence.") |
|
model, spacy_model, spacy_blank_model = load_models() |
|
|
|
|
|
spacy_doc = spacy_model(text) |
|
if len(list(spacy_doc.sents)) >= 2: |
|
last_two_sentences = ' '.join( |
|
[item.text for item in list(spacy_doc.sents)[-2::]]) |
|
else: |
|
last_two_sentences = ' '.join( |
|
[item.text for item in list(spacy_doc.sents)[-1::]]) |
|
to_skip_count = len(last_two_sentences.split()) |
|
augmented_sent = last_two_sentences + " " + text |
|
|
|
if st.button("Get Named Entities"): |
|
predictions, raw_outputs = model.predict([augmented_sent], |
|
split_on_space=True) |
|
transformer_predictions = [ |
|
list(val.values())[0] for val in predictions[0] |
|
] |
|
transformer_predictions = transformer_predictions[to_skip_count::] |
|
biluo_tags = iob_to_biluo(transformer_predictions) |
|
doc = spacy_blank_model.make_doc(text) |
|
entities = biluo_tags_to_offsets(doc, biluo_tags) |
|
entities_formatted = [] |
|
for tag in entities: |
|
entities_formatted.append({ |
|
"start": tag[0], |
|
"end": tag[1], |
|
"label": tag[2], |
|
"score": 1.0 |
|
}) |
|
ner_for_display = [{ |
|
"text": doc.text, |
|
"ents": entities_formatted, |
|
"title": None |
|
}] |
|
st.title("Named Entity Results") |
|
html_ner = displacy.render(ner_for_display, style="ent", manual=True) |
|
html_ner = html_ner.replace("\n", " ") |
|
st.write(HTML_WRAPPER.format(html_ner), unsafe_allow_html=True) |
|
|
|
|
|
if __name__ == '__main__': |
|
main() |
|
|