app.py

import streamlit as st
import transformers
from annotated_text import annotated_text

ENTITY_TO_COLOR = {
    'PER': '#8ef',
    'LOC': '#faa',
    'ORG': '#afa',
    'MISC': '#fea',
}


@st.cache(allow_output_mutation=True, show_spinner=False)
def get_pipe(model_name):
    model = transformers.AutoModelForTokenClassification.from_pretrained(model_name)
    tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
    pipe = transformers.pipeline("token-classification",
                                 model=model,
                                 tokenizer=tokenizer,
                                 aggregation_strategy="simple")
    return pipe


def parse_text(text, prediction):
    start = 0
    parsed_text = []
    for p in prediction:
        parsed_text.append(text[start:p["start"]])
        parsed_text.append((p["word"], p["entity_group"], ENTITY_TO_COLOR[p["entity_group"]]))
        start = p["end"]
    parsed_text.append(text[start:])
    return parsed_text


st.set_page_config(page_title="Named Entity Recognition")
st.title("Named Entity Recognition")
st.write("Type text into the text box and then press 'Predict' to get the named entities.")

option = st.selectbox('Select model',
                      ("dslim/bert-base-NER",
                       "dslim/bert-large-NER",
                       "Davlan/bert-base-multilingual-cased-ner-hrl"))
default_text = """Mark Zuckerberg has unveiled a new VR headset, the Quest Pro, at an online event held for developers. With a price tag of $1,499, Quest Pro is almost four times the price of Meta's current headset, the Quest 2, which starts at $399. It boasts thinner lenses, a curved battery around the head strap at the back, and controllers which self-track. The headset also allows users to see their real environment around the periphery of the screen. Quest Pro has mixed reality capabilities, unlike its predecessor - meaning digital content can be viewed overlaid on the real world.
Meta boss Mark Zuckerberg said mixed reality was "the next major step for VR".
Gartner analyst Tuong Nguyen said the high launch price made it a device more suited to "high-end, enthusiast and potentially enterprise users" than mass-market. Microsoft CEO Satya Nadella also announced that the office platform Windows 365 will be available on it, as Meta seeks to position mixed and virtual reality as a work tool as well as a form of entertainment."""

text = st.text_area('Enter text here:', value=default_text)
st.write('Model used for prediction:', option)
submit = st.button('Predict')

with st.spinner("Loading model..."):
    pipe = get_pipe(model_name=option)

if (submit and len(text.strip()) > 0) or len(text.strip()) > 0:
    prediction = pipe(text)

    parsed_text = parse_text(text, prediction)

    st.header("Prediction:")
    annotated_text(*parsed_text)