EverstreamAnalyticsAI/POS

1# code is adopted from https://huggingface.co/spaces/flair/model_demo/blob/main/app.py and modfied to fit the purpose of this project.2import spacy.displacy3import streamlit as st4from flair.nn import Classifier5from flair.splitter import SegtokSentenceSplitter6from colorhash import ColorHash7 8# st.title("Flair NER Demo")9st.set_page_config(layout="centered")10 11# Block 2: Users can input text12st.subheader("Input your text here")13input_text = st.text_area('Write or Paste Text Below',14                          value='May visited the Eiffel Tower in Paris last May.\n\n'15                                'There she ran across a sign in German that read: "Dirk liebt den Eiffelturm"',16                          height=128,17                          max_chars=None,18                          label_visibility="collapsed")19 20 21@st.cache_resource22def get_model():23    return Classifier.load('pos')24 25 26def get_html(html: str):27    WRAPPER = """<div style="overflow-x: auto; border: 1px solid #e6e9ef; border-radius: 0.25rem; padding: 1rem; margin-bottom: 2.5rem">{}</div>"""28    html = html.replace("\n", " ")29    return WRAPPER.format(html)30 31 32def color_variant(hex_color, brightness_offset=1):33    """ takes a color like #87c95f and produces a lighter or darker variant34    taken from: https://chase-seibert.github.io/blog/2011/07/29/python-calculate-lighterdarker-rgb-colors.html35    """36    if len(hex_color) != 7:37        raise Exception("Passed %s into color_variant(), needs to be in #87c95f format." % hex_color)38    rgb_hex = [hex_color[x:x + 2] for x in [1, 3, 5]]39    new_rgb_int = [int(hex_value, 16) + brightness_offset for hex_value in rgb_hex]40    new_rgb_int = [min([255, max([0, i])]) for i in new_rgb_int]  # make sure new values are between 0 and 25541    # hex() produces "0x88", we want just "88"42    return "#" + "".join([hex(i)[2:] for i in new_rgb_int])43 44 45# Block 3: Output is displayed46button_clicked = st.button("**Click here** to tag the input text", key=None)47 48if button_clicked:49    splitter = SegtokSentenceSplitter()50    # TODO: perhaps truncate input_text51    sentences = splitter.split(input_text)52 53    # get the model and predict54    model = get_model()55    model.predict(sentences)56 57    spacy_display = {"ents": [], "text": input_text, "title": None}58 59    predicted_labels = set()60    for sentence in sentences:61        for prediction in sentence.get_labels():62            entity_fields = {63                "start": prediction.data_point.start_position + sentence.start_position,64                "end": prediction.data_point.end_position + sentence.start_position,65                "label": prediction.value,66            }67 68            spacy_display["ents"].append(entity_fields)69            predicted_labels.add(entity_fields["label"])70 71    # create colors for each label72    colors = {}73    for label in predicted_labels:74        colors[label] = color_variant(ColorHash(label).hex, brightness_offset=85)75 76    # use displacy to render77    html = spacy.displacy.render(spacy_display,78                                 style="ent",79                                 minify=True,80                                 manual=True,81                                 options={82                                     "colors": colors,83                                 },84                                 )85    style = "<style>mark.entity { display: inline-block }</style>"86    st.subheader("Tagged text")87    st.write(f"{style}{get_html(html)}", unsafe_allow_html=True)