Update app.py

app.py

@@ -197,29 +197,44 @@ if st.button("Predict Personality"):
     if not displayed:
         st.write("No predictions exceed the confidence threshold.")
         """
+import streamlit as st
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
 import torch
 
-# Check if CUDA is available, otherwise use CPU
+# Initialize the tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained('facebook/bart-large-mnli')
+model = AutoModelForSequenceClassification.from_pretrained('facebook/bart-large-mnli')
+
+# Move model to appropriate device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
 
-# Load the model and tokenizer
-nli_model = AutoModelForSequenceClassification.from_pretrained('facebook/bart-large-mnli').to(device)
-tokenizer = AutoTokenizer.from_pretrained('facebook/bart-large-mnli')
+# Streamlit app
+def main():
+    st.title("Natural Language Inference with BART")
+
+    # Text input for premise and hypothesis
+    premise = st.text_area("Enter the premise:", value="", height=150)
+    hypothesis = st.text_input("Enter the hypothesis:")
 
-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 = 'This response is Helper personality.'
+    if st.button("Analyze"):
+        if premise and hypothesis:
+            # Tokenize and encode the premise and hypothesis
+            inputs = tokenizer.encode(premise, hypothesis, return_tensors='pt', truncation_strategy='only_first').to(device)
 
-# Tokenize the input text pair
-inputs = tokenizer.encode(premise, hypothesis, return_tensors='pt', truncation_strategy='only_first').to(device)
+            # Model inference
+            with torch.no_grad():
+                logits = model(inputs)[0]
 
-# Perform inference
-logits = nli_model(inputs)[0]
+            # Calculate probabilities
+            entail_contradiction_logits = logits[:, [0, 2]]
+            probs = entail_contradiction_logits.softmax(dim=1)
+            prob_label_is_true = probs[:, 1].item()
 
-# Process logits to get probabilities
-entail_contradiction_logits = logits[:, [0, 2]]
-probs = entail_contradiction_logits.softmax(dim=1)
-prob_label_is_true = probs[:, 1]
+            # Display results
+            st.write(f"Probability of Entailment: {prob_label_is_true:.4f}")
+        else:
+            st.error("Please enter both a premise and a hypothesis.")
 
-# Print the probability that the label is true
-print(prob_label_is_true)
+if __name__ == "__main__":
+    main()