Create app.py

app.py

@@ -0,0 +1,60 @@
+import torch
+import gradio as gr
+import numpy as np
+import matplotlib.pyplot as plt
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+
+torch.set_num_threads(torch.get_num_threads())
+
+# 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
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+
+# Define toxicity labels
+labels = ["toxic", "severe_toxic", "obscene", "threat", "insult", "identity_hate"]
+
+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 = {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]
+    
+    return {labels[i]: float(probabilities[i]) for i in range(len(labels))}
+
+def visualize_toxicity(comment):
+    """Generates a bar chart showing toxicity levels."""
+    scores = predict_toxicity(comment)
+    
+    # Create bar chart
+    plt.figure(figsize=(6, 4))
+    plt.bar(scores.keys(), scores.values(), color=['blue', 'red', 'green', 'purple', 'orange', 'brown'])
+    plt.ylim(0, 1)
+    plt.ylabel("Toxicity Score")
+    plt.title("Toxicity Analysis")
+    plt.xticks(rotation=45)
+    plt.grid(axis='y', linestyle='--', alpha=0.7)
+    
+    # Save plot to display in Gradio
+    plt.savefig("toxicity_plot.png")
+    plt.close()
+    
+    return "toxicity_plot.png"
+
+# Gradio interface
+demo = gr.Interface(
+    fn=visualize_toxicity,
+    inputs=gr.Textbox(label="Enter a comment:"),
+    outputs=gr.Image(type="file", label="Toxicity Scores"),
+    title="Toxicity Detection with RogueBERT",
+    description="Enter a comment to analyze its toxicity levels. The results will be displayed as a bar chart."
+)
+
+demo.launch()