app.py

# Copyright 2023 Dmitry Ustalov
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

__author__ = 'Dmitry Ustalov'
__license__ = 'Apache 2.0'

import csv
import re
import subprocess
from dataclasses import dataclass
from tempfile import NamedTemporaryFile
from typing import Dict, IO, List, cast, Tuple, Optional

import gradio as gr
import matplotlib.pyplot as plt
import networkx as nx
import pandas as pd


@dataclass
class Algorithm:
    name: str
    mode: Optional[str] = None
    local_name: Optional[str] = None
    local_params: Optional[str] = None
    global_name: Optional[str] = None
    global_params: Optional[str] = None

    def args_clustering(self) -> List[str]:
        args = [self.name]

        if self.mode:
            args.extend(['--mode', self.mode])

        args.extend(self.args_graph())

        if self.global_name:
            args.extend(['--global', self.global_name])

        if self.global_params:
            args.extend(['--global-params', self.global_params])

        return args

    def args_graph(self) -> List[str]:
        args = []

        if self.local_name:
            args.extend(['--local', self.local_name])

        if self.local_params:
            args.extend(['--local-params', self.local_params])

        return args


ALGORITHMS: Dict[str, Algorithm] = {
    'Watset[CW_top, CW_top]': Algorithm('watset', None, 'cw', 'mode=top', 'cw', 'mode=top'),
    'Watset[CW_lin, CW_top]': Algorithm('watset', None, 'cw', 'mode=lin', 'cw', 'mode=top'),
    'Watset[CW_log, CW_top]': Algorithm('watset', None, 'cw', 'mode=log', 'cw', 'mode=top'),
    'Watset[MCL, CW_top]': Algorithm('watset', None, 'mcl', None, 'cw', 'mode=top'),
    'Watset[CW_top, CW_lin]': Algorithm('watset', None, 'cw', 'mode=top', 'cw', 'mode=lin'),
    'Watset[CW_lin, CW_lin]': Algorithm('watset', None, 'cw', 'mode=lin', 'cw', 'mode=lin'),
    'Watset[CW_log, CW_lin]': Algorithm('watset', None, 'cw', 'mode=log', 'cw', 'mode=lin'),
    'Watset[MCL, CW_lin]': Algorithm('watset', None, 'mcl', None, 'cw', 'mode=lin'),
    'Watset[CW_top, CW_log]': Algorithm('watset', None, 'cw', 'mode=top', 'cw', 'mode=log'),
    'Watset[CW_lin, CW_log]': Algorithm('watset', None, 'cw', 'mode=lin', 'cw', 'mode=log'),
    'Watset[CW_log, CW_log]': Algorithm('watset', None, 'cw', 'mode=log', 'cw', 'mode=log'),
    'Watset[MCL, CW_log]': Algorithm('watset', None, 'mcl', None, 'cw', 'mode=log'),
    'CW_top': Algorithm('cw', 'top'),
    'CW_lin': Algorithm('cw', 'lin'),
    'CW_log': Algorithm('cw', 'log'),
    'MaxMax': Algorithm('maxmax')
}

SENSE = re.compile(r'^(?P<item>\d+)#(?P<sense>\d+)$')


def visualize(G: nx.Graph, seed: int = 0) -> plt.Figure:
    pos = nx.spring_layout(G, seed=seed)

    fig = plt.figure(dpi=240)
    plt.axis('off')
    nx.draw_networkx_edges(G, pos, alpha=.15)
    nx.draw_networkx_labels(G, pos)

    return fig


