Update app.py

app.py

@@ -1,43 +1,19 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModel
-from torch.nn import functional as F
-import pandas as pd
+from transformers import pipeline
 
-tokenizer = AutoTokenizer.from_pretrained('deepset/sentence_bert')
-model = AutoModel.from_pretrained('deepset/sentence_bert')
+# Initialize the zero-shot classification pipeline
+classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
 
-st.title('Semantic Similarity Checker')
+# Define the candidate labels according to the Enneagram types
+labels = ["Peacemaker", "Loyalist", "Achiever", "Reformer", "Individualist", "Helper", "Challenger", "Investigator", "Enthusiast"]
 
-
-sentence = st.text_area('Enter')
-user_labels = st.text_area('Enter label')
-labels = [label.strip() for label in user_labels.split(',')]
-
-# run inputs through model and mean-pool over the sequence
-# dimension to get sequence-level representations
-
-
-if st.button("Calculate Similarities"):
-    if sentence and labels:
-        inputs = tokenizer.batch_encode_plus([sentence] + labels,
-                                     return_tensors='pt',
-                                     pad_to_max_length=True)
-        input_ids = inputs['input_ids']
-        attention_mask = inputs['attention_mask']
-        output = model(input_ids, attention_mask=attention_mask)[0]
-        sentence_rep = output[:1].mean(dim=1)
-        label_reps = output[1:].mean(dim=1)
-        
-        similarities = F.cosine_similarity(sentence_rep, label_reps)
-        similarities = similarities.cpu().detach().numpy()  # Convert to numpy array for easier handling
-        
-        # Sorting indices for displaying in order
-        sorted_indices = similarities.argsort()[::-1]
-        sorted_labels = [labels[idx] for idx in sorted_indices]
-        sorted_similarities = similarities[sorted_indices]
-        
-        # Display results in bar chart
-        df = pd.DataFrame({'Label': sorted_labels, 'Similarity': sorted_similarities})
-        st.bar_chart(df.set_index('Label'))
-    else:
-        st.error("Please enter both a sentence and some labels.")
+# Streamlit interface
+st.title("Resume-based Personality Prediction")
+resume_text = st.text_area("Enter Resume Text Here", height=300)
+if st.button("Predict Personality"):
+    # Make prediction
+    result = classifier(resume_text, labels)
+    # Display the results
+    st.write("Predictions:")
+    for label, score in zip(result['labels'], result['scores']):
+        st.write(f"{label}: {score*100:.2f}%")