hs-knowledge/ner_app

1# Streamlit app to highlight NER entities2import random3import streamlit as st4from datasets import load_dataset5from annotated_text import annotated_text6 7 8# Show highlighted ner entities in a tweet9def display_ner(example):10    ner_output = example["ner_output"]11    chunks = []12    current_chunk = ""13    current_type = None14 15    # Check if there are two labels repeated16    previous_label = None17 18    for label in ner_output["labels"]:19        if (20            label21            and previous_label22            and previous_label == label23            and label != "O"24            and not label.startswith("I-")25            and not label.startswith("B-")26        ):27            pass28        previous_label = label29 30    for token, label in zip(ner_output["tokens"], ner_output["labels"]):31        if label is None:32            # Perhaps it is too long33            continue34        if label == "O":35            if current_type is not None:36                # Add previous entity37                chunks.append((current_chunk.strip(), current_type))38                current_chunk = token + " "39                current_type = None40            else:41                current_chunk += token + " "42                current_type = None43        elif label.startswith("B-"):44            if current_chunk:45                chunks.append((current_chunk.strip(), current_type))46            current_chunk = token + " "47            current_type = label[2:]48        elif label.startswith("I-"):49            current_chunk += token + " "50            current_type = label[2:]51        else:52            # It doesn't start with B- or I- => add single token53            if label != current_type:54                chunks.append((current_chunk.strip(), current_type))55                current_chunk = token + " "56                current_type = label57            else:58                current_chunk += token + " "59                current_type = label60 61    if current_chunk:62        chunks.append((current_chunk.strip(), current_type))63 64    # Display text65    chunks = [(c, t) if t is not None else c for c, t in chunks]66    annotated_text(*chunks)67 68 69def display_text(example, text_column):70    # Use annotated_text to show entities71    text = example[text_column]72 73    # Sort entities by start74    entities = sorted(example["entities"], key=lambda x: x["start"])75 76    for entity in entities:77        entity_text = entity["text"]78        # find in text79        start = text.find(entity_text)80        end = start + len(entity_text)81        entity["start"] = start82        entity["end"] = end83    # Chunk text84 85    if len(entities) == 0:86        annotated_text(*[text])87        return88 89    chunks = []90    last_index = 091    for i in range(len(entities)):92        entity = entities[i]93        start, end = entity["start"], entity["end"]94 95        if last_index < start:96            chunk_before_entity = text[last_index : entity["start"]]97            chunks.append((chunk_before_entity, None))98        chunks.append((entity["text"], entity["type"]))99 100        last_index = end101 102    if last_index < len(text):103        chunks.append((text[last_index:], None))104 105    # description = entity["kg_result"]["detailedDescription"]["articleBody"]106    chunks = [(c, t) if t is not None else c for c, t in chunks]107    annotated_text(*chunks)108 109 110# selectbox to choose dataset111 112selected_dataset = st.sidebar.selectbox(113    "Select dataset", ["hateval_enriched", "sbf-enriched", "hatecheck-enriched"]114)115 116# Load data117ds = load_dataset(f"hs-knowledge/{selected_dataset}")118 119text_column = {120    "hateval_enriched": "text",121    "sbf-enriched": "post",122    "hatecheck-enriched": "test_case",123}124 125elements = random.choices(range(len(ds["train"])), k=50)126ds["train"] = ds["train"].select(elements)127 128for ex in ds["train"]:129    # display_text(ex)130    st.markdown("---")131    display_ner(ex)132    with st.expander("Show entities"):133        for ent in ex["entities"]:134            entity_name = ent["text"]135            entity_type = ent["type"]136            entity_description = ent["kg_result"]["detailedDescription"]["articleBody"]137            st.write(f"{entity_name} ({entity_type}): {entity_description}")138