update

app.py

@@ -10,52 +10,8 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 # Load model directly
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
-tokenizer = AutoTokenizer.from_pretrained("mavinsao/mi-roberta-mental-illness")
-model = AutoModelForSequenceClassification.from_pretrained("mavinsao/mi-roberta-mental-illness")
-
-# Create a common label map
-common_label_map = {'ADHD': 0, 'Anxiety': 1, 'bipolar': 2, 'BPD': 3, 'depression': 4, 'OCD': 5, 'ptsd': 6, 'none': 7}
-num_classes = 8
-
-
-def predict_labels(sentence, tokenizer, model, device, threshold=0.5, top_n=5):
-    # Tokenize the sentence and create attention mask
-    tokenized_input = tokenizer(
-        sentence,
-        add_special_tokens=True,
-        max_length=512,
-        padding="max_length",
-        truncation=True,
-        return_tensors="pt"
-    )
-
-    # Move the input tensors to the device
-    input_ids = tokenized_input['input_ids'].to(device)
-    attention_mask = tokenized_input['attention_mask'].to(device)
-
-    # Set the model to evaluation mode
-    model.eval()
-
-    # Make a prediction
-    with torch.no_grad():
-        output = model(input_ids, attention_mask)
-
-    # Apply thresholding to the logits to obtain predicted labels
-    logits = output.logits
-    sigmoid_output = torch.sigmoid(logits.squeeze(dim=0))
-    indices_above_threshold = torch.arange(logits.shape[-1], device=device)[sigmoid_output > threshold]
-
-    # Sort the indices by their sigmoid values
-    sorted_indices = indices_above_threshold[torch.argsort(sigmoid_output[indices_above_threshold], descending=True)]
-
-    # Map the predicted label indices back to the original class labels using the common label map
-    predicted_labels_with_score = [{"label": list(common_label_map.keys())[index], "score": sigmoid_output[index].item()} for index in sorted_indices[:top_n]]
-
-    # Create a JSON object with labels, scores, and short forms
-    json_result = [{"label": entry["label"], "score": entry["score"]} for entry in predicted_labels_with_score]
-
-    return json.dumps(json_result, indent=4)
-
+tokenizer = AutoTokenizer.from_pretrained("mavinsao/mi-roberta-classification")
+model = AutoModelForSequenceClassification.from_pretrained("mavinsao/mi-roberta-classification")
 
 # Streamlit app
 st.title('Mental Illness Prediction')
@@ -63,8 +19,14 @@ 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
-    predicted_response = predict_labels(sentence, tokenizer, model, device)
-    st.json(predicted_response)
+
+    with torch.no_grad():
+    logits = model(sentence).logits
+
+    predicted_class_ids = torch.arange(0, logits.shape[-1])[torch.sigmoid(logits).squeeze(dim=0) > 0.5]
+
+    st.json(predicted_class_ids)