Update app.py

app.py

@@ -198,19 +198,28 @@ if st.button("Predict Personality"):
         st.write("No predictions exceed the confidence threshold.")
         """
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
-nli_model = AutoModelForSequenceClassification.from_pretrained('facebook/bart-large-mnli')
+import torch
+
+# Check if CUDA is available, otherwise use CPU
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Load the model and tokenizer
+nli_model = AutoModelForSequenceClassification.from_pretrained('facebook/bart-large-mnli').to(device)
 tokenizer = AutoTokenizer.from_pretrained('facebook/bart-large-mnli')
 
 premise = 'A few years ago, I was juggling a demanding job, volunteer commitments, and personal relationships, all while trying to manage chronic health issues. The challenge was overwhelming at times, but I approached it by prioritizing open communication with my employer and loved ones about my limits. I learned to delegate and accept help, which was difficult for me as I usually prefer to keep the peace by handling things myself. This experience taught me the importance of setting boundaries and the strength in vulnerability.'
-hypothesis = f'This example is Helper.'
+hypothesis = 'This example is Helper.'
 
-# run through model pre-trained on MNLI
-x = tokenizer.encode(premise, hypothesis, return_tensors='pt',
-                     truncation_strategy='only_first')
-logits = nli_model(x.to(device))[0]
+# Tokenize the input text pair
+inputs = tokenizer.encode(premise, hypothesis, return_tensors='pt', truncation_strategy='only_first').to(device)
 
-# we throw away "neutral" (dim 1) and take the probability of
-# "entailment" (2) as the probability of the label being true 
-entail_contradiction_logits = logits[:,[0,2]]
+# Perform inference
+logits = nli_model(inputs)[0]
+
+# Process logits to get probabilities
+entail_contradiction_logits = logits[:, [0, 2]]
 probs = entail_contradiction_logits.softmax(dim=1)
-prob_label_is_true = probs[:,1]
+prob_label_is_true = probs[:, 1]
+
+# Print the probability that the label is true
+print(prob_label_is_true)