Update app.py

app.py

@@ -1,10 +1,9 @@
-import torch
 from sklearn.model_selection import train_test_split
+import torch
 from transformers import BertTokenizer, BertForSequenceClassification, Trainer, TrainingArguments
 
-# Load multilingual BERT tokenizer and model
+# Load multilingual BERT tokenizer
 tokenizer = BertTokenizer.from_pretrained("bert-base-multilingual-cased")
-model = BertForSequenceClassification.from_pretrained("bert-base-multilingual-cased", num_labels=2)
 
 # Example dataset in Hindi
 dataset = [
@@ -20,10 +19,13 @@ dataset = [
 # Split dataset into training and evaluation sets
 train_data, eval_data = train_test_split(dataset, test_size=0.2)
 
-# Tokenization
+# Tokenizer function that also keeps the label in the dataset
 def tokenize_function(example):
-    return tokenizer(example['customer_input'], example['agent_response'], padding='max_length', truncation=True, max_length=512)
+    tokenized_example = tokenizer(example['customer_input'], example['agent_response'], padding='max_length', truncation=True, max_length=512)
+    tokenized_example['label'] = 1 if example['label'] == 'compliant' else 0  # Convert 'compliant' to 1 and 'non-compliant' to 0
+    return tokenized_example
 
+# Apply tokenization to the entire dataset
 train_data = [tokenize_function(x) for x in train_data]
 eval_data = [tokenize_function(x) for x in eval_data]
 
@@ -31,7 +33,7 @@ eval_data = [tokenize_function(x) for x in eval_data]
 class DialogueDataset(torch.utils.data.Dataset):
     def __init__(self, data):
         self.data = data
-        self.labels = [1 if item["label"] == "compliant" else 0 for item in data]
+        self.labels = [item['label'] for item in data]
 
     def __len__(self):
         return len(self.data)
@@ -40,19 +42,25 @@ class DialogueDataset(torch.utils.data.Dataset):
         item = self.data[idx]
         input_ids = torch.tensor(item['input_ids'])
         attention_mask = torch.tensor(item['attention_mask'])
-        label = torch.tensor(self.labels[idx])
+        label = torch.tensor(item['label'])
         return {"input_ids": input_ids, "attention_mask": attention_mask, "labels": label}
 
+
+# Create PyTorch datasets
 train_dataset = DialogueDataset(train_data)
 eval_dataset = DialogueDataset(eval_data)
 
+
+# Load multilingual BERT model for sequence classification
+model = BertForSequenceClassification.from_pretrained("bert-base-multilingual-cased", num_labels=2)
+
 # Training arguments
 training_args = TrainingArguments(
     output_dir="./results",
     evaluation_strategy="epoch",  # Evaluate every epoch
     per_device_train_batch_size=8,
     per_device_eval_batch_size=8,
-    num_train_epochs=2,
+    num_train_epochs=3,
     weight_decay=0.01,
     logging_dir='./logs',
 )
@@ -72,7 +80,7 @@ trainer.train()
 eval_results = trainer.evaluate()
 print("Evaluation Results:", eval_results)
 
-# Inference function
+
 def check_compliance(customer_input, agent_response):
     inputs = tokenizer(customer_input, agent_response, return_tensors="pt", padding=True, truncation=True, max_length=512)
     with torch.no_grad():
@@ -89,4 +97,4 @@ def check_compliance(customer_input, agent_response):
 test_customer_input = "मेरे पास अकाउंट एक्सेस नहीं हो रहा है। क्या आप मेरी मदद कर सकते हैं?"
 test_agent_response = "मुझे खेद है। कृपया अपना उपयोगकर्ता नाम साझा करें, ताकि मैं आपकी सहायता कर सकूं।"
 result = check_compliance(test_customer_input, test_agent_response)
-print(result)
+print(result)