Update app.py

app.py

@@ -2,13 +2,26 @@ import gradio as gr
 from transformers import RobertaTokenizer, RobertaForSequenceClassification
 import torch
 
-# Load the model and tokenizer from the specified directory
-model_path = 'GoalZero/aidetection-ada-v0.1'
-tokenizer = RobertaTokenizer.from_pretrained(model_path)
-model = RobertaForSequenceClassification.from_pretrained(model_path)
+# Define available models
+model_options = {
+    "GoalZero/aidetection-ada-v0.2": "GoalZero/aidetection-ada-v0.2",
+    "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)
 
 # Define the prediction function
-def classify_text(text):
+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', ' ')
     
@@ -30,12 +43,15 @@ def classify_text(text):
 # Create the Gradio interface
 iface = gr.Interface(
     fn=classify_text,
-    inputs=gr.Textbox(lines=2, placeholder="Enter text here..."),
+    inputs=[
+        gr.Textbox(lines=2, placeholder="Enter text here..."),
+        gr.Dropdown(choices=list(model_options.keys()), value=default_model, label="Select Model")
+    ],
     outputs="json",
-    title="GoalZero Ada v0.1 Demo",
-    description="Enter some text and get the probability of the text being written by AI. Full checkpoints of the model will be released soon.",
+    title="GoalZero Ada Model Selector",
+    description="Enter text to get the probability of it being AI-written. Select a model version to use.",
 )
 
 # Launch the app
 if __name__ == "__main__":
-    iface.launch(share=True)
+    iface.launch(share=True)