def watset(G: nx.Graph, algorithm: str, seed: int = 0,
           jar: str = 'watset.jar', timeout: int = 10) -> Tuple[pd.DataFrame, Optional[nx.Graph]]:
    with (NamedTemporaryFile() as graph,
          NamedTemporaryFile(mode='rb') as clusters,
          NamedTemporaryFile(mode='rb') as senses):
        nx.write_edgelist(G, graph.name, delimiter='\t', data=['weight'])

        try:
            result = subprocess.run(['java', '-jar', jar,
                                     '--input', graph.name, '--output', clusters.name, '--seed', str(seed),
                                     *ALGORITHMS[algorithm].args_clustering()],
                                    capture_output=True, text=True, timeout=timeout)

            if result.returncode != 0:
                raise gr.Error(f'Backend error (code {result.returncode}): {result.stderr}')
        except subprocess.SubprocessError as e:
            raise gr.Error(f'Backend error: {e}')

        df_clusters = pd.read_csv(clusters, sep='\t', names=('cluster', 'size', 'items'),
                                  dtype={'cluster': int, 'size': int, 'items': str})

        df_clusters['items'] = df_clusters['items'].str.split(', ')

        if ALGORITHMS[algorithm].name == 'watset':
            try:
                result = subprocess.run(['java', '-jar', jar,
                                         '--input', graph.name, '--output', senses.name, '--seed', str(seed),
                                         'graph', *ALGORITHMS[algorithm].args_graph()],
                                        capture_output=True, text=True, timeout=timeout)

                if result.returncode != 0:
                    raise gr.Error(f'Backend error (code {result.returncode}): {result.stderr}')
            except subprocess.SubprocessError as e:
                raise gr.Error(f'Backend error: {e}')

            G_senses = nx.read_edgelist(senses.name, delimiter='\t', comments='\n', data=[('weight', float)])

            return df_clusters, G_senses

        return df_clusters, None


def handler(file: IO[bytes], algorithm: str, seed: int) -> Tuple[pd.DataFrame, plt.Figure]:
    if file is None:
        raise gr.Error('File must be uploaded')

    if algorithm not in ALGORITHMS:
        raise gr.Error(f'Unknown algorithm: {algorithm}')

    with open(file.name) as f:
        try:
            dialect = csv.Sniffer().sniff(f.readline(4096))
            delimiter = dialect.delimiter
        except csv.Error:
            delimiter = ','

    G: nx.Graph = nx.read_edgelist(file.name, delimiter=delimiter, comments='\n', data=[('weight', float)])

    mapping, reverse = {}, {}

    for i, node in enumerate(G):
        mapping[node] = i
        reverse[i] = node

    nx.relabel_nodes(G, mapping, copy=False)

    df_clusters, G_senses = watset(G, algorithm=algorithm, seed=seed)

    nx.relabel_nodes(G, reverse, copy=False)

    df_clusters['items'] = df_clusters['items'].apply(lambda items: sorted(reverse[int(item)] for item in items))

    if G_senses is None:
        fig = visualize(G, seed=seed)
    else:
        sense_mapping = {node: f'{reverse[int(match["item"])]}#{match["sense"]}'  # type: ignore
                         for node in G_senses for match in (SENSE.match(node),)}

        nx.relabel_nodes(G_senses, sense_mapping, copy=False)

        fig = visualize(G_senses, seed=seed)

    return df_clusters, fig


def main() -> None:
    iface = gr.Interface(
        fn=handler,
        inputs=[
            gr.File(
                file_types=['.tsv', '.csv'],
                label='Graph'
            ),
            gr.Dropdown(
                choices=cast(List[str], ALGORITHMS),
                value='Watset[MCL, CW_lin]',
                label='Algorithm'
            ),
            gr.Number(
                label='Seed',
                precision=0
            )
        ],
        outputs=[
            gr.Dataframe(
                headers=['cluster', 'size', 'items'],
                label='Clustering'
            ),
            gr.Plot(
                label='Graph'
            )
        ],
        examples=[
            ['java.tsv', 'Watset[MCL, CW_lin]', 0],
            ['java.tsv', 'MaxMax', 0]
        ],
        title='Structure Discovery with Watset',
        description='''
**Watset** is a powerful algorithm for structure discovery in graphs.

By capturing the ambiguity of nodes in a graph, Watset efficiently finds clusters in the input data.

Whether you're working with linguistic data or other networks, Watset is the go-to solution for unlocking hidden patterns and structures.
        ''',
        article='''
**More Watset:**

- Paper: <https://doi.org/10.1162/COLI_a_00354> ([arXiv](https://arxiv.org/abs/1808.06696))
- Implementation: <https://github.com/nlpub/watset-java>
- Maven Central: <https://search.maven.org/artifact/org.nlpub/watset>
- conda-forge: <https://anaconda.org/conda-forge/watset>
        ''',
        allow_flagging='never'
    )

    iface.launch()


if __name__ == '__main__':
    main()