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("Compliance", divider="gray")
+st.subheader("Legal", divider="gray")
 st.link_button("by nlpblogs", "https://nlpblogs.com", type="tertiary")
 
 expander = st.expander("**Important notes on the ProductTag**")
@@ -122,114 +122,100 @@ if not comet_initialized:
 
 # --- Label Definitions ---
 labels = [
-    "medical_record_number",
-    "date_of_birth",
-    "ssn",
-    "date",
-    "first_name",
-    "email",
-    "last_name",
-    "customer_id",
-    "employee_id",
-    "name",
-    "street_address",
-    "phone_number",
-    "ipv4",
-    "credit_card_number",
-    "license_plate",
-    "address",
-    "user_name",
-    "device_identifier",
-    "bank_routing_number",
-    "date_time",
-    "company_name",
-    "unique_identifier",
-    "biometric_identifier",
-    "account_number",
-    "city",
-    "certificate_license_number",
-    "time",
-    "postcode",
-    "vehicle_identifier",
-    "coordinate",
-    "country",
-    "api_key",
-    "ipv6",
-    "password",
-    "health_plan_beneficiary_number",
-    "national_id",
-    "tax_id",
-    "url",
-    "state",
-    "swift_bic",
-    "cvv",
-    "pin"
+"Plaintiff",
+"Petitioner",
+
+"Defendant",
+"Respondent",
+
+"Appellant",
+"Appellee",
+
+"Debtor",
+"Creditor",
+
+"Signer",
+"Witness",
+"Courts",
+"Judges",
+"Lawyers",
+"Attorneys",
+"Statutes",
+"Laws",
+"Provisions",
+"Case Citations",
+"Legal Documents"
+
 ]
+    
 
 # Create a mapping dictionary for labels to categories
 category_mapping = {
-    "Personal Identifiers": [
-        "date_of_birth",
-        "first_name",
-        "last_name",
-        "name",
-        "biometric_identifier",
-        "user_name",
-        "password"
+    "Parties": [
+        "Plaintiff",
+        "Petitioner",
+        "Defendant",
+        "Respondent",
+        "Appellant",
+        "Appellee",
+        "Debtor",
+        "Creditor",
+        "Signer",
+        "Witness"
     ],
     
-    "Contact & Location Information": [
-        "email",
-        "street_address",
-        "phone_number",
-        "address",
-        "city",
-        "postcode",
-        "coordinate",
-        "country",
-        "state",
-        "url"
+    "Judicial & Governmental Bodies": [
+        "Courts",
+        "Judges",
+        "Lawyers",
+        "Attorneys"
+
     ],
     
-    "Financial & Business Data": [
-        "credit_card_number",
-        "bank_routing_number",
-        "account_number",
-        "swift_bic",
-        "cvv",
-        "pin",
-        "company_name",
-        "api_key"
+    "Legal Instruments & Concepts": [
+        "Statutes",
+        "Laws",
+        "Provisions",
+        "Case_citations",
+        "Legal_documents"
+
     ],
     
-    "Government & Official IDs": [
-        "ssn",
-        "license_plate",
-        "certificate_license_number",
-        "national_id",
-        "tax_id",
-        "medical_record_number",
-        "health_plan_beneficiary_number"
+    "Dates & Timeframes:": [
+        "Effective_dates",
+        "Execution_dates",
+        "Expiration_dates"
+
+
     ],
     
-    "Technical & System Data ": [
-        "ipv4",
-        "device_identifier",
-        "ipv6"
+    "Financial & Monetary Entities:": [
+        "Money",
+        "Amounts"
+        
     ],
     
-    "Unique Identifiers & Registration Numbers": [
-        "customer_id",
-        "employee_id",
-        "unique_identifier",
-        "vehicle_identifier"
+    "Contracts": [
+        "Contract_terms"
+        
     ],
     
-    "Date & Time Stamps": [
-        "date",
-        "date_time",
-        "time"
+    "Court Judgments:": [
+        "Case_number",
+        "Witnesses",
+       
     ]
+
+
+
+     "Criminal Law:": [
+        "Crimes",
+        "Offenses",
+        "Victims"
+       
+    ]
+
+    
 }
 
 
@@ -239,7 +225,7 @@ category_mapping = {
 def load_ner_model():
     """Loads the GLiNER model and caches it."""
     try:
-        return GLiNER.from_pretrained("gretelai/gretel-gliner-bi-large-v1.0", nested_ner=True, num_gen_sequences=2, gen_constraints= labels)
+        return GLiNER.from_pretrained("knowledgator/gliner-decoder-large-v1.0", nested_ner=True, num_gen_sequences=2, gen_constraints= labels)
     except Exception as e:
         st.error(f"Failed to load NER model. Please check your internet connection or model availability: {e}")
         st.stop()