change to model

app.py

@@ -5,67 +5,20 @@ import json
 import streamlit as st
 
 # Set device (GPU if available, otherwise CPU)
-device = torch.device("cpu")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Define the path to the model and tokenizer
-model_path = 'mi-roberta-base'  # pre-trained model
-tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
+# 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
 
 
-# MiRoBERTa
-class MIRobertaClassifier(nn.Module):
-    def __init__(self, num_classes, dropout_prob=0.3):
-        super(MIRobertaClassifier, self).__init__()
-        self.roberta = RobertaModel.from_pretrained(model_path)
-        self.dropout = nn.Dropout(dropout_prob)
-        self.fc = nn.Linear(self.roberta.config.hidden_size, num_classes)
-
-    def forward(self, input_ids, attention_mask):
-        outputs = self.roberta(input_ids=input_ids, attention_mask=attention_mask)
-        last_hidden_state = outputs.last_hidden_state[:, 0, :]
-        x = self.dropout(last_hidden_state)
-        logits = self.fc(x)
-        return logits
-
-
-# RoBERTa
-class RobertaClassifier(nn.Module):
-    def __init__(self, num_classes, dropout_prob=0.3):
-        super(RobertaClassifier, self).__init__()
-        self.roberta = RobertaModel.from_pretrained('roberta-base')
-        self.dropout = nn.Dropout(dropout_prob)
-        self.fc = nn.Linear(self.roberta.config.hidden_size, num_classes)
-
-    def forward(self, input_ids, attention_mask):
-        outputs = self.roberta(input_ids=input_ids, attention_mask=attention_mask)
-        last_hidden_state = outputs.last_hidden_state[:, 0, :]
-        x = self.dropout(last_hidden_state)
-        logits = self.fc(x)
-        return logits
-
-
-# Load the state dictionary into the model
-roberta_loaded_model_state = torch.load('reddit_roberta_state.pth', map_location=device)
-# Create an instance of your model
-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
-mi_model = MIRobertaClassifier(num_classes=num_classes).to(device)
-
-# Load the state dictionary into the model
-mi_model.load_state_dict(mi_loaded_model_state['state_dict'])
-
-
-def predict_label(sentence, tokenizer, model1, model2, device):
+def predict_labels(sentence, tokenizer, model, device, threshold=0.5, top_n=5):
     # Tokenize the sentence and create attention mask
     tokenized_input = tokenizer(
         sentence,
@@ -81,36 +34,27 @@ def predict_label(sentence, tokenizer, model1, model2, device):
     attention_mask = tokenized_input['attention_mask'].to(device)
 
     # Set the model to evaluation mode
-    mi_model.eval()
-    roberta_model.eval()
+    model.eval()
 
     # Make a prediction
     with torch.no_grad():
-        outputs1 = mi_model(input_ids, attention_mask)
-        outputs2 = roberta_model(input_ids, attention_mask)
-
-    # Ensemble predictions: averaging logits from both models
-    ensemble_outputs = (outputs1 + outputs2) / 2
-
-    # Apply softmax to get probabilities
-    probabilities = torch.softmax(ensemble_outputs, dim=1)[0].tolist()
+        output = model(input_ids, attention_mask)
 
-    # Map the predicted index back to the original class label using class_names
-    class_names = list(common_label_map.keys())
+    # 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]
 
-    # 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 the indices by their sigmoid values
+    sorted_indices = indices_above_threshold[torch.argsort(sigmoid_output[indices_above_threshold], descending=True)]
 
-    # 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)}
+    # 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]]
 
-    # Get the predicted label
-    predicted_index = torch.argmax(ensemble_outputs, dim=1)
+    # 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 sorted_label_scores
+    return json.dumps(json_result, indent=4)
 
 
 # Streamlit app
@@ -122,5 +66,5 @@ sentence = st.text_area("Enter the long sentence to predict your mental illness
 # Prediction button
 if st.button('Predict'):
     # Predict label
-    predicted_response = predict_label(sentence, tokenizer, mi_model, roberta_model, device)
+    predicted_response = predict_labels(sentence, tokenizer, model, device)
     st.json(predicted_response)