Update app.py

app.py

@@ -8,50 +8,93 @@ model_options = {
     "GoalZero/aidetection-ada-v0.1": "GoalZero/aidetection-ada-v0.1"
 }
 
-# Initialize tokenizer and model with the default model
-default_model = model_options["GoalZero/aidetection-ada-v0.2"]
-tokenizer = RobertaTokenizer.from_pretrained(default_model)
-model = RobertaForSequenceClassification.from_pretrained(default_model)
+# Initialize global variables for model and tokenizer
+model = None
+tokenizer = None
+
+def load_model(model_name):
+    """Helper function to load model and tokenizer"""
+    try:
+        return (
+            RobertaForSequenceClassification.from_pretrained(model_name),
+            RobertaTokenizer.from_pretrained(model_name)
+        )
+    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)
+except Exception as e:
+    print(f"Error loading default model: {str(e)}")
 
-# Define the prediction function
 def classify_text(text, model_choice):
-    global model, tokenizer  # Access the global model and tokenizer variables
-    
-    # Check if the model needs to be changed
-    if model_choice != model.name_or_path:
-        model = RobertaForSequenceClassification.from_pretrained(model_choice)
-        tokenizer = RobertaTokenizer.from_pretrained(model_choice)
-    
-    # Remove periods and new lines from the input text
-    cleaned_text = text.replace('.', '').replace('\n', ' ')
-    
-    # Tokenize the cleaned input text
-    inputs = tokenizer(cleaned_text, return_tensors='pt', padding=True, truncation=True, max_length=128)
-    
-    # Get the model's prediction
-    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 the class '1'
-    prob_1 = probabilities[0][1].item()
+    global model, tokenizer
     
-    return {"Probability of being AI": prob_1}
+    try:
+        # Check if we need to change the model
+        if model is None or model_choice != model.name_or_path:
+            model, tokenizer = load_model(model_choice)
+        
+        # Clean the input text
+        cleaned_text = text.replace('.', '').replace('\n', ' ')
+        
+        # Tokenize the cleaned input text
+        inputs = tokenizer(
+            cleaned_text,
+            return_tensors='pt',
+            padding=True,
+            truncation=True,
+            max_length=128
+        )
+        
+        # Get the model's prediction
+        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()
+            
+            return {
+                "AI Probability": round(prob_1 * 100, 2),
+                "Model used": model_choice
+            }
+            
+    except Exception as e:
+        return {
+            "error": f"An error occurred: {str(e)}",
+            "Model used": model_choice
+        }
 
 # Create the Gradio interface
 iface = gr.Interface(
     fn=classify_text,
     inputs=[
-        gr.Textbox(lines=2, placeholder="Enter text here..."),
-        gr.Dropdown(choices=list(model_options.keys()), value=default_model, label="Select Model")
+        gr.Textbox(
+            lines=2,
+            placeholder="Enter text here...",
+            label="Input Text"
+        ),
+        gr.Dropdown(
+            choices=list(model_options.keys()),
+            value="GoalZero/aidetection-ada-v0.2",
+            label="Select Model Version"
+        )
     ],
-    outputs="json",
-    title="GoalZero Ada Model Selector",
+    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=[
+        ["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"]
+    ]
 )
 
 # Launch the app
 if __name__ == "__main__":
-    iface.launch(share=True)
+    iface.launch(share=True)