Update app.py

app.py

@@ -1,6 +1,6 @@
 import streamlit as st
 import pandas as pd
-from transformers import BertModel, BertTokenizer, RobertaModel, RobertaTokenizer
+from transformers import AutoModel, AutoTokenizer
 from sklearn.decomposition import PCA
 import plotly.graph_objs as go
 import numpy as np
@@ -9,11 +9,14 @@ from database_utils import init_db, save_embeddings_to_db, get_all_embeddings, c
 @st.cache_resource
 def load_model(model_name):
     if model_name == "BERT":
-        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
-        model = BertModel.from_pretrained('bert-base-uncased')
+        model_path = 'bert-base-uncased'
     elif model_name == "RoBERTa":
-        tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
-        model = RobertaModel.from_pretrained('roberta-base')
+        model_path = 'roberta-base'
+    else:
+        raise ValueError(f"Unsupported model: {model_name}")
+    
+    tokenizer = AutoTokenizer.from_pretrained(model_path)
+    model = AutoModel.from_pretrained(model_path)
     return tokenizer, model
 
 def get_embeddings(phrases, tokenizer, model):