Update src/streamlit_app.py

src/streamlit_app.py

@@ -69,7 +69,7 @@ st.markdown(
 # --- Page Configuration and UI Elements ---
 st.set_page_config(layout="wide", page_title="Named Entity Recognition App")
 
-st.subheader("FOODS", divider="gray")
+st.subheader("newS", divider="gray")
 st.link_button("by nlpblogs", "https://nlpblogs.com", type="tertiary")
 
 expander = st.expander("**Important notes on the ProductTag**")
@@ -119,29 +119,33 @@ if not comet_initialized:
 
 # --- Label Definitions ---
 labels = [
-    "Fruits Vegetables",
-    "Lactose, Diary, Eggs, Cheese, Yoghurt",
-    "Meat, Fish, Seafood",
-    "Frozen, Prepared Meals",
-    "Baking, Cooking",
-    "Cereals, Grains, Canned, Seeds",
-    "Breads",
-    "Snacks, Pastries, Treats",
-    "Frozen Desserts",
-    "Hot Drinks, Chilled Drinks",
-    "Alcoholic Drinks",
-    "Spices, Sauces",
-    "World Foods",
-    "Dietary Restrictions, Health, Allergens, Lifestyle"
+"People", 
+"Locations",
+
+"Organizations",
+
+"Dates",
+
+"Events",
+
+"Facilities",
+
+"Vehicles",
+
+"Numbers"
   ]
 # Create a mapping dictionary for labels to categories
 category_mapping = {
     
-    "Fresh & Primary Ingredients": ["Fruits Vegetables", "Meat, Fish, Seafood", "Lactose, Diary, Eggs, Cheese, Yoghurt"],
-    "Pantry & Staple Goods" :["Cereals, Grains, Canned, Seeds", "Breads", "Spices, Sauces", "Baking, Cooking"],
-    "Prepared & Convenience Foods": ["Frozen, Prepared Meals", "Snacks, Pastries, Treats", "Frozen Desserts", "World Foods"],
-    "Beverages": ["Alcoholic Drinks", "Alcoholic Drinks"],
-    "Meta-Labels & Lifestyle": ["Dietary Restrictions, Health, Allergens, Lifestyle"]
+    "People": ["People"],
+    "Locations" :["Locations"],
+    "Organizations": ["Organizations"],
+    "Dates": ["Dates"],
+    "Events": ["Events"],
+    "Facilities": ["Facilities"],
+    "Vehicles": ["Vehicles"],
+    "Numbers": ["Numbers"],
+    
    
 }
 
@@ -150,7 +154,7 @@ category_mapping = {
 def load_ner_model():
     """Loads the GLiNER model and caches it."""
     try:
-        return GLiNER.from_pretrained("empathyai/gliner_large-v2.5-groceries", nested_ner=True, num_gen_sequences=2, gen_constraints= labels, threshold = 0.70)
+        return GLiNER.from_pretrained("EmergentMethods/gliner_small_news-v2.1", nested_ner=True, num_gen_sequences=2, gen_constraints= labels, threshold = 0.70)
     except Exception as e:
         st.error(f"Failed to load NER model. Please check your internet connection or model availability: {e}")
         st.stop()