Enhance Dockerfile and main.py: add non-root user, improve PII masking tool UI with label selection

Dockerfile

@@ -2,12 +2,23 @@ FROM python:3.10-slim
 
 WORKDIR /app
 
-RUN apt-get update && apt-get install -y redis-server && rm -rf /var/lib/apt/lists/*
-
-COPY requirements.txt /app
+# Install system packages as root
+RUN apt-get update && \
+    apt-get install -y redis-server && \
+    rm -rf /var/lib/apt/lists/*
 
+# Copy and install Python requirements
+COPY requirements.txt /app/
 RUN pip install --no-cache-dir -r requirements.txt
 
-COPY . /app
+# Copy application code
+COPY . /app/
+
+# Create non-root user and set permissions
+RUN useradd -m -u 1000 user && \
+    chown -R user:user /app
+
+# Switch to non-root user
+USER user
 
 CMD ["/bin/bash", "-c", "redis-server --daemonize yes && streamlit run main.py --server.port=7860 --server.address=0.0.0.0"]

main.py

@@ -1,16 +1,35 @@
 import math
 import streamlit as st
 
-from utils.model import pii_masking_pipeline
+from utils.model import pii_masking_pipeline, valid_labels
 from utils.redis import redis_client
 
 st.set_page_config(page_title="PII Masking Tool", page_icon="🔒")
 
 st.title("PII Masking Tool")
 
+if 'selected_labels' not in st.session_state:
+    st.session_state.selected_labels = ["name", "nric", "phone number", "address", "email"]
+
+def select_all_labels():
+    st.session_state.selected_labels = valid_labels
+
 # Text input
 text_to_mask = st.text_area("Enter text to mask PII:", height=200)
 
+col1, col2 = st.columns([3, 1])
+with col1:
+    selected_labels = st.multiselect(
+        "Select PII types to mask:",
+        options=valid_labels,
+        key='selected_labels', 
+        help="Choose the types of PII you want to redact from the text"
+    )
+
+with col2:
+    st.button("Select All", on_click=select_all_labels)
+
+
 if st.button("Mask Text"):
     if text_to_mask:
         masked_text = pii_masking_pipeline(text_to_mask)