import gradio as gr
from transformers import BertTokenizer, BertForSequenceClassification
import torch

# Load the model and tokenizer
tokenizer = BertTokenizer.from_pretrained("Minej/bert-base-personality")
model = BertForSequenceClassification.from_pretrained("Minej/bert-base-personality")

# Define the personality detection function
def personality_detection(text):
    inputs = tokenizer(text, truncation=True, padding=True, return_tensors="pt")
    outputs = model(**inputs)
    predictions = outputs.logits.squeeze().detach().numpy()

    label_names = ['Extroversion', 'Neuroticism', 'Agreeableness', 'Conscientiousness', 'Openness']
    result = {label_names[i]: predictions[i] for i in range(len(label_names))}
    return result

# Set up Gradio Interface
interface = gr.Interface(
    fn=personality_detection,
    inputs=gr.Textbox(lines=5, placeholder="Enter text for personality detection..."),
    outputs=gr.Label(num_top_classes=5),
    title="Personality Detection from Text",
    description="This app detects personality traits based on the input text using a fine-tuned BERT model."
)

# Launch the app
interface.launch(share=True)