Update app.py

app.py

@@ -1,25 +1,33 @@
 import gradio as gr
-from transformers import RobertaTokenizer, RobertaForSequenceClassification
+from transformers import (
+    RobertaTokenizer, RobertaForSequenceClassification,
+    AutoTokenizer, AutoModelForSequenceClassification
+)
 import torch
 
-# Define available models
+# Define available models including DeBERTa
 model_options = {
     "GoalZero/aidetection-ada-v0.2": "GoalZero/aidetection-ada-v0.2",
     "GoalZero/aidetection-ada-v0.1": "GoalZero/aidetection-ada-v0.1",
-    "GoalZero/babbage-mini-v0.1": "GoalZero/babbage-mini-v0.1"
+    "GoalZero/babbage-mini-v0.1": "GoalZero/babbage-mini-v0.1",
+    "GoalZero/ada-2534": "GoalZero/ada-2534"
 }
 
-# Initialize global variables for model and tokenizer
+# Initialize global variables
 model = None
 tokenizer = None
+current_model_name = None
 
 def load_model(model_name):
-    """Helper function to load model and tokenizer"""
+    """Load model and tokenizer, handling both RoBERTa and DeBERTa"""
     try:
-        return (
-            RobertaForSequenceClassification.from_pretrained(model_name),
-            RobertaTokenizer.from_pretrained(model_name)
-        )
+        if "deberta" in model_name.lower() or "ada-2534" in model_name.lower():
+            model = AutoModelForSequenceClassification.from_pretrained(model_name)
+            tokenizer = AutoTokenizer.from_pretrained(model_name)
+        else:
+            model = RobertaForSequenceClassification.from_pretrained(model_name)
+            tokenizer = RobertaTokenizer.from_pretrained(model_name)
+        return model, tokenizer
     except Exception as e:
         raise Exception(f"Failed to load model {model_name}: {str(e)}")
 
@@ -27,21 +35,23 @@ def load_model(model_name):
 try:
     default_model = "GoalZero/aidetection-ada-v0.2"
     model, tokenizer = load_model(default_model)
+    current_model_name = default_model
 except Exception as e:
     print(f"Error loading default model: {str(e)}")
 
 def classify_text(text, model_choice):
-    global model, tokenizer
+    global model, tokenizer, current_model_name
     
     try:
-        # Check if we need to change the model
-        if model is None or model_choice != model.name_or_path:
+        # Reload model if needed
+        if model is None or model_choice != current_model_name:
             model, tokenizer = load_model(model_choice)
+            current_model_name = model_choice
         
-        # Clean the input text
+        # Clean input
         cleaned_text = text.replace('.', '').replace('\n', ' ')
         
-        # Tokenize the cleaned input text
+        # Tokenize
         inputs = tokenizer(
             cleaned_text,
             return_tensors='pt',
@@ -50,21 +60,16 @@ def classify_text(text, model_choice):
             max_length=128
         )
         
-        # Get the model's prediction
+        # Predict
         with torch.no_grad():
             outputs = model(**inputs)
-            
-            # Apply softmax to get probabilities
             probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
-            
-            # Get the probability of class '1'
-            prob_1 = probabilities[0][1].item()
+            prob_ai = probabilities[0][1].item()
             
             return {
-                "AI Probability": round(prob_1 * 100, 10),
+                "AI Probability": round(prob_ai * 100, 10),
                 "Model used": model_choice
             }
-            
     except Exception as e:
         return {
             "error": f"An error occurred: {str(e)}",
@@ -92,10 +97,11 @@ iface = gr.Interface(
     examples=[
         ["Waymo is an American autonomous driving technology company that originated as the Google Self-Driving Car Project in 2009. It is now a subsidiary of Alphabet Inc., headquartered in Mountain View, California. The name \"Waymo\" was adopted in December 2016 when the project was rebranded and spun out of Google to focus on developing fully autonomous vehicles aimed at improving transportation safety and convenience", "GoalZero/babbage-mini-v0.1"],
         ["WWII demonstrated the importance of alliances in global conflicts. The Axis and Allied powers were formed as countries sought to protect their interests and expand their influence. This lesson underscores the potential for future global conflicts to involve complex alliances, similar to the Cold War era’s NATO and Warsaw Pact alignments.", "GoalZero/aidetection-ada-v0.2"],
-        ["Eustace was a thorough gentleman. There was candor in his quack, and affability in his waddle; and underneath his snowy down beat a pure and sympathetic heart. In short, he was a most exemplary duck.", "GoalZero/aidetection-ada-v0.1"]
+        ["Eustace was a thorough gentleman. There was candor in his quack, and affability in his waddle; and underneath his snowy down beat a pure and sympathetic heart. In short, he was a most exemplary duck.", "GoalZero/aidetection-ada-v0.1"],
+        ["This is an example of AI-written text using the DeBERTa model for testing purposes.", "GoalZero/ada-2534"]
     ]
 )
 
 # Launch the app
 if __name__ == "__main__":
-    iface.launch(share=True)
+    iface.launch(share=True)