Update pages/21_NLP_Transformer.py

pages/21_NLP_Transformer.py

@@ -1,90 +1,197 @@
+import pandas as pd
+from sklearn.model_selection import train_test_split
 import torch
-from transformers import BertTokenizer, BertForSequenceClassification, AdamW, get_scheduler
-from datasets import load_dataset
-from torch.utils.data import DataLoader
+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
 import streamlit as st
-import matplotlib.pyplot as plt
-from tqdm.auto import tqdm
 
-# Load pre-trained model and tokenizer from Hugging Face
-tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
-model = BertForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=2)
-device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
-model.to(device)
+# 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)
 
-# Streamlit interface
-st.title("Sentiment Analysis with BERT")
+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)
+
+class SentimentDataset(Dataset):
+    def __init__(self, dataframe, tokenizer, max_len):
+        self.tokenizer = tokenizer
+        self.data = dataframe
+        self.max_len = max_len
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, index):
+        review = str(self.data.iloc[index, 0])
+        label = self.data.iloc[index, 1]
+
+        encoding = self.tokenizer.encode_plus(
+            review,
+            add_special_tokens=True,
+            max_length=self.max_len,
+            return_token_type_ids=False,
+            pad_to_max_length=True,
+            return_attention_mask=True,
+            return_tensors='pt',
+        )
+
+        return {
+            'review_text': review,
+            'input_ids': encoding['input_ids'].flatten(),
+            'attention_mask': encoding['attention_mask'].flatten(),
+            'labels': torch.tensor(label, dtype=torch.long)
+        }
+
+def train_epoch(model, data_loader, loss_fn, optimizer, device, scheduler, n_examples):
+    model = model.train()
+    losses = []
+    correct_predictions = 0
+
+    for d in data_loader:
+        input_ids = d["input_ids"].to(device)
+        attention_mask = d["attention_mask"].to(device)
+        labels = d["labels"].to(device)
+
+        outputs = model(
+            input_ids=input_ids,
+            attention_mask=attention_mask
+        )
 
-# Training setup
-st.sidebar.title("Training Setup")
-num_epochs = st.sidebar.slider("Number of Epochs", 1, 5, 3)
-batch_size = st.sidebar.slider("Batch Size", 4, 32, 8)
-learning_rate = st.sidebar.slider("Learning Rate", 1e-6, 1e-3, 5e-5, format="%.6f")
+        loss = loss_fn(outputs.logits, labels)
+        correct_predictions += torch.sum(torch.argmax(outputs.logits, dim=1) == labels)
+        losses.append(loss.item())
 
-# Define a custom hash function for AddedToken type
-@st.cache_data(hash_funcs={tokenizer.__class__: id})
-def load_and_preprocess_data():
-    dataset = load_dataset("imdb", split="train[:1%]")
-    def preprocess_function(examples):
-        return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=128)
-    encoded_dataset = dataset.map(preprocess_function, batched=True)
-    encoded_dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'label'])
-    encoded_dataset = encoded_dataset.rename_column("label", "labels")  # Rename the column to 'labels'
-    return DataLoader(encoded_dataset, shuffle=True, batch_size=batch_size)
+        loss.backward()
+        optimizer.step()
+        scheduler.step()
+        optimizer.zero_grad()
 
-train_dataloader = load_and_preprocess_data()
+    return correct_predictions.double() / n_examples, np.mean(losses)
 
-# Initialize training status
-training_completed = st.sidebar.empty()
+def eval_model(model, data_loader, loss_fn, device, n_examples):
+    model = model.eval()
+    losses = []
+    correct_predictions = 0
+
+    with torch.no_grad():
+        for d in data_loader:
+            input_ids = d["input_ids"].to(device)
+            attention_mask = d["attention_mask"].to(device)
+            labels = d["labels"].to(device)
+
+            outputs = model(
+                input_ids=input_ids,
+                attention_mask=attention_mask
+            )
+
+            loss = loss_fn(outputs.logits, labels)
+            correct_predictions += torch.sum(torch.argmax(outputs.logits, dim=1) == labels)
+            losses.append(loss.item())
+
+    return correct_predictions.double() / n_examples, np.mean(losses)
+
+def create_data_loader(df, tokenizer, max_len, batch_size):
+    ds = SentimentDataset(
+        dataframe=df,
+        tokenizer=tokenizer,
+        max_len=max_len
+    )
+
+    return DataLoader(
+        ds,
+        batch_size=batch_size,
+        num_workers=4
+    )
+
+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
+
+train_data_loader = create_data_loader(train_df, tokenizer, MAX_LEN, BATCH_SIZE)
+test_data_loader = create_data_loader(test_df, tokenizer, MAX_LEN, BATCH_SIZE)
+
+EPOCHS = 2
+optimizer = AdamW(model.parameters(), lr=2e-5, correct_bias=False)
+total_steps = len(train_data_loader) * EPOCHS
+scheduler = get_linear_schedule_with_warmup(
+    optimizer,
+    num_warmup_steps=0,
+    num_training_steps=total_steps
+)
+loss_fn = torch.nn.CrossEntropyLoss().to(device)
+model = model.to(device)
 
 # Training loop
