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

# Load the model and tokenizer from the specified directory
model_path = './finetuned_roberta'
tokenizer = RobertaTokenizer.from_pretrained(model_path)
model = RobertaForSequenceClassification.from_pretrained(model_path)

# Define the prediction function
def classify_text(text):
    # Tokenize the input text
    inputs = tokenizer(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 1": prob_1}

# Create the Gradio interface
iface = gr.Interface(
    fn=classify_text,
    inputs=gr.Textbox(lines=2, placeholder="Enter text here..."),
    outputs="json",
    title="Text Classification with RoBERTa",
    description="Enter some text and get the probability of the text being classified as class 1.",
    enable_queue=True,  # Ensure the API is enabled
)

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