Could not use `streamlit` to annotate a dataset of multiple labels

I am trying to build an annotation interface using streamlit.

In my dataset, each data point may have multiple labels (i.e. labels in the code below). However, I could only select one label using st.multiselect() rather than the expected "multiple select". Specifically, every time I click the one of the choices, the page will be updated and the next data point pops up.

I am not sure what went wrong after getting stuck in this for hours. Could anyone provide any pointers for me?

import pandas as pd
import streamlit as st

df = pd.read_pickle("unlabeled.pkl")
records = df.to_dict("records")

if "annotations" not in st.session_state:
    st.session_state.records = records
    st.session_state.current_record = records[0]

annotated_data = list()

if st.session_state.records:
    labels = st.session_state.current_record["labels"]
    example = st.session_state.current_record["example"]
    text = st.session_state.current_record["text"]

    demo = "\n".join(["- {}".format(ee) for ee in example])
    text = "- {}".format(text)

    st.write(f"# Example\n{demo}\n# Output\n{text}")

    labels = st.multiselect(
        label="Select Labels",
        options=labels
    )

    st.write('You Selected:', labels)

    if st.button("Save"):
        st.session_state.records.remove(st.session_state.current_record)
        st.session_state.current_record = st.session_state.records[0]

    annotated_data.append(
        {
            **st.session_state.current_record,
            "label": labels
        }
    )

    if len(annotated_data) % 50 == 0:
        save_data(annotated_data)

save_data(annotated_data)

You can use a form to prevent the app from rerunning after interacting with the multiselect widget (instead, the app won't rerun until you hit "Submit").