import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

# Load the model and tokenizer
model_name = "KevSun/Personality_LM"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

# Function to predict personality traits
def predict_personality(text):
    inputs = tokenizer(text, return_tensors="pt")
    outputs = model(**inputs)
    probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
    labels = ["Introverted", "Extroverted", "Open", "Agreeable", "Conscientious", "Neurotic"]
    predictions = {label: prob.item() for label, prob in zip(labels, probs[0])}
    return predictions

# Create the Gradio interface
interface = gr.Interface(
    fn=predict_personality,
    inputs=gr.Textbox(lines=2, placeholder="Enter a sentence here..."),
    outputs=gr.Label(),
    title="Personality Analyzer",
    description="Enter a sentence and get a prediction of personality traits."
)

# Launch the Gradio app
interface.launch()