Embedding similarity filtering of entity descriptions

filter_explanations_dash.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import json
from base64 import b64decode

import dash_bootstrap_components as dbc
import numpy as np
from dash import Dash, Input, Output, State, callback, dash_table, dcc, html
from dash.exceptions import PreventUpdate
from sentence_transformers import SentenceTransformer
from sentence_transformers.util import cos_sim

model = SentenceTransformer("distiluse-base-multilingual-cased-v1")


@callback(
    Output("original_data_table", "data"),
    Output("embedding-store", "data"),
    Input("upload-data", "contents"),
)
def load_explanation_data(data):
    if not data:
        raise PreventUpdate
    _, data = data.split(",")
    data = json.loads(b64decode(data).decode("utf-8"))
    entity_explanations = [{"name": k, "explanation": v} for k, v in data.items()]
    return entity_explanations, model.encode(
        [item["explanation"] for item in entity_explanations]
    )


@callback(
    Output("filtered_data_table", "data", allow_duplicate=True),
    Input("original_data_table", "data"),
    prevent_initial_call=True,
)
def load_filtered_data(data):
    # This function is used to populate filtered data in case of a new upload
    if not data:
        raise PreventUpdate
    return data


@callback(
    Output("filtered_data_table", "data"),
    Output("filtered_explanations", "children"),
    Input("threshold", "value"),
    Input("original_data_table", "selected_rows"),
    State("original_data_table", "data"),
    State("embedding-store", "data"),
    prevent_initial_call=True,
)
def filter_explanations(threshold, selected_rows, original_data, embeddings):
    if not selected_rows:
        raise PreventUpdate

    embeddings = np.array(embeddings)
    sims = cos_sim(embeddings, embeddings[selected_rows])
    reduced_sims = (sims > threshold).sum(axis=1)
    # return any item that has no similarity to any bad explanation
    return [
        item for i, item in enumerate(original_data) if reduced_sims[i] == 0
    ], f"Number of filtered explanations: {(reduced_sims > 0).sum()}"


@callback(
    Output("export", "data"),
    Input("export-btn", "n_clicks"),
    State("filtered_data_table", "data"),
    prevent_initial_call=True,
)
def export_data(n_clicks, data):
    export_data = {item["name"]: item["explanation"] for item in data}
    return {
        "content": json.dumps(export_data, indent=2),
        "filename": "filtered_explanations.json",
    }


upload_button = html.Div(
    [
        dcc.Upload(
            id="upload-data", children=dbc.Button("Upload explanations"), multiple=False
        )
    ]
)
export_button = html.Div(
    [
        dbc.Button("Export filtered explanations", id="export-btn"),
        dcc.Download(id="export"),
    ]
)
threshold_slider = html.Div(
    [
        dbc.Label("Select similarity threshold"),
        dcc.Slider(id="threshold", min=0, max=1, step=0.1, value=0.5),
    ]
)
filtered_explanations = dcc.Markdown(
    id="filtered_explanations", children="Number of filtered explanations: 0"
)

top_row = dbc.Row(
    [
        dbc.Col(upload_button),
        dbc.Col(threshold_slider),
        dbc.Col(filtered_explanations),
        dbc.Col(export_button),
    ],
    class_name="py-2",
)

original_data_table = dash_table.DataTable(
    id="original_data_table",
    row_selectable="multi",
    editable=False,
    page_size=10,
    style_data={"whiteSpace": "normal", "height": "auto"},
)

filtered_data_table = dash_table.DataTable(
    id="filtered_data_table",
    editable=False,
    page_size=10,
    style_data={"whiteSpace": "normal", "height": "auto"},
)

mid_row = dbc.Row(
    [dbc.Col(original_data_table, width=6), dbc.Col(filtered_data_table, width=6)],
    class_name="py-2",
)

app = Dash(external_stylesheets=[dbc.themes.BOOTSTRAP])
app.layout = dbc.Container(
    [
        dcc.Store(id="embedding-store"),
        html.H1("Filter entities based on embedding similarity"),
        html.Hr(),
        top_row,
        mid_row,
    ]
)

if __name__ == "__main__":
    app.run_server(debug=True)

This demo takes a json dictionary {name: explanation} as input, and allows you to select entities to filter out.

The filtering happens based on cosine similarity of explanation embeddings using distiluse-base-multilingual-cased-v1. Embeddings of just the explanation string are currently used, the entity name is not taken into account.