Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModel, AutoModelForSequenceClassification
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
 import spacy
 import subprocess
@@ -24,19 +24,18 @@ word_vectors = api.load("glove-wiki-gigaword-50")
 # Check for GPU and set the device accordingly
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Load GPT-J-6B model and tokenizer from Hugging Face
-tokenizer_ai = AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B")
-model_ai = AutoModel.from_pretrained("EleutherAI/gpt-j-6B").to(device)
+# Load AI Detector model and tokenizer from Hugging Face (roberta-base-openai-detector)
+tokenizer_ai = AutoTokenizer.from_pretrained("roberta-base-openai-detector")
+model_ai = AutoModelForSequenceClassification.from_pretrained("roberta-base-openai-detector").to(device)
 
-# AI detection function using GPT-J-6B-based model
+# AI detection function using the RoBERTa-based model
 def detect_ai_generated(text):
     inputs = tokenizer_ai(text, return_tensors="pt", truncation=True, max_length=512).to(device)
     with torch.no_grad():
         outputs = model_ai(**inputs)
-        # Process the hidden states to generate a score (example logic, adjust as needed)
-        hidden_state = outputs.last_hidden_state[:, 0, :]  # Use the first token's representation
-        score = torch.mean(hidden_state).item()
-    return f"AI-Generated Content Score: {score:.2f}"
+        probabilities = torch.softmax(outputs.logits, dim=1)
+    ai_probability = probabilities[0][1].item() * 100  # Probability of being AI-generated
+    return f"AI-Generated Content Probability: {ai_probability:.2f}%"
 
 # Function to get synonyms using NLTK WordNet
 def get_synonyms_nltk(word, pos):