added pipeline for classification with bert

app.py

@@ -2,6 +2,7 @@ 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')
 
@@ -11,6 +12,19 @@ for split in dataset.keys():
     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."
 
@@ -22,9 +36,18 @@ if st.button('Encontre exemplos!'):
     get_stopwords()
     df['preprocessed_review'] = df['review'].copy()
     df['preprocessed_review'] = df['preprocessed_review'].apply(preprocess_text)
-    cols = ['review','preprocessed_review','sentiment']    
-
-
+    
+    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])