Update the model loading module

app.py

@@ -4,7 +4,7 @@ import torch
 import pandas as pd
 import streamlit as st
 import matplotlib.pyplot as plt
-from transformers import BertForSequenceClassification, AutoTokenizer
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
 from constants import DIALECTS
 
 import altair as alt
@@ -83,7 +83,7 @@ def convert_df(df):
 
 @st.cache_resource
 def load_model(model_name):
-    model = BertForSequenceClassification.from_pretrained(model_name)
+    model = AutoModelForSequenceClassification.from_pretrained(model_name)
     return model