import streamlit as st
import pandas as pd
from preprocess_data import preprocess_text,get_stopwords
from datasets import load_dataset

dataset = load_dataset('danielcd99/imdb')

dataframes = {}
for split in dataset.keys():
    # Convert the dataset split to a pandas DataFrame
    df = dataset[split].to_pandas()
    dataframes[split] = df

TITLE_TEXT = f"IMDB reviews"
DESCRIPTION_TEXT = f"Esta é uma aplicação para o trabalho de NLP. Utilizamos a base de dados de reviews do IMDb com 50.000 comentários entre positivos e negativos (a base está balanceada). Por meio desta interface é possível visualizar como os exemplos da nossa base de teste foram classificados com um BERT treinado para esta task."

st.title(TITLE_TEXT)
st.write(DESCRIPTION_TEXT)

if st.button('Encontre exemplos!'):
    df = df.sample(5)
    get_stopwords()
    df['preprocessed_review'] = df['review'].copy()
    df['preprocessed_review'] = df['preprocessed_review'].apply(preprocess_text)
    cols = ['review','preprocessed_review','sentiment']    



    st.table(df[cols])