Update app.py

app.py

@@ -9,14 +9,18 @@ st.sidebar.image("logo-wordlift.png")
 model_options = {"aida_model", "wikipedia_model_with_numbers"}
 selected_model_name = st.sidebar.selectbox("Select the Model", list(model_options))
 
+# Select entity_set
+entity_set_options = {"wikidata", "wikipedia"}
+selected_entity_set = st.sidebar.selectbox("Select the Entity Set", list(entity_set_options))
+
 @st.cache_resource  # 👈 Add the caching decorator
-def load_model(model_name):
+def load_model(model_name, entity_set):
     # Load the pretrained model
-    refined_model = Refined.from_pretrained(model_name=model_name, entity_set="wikipedia")
+    refined_model = Refined.from_pretrained(model_name=model_name, entity_set=entity_set)
     return refined_model
 
 # Use the cached model
-refined_model = load_model(selected_model_name)
+refined_model = load_model(selected_model_name, selected_entity_set)
 
 # Helper functions
 def get_wikidata_id(entity_string):