Update app.py

app.py

@@ -1,16 +1,17 @@
 import streamlit as st
 import joblib
 from sentence_transformers import SentenceTransformer
+import numpy as np
 
 # Load the pre-trained embedding model
 @st.cache_resource  # Cache the embedding model to save loading time
 def load_embedding_model():
     return SentenceTransformer('neuml/pubmedbert-base-embeddings')
 
-# Load the MLP model
+# Load the multilabel classification model
 @st.cache_resource  # Cache the loaded model
-def load_mlp_model():
-    with open("MLP.pkl", "rb") as file:
+def load_multilabel_model():
+    with open("multilabel_model.pkl", "rb") as file:
         return joblib.load(file)
 
 # Embed text
@@ -19,9 +20,15 @@ def get_embeddings(title, abstract, embedding_model):
     combined_text = title + " " + abstract
     return embedding_model.encode(combined_text)
 
+# Map predicted binary outputs to labels
+LABELS = ["device", "screening", "drug", "surgery", "imaging", "telemedicine"]
+
+def decode_predictions(predictions):
+    return [label for label, pred in zip(LABELS, predictions) if pred == 1]
+
 # Main Streamlit app
 def main():
-    st.title("MLP Predictor for Titles and Abstracts")
+    st.title("Multilabel Classifier for Titles and Abstracts")
     
     # Input fields
     title = st.text_input("Enter the Title:")
@@ -29,10 +36,10 @@ def main():
     
     # Load models
     embedding_model = load_embedding_model()
-    mlp_model = load_mlp_model()
+    multilabel_model = load_multilabel_model()
     
     # Predict button
-    if st.button("Predict Label"):
+    if st.button("Predict Labels"):
         if title.strip() == "" or abstract.strip() == "":
             st.error("Both Title and Abstract are required!")
         else:
@@ -40,10 +47,17 @@ def main():
             embeddings = get_embeddings(title, abstract, embedding_model)
             
             # Make prediction
-            prediction = mlp_model.predict([embeddings])[0]  # Input should be a 2D array
+            predictions = multilabel_model.predict([embeddings])[0]  # Input should be a 2D array
             
-            # Display result
-            st.success(f"The predicted label is: {prediction}")
+            # Decode predictions
+            predicted_labels = decode_predictions(predictions)
+            
+            # Display results
+            if predicted_labels:
+                st.success(f"The predicted labels are: {', '.join(predicted_labels)}")
+            else:
+                st.warning("No relevant labels were predicted.")
 
 if __name__ == "__main__":
     main()
+