Update app.py

app.py

@@ -10,10 +10,12 @@ from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay
 from sklearn.model_selection import train_test_split
 from sklearn.metrics import classification_report
 from transformers import AutoTokenizer, AutoModelForTokenClassification, AutoModelForSequenceClassification
+
 # Load BERT model and tokenizer via HuggingFace Transformers
 bert = XLMRobertaModel.from_pretrained('castorini/afriberta_large')
 tokenizer = XLMRobertaTokenizer.from_pretrained('castorini/afriberta_large')
-# Define the model architecture
+
+# Define the model architecture for three classes
 class BERT_Arch(nn.Module):
     def __init__(self, bert):
         super(BERT_Arch, self).__init__()
@@ -21,7 +23,7 @@ class BERT_Arch(nn.Module):
         self.dropout = nn.Dropout(0.1)          # Dropout layer
         self.relu = nn.ReLU()                   # ReLU activation function
         self.fc1 = nn.Linear(768, 512)          # Dense layer 1
-        self.fc2 = nn.Linear(512, 2)            # Dense layer 2 (Output layer)
+        self.fc2 = nn.Linear(512, 3)            # Dense layer 2 (Output layer for 3 classes)
         self.softmax = nn.LogSoftmax(dim=1)     # Softmax activation function
 
     def forward(self, sent_id, mask):           # Define the forward pass
@@ -39,19 +41,21 @@ fake_news_model_path = "Hate_Speech_model.pt"
 fake_news_model = torch.load(fake_news_model_path, map_location=torch.device('cpu'))
 fake_news_model.eval()
 
-# Function to detect fake news
+# Mapping labels to classes
+LABELS = {0: "Free", 1: "Hate", 2: "Offensive"}
+
+# Function to detect hate speech
 def detect_fake_news(text):
     inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=128)
     with torch.no_grad():
         outputs = fake_news_model(inputs['input_ids'], inputs['attention_mask'])
         label = torch.argmax(outputs, dim=1).item()
-    fake_news_result = "Hate" if label == 1 else "Real"
-    return fake_news_result
+    return LABELS[label]
 
 # Function to handle post logic
 def post_text(text, fake_news_result):
-    if fake_news_result == "Hate":
-        return "Your message contains Hate Speech and cannot be posted.", ""
+    if fake_news_result in ["Hate", "Offensive"]:
+        return f"Your message contains {fake_news_result} Speech and cannot be posted.", ""
     else:
         return "The text is safe to post.", text
 
@@ -64,7 +68,7 @@ with interface:
     with gr.Row():
         detect_fake_button = gr.Button("Detect Hate Speech")
     with gr.Row():
-        fake_news_result_box = gr.Textbox(label="Hate speech Detection Result", interactive=False)
+        fake_news_result_box = gr.Textbox(label="Hate Speech Detection Result", interactive=False)
     with gr.Row():
         post_button = gr.Button("Post Text")
     with gr.Row():