app.py

import sentencepiece
import streamlit as st
import pandas as pd
import spacy
import roner

example_list = [
    "Ana merge în București.",
    """Ana merge în București. Ana merge în București. Ana merge în București. Ana merge în București. Ana merge în București. Ana merge în București."""
]

st.set_page_config(layout="wide")

st.title("Demo for Romanian NER")

model_list = ['dumitrescustefan/bert-base-romanian-ner']

st.sidebar.header("Select NER Model")
model_checkpoint = st.sidebar.radio("", model_list)

st.sidebar.write("For details of models: 'https://huggingface.co/dumitrescustefan/")
st.sidebar.write("")

xlm_agg_strategy_info = "'aggregation_strategy' can be selected as 'simple' or 'none' for 'xlm-roberta' because of the RoBERTa model's tokenization approach."

st.sidebar.header("Select Aggregation Strategy Type")
if model_checkpoint == "akdeniz27/xlm-roberta-base-turkish-ner":
    aggregation = st.sidebar.radio("", ('simple', 'none'))
    st.sidebar.write(xlm_agg_strategy_info)
elif model_checkpoint == "xlm-roberta-large-finetuned-conll03-english":
    aggregation = st.sidebar.radio("", ('simple', 'none'))
    st.sidebar.write(xlm_agg_strategy_info)
    st.sidebar.write("")
    st.sidebar.write("This English NER model is included just to show the zero-shot transfer learning capability of XLM-Roberta.")
else:
    aggregation = st.sidebar.radio("", ('first', 'simple', 'average', 'max', 'none'))
    
st.sidebar.write("Please refer 'https://huggingface.co/transformers/_modules/transformers/pipelines/token_classification.html' for entity grouping with aggregation_strategy parameter.")

st.subheader("Select Text Input Method")
input_method = st.radio("", ('Select from Examples', 'Write or Paste New Text'))
if input_method == 'Select from Examples':
    selected_text = st.selectbox('Select Text from List', example_list, index=0, key=1)
    st.subheader("Text to Run")
    input_text = st.text_area("Selected Text", selected_text, height=128, max_chars=None, key=2)
elif input_method == "Write or Paste New Text":
    st.subheader("Text to Run")
    input_text = st.text_area('Write or Paste Text Below', value="", height=128, max_chars=None, key=2)

@st.cache(allow_output_mutation=True)
def setModel(named_persons_only):
    ner = roner.NER(named_persons_only=named_persons_only)
    return ner

@st.cache(allow_output_mutation=True)
def get_html(html: str):
    WRAPPER = """<div style="overflow-x: auto; border: 1px solid #e6e9ef; border-radius: 0.25rem; padding: 1rem; margin-bottom: 2.5rem">{}</div>"""
    html = html.replace("\n", " ")
    return WRAPPER.format(html)
        
Run_Button = st.button("Run", key=None)
if Run_Button == True:
    
    ner = setModel(named_persons_only = False)
    output = ner(input_text)[0]

    df = pd.DataFrame.from_dict(output)

    st.subheader("Recognized Entities")
    st.dataframe(df)

    """
    st.subheader("Spacy Style Display")
    spacy_display = {}
    spacy_display["ents"] = []
    spacy_display["text"] = input_text
    spacy_display["title"] = None

    for entity in output:
        if aggregation != "none":
            spacy_display["ents"].append({"start": entity["start"], "end": entity["end"], "label": entity["entity_group"]})
        else:
            spacy_display["ents"].append({"start": entity["start"], "end": entity["end"], "label": entity["entity"]})
    
    entity_list = ["PER", "LOC", "ORG", "MISC"]
    colors = {'PER': '#85DCDF', 'LOC': '#DF85DC', 'ORG': '#DCDF85', 'MISC': '#85ABDF',}
    html = spacy.displacy.render(spacy_display, style="ent", minify=True, manual=True, options={"ents": entity_list, "colors": colors})
    style = "<style>mark.entity { display: inline-block }</style>"
    st.write(f"{style}{get_html(html)}", unsafe_allow_html=True)
    """