Update app.py

app.py

@@ -1,15 +1,21 @@
-from sklearn.model_selection import train_test_split
 import torch
-from transformers import BertTokenizer, BertForSequenceClassification, Trainer, TrainingArguments
+from torch.utils.data import Dataset, DataLoader  # Import Dataset here
+from transformers import BertTokenizer, BertForSequenceClassification
+from torch.optim import AdamW
+from sklearn.model_selection import train_test_split
+from sklearn.metrics import classification_report
 import pandas as pd
+from tqdm import tqdm
 
-# Load the CSV file into a pandas DataFrame
-df = pd.read_csv('conversation_data.csv')
-
-# Ensure that the columns are named correctly
-print(df.head())
+# Initialize the tokenizer and model for BERT
+tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=3)  # Multi-label classification
 
+# Ensure the model is on the right device (GPU or CPU)
+device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+model.to(device)
 
+# Custom Dataset Class for Conversation Data
 class ConversationDataset(Dataset):
     def __init__(self, agent_utterances, customer_utterances, rule_followed, question_asked, question_answered):
         self.agent_utterances = agent_utterances
@@ -35,17 +41,10 @@ class ConversationDataset(Dataset):
         labels = torch.tensor([self.rule_followed[idx], self.question_asked[idx], self.question_answered[idx]], dtype=torch.long)
         return {**inputs, 'labels': labels}
 
+# Example of reading the data from CSV
+df = pd.read_csv('conversation_data.csv')
 
-# Initialize the tokenizer and model for BERT
-tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
-model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=3)  # Multi-label classification
-
-# Check if GPU is available
-device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
-model.to(device)
-
-
-# Extract agent utterances, customer utterances, and labels (rule_followed, question_asked, question_answered)
+# Extracting the agent and customer utterances along with labels
 agent_utterances = df['Agent Utterance'].tolist()
 customer_utterances = df['Customer Utterance'].tolist()
 rule_followed = df['Rule Followed'].tolist()
@@ -67,14 +66,17 @@ val_agent, val_customer = zip(*X_val)
 val_rule, val_question, val_answer = zip(*y_val)
 
 # Create dataset objects for training and validation
-train_dataset = ConversationDataset(train_agent, train_customer, train_rule, train_question, train_answer)
-val_dataset = ConversationDataset(val_agent, val_customer, val_rule, val_question, val_answer)
+train_dataset = ConversationDataset(
+    train_agent, train_customer, train_rule, train_question, train_answer
+)
+val_dataset = ConversationDataset(
+    val_agent, val_customer, val_rule, val_question, val_answer
+)
 
 # Create DataLoader for training and validation
 train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True)
 val_loader = DataLoader(val_dataset, batch_size=8, shuffle=False)
 
-
 # Set optimizer
 optimizer = AdamW(model.parameters(), lr=1e-5)
 
@@ -100,38 +102,41 @@ for epoch in range(3):  # Training for 3 epochs
         
         loop.set_postfix(loss=loss.item())
 
-
-
-# Evaluate the model
+# After training, evaluate the model on the validation set
 model.eval()
 all_preds = []
 all_labels = []
 
-with torch.no_grad():
+with torch.no_grad():  # Disable gradient calculation for inference
     for batch in val_loader:
+        # Move the batch to the device (GPU/CPU)
         batch = {key: value.squeeze(1).to(device) for key, value in batch.items()}
+
+        # Forward pass: Get logits
         outputs = model(**batch)
         
-        # Get predictions (we are predicting 3 labels)
+        # Get predictions (highest probability)
         logits = outputs.logits
         preds = torch.argmax(logits, dim=-1)
         
+        # Append predictions and true labels
         all_preds.append(preds.cpu().numpy())
         all_labels.append(batch['labels'].cpu().numpy())
 
-# Flatten the lists
+# Flatten lists of predictions and labels
 all_preds = [item for sublist in all_preds for item in sublist]
 all_labels = [item for sublist in all_labels for item in sublist]
 
-# Convert predictions and labels into multi-label format (for classification report)
+# Print classification report for each task
 print(classification_report(all_labels, all_preds, target_names=["Rule Followed", "Question Asked", "Question Answered"]))
 
 
 
+
 # Test the model with new data
 # Example single prediction
 test_agent_utterance = "What is your account number?"
-test_customer_utterance = "12345"
+test_customer_utterance = "888888"
 
 # Combine agent and customer utterance
 input_text = test_agent_utterance + " [SEP] " + test_customer_utterance