app.py

import json
import os
from collections import Counter
import pydeck as pdk
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
import plotly.io as pio
import streamlit as st
from safetensors import safe_open
from sentence_transformers import SentenceTransformer

from semantic_search import predict

HF_TOKEN = os.environ.get("HF_TOKEN")
CITIES_ENRICHED = os.path.join("data", "cities_enriched_manually.csv")
DATA = os.path.join("2025-06-13_musterdatenkatalog.json")

TAXONOMY = os.path.join("taxonomy_processed_v3.json")
MAP_PATH = os.path.join("data", "map_data.csv")  # this is only for saving
MAP_PATH_WITH_COORD = os.path.join(
    "data", "map_data_with_coord.csv"
)  # this is for saving the data with coordinates
# and local testing


def get_tree_map_data(
    data: dict,
    countings_parents: dict,
    countings_labels: dict,
    root: str = " ",
) -> tuple:
    names: list = [""]
    parents: list = [root]
    values: list = ["0"]

    for group, labels in data.items():
        parents.append(root)
        if group in countings_parents:
            values.append(str(countings_parents[group]))
            group_name_with_count = (
                group
                + "<br>"
                + "Anzahl Datensätze:"
                + " "
                + str(countings_parents[group])
            )
            names.append(group_name_with_count)
        else:
            values.append("0")
            group_name_with_count = group + "<br>" + "Anzahl Datensätze:" + " " + "0"
            names.append(group_name_with_count)
        for label in labels:
            if "-" in label:
                label = label.split("-")
                label = label[0] + "<br> -" + label[1]
            if label in countings_labels:
                label_name_with_count = (
                    label
                    + "<br>"
                    + "<br>"
                    + "Anzahl Datensätze:"
                    + "<br>"
                    + ""
                    + str(countings_labels[label])
                )
                names.append(label_name_with_count)
                parents.append(group_name_with_count)
                values.append(str(countings_labels[label]))
            if label not in countings_labels:
                if "<br>" in label:
                    if (
                        label.split("<br>")[0].strip() + label.split("<br>")[-1]
                        in countings_labels
                    ):
                        label_name_with_count = (
                            label
                            + "<br>"
                            + "<br>"
                            + "Anzahl Datensätze:"
                            + "<br>"
                            + ""
                            + str(
                                countings_labels[
                                    label.split("<br>")[0].strip()
                                    + label.split("<br>")[-1]
                                ]
                            )
                        )
                else:
                    print(label)
                    label_name_with_count = (
                        label
                        + "<br>"
                        + "<br>"
                        + "Anzahl Datensätze:"
                        + "<br>"
                        + ""
                        + "0"
                    )
                    names.append(label_name_with_count)
                    parents.append(group_name_with_count)
                    values.append("0")
    return parents, names, values


def load_json(path: str) -> dict:
    with open(path, "r") as fp:
        return json.load(fp)


# Load Data
data = load_json(DATA)
taxonomy = load_json(TAXONOMY)
taxonomy_labels = [el["group"] + " - " + el["label"] for el in taxonomy]

theme_counts = dict(Counter([el["THEMA"] for el in data]))
labels_counts = dict(Counter([el["BEZEICHNUNG"] for el in data]))

names = [""]
parents = ["Musterdatenkatalog"]

taxonomy_group_label_mapper: dict = {el["group"]: [] for el in taxonomy}

for el in taxonomy:
    if el["group"] != "Sonstiges":
        taxonomy_group_label_mapper[el["group"]].append(el["label"])
    else:
        taxonomy_group_label_mapper[el["group"]].append("Sonstiges ")

del taxonomy_group_label_mapper["Sonstiges"]

parents, names, values = get_tree_map_data(
    data=taxonomy_group_label_mapper,
    countings_parents=theme_counts,
    countings_labels=labels_counts,
    root="Musterdatenkatalog",
)

df = pd.DataFrame(data={"thema": parents, "bezeichnung": names, "value": values})
df["value"] = df["value"].astype(str)
df["bezeichnung"] = df["bezeichnung"]

