Update app.py

app.py

@@ -1,61 +1,75 @@
 import os
 os.system("pip install torch transformers gradio matplotlib")
 
+# Install required packages
+# !pip install torch transformers gradio matplotlib
 
 import torch
 import gradio as gr
-import numpy as np
 import matplotlib.pyplot as plt
-import pandas as pd
+import numpy as np
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
-torch.set_num_threads(torch.get_num_threads())
+# Load model and tokenizer from Hugging Face Hub
+model_name = "HyperX-Sentience/RogueBERT-Toxicity-85K"
+model = AutoModelForSequenceClassification.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-# Load the trained model and tokenizer from Hugging Face Hub
-model_path = "HyperX-Sentience/RogueBERT-Toxicity-85K"
-model = AutoModelForSequenceClassification.from_pretrained(model_path)
-tokenizer = AutoTokenizer.from_pretrained(model_path)
-
-# Move the model to CUDA if available
+# Move model to CUDA if available
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 
-# Define toxicity labels
+# Toxicity category labels
 labels = ["toxic", "severe_toxic", "obscene", "threat", "insult", "identity_hate"]
 
+# Function to predict toxicity
 def predict_toxicity(comment):
-    """Predicts the toxicity levels of a given comment."""
-    inputs = tokenizer(comment, truncation=True, padding="max_length", max_length=128, return_tensors="pt")
+    inputs = tokenizer([comment], truncation=True, padding="max_length", max_length=128, return_tensors="pt")
     inputs = {key: val.to(device) for key, val in inputs.items()}
     
     with torch.no_grad():
         outputs = model(**inputs)
-        probabilities = torch.sigmoid(outputs.logits).cpu().numpy()[0]
+        logits = outputs.logits
+        probabilities = torch.sigmoid(logits).cpu().numpy()[0]
     
-    return {labels[i]: float(probabilities[i]) for i in range(len(labels))}
+    toxicity_scores = {label: float(probabilities[i]) for i, label in enumerate(labels)}
+    return toxicity_scores
 
-def format_toxicity_data(comment):
-    """Formats the toxicity scores for a modern bar graph."""
-    scores = predict_toxicity(comment)
-    df = pd.DataFrame({"Category": list(scores.keys()), "Score": list(scores.values())})
-    return df
+# Function to create a bar chart
+def plot_toxicity(comment):
+    toxicity_scores = predict_toxicity(comment)
+    categories = list(toxicity_scores.keys())
+    scores = list(toxicity_scores.values())
+    
+    plt.figure(figsize=(8, 5), facecolor='black')
+    ax = plt.gca()
+    ax.set_facecolor('black')
+    bars = plt.bar(categories, scores, color='#20B2AA', edgecolor='white')  # Sea green
+    
+    plt.xticks(color='white', fontsize=12)
+    plt.yticks(color='white', fontsize=12)
+    plt.title("Toxicity Score Analysis", color='white', fontsize=14)
+    plt.ylim(0, 1)
+    
+    for bar in bars:
+        yval = bar.get_height()
+        plt.text(bar.get_x() + bar.get_width()/2, yval + 0.02, f'{yval:.2f}', ha='center', color='white', fontsize=10)
+    
+    plt.tight_layout()
+    plt.savefig("toxicity_chart.png", facecolor='black')
+    plt.close()
+    
+    return "toxicity_chart.png"
 
-# Gradio interface
+# Gradio UI
 demo = gr.Interface(
-    fn=format_toxicity_data,
-    inputs=gr.Textbox(label="Enter a comment:"),
-    outputs=gr.BarPlot(
-        value=None,
-        x="Category",
-        y="Score",
-        title="Toxicity Analysis",
-        y_lim=[0, 1],
-        color="blue",
-        label="Toxicity Scores",
-        interactive=False
-    ),
-    title="Toxicity Detection with RogueBERT",
-    description="Enter a comment to analyze its toxicity levels. The results will be displayed as a modern bar chart."
+    fn=plot_toxicity,
+    inputs=gr.Textbox(label="Enter a comment"),
+    outputs=gr.Image(type="file", label="Toxicity Analysis"),
+    title="Toxicity Detector",
+    description="Enter a comment to analyze its toxicity scores across different categories.",
 )
 
-demo.launch()
+# Launch the Gradio app
+if __name__ == "__main__":
+    demo.launch()