Demo.py

import streamlit as st
import sparknlp
import os
import pandas as pd

from sparknlp.base import *
from sparknlp.annotator import *
from pyspark.ml import Pipeline
from sparknlp.pretrained import PretrainedPipeline
from annotated_text import annotated_text

# Page configuration
st.set_page_config(
    layout="wide", 
    initial_sidebar_state="auto"
)

# CSS for styling
st.markdown("""

    <style>

        .main-title {

            font-size: 36px;

            color: #4A90E2;

            font-weight: bold;

            text-align: center;

        }

        .section p, .section ul {

            color: #666666;

        }

    </style>

""", unsafe_allow_html=True)

@st.cache_resource
def init_spark():
    return sparknlp.start()

@st.cache_resource
def create_pipeline(model):
    document_assembler = DocumentAssembler() \
        .setInputCol("text") \
        .setOutputCol("document")

    tokenizer = Tokenizer() \
        .setInputCols(["document"]) \
        .setOutputCol("token")

    if model == "ner_wikiner_glove_840B_300":
        embeddings = WordEmbeddingsModel.pretrained('glove_840B_300', lang='xx') \
            .setInputCols(['document', 'token']) \
            .setOutputCol('embeddings')
    elif model == "wikiner_6B_300":
        embeddings = WordEmbeddingsModel.pretrained('glove_6B_300', lang='xx') \
            .setInputCols(['document', 'token']) \
            .setOutputCol('embeddings')
    elif model == "wikiner_6B_100":
        embeddings = WordEmbeddingsModel.pretrained('glove_100d') \
            .setInputCols(['document', 'token']) \
            .setOutputCol('embeddings')

    lang = 'xx' if model == "ner_wikiner_glove_840B_300" else 'pt'
    ner_model = NerDLModel.pretrained(model, lang) \
        .setInputCols(["document", "token", "embeddings"]) \
        .setOutputCol("ner")

    ner_converter = NerConverter() \
        .setInputCols(["document", "token", "ner"]) \
        .setOutputCol("ner_chunk")

    pipeline = Pipeline(stages=[
        document_assembler,
        tokenizer,
        embeddings,
        ner_model,
        ner_converter
    ])

    return pipeline

def fit_data(pipeline, data):
  empty_df = spark.createDataFrame([['']]).toDF('text')
  pipeline_model = pipeline.fit(empty_df)
  model = LightPipeline(pipeline_model)
  result = model.fullAnnotate(data)
  return result

def annotate(data):
    document, chunks, labels = data["Document"], data["NER Chunk"], data["NER Label"]
    annotated_words = []
    for chunk, label in zip(chunks, labels):
        parts = document.split(chunk, 1)
        if parts[0]:
            annotated_words.append(parts[0])
        annotated_words.append((chunk, label))
        document = parts[1]
    if document:
        annotated_words.append(document)
    annotated_text(*annotated_words)

# Set up the page layout
st.markdown('<div class="main-title">Reconhecer pessoas, lugares, organizações e outras entidades</div>', unsafe_allow_html=True)

# Sidebar content
model = st.sidebar.selectbox(
    "Choose the pretrained model",
    ["ner_wikiner_glove_840B_300", "wikiner_6B_100", "wikiner_6B_300"],
    help="For more info about the models visit: https://sparknlp.org/models"
)

# Reference notebook link in sidebar
link = """

<a href="https://colab.research.google.com/github/JohnSnowLabs/spark-nlp-workshop/blob/master/tutorials/streamlit_notebooks/NER_PT.ipynb">

    <img src="https://colab.research.google.com/assets/colab-badge.svg" style="zoom: 1.3" alt="Open In Colab"/>

</a>

"""
st.sidebar.markdown('Reference notebook:')
st.sidebar.markdown(link, unsafe_allow_html=True)

