import gradio as gr
from transformers import RobertaTokenizer, RobertaForSequenceClassification
import torch

# 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, 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()
    
    return {"Probability of being AI": prob_1}

# 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")
    ],
    outputs="json",
    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)