Fix the model class

app.py

@@ -7,7 +7,7 @@ import streamlit as st
 
 import torch
 import datasets
-from transformers import AutoModel, AutoTokenizer
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
 
 model_name = st.text_input("Enter a model's name on HF")
 # MODEL_NAME = "AMR-KELEG/NADI2024-baseline"
@@ -33,7 +33,7 @@ DIALECTS = ["Algeria",
 assert len(DIALECTS) == 18
 
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModel.from_pretrained(model_name)
+model = AutoModelForSequenceClassification.from_pretrained(model_name)
 
 def predict_top_p(text, P=0.9):
     """Predict the top dialects with an accumulative confidence of at least P."""