Update app.py

app.py

@@ -1,17 +1,13 @@
 import torch
-from transformers import XLNetTokenizer, XLNetForSequenceClassification
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
 import gradio as gr
 
+# Path to the saved model in Hugging Face Spaces
+model_path = 'https://huggingface.co/spaces/AliArshad/SeverityPrediction/blob/main/XLNet_model_project_Core.pt'
 
-
-# Path to the saved model
-model_path = 'https://huggingface.co/spaces/AliArshad/SeverityPrediction/blob/main/XLNet_model_project_Core.pt'  # Update with your model path
-
-# Load the saved model
-tokenizer = XLNetTokenizer.from_pretrained('xlnet-base-cased')
-model = XLNetForSequenceClassification.from_pretrained('xlnet-base-cased', num_labels=2)
-model.load_state_dict(torch.load(model_path))
-model.eval()
+# Load the model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained('xlnet-base-cased')
+model = AutoModelForSequenceClassification.from_pretrained(model_path)
 
 # Function for prediction
 def xl_net_predict(text):