app.py

# Streamlit app to highlight NER entities
import random
import streamlit as st
from datasets import load_dataset
from annotated_text import annotated_text

# Load data
ds = load_dataset("hs-knowledge/hateval_enriched")

# Show highlighted ner entities in a tweet


def display_text(example):
    # Use annotated_text to show entities
    text = example["text"]

    # Sort entities by start
    entities = sorted(example["entities"], key=lambda x: x["start"])

    for entity in entities:
        entity_text = entity["text"]
        # find in text
        start = text.find(entity_text)
        end = start + len(entity_text)
        entity["start"] = start
        entity["end"] = end
    # Chunk text

    if len(entities) == 0:
        annotated_text(*[text])
        return

    chunks = []
    last_index = 0
    for i in range(len(entities)):
        entity = entities[i]
        start, end = entity["start"], entity["end"]

        if last_index < start:
            chunk_before_entity = text[last_index : entity["start"]]
            chunks.append((chunk_before_entity, None))
        chunks.append((entity["text"], entity["type"]))

        last_index = end

    if last_index < len(text):
        chunks.append((text[last_index:], None))

    # description = entity["kg_result"]["detailedDescription"]["articleBody"]
    chunks = [(c, t) if t is not None else c for c, t in chunks]
    annotated_text(*chunks)


# Get first 1000 examples

elements = random.choices(range(len(ds["train"])), k=50)
ds["train"] = ds["train"].select(elements)

for ex in ds["train"]:
    st.write("=" * 80)
    display_text(ex)
    with st.expander("Show entities"):
        for ent in ex["entities"]:
            entity_name = ent["text"]
            entity_type = ent["type"]
            entity_description = ent["kg_result"]["detailedDescription"]["articleBody"]
            st.write(f"{entity_name} ({entity_type}): {entity_description}")