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 the prediction function
def classify_text(text):
    # 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..."),
    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.",
)

# Launch the app
if __name__ == "__main__":
    iface.launch(share=True)