# Load examples
examples = [
    "William Henry Gates III (nascido em 28 de outubro de 1955) é um magnata dos negócios americano, desenvolvedor de software, investidor e filantropo. Ele é mais conhecido como co-fundador da Microsoft Corporation. Durante sua carreira na Microsoft, Gates ocupou os cargos de presidente, diretor executivo (CEO), presidente e arquiteto-chefe de software, sendo também o maior acionista individual até maio de 2014. Ele é um dos empreendedores mais conhecidos e pioneiros da revolução dos microcomputadores das décadas de 1970 e 1980. Nascido e criado em Seattle, Washington, Gates co-fundou a Microsoft com seu amigo de infância Paul Allen em 1975, em Albuquerque, Novo México; tornou-se a maior empresa de software para computadores pessoais do mundo. Gates liderou a empresa como presidente e CEO até renunciar como CEO em janeiro de 2000, mas permaneceu como presidente e tornou-se arquiteto-chefe de software. No final dos anos 1990, Gates foi criticado por suas táticas comerciais, consideradas anticompetitivas. Esta opinião foi confirmada por várias decisões judiciais. Em junho de 2006, Gates anunciou que passaria para um papel de meio período na Microsoft e trabalho em tempo integral na Fundação Bill & Melinda Gates, a fundação de caridade privada que ele e sua esposa Melinda Gates estabeleceram em 2000. Ele transferiu gradualmente suas responsabilidades para Ray Ozzie e Craig Mundie. Ele renunciou como presidente da Microsoft em fevereiro de 2014 e assumiu um novo cargo como conselheiro de tecnologia para apoiar o recém-nomeado CEO Satya Nadella.",
    "A Mona Lisa é uma pintura a óleo do século XVI criada por Leonardo. Está exposta no Louvre, em Paris.",
    "Quando Sebastian Thrun começou a trabalhar em carros autônomos na Google em 2007, poucas pessoas fora da empresa o levaram a sério. 'Posso dizer que CEOs muito importantes de grandes empresas automobilísticas americanas apertavam minha mão e se afastavam porque não valia a pena conversar comigo', disse Thrun, agora co-fundador e CEO da startup de educação superior online Udacity, em uma entrevista com o Recode no início desta semana.",
    "O Facebook é um serviço de rede social lançado como TheFacebook em 4 de fevereiro de 2004. Foi fundado por Mark Zuckerberg com seus colegas de quarto e colegas de Harvard Eduardo Saverin, Andrew McCollum, Dustin Moskovitz e Chris Hughes. A adesão ao site inicialmente era limitada aos estudantes de Harvard, mas foi expandida para outras faculdades na área de Boston, a Ivy League e gradualmente para a maioria das universidades nos Estados Unidos e Canadá.",
    "A história do processamento de linguagem natural geralmente começou na década de 1950, embora o trabalho possa ser encontrado em períodos anteriores. Em 1950, Alan Turing publicou um artigo intitulado 'Computing Machinery and Intelligence', que propunha o que agora é chamado de teste de Turing como critério de inteligência.",
    "Geoffrey Everest Hinton é um psicólogo cognitivo e cientista da computação inglês-canadense, mais conhecido por seu trabalho em redes neurais artificiais. Desde 2013, ele divide seu tempo trabalhando para o Google e a Universidade de Toronto. Em 2017, ele co-fundou e tornou-se o Conselheiro Científico Chefe do Instituto Vector em Toronto.",
    "Quando eu disse a John que queria me mudar para o Alasca, ele me avisou que eu teria dificuldade em encontrar um Starbucks lá.",
    "Steven Paul Jobs foi um magnata dos negócios americano, designer industrial, investidor e proprietário de mídia. Ele foi presidente, diretor executivo (CEO) e co-fundador da Apple Inc., presidente e acionista majoritário da Pixar, membro do conselho administrativo da The Walt Disney Company após sua aquisição da Pixar, e fundador, presidente e CEO da NeXT. Jobs é amplamente reconhecido como um pioneiro da revolução do computador pessoal das décadas de 1970 e 1980, junto com o co-fundador da Apple, Steve Wozniak. Jobs nasceu em São Francisco, Califórnia, e foi colocado para adoção. Ele cresceu na área da Baía de São Francisco. Ele frequentou a Reed College em 1972 antes de abandonar o mesmo ano e viajar pela Índia em 1974 em busca de iluminação e estudar o budismo zen.",
    "Titanic é um filme épico americano de romance e desastre de 1997 dirigido, escrito, co-produzido e co-editado por James Cameron. Incorporando aspectos tanto históricos quanto ficcionais, é baseado em relatos do naufrágio do RMS Titanic, e estrelado por Leonardo DiCaprio e Kate Winslet como membros de diferentes classes sociais que se apaixonam a bordo do navio durante sua malfadada viagem inaugural.",
    "Além de ser o rei do norte, John Snow é um médico inglês e um líder no desenvolvimento da anestesia e higiene médica. Ele é considerado o primeiro a usar dados para curar o surto de cólera em 1834."
]

selected_text = st.selectbox("Select an example", examples)
custom_input = st.text_input("Try it with your own Sentence!")

text_to_analyze = custom_input if custom_input else selected_text

st.subheader('Full example text')
HTML_WRAPPER = """<div class="scroll entities" style="overflow-x: auto; border: 1px solid #e6e9ef; border-radius: 0.25rem; padding: 1rem; margin-bottom: 2.5rem; white-space:pre-wrap">{}</div>"""
st.markdown(HTML_WRAPPER.format(text_to_analyze), unsafe_allow_html=True)

# Initialize Spark and create pipeline
spark = init_spark()
pipeline = create_pipeline(model)
output = fit_data(pipeline, text_to_analyze)

# Display matched sentence
st.subheader("Processed output:")

results = {
    'Document': output[0]['document'][0].result,
    'NER Chunk': [n.result for n in output[0]['ner_chunk']],
    "NER Label": [n.metadata['entity'] for n in output[0]['ner_chunk']]
}

annotate(results)

with st.expander("View DataFrame"):
    df = pd.DataFrame({'NER Chunk': results['NER Chunk'], 'NER Label': results['NER Label']})
    df.index += 1
    st.dataframe(df)