Update app.py

app.py

@@ -1,42 +1,58 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
-import torch  # Import torch
+import torch
+import numpy as np  # Import numpy
 
-# Load the model and tokenizer (same as your original code)
-model_name = "frameai/PersianSentiment"
+
+# Check for GPU availability and set device
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device}")
+
+# Load the model and tokenizer
+model_name = "explorewithai/PersianSwear-Detector"  # Corrected model name
+loaded_model = AutoModelForSequenceClassification.from_pretrained(model_name).to(device)  # Move model to device
 loaded_tokenizer = AutoTokenizer.from_pretrained(model_name)
-loaded_model = AutoModelForSequenceClassification.from_pretrained(model_name)
+
 
 
 def predict_sentiment(text):
-    """Predicts the sentiment of a given text."""
-    inputs = loaded_tokenizer(text, return_tensors="pt", padding=True, truncation=True)  # Add padding and truncation
-    outputs = loaded_model(**inputs)
-    # Use softmax to get probabilities and argmax to get the predicted class
-    probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
-    predictions = torch.argmax(probabilities, dim=-1).item()
-
-    if predictions == 0:
-        sentiment = "Negative"
-    elif predictions == 1:
-        sentiment = "Positive"
+    """Predicts the sentiment (Bad Word, Good Word, Neutral Word) of a given text."""
+    inputs = loaded_tokenizer(text, return_tensors="pt", padding=True, truncation=True).to(device)  # Move inputs to GPU
+    with torch.no_grad():  # Ensure no gradients are calculated
+        outputs = loaded_model(**inputs)
+    logits = outputs.logits
+    probabilities = torch.nn.functional.softmax(logits, dim=-1) # Get probabilities
+    prediction = torch.argmax(logits, dim=-1).item()
+
+    # Map numeric labels to meaningful strings and get probabilities
+    if prediction == 4:
+        sentiment = "Bad Word"
+    elif prediction == 0:
+        sentiment = "Good Word"
+    elif prediction == 3:
+        sentiment = "Neutral Word"
     else:
-        sentiment = "Neutral"
+        sentiment = "Unknown"  # Should not happen, but good practice
+
+    # Create a dictionary for the probabilities
+    prob_dict = {}
+    if "Bad Word" in ["Bad Word", "Good Word", "Neutral Word"]:
+        prob_dict["Bad Word"] = float(probabilities[0][4]) if 4 < probabilities.shape[1] else 0.0
+    if "Good Word" in ["Bad Word", "Good Word", "Neutral Word"]:
+        prob_dict["Good Word"] = float(probabilities[0][0]) if 0 < probabilities.shape[1] else 0.0
+    if "Neutral Word" in ["Bad Word", "Good Word", "Neutral Word"]:
+        prob_dict["Neutral Word"] = float(probabilities[0][3]) if 3 < probabilities.shape[1] else 0.0
+
+    return prob_dict, sentiment
 
-    #  Return probabilities as well for a more informative output
-    return {
-        "Negative": float(probabilities[0][0]),
-        "Positive": float(probabilities[0][1]),
-        "Neutral": float(probabilities[0][2]),
-    }, sentiment
 
 # Create example sentences
 examples = [
-    ["این فیلم عالی بود!"],  # Positive example
-    ["من این غذا را دوست نداشتم."],  # Negative example
-    ["هوا خوب است."],  # Neutral (could be slightly positive, depends on context)
-    ["کتاب جالبی بود اما کمی خسته کننده هم بود."] , # Mixed/Neutral
-      ["اصلا راضی نبودم."] #negative
+    ["چه کت و شلوار زیبایی"],  # Good word example
+    ["این فیلم خیلی زیبا بود"],  # Good word example
+    ["میز"],  # Neutral word example
+    ["کثافت"],  # Bad word example
+    ["هوا خوب است."] #neutral
 ]
 
 
@@ -46,15 +62,13 @@ iface = gr.Interface(
     inputs=gr.Textbox(label="Enter Persian Text", lines=5, placeholder="Type your text here..."),
     outputs=[
         gr.Label(label="Sentiment Probabilities"),
-        gr.Textbox(label="Predicted Sentiment") # Add output component for the sentiment string
-
+        gr.Textbox(label="Predicted Sentiment")  # Output component for the sentiment string
     ],
-    title="Persian Sentiment Analysis",
-    description="Enter a Persian sentence and get its sentiment (Positive, Negative, or Neutral).",
+    title="Persian Swear Word Detection",
+    description="Enter a Persian sentence and get its sentiment (Good Word, Bad Word, or Neutral Word).",
     examples=examples,
-    live=False  # set to True for automatic updates as you type
+    live=False  # Set to True for automatic updates as you type
 )
 
-
 if __name__ == "__main__":
     iface.launch()