Update pages/21_NLP_Transformer.py

pages/21_NLP_Transformer.py

@@ -1,22 +1,19 @@
-import pandas as pd
-from sklearn.model_selection import train_test_split
 import torch
 from torch.utils.data import DataLoader, Dataset
 from transformers import BertTokenizer, BertForSequenceClassification, AdamW
 from transformers import get_linear_schedule_with_warmup
 import numpy as np
-from sklearn.metrics import accuracy_score, classification_report
+from datasets import load_dataset
 import streamlit as st
 
-# Load and preprocess the IMDb dataset
-data_url = "https://ai.stanford.edu/~amaas/data/sentiment/aclImdb_v1.tar.gz"
-df = pd.read_csv(data_url)
+# Load IMDb dataset
+dataset = load_dataset('imdb')
+train_df = dataset['train'].to_pandas()
+test_df = dataset['test'].to_pandas()
 
-df['label'] = df['sentiment'].map({'positive': 1, 'negative': 0})
-train_df, test_df = train_test_split(df, test_size=0.2, random_state=42)
-
-train_df.to_csv('train.csv', index=False)
-test_df.to_csv('test.csv', index=False)
+# Preprocess the data
+train_df = train_df[['text', 'label']]
+test_df = test_df[['text', 'label']]
 
 class SentimentDataset(Dataset):
     def __init__(self, dataframe, tokenizer, max_len):
@@ -113,10 +110,6 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
 model = BertForSequenceClassification.from_pretrained('bert-base-uncased')
 
-# Load data
-train_df = pd.read_csv('train.csv')
-test_df = pd.read_csv('test.csv')
-
 # Create data loaders
 BATCH_SIZE = 16
 MAX_LEN = 128