-if st.sidebar.button("Train"):
-    optimizer = AdamW(model.parameters(), lr=learning_rate)
-    num_training_steps = num_epochs * len(train_dataloader)
-    lr_scheduler = get_scheduler(
-        name="linear", optimizer=optimizer, num_warmup_steps=0, num_training_steps=num_training_steps
+for epoch in range(EPOCHS):
+    train_acc, train_loss = train_epoch(
+        model,
+        train_data_loader,
+        loss_fn,
+        optimizer,
+        device,
+        scheduler,
+        len(train_df)
     )
 
-    progress_bar = tqdm(range(num_training_steps))
-    loss_values = []
-
-    model.train()
-    for epoch in range(num_epochs):
-        for batch in train_dataloader:
-            batch = {k: v.to(device) for k, v in batch.items()}
-            outputs = model(**batch)
-            loss = outputs.loss
-            loss.backward()
-
-            optimizer.step()
-            lr_scheduler.step()
-            optimizer.zero_grad()
-            progress_bar.update(1)
-            loss_values.append(loss.item())
-
-    training_completed.success("Training completed")
-
-    # Plot loss values
-    st.write("### Training Loss")
-    plt.figure(figsize=(10, 6))
-    plt.plot(loss_values, label="Training Loss")
-    plt.xlabel("Training Steps")
-    plt.ylabel("Loss")
-    plt.legend()
-    st.pyplot(plt)
-
-# Text input for prediction
-st.write("### Predict Sentiment")
-user_input = st.text_area("Enter text:", "I loved this movie!")
-
-if user_input:
-    inputs = tokenizer(user_input, padding="max_length", truncation=True, max_length=128, return_tensors="pt")
-    inputs = {k: v.to(device) for k, v in inputs.items()}
-    
-    model.eval()
+    print(f'Epoch {epoch + 1}/{EPOCHS}')
+    print(f'Train loss {train_loss} accuracy {train_acc}')
+
+    val_acc, val_loss = eval_model(
+        model,
+        test_data_loader,
+        loss_fn,
+        device,
+        len(test_df)
+    )
+
+    print(f'Val loss {val_loss} accuracy {val_acc}')
+
+# Save the model
+model.save_pretrained('bert-sentiment-model')
+tokenizer.save_pretrained('bert-sentiment-model')
+
+# Streamlit app
+model = BertForSequenceClassification.from_pretrained('bert-sentiment-model')
+tokenizer = BertTokenizer.from_pretrained('bert-sentiment-model')
+model = model.eval()
+
+def predict_sentiment(text):
+    encoding = tokenizer.encode_plus(
+        text,
+        add_special_tokens=True,
+        max_length=128,
+        return_token_type_ids=False,
+        pad_to_max_length=True,
+        return_attention_mask=True,
+        return_tensors='pt',
+    )
+    input_ids = encoding['input_ids']
+    attention_mask = encoding['attention_mask']
+
     with torch.no_grad():
-        outputs = model(**inputs)
-        prediction = outputs.logits.argmax(dim=-1).item()
-        sentiment = "Positive" if prediction == 1 else "Negative"
-        
-    st.write(f"Sentiment: **{sentiment}**")
+        outputs = model(input_ids, attention_mask=attention_mask)
+        probabilities = torch.nn.functional.softmax(outputs.logits, dim=1)
+        predicted_class = torch.argmax(probabilities, dim=1).item()
+    
+    return 'positive' if predicted_class == 1 else 'negative'
+
+st.title("Sentiment Analysis with BERT")
+user_input = st.text_area("Enter a movie review:")
+
+if st.button("Analyze"):
+    sentiment = predict_sentiment(user_input)
+    st.write(f'The sentiment of the review is: **{sentiment}**')