Update src/streamlit_app.py

src/streamlit_app.py

@@ -118,33 +118,40 @@ if not comet_initialized:
 
 
 # --- Label Definitions ---
-labels = [
-"People", 
-"Locations",
-
-"Organizations",
 
-"Dates",
 
-"Events",
+labels = [
+    'person',
+    'organization',
+    'location',
+    'date',
+    'time',
+    'event',
+    'title',
+    'product',
+    'law',
+    'policy',
+    'work of art',
+    'geopolitical entity',
+    'number',
+    'cause of death',
+    'weapon',
+    'vehicle',
+    'facility',
+    'temporal expression',
+]
 
-"Facilities",
 
-"Vehicles",
 
-"Numbers"
-  ]
 # Create a mapping dictionary for labels to categories
 category_mapping = {
     
-    "People": ["People"],
-    "Locations" :["Locations"],
-    "Organizations": ["Organizations"],
-    "Dates": ["Dates"],
-    "Events": ["Events"],
-    "Facilities": ["Facilities"],
-    "Vehicles": ["Vehicles"],
-    "Numbers": ["Numbers"],
+    "People & Groups": ["Person", "Organization", "Title"],
+    "Topics & Objects":["Event", "Product", "Law", "Policy", "Work of Art", "Weapon", "Vehicle"],
+    "Temporal": ["Date", "Time", "Temporal expression"],
+    "Locations": ["Location", "Geopolitical entity", "Facility"],
+    "Quantitative & Contextual": ["Number", "Cause of death"]
+    
     
    
 }
@@ -154,7 +161,7 @@ category_mapping = {
 def load_ner_model():
     """Loads the GLiNER model and caches it."""
     try:
-        return GLiNER.from_pretrained("EmergentMethods/gliner_small_news-v2.1", nested_ner=True, num_gen_sequences=2, gen_constraints= labels, threshold = 0.70)
+        return GLiNER.from_pretrained("EmergentMethods/gliner_medium_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()