Update app.py

app.py

@@ -9,6 +9,7 @@ from sklearn.preprocessing import MultiLabelBinarizer
 # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # Load the trained model and tokenizer
+
 @st.cache_resource
 def load_model():
     model = AutoModelForSequenceClassification.from_pretrained(
@@ -16,13 +17,13 @@ def load_model():
         num_labels=8,  # Adjust based on your label count
         problem_type="multi_label_classification"
     )
-    model.load_state_dict(torch.load('best_model_v2.pth'))
+    # Map the model to the appropriate device
+    model.load_state_dict(torch.load('best_model_v2.pth', map_location=torch.device('cpu')))
     model.eval()
     tokenizer = AutoTokenizer.from_pretrained("microsoft/BiomedNLP-BiomedBERT-base-uncased-abstract")
-    model = model.to(device)  # Move the model to the correct device
-
     return model, tokenizer
 
+
 @st.cache_resource
 def load_mlb():
     # Define the classes based on your label set
@@ -62,7 +63,9 @@ if st.button('Predict'):
         inputs = tokenizer(clinical_note, truncation=True, padding="max_length", max_length=512, return_tensors='pt')
         
         # Move inputs to the GPU if available
-        inputs = {key: val.to(device) for key, val in inputs.items()}
+        # inputs = {key: val.to(device) for key, val in inputs.items()}
+        inputs = {key: val.to(torch.device('cpu')) for key, val in inputs.items()}
+
         
         # Model inference
         with torch.no_grad():
@@ -91,30 +94,4 @@ if st.button('Predict'):
     #     st.write("Please enter clinical notes for prediction.")
 
 
-# # Prediction button
-# if st.button('Predict'):
-#     if clinical_note:
-#         # Tokenize the input clinical note
-#         inputs = tokenizer(clinical_note, truncation=True, padding="max_length", max_length=512, return_tensors='pt')
-        
-#         # Move inputs to the GPU if available
-#         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-#         inputs = {key: val.to(device) for key, val in inputs.items()}
-        
-#         # Model inference
-#         with torch.no_grad():
-#             outputs = model(**inputs)
-#             logits = outputs.logits
-        
-#         # Apply sigmoid and threshold the output (0.5 for multi-label classification)
-#         pred_labels = (torch.sigmoid(logits) > 0.5).cpu().numpy()
-        
-#         # Get the predicted ICD and CPT codes
-#         predicted_codes = mlb.inverse_transform(pred_labels)
-        
-#         # Show the results
-#         st.write("Predicted ICD and CPT Codes:")
-#         st.write(predicted_codes)
-        
-#     else:
-#         st.write("Please enter clinical notes for prediction.")
+