Update app.py

app.py

@@ -142,7 +142,6 @@
 
 import streamlit as st
 import transformers
-import torch
 
 # Load the model and tokenizer
 model = transformers.AutoModelForSequenceClassification.from_pretrained("DeeeTeeee01/mytest_trainer_roberta-base")
@@ -158,7 +157,10 @@ def predict_sentiment(text):
     sentiment = prediction["label"]
     score = prediction["score"]
 
-    return sentiment, score
+    # Convert scores to a dictionary with labels as keys and percentages as values
+    scores_dict = {label: score * 100 for label, score in zip(prediction["labels"], prediction["scores"])}
+    
+    return sentiment, scores_dict
 
 # Setting the page configurations
 st.set_page_config(
@@ -182,14 +184,13 @@ predict_button = st.button("Predict")
 
 # Show sentiment output
 if predict_button and text:
-    sentiment, score = predict_sentiment(text)
-    st.write(f"The sentiment is {sentiment} with a score of {score*100:.2f}% for each category.")
+    sentiment, scores_dict = predict_sentiment(text)
+    st.write(f"The sentiment is {sentiment} with a score of {scores_dict[sentiment]:.2f}% for each category.")
     
     # Display individual percentages
     st.write("Sentiment Breakdown:")
-    st.write(f"- Negative: {score[0]*100:.2f}%")
-    st.write(f"- Positive: {score[1]*100:.2f}%")
-    st.write(f"- Neutral: {score[2]*100:.2f}%")
+    for label, score in scores_dict.items():
+        st.write(f"- {label}: {score:.2f}%")
 
 # Define the CSS style for the app
 st.markdown(
@@ -206,3 +207,4 @@ h1 {
 """,
 unsafe_allow_html=True
 )
+