fix

app.py

@@ -1,17 +1,13 @@
 import torch
-import torch.nn as nn
-from transformers import RobertaTokenizer, RobertaModel
-import json
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import streamlit as st
 
 # Set device (GPU if available, otherwise CPU)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Load model directly
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
-
+# Load tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained("mavinsao/mi-roberta-classification")
-model = AutoModelForSequenceClassification.from_pretrained("mavinsao/mi-roberta-classification")
+model = AutoModelForSequenceClassification.from_pretrained("mavinsao/mi-roberta-classification").to(device)
 
 # Streamlit app
 st.title('Mental Illness Prediction')
@@ -19,14 +15,19 @@ st.title('Mental Illness Prediction')
 # Input text area for user input
 sentence = st.text_area("Enter the long sentence to predict your mental illness state:")
 
-
 # Prediction button
 if st.button('Predict'):
-    # Predict label
+    # Tokenize the input sentence
+    inputs = tokenizer(sentence, return_tensors="pt", padding=True, truncation=True).to(device)
 
+    # Forward pass
     with torch.no_grad():
-      logits = model(sentence).logits
+        outputs = model(**inputs)
+
+    # Get predicted probabilities
+    probabilities = torch.sigmoid(outputs.logits).squeeze(dim=0)
 
-    predicted_class_ids = torch.arange(0, logits.shape[-1])[torch.sigmoid(logits).squeeze(dim=0) > 0.5]
+    # Get predicted labels with probability greater than 0.5
+    predicted_labels = [label for i, label in enumerate(tokenizer.labels) if probabilities[i] > 0.5]
 
-    st.json(predicted_class_ids)
+    st.write("Predicted labels:", predicted_labels)