Spaces:
Sleeping
Sleeping
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) |