akash-43's picture
Updated app.py
512caf5 verified
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)