Updated model's, and added staff

src/streamlit_app.py

@@ -68,14 +68,17 @@ st.markdown("""
 """, unsafe_allow_html=True)
 
 # ==================== MODEL LOADING ====================
+MODEL_OPTIONS = {
+    "Metadata Baseline": "liaad/Citilink-BERTimbau-large-metadata-pt-baseline",
+    "Metadata Augmented": "liaad/Citilink-BERTimbau-large-Metadata-pt",
+}
+
 @st.cache_resource
-def load_model():
+def load_model(model_name: str):
     """Loads the model and tokenizer (cached to avoid reloading)"""
-    MODEL_NAME = "liaad/Citilink-BERTimbau-large-metadata-pt-baseline"
-    
     try:
-        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-        model = AutoModelForTokenClassification.from_pretrained(MODEL_NAME)
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModelForTokenClassification.from_pretrained(model_name)
         model.eval()
         return tokenizer, model
     except Exception as e:
@@ -216,6 +219,9 @@ def translate_entity_label(label):
         "PARTICIPANTE-VEREADOR-PRESENTE": ("👥", "COUNCILOR - PRESENT"),
         "PARTICIPANTE-VEREADOR-AUSENTE": ("👥", "COUNCILOR - ABSENT"),
         "PARTICIPANTE-VEREADOR-SUBSTITUIDO": ("👥", "COUNCILOR - SUBSTITUTED"),
+        "PARTICIPANTE-FUNCIONARIO-PRESENTE": ("🗂️", "STAFF - PRESENT"),
+        "PARTICIPANTE-FUNCIONARIO-AUSENTE": ("🗂️", "STAFF - ABSENT"),
+        "PARTICIPANTE-FUNCIONARIO-SUBSTITUIDO": ("🗂️", "STAFF - SUBSTITUTED"),
     }
     
     # Check if exact translation exists
@@ -239,7 +245,8 @@ def get_entity_style(label):
         "TIPO-REUNIAO": "#E67E22",         # Orange
         "HORARIO": "#F39C12",              # Yellow-orange
         "PARTICIPANTE-PRESIDENTE": "#8E44AD",  # Purple for President
-        "PARTICIPANTE-VEREADOR": "#27AE60"     # Green for Councilors
+        "PARTICIPANTE-VEREADOR": "#27AE60",    # Green for Councilors
+        "PARTICIPANTE-FUNCIONARIO": "#1A6B8A"  # Steel blue for Staff
     }
     
     # Border styles for different attendance states
@@ -273,6 +280,15 @@ def get_entity_style(label):
             border_style = border_styles["AUSENTE"]
         elif "SUBSTITUIDO" in label:
             border_style = border_styles["SUBSTITUIDO"]
+    elif "PARTICIPANTE-FUNCIONARIO" in label:
+        color = colors["PARTICIPANTE-FUNCIONARIO"]
+        # Determine border style
+        if "PRESENTE" in label:
+            border_style = border_styles["PRESENTE"]
+        elif "AUSENTE" in label:
+            border_style = border_styles["AUSENTE"]
+        elif "SUBSTITUIDO" in label:
+            border_style = border_styles["SUBSTITUIDO"]
     else:
         # For other entities, use specific colors
         base_category = label.split("-")[0]
@@ -306,6 +322,16 @@ def display_entities_compact(entities_dict):
 # ==================== SIDEBAR ====================
 with st.sidebar:
     st.markdown("### ⚙️ Configuration")
+    
+    # Model selection
+    selected_model_label = st.selectbox(
+        "🤖 Select Model",
+        options=list(MODEL_OPTIONS.keys())
+    )
+    selected_model_name = MODEL_OPTIONS[selected_model_label]
+    
+    st.markdown("---")
+    
     st.markdown("Choose an example or enter your own text:")
     
     selected_example = st.selectbox(
@@ -317,10 +343,10 @@ with st.sidebar:
     
     # About Section
     st.markdown("### 📋 About")
-    st.markdown("""
+    st.markdown(f"""
     **MiNER Stage 2** uses Named Entity Recognition models to automatically extract metadata from meeting minutes.
 
-    - **Model**: BERTimbau fine-tuned
+    - **Model**: {selected_model_label}
     - **Languages**: Portuguese
     - **Method**: Token Classification (NER) with BIO tagging
     """)
@@ -329,9 +355,9 @@ with st.sidebar:
     
     # Resources Section
     st.markdown("### 🔗 Resources")
-    st.markdown("""
+    st.markdown(f"""
     - [GitHub Repository](https://github.com/LIAAD/MiNER)
-    - [Model](https://huggingface.co/liaad/Citilink-BERTimbau-large-metadata-pt-baseline)
+    - [Model on HuggingFace](https://huggingface.co/{selected_model_name})
     """)
 
 # ==================== MAIN CONTENT ====================
@@ -412,8 +438,8 @@ No Edifício Sede do Município de Alandroal, o Senhor Presidente da Câmara Mun
     """)
 
 # Load model
-with st.spinner("Loading model..."):
-    tokenizer, model = load_model()
+with st.spinner(f"Loading model **{selected_model_label}**..."):
+    tokenizer, model = load_model(selected_model_name)
 
 if tokenizer is None or model is None:
     st.error("❌ Failed to load model. Please check if the model path is correct.")