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

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


MODEL_PATH = 'danielcd99/BERT_imdb'

def load_pipeline():
    pipe=pipeline(
    "text-classification",
    model=MODEL_PATH
    )
    return pipe

pipe = load_pipeline()


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)
    
    predictions = []
    for review in df['preprocessed_reviews']:
        label = pipeline(review)[0]['label']
        if label == 'LABEL_0':
            predictions.append('Negative')
        else:
            predictions.append('Positive')

    df['predictions'] = predictions

    cols = ['review','sentiment', 'predictions']    

    st.table(df[cols])