commit

app.py

@@ -1,24 +1,30 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from transformers import BertTokenizer, BertForSequenceClassification
 import torch
 
+# Function to load model and tokenizer
+def load_model():
+    tokenizer = BertTokenizer.from_pretrained("Minej/bert-base-personality")
+    model = BertForSequenceClassification.from_pretrained("Minej/bert-base-personality")
+    return tokenizer, model
+
 # Load the model and tokenizer
-model_name = "KevSun/Personality_LM"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSequenceClassification.from_pretrained(model_name)
+tokenizer, model = load_model()
 
 # 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
+def personality_detection(text):
+    inputs = tokenizer(text, truncation=True, padding=True, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model(**inputs)
+        predictions = torch.nn.functional.softmax(outputs.logits, dim=-1).squeeze().numpy()
+
+    label_names = ['Extroversion', 'Neuroticism', 'Agreeableness', 'Conscientiousness', 'Openness']
+    result = {label_names[i]: predictions[i] for i in range(len(label_names))}
+    return result
 
 # Create the Gradio interface
 interface = gr.Interface(
-    fn=predict_personality,
+    fn=personality_detection,
     inputs=gr.Textbox(lines=2, placeholder="Enter a sentence here..."),
     outputs=gr.Label(),
     title="Personality Analyzer",
@@ -26,4 +32,4 @@ interface = gr.Interface(
 )
 
 # Launch the Gradio app on a specific port
-interface.launch(server_port=7862)  # You can change 7861 to another port if necessary
+interface.launch(server_port=7861)  # You can change 7861 to another port if necessary