fig = go.Figure(
    go.Treemap(
        labels=df["bezeichnung"],
        parents=df["thema"],
        textinfo="label",
    )
)
fig.update_layout(margin=dict(t=50, l=25, r=25, b=25))
fig.update_layout(height=1000, width=1000, template="plotly")


# load data ready to plot for local testing

# germany.drop(columns=["lat", "lon"], inplace=True)

# # or generate it directly in this script
# map_data = load_data()
# map_data = merge_geoemtry(map_data, pd.read_csv(filepath_or_buffer=CITIES_ENRICHED))
# germany = add_coor(map_data)
# germany.to_csv(MAP_PATH_WITH_COORD, index=False)

# # germany need columns with lat and lon as well as hover data


tensors = {}
with safe_open("corpus_embeddings.pt", framework="pt", device="cpu") as f:
    for k in f.keys():
        tensors[k] = f.get_tensor(k)

model = SentenceTransformer(
    model_name_or_path="and-effect/musterdatenkatalog_clf",
    device="cpu",
    use_auth_token=HF_TOKEN,
)


st.set_page_config(layout="wide")


@st.cache_data
def load_data() -> pd.DataFrame:
    germany = pd.read_csv(MAP_PATH)
    return germany


germany = load_data()
germany["lat"] = pd.to_numeric(germany["lat"])
germany["lon"] = pd.to_numeric(germany["lon"])
germany = germany[["ORG", "lat", "lon", "Count"]]


# fig_map = px.scatter_mapbox(
#     germany,
#     lat="lat",
#     lon="lon",
#     hover_name="ORG",
#     custom_data=["Count"],
#     # scope="europe",
#     height=700,
#     zoom=5,
# )
# # Custom hover template
# fig_map.update_traces(
#     hovertemplate="<br>".join(
#         [
#             "Kommune: %{hovertext}",  # Use hover_name as hovertext
#             "Count: %{customdata[0]}",  # Access elements in custom_data
#         ]
#     )
# )
# fig_map.update_layout(
#     # geo=dict(
#     #     showland=True,
#     #     landcolor="LightGray",
#     #     showocean=True,
#     #     oceancolor="LightBlue",
#     #     # showcountries=True,
#     #     # countrycolor="Gray",
#     #     showsubunits=True,
#     #     # subunitcolor="Gray",
#     #     fitbounds="locations",  # Fit the map bounds to the locations
#     #     lataxis=dict(range=[47, 55]),  # Approximate latitude range for Germany
#     #     lonaxis=dict(range=[5, 16]),  # Approximate longitude range for Germany
#     # ),
#     mapbox_style="carto-positron",
#     # height=700,
# )

# Define the layer
layer = pdk.Layer(
    "ScatterplotLayer",
    data=germany,
    get_position="[lon, lat]",
    get_radius=10000,  # or use 'Count' to scale
    get_color=[0, 0, 255, 160],
    pickable=True,
)

# Define the view
view_state = pdk.ViewState(
    latitude=51.1,
    longitude=10.5,
    zoom=4.5,
    pitch=0,
)

# Render the deck
r = pdk.Deck(
    layers=[layer],
    initial_view_state=view_state,
    tooltip={"text": "Kommune: {ORG}\nCount: {Count}"},
    map_style=None,
    height=700,
)

st.title("Musterdatenkatalog (MDK)")

st.markdown(
    """
<style>
.font {
    font-size:20px !important;
}
</style>
""",
    unsafe_allow_html=True,
)

st.markdown(
    """
<style>
.prediction {
    font-size:10px !important;
}
</style>
""",
    unsafe_allow_html=True,
)


st.markdown(
    '<p class="font">This demo showcases the algorithm of Musterdatenkatalog (MDK) of the Bertelsmann Stiftung. The MDK is a taxonomy of Open Data in municipalities in Germany. It is intended to help municipalities in Germany, as well as data analysts and journalists, to get an overview of the topics and the extent to which cities have already published data sets.</p>',
    unsafe_allow_html=True,
)

