Update app.py

app.py

@@ -1,107 +1,110 @@
 import gradio as gr
-from transformers import (
-    RobertaTokenizer, RobertaForSequenceClassification,
-    AutoTokenizer, AutoModelForSequenceClassification
-)
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, GPT2LMHeadModel
 import torch
+import math
+
+# -----------------------------
+# Device setup
+# -----------------------------
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device}")
+
+# -----------------------------
+# GPT-2 for perplexity
+# -----------------------------
+gpt2_tokenizer = AutoTokenizer.from_pretrained("gpt2")
+gpt2_model = GPT2LMHeadModel.from_pretrained("gpt2").to(device)
+gpt2_model.eval()
 
-# Define available models including DeBERTa
+def compute_perplexity(sentence):
+    tokens = gpt2_tokenizer(sentence, return_tensors='pt').to(device)
+    with torch.no_grad():
+        outputs = gpt2_model(**tokens, labels=tokens["input_ids"])
+        loss = outputs.loss.item()
+        ppl = math.exp(loss)
+    # Normalize and round
+    return round(ppl / 100, 2)
+
+# -----------------------------
+# AI detection models
+# -----------------------------
 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/ada-2534": "GoalZero/ada-2534"
 }
 
-# Initialize global variables
+# Initialize globals
 model = None
 tokenizer = None
 current_model_name = None
 
 def load_model(model_name):
-    """Load model and tokenizer, handling both RoBERTa and DeBERTa"""
     try:
-        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)
+        model = AutoModelForSequenceClassification.from_pretrained(model_name).to(device)
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
         return model, tokenizer
     except Exception as e:
         raise Exception(f"Failed to load model {model_name}: {str(e)}")
 
 # Load default model
-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)}")
+default_model = "GoalZero/ada-2534"
+model, tokenizer = load_model(default_model)
+current_model_name = default_model
 
+# -----------------------------
+# Classification function
+# -----------------------------
 def classify_text(text, model_choice):
     global model, tokenizer, current_model_name
-    
     try:
         # 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 input
-        cleaned_text = text.replace('.', '').replace('\n', ' ')
-        
-        # Tokenize
+
+        # Compute perplexity
+        ppl_score = compute_perplexity(text)
+        # Prepend perplexity to text
+        text_with_ppl = f"{ppl_score} {text}"
+
+        # Tokenize and predict
         inputs = tokenizer(
-            cleaned_text,
+            text_with_ppl,
             return_tensors='pt',
             padding=True,
             truncation=True,
             max_length=128
-        )
-        
-        # Predict
+        ).to(device)
+
         with torch.no_grad():
             outputs = model(**inputs)
             probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
             prob_ai = probabilities[0][1].item()
-            
-            return {
-                "AI Probability": round(prob_ai * 100, 10),
-                "Model used": model_choice
-            }
-    except Exception as e:
+
         return {
-            "error": f"An error occurred: {str(e)}",
+            "AI Probability": round(prob_ai * 100, 10),
+            "Perplexity": ppl_score,
             "Model used": model_choice
         }
+    except Exception as e:
+        return {"error": str(e), "Model used": model_choice}
 
-# Create the Gradio interface
+# -----------------------------
+# Gradio interface
+# -----------------------------
 iface = gr.Interface(
     fn=classify_text,
     inputs=[
-        gr.Textbox(
-            lines=2,
-            placeholder="Enter text here...",
-            label="Input Text"
-        ),
+        gr.Textbox(lines=2, placeholder="Enter text here...", label="Input Text"),
         gr.Dropdown(
             choices=list(model_options.keys()),
-            value="GoalZero/aidetection-ada-v0.2",
+            value=default_model,
             label="Select Model Version"
         )
     ],
     outputs=gr.JSON(label="Results"),
-    title="GoalZero Ada AI Detection",
-    description="Enter text to get the probability of it being AI-written. Select a model version to use.",
-    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"],
-        ["This is an example of AI-written text using the DeBERTa model for testing purposes.", "GoalZero/ada-2534"]
-    ]
+    title="GoalZero Ada AI Detection with Perplexity",
+    description="Enter text to get the probability of it being AI-written. The perplexity score is also computed and prepended."
 )
 
-# Launch the app
 if __name__ == "__main__":
     iface.launch(share=True)