Update app.py

app.py

@@ -1,6 +1,7 @@
 import torch
+import torch
 import torch.nn as nn
-from transformers import ViTFeatureExtractor, ViTModel, BertTokenizerFast, BertModel
+from transformers import ViTImageProcessor, ViTModel, BertTokenizerFast, BertModel
 from PIL import Image
 import gradio as gr
 
@@ -32,13 +33,13 @@ class VisionLanguageModel(nn.Module):
         logits = self.classifier(combined_features)
         return logits
 
-# Load the model checkpoint
+# Load the model checkpoint with safer loading
 model = VisionLanguageModel()
-model.load_state_dict(torch.load('best_model.pth', map_location=torch.device('cpu')))
+model.load_state_dict(torch.load('best_model.pth', map_location=torch.device('cpu'), weights_only=True))
 model.eval()
 
 tokenizer = BertTokenizerFast.from_pretrained('bert-base-uncased')
-feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224-in21k')
+feature_extractor = ViTImageProcessor.from_pretrained('google/vit-base-patch16-224-in21k')
 
 def predict(image, text_input):
     # Preprocess the image
@@ -64,12 +65,12 @@ def predict(image, text_input):
     _, prediction = torch.max(outputs, dim=1)
     return "Malignant" if prediction.item() == 1 else "Benign"
 
-# Define Gradio interface
+# Define Gradio interface with updated component syntax
 iface = gr.Interface(
     fn=predict,
     inputs=[
-        gr.inputs.Image(type="pil", label="Upload Skin Lesion Image"),
-        gr.inputs.Textbox(label="Clinical Information (e.g., patient age, symptoms)")
+        gr.Image(type="pil", label="Upload Skin Lesion Image"),
+        gr.Textbox(label="Clinical Information (e.g., patient age, symptoms)")
     ],
     outputs="text",
     title="Skin Lesion Classification Demo",
@@ -77,3 +78,4 @@ iface = gr.Interface(
 )
 
 iface.launch()
+