st.markdown(
    '<p class="font"> For more details checkout the <a href=https://www.bertelsmann-stiftung.de/de/unsere-projekte/smart-country/musterdatenkatalog> Musterdatenkatalog </a>.</p>',
    unsafe_allow_html=True,
)


col1, col2, col3 = st.columns(3)
col1.metric("Datensätze", len(data))
col2.metric("Themen", len(theme_counts))
col3.metric("Bezeichnungen", len(labels_counts))

st.header("Explore the MDK-Classifier")

st.markdown(
    '<p class="font"> This section allows you to predict a label from the MDK Taxonomy for a title of a dataset from municipalities. You can either enter your own dataset title or click on one of the examples. Checkout also <a href=https://www.govdata.de/> GOVDATA </a> for more dataset title examples. \
    \
    If you click on predict, the model will predict the most likely label for the dataset title. You can also change the number of labels that should be predicted. For example, if you change the Top Results to 3, the model will predict the 3 most likely labels for the dataset title in descending order. </p>',
    unsafe_allow_html=True,
)

st.markdown(
    """
<style>
/* Style columns */
[data-testid="column"] {
      border-radius: 15px;
         background-color: white;
         box-shadow: 0 0 10px #eee;
         border: 1px solid #ddd;
         padding: 1rem;;
} 

/* Style containers */
[data-testid="stVerticalBlock"] > [style*="flex-direction: column;"] > [data-testid="stVerticalBlock"] {
      border-radius: 15px;
         background-color: white;
         box-shadow: 0 0 10px #eee;
         border: 1px solid #ddd;
         padding: 1rem;;
}
</style>
""",
    unsafe_allow_html=True,
)


col1, col2 = st.columns([1.2, 1])

st.markdown(
    """
<style>
.example {
    font-size:24px !important;
}
</style>
""",
    unsafe_allow_html=True,
)


with col2:
    st.markdown(
        '<p class="example">Example Titles of Datasets</p>',
        unsafe_allow_html=True,
    )
    examples = [
        "Spielplätze",
        "Berliner Weihnachtsmärkte 2022",
        "Hochschulwechslerquoten zum Masterstudium nach Bundesländern",
        "Umringe der Bebauungspläne von Etgert",
    ]

    for example in examples:
        if st.button(example):
            if "key" not in st.session_state:
                st.session_state["query"] = example


with col1:
    tabs_font_css = """
    <style>
    div[class*="stTextInput"] label p {
    font-size: 2px;
    }
    </style>    
    """

    st.write(tabs_font_css, unsafe_allow_html=True)

    st.markdown(
        '<p class="example">Enter a dataset title</p>',
        unsafe_allow_html=True,
    )

    if "query" not in st.session_state:
        query = st.text_input("")
    if "query" in st.session_state and st.session_state.query in examples:
        query = st.text_input("Enter a dataset title", value=st.session_state.query)
    if "query" in st.session_state and st.session_state.query not in examples:
        del st.session_state["query"]
        query = st.text_input("Enter a dataset title")

    top_k = st.select_slider("Top Results", options=[1, 2, 3, 4, 5], value=1)

    predictions = predict(
        query=query,
        corpus_embeddings=tensors["corpus_embeddings"],
        corpus_labels=taxonomy_labels,
        top_k=top_k,
        model=model,
    )

    if st.button("Predict"):
        for prediction in predictions:
            st.markdown(f'<p class="font"> {prediction} <p>', unsafe_allow_html=True)


st.header("Musterdatenkatalog Taxonomy")
st.write("Data as of 13.06.2025")
st.plotly_chart(fig)

st.header("Locations with Musterdatensätzen")
st.write("Data as of 13.06.2026")
st.markdown(
    """<p class="font">Hover over the map to see how many datasets are available
            for this location. </p>
   """,
    unsafe_allow_html=True,
)
st.pydeck_chart(r)