fix app

app.py

@@ -1,8 +1,8 @@
-import streamlit as st
 import torch
 import torch.nn as nn
 from transformers import RobertaTokenizer, RobertaModel
 import json
+import streamlit as st
 
 # Set device (GPU if available, otherwise CPU)
 device = torch.device("cpu")
@@ -56,8 +56,6 @@ roberta_model = MIRobertaClassifier(num_classes=num_classes).to(device)
 # Load the state dictionary into the model
 roberta_model.load_state_dict(roberta_loaded_model_state['state_dict'])
 
-
-        
 # Load the state dictionary into the model
 mi_loaded_model_state = torch.load('reddit_miroberta_state.pth', map_location=device)
 # Create an instance of your model
@@ -67,7 +65,6 @@ mi_model = MIRobertaClassifier(num_classes=num_classes).to(device)
 mi_model.load_state_dict(mi_loaded_model_state['state_dict'])
 
 
-
 def predict_label(sentence, tokenizer, model1, model2, device):
     # Tokenize the sentence and create attention mask
     tokenized_input = tokenizer(
@@ -96,28 +93,31 @@ def predict_label(sentence, tokenizer, model1, model2, device):
     ensemble_outputs = (outputs1 + outputs2) / 2
 
     # Apply softmax to get probabilities
-    probabilities = torch.max(ensemble_outputs, dim=1)[0].tolist()
+    probabilities = torch.softmax(ensemble_outputs, dim=1)[0].tolist()
 
     # Map the predicted index back to the original class label using class_names
     class_names = list(common_label_map.keys())
-    
-    predicted_index = torch.argmax(ensemble_outputs, dim=1)
-    predicted_label = class_names[predicted_index.item()]
 
-    # Create JSON response
-    response = {
-        "label": predicted_label,
-        "score": probabilities[common_label_map[predicted_label]]
-    }
+    # Get predicted index and score for each label
+    label_scores = {}
+    for i, label in enumerate(class_names):
+        label_index = common_label_map[label]
+        label_scores[label] = probabilities[label_index]
+
+    # Sort label scores by score values in descending order
+    sorted_label_scores = {k: v for k, v in sorted(label_scores.items(), key=lambda item: item[1], reverse=True)}
+
+    # Get the predicted label
+    predicted_index = torch.argmax(ensemble_outputs, dim=1)
 
-    return response
+    return sorted_label_scores
 
 
 # Streamlit app
 st.title('Mental Illness Prediction')
 
 # Input text area for user input
-sentence = st.text_area("Enter the sentence to predict your mental illness state:")
+sentence = st.text_area("Enter the long sentence to predict your mental illness state:")
 
 # Prediction button
 if st.button('Predict'):