first commit

Dockerfile

@@ -0,0 +1,24 @@
+FROM python:3.9
+
+RUN --mount=target=/root/packages.txt,source=packages.txt 	apt-get update &&  xargs -r -a /root/packages.txt apt-get install -y
+
+RUN pip install --no-cache-dir pip==22.3.1 &&     pip install --no-cache-dir         datasets         "huggingface-hub>=0.12.1" "protobuf<4" "click<8.1"
+
+WORKDIR /home/user/app
+
+RUN apt-get install -y 	git 	git-lfs 	ffmpeg 	libsm6 	libxext6 	cmake 	libgl1-mesa-glx && git lfs install
+
+RUN pip install --no-cache-dir Cython "gradio==3.37.0" "torch==1.10.1"  jieba subword-nmt sacremoses transformers
+
+RUN git clone --recurse-submodules https://github.com/thu-coai/DA-Transformer.git && cd DA-Transformer && pip install -e . && cd dag_search && python3 setup.py build_ext --inplace && pip install -e . && cd ../..
+
+RUN mkdir -p /home/user && chmod 777 /home/user
+
+USER user
+
+ENV HOME=/home/user \
+	PATH=/home/user/.local/bin:$PATH
+
+COPY . .
+
+CMD ["python3", "app.py"]

app.py

@@ -0,0 +1,554 @@
+import argparse
+from collections import defaultdict
+import datetime
+import json
+import os, sys
+import time
+import concurrent
+
+import math
+import gradio as gr
+import requests
+import logging
+import numpy as np
+import matplotlib.pyplot as plt
+import fairseq
+
+logger = logging.getLogger(__name__)
+
+fairseq_path = os.path.dirname(os.path.dirname(fairseq.__file__))
+
+sys.path.insert(1, f"{fairseq_path}")
+from fs_plugins.models.glat_decomposed_with_link import GlatDecomposedLink
+
+sys.path.insert(1, f"{fairseq_path}/examples")
+from mass.s2s_model import TransformerMASSModel
+from transformer.hub_interface import TransformerHubInterface
+
+notice_markdown = ("""
+# ⚡ Directed Acyclic Transformer: A Non-Autoregressive Sequence-to-Sequence Model designed for Parallel Text Generation.
+- **Fast Generation**: DA-Transformer offers faster inference compared to autoregressive Transformers (with fairseq implementation), with a reduction in latency by 7~14x and an increase in throughput by ~20x.
+- **High Quality**: DA-Transformer performs competitively with autoregressive Transformers, even with pre-trained models like BART, in a variety of text generation tasks.
+- **Easy Training**: DA-Transformer can be trained end-to-end without requiring knowledge distillation, making it simple and straightforward to train.
+
+## Resources
+
+- Codes: [[Github]](https://github.com/thu-coai/DA-Transformer)
+- Papers: [[Machine Translation]](https://proceedings.mlr.press/v162/huang22m/huang22m.pdf) [[Pre-training]](https://arxiv.org/pdf/2304.11791.pdf)
+
+## Terms of use
+By using this service, users are required to agree to the following terms: The service is a research preview intended for non-commercial use only. It does not gaurantee the correctness of the output text. The service may collect user data for future research.
+
+## This demo contains models for
+- [Zh-En Translation](https://huggingface.co/thu-coai/dat_base_translation_zhen)
+- [En-De Translation](https://huggingface.co/thu-coai/dat_base_translation_ende)
+- [Question Generation](https://huggingface.co/thu-coai/dat_uncased_squad)
+""")
+
+learn_more_markdown = ("""
+""")
+
+
+css = """
+pre {
+    white-space: pre-wrap;       /* Since CSS 2.1 */
+    white-space: -moz-pre-wrap;  /* Mozilla, since 1999 */
+    white-space: -pre-wrap;      /* Opera 4-6 */
+    white-space: -o-pre-wrap;    /* Opera 7 */
+    word-wrap: break-word;       /* Internet Explorer 5.5+ */
+}
+"""
+
+available_models = {
+    "dat_base_translation_ende": {
+        "class": GlatDecomposedLink,
+        "args":{
+            "model_name_or_path": "hfhub://thu-coai/dat_base_translation_ende",
+            "decode_strategy": "beamsearch",
+            "decode_max_workers": 1,
+            "decode_threads_per_worker": 4,
+            "decode_dedup": True,
+            "decode_alpha": 1.1,
+            "decode_gamma": 0,
+            "decode_beam_size": 200,
+            "decode_batch_size": 1,
+            "decode_top_cand": 5,
+            "decode_max_beam_per_length": 10,
+            "max_decoder_batch_tokens": 2048
+        },
+        "examples": ["I am a fast translation model."],
+        "expected_load_time": 17
+    },
+    "dat_base_translation_zhen": {
+        "class": GlatDecomposedLink,
+        "args":{
+            "model_name_or_path": "hfhub://thu-coai/dat_base_translation_zhen",
+            "decode_strategy": "beamsearch",
+            "decode_max_workers": 1,
+            "decode_threads_per_worker": 4,
+            "decode_dedup": True,
+            "decode_alpha": 1.1,
+            "decode_gamma": 0,
+            "decode_beam_size": 200,
+            "decode_batch_size": 1,
+            "decode_top_cand": 5,
+            "decode_max_beam_per_length": 10,
+            "max_decoder_batch_tokens": 2048
+        },
+        "examples": ["我是一个高速的机器翻译模型。"],
+        "expected_load_time": 17
+    },
+    "dat_uncased_squad": {
+        "class": GlatDecomposedLink,
+        "args":{
+            "model_name_or_path": "hfhub://thu-coai/dat_uncased_squad",
+            "decode_strategy": "beamsearch",
+            "decode_max_workers": 1,
+            "decode_threads_per_worker": 4,
+            "decode_gamma": 0,
+            "decode_beam_size": 200,
+            "decode_batch_size": 1,
+            "decode_top_cand": 5,
+            "decode_no_consecutive_repeated_tokens": 3,
+            "decode_no_repeated_tokens": 2,
+            "decode_max_beam_per_length": 10,
+            "max_decoder_batch_tokens": 2048
+        },
+        "examples": ["Two [SEP] Two additional teams of 40 attendants each will accompany the flame on its mainland China route."],
+        "expected_load_time": 20
+    },
+    "mass_uncased_squad": {
+        "class": TransformerMASSModel,
+        "args":{
+            "model_name_or_path": "hfhub://thu-coai/mass_uncased_squad"
+        },
+        "examples": ["Two [SEP] Two additional teams of 40 attendants each will accompany the flame on its mainland China route."],
+        "expected_load_time": 10
+    },
+    "transformer_base_translation_ende": {
+        "class": TransformerHubInterface,
+        "args":{
+            "model_name_or_path": "hfhub://thu-coai/transformer_base_translation_ende"
+        },
+        "examples": ["I am a fast translation model."],
+        "expected_load_time": 10
+    },
+    "transformer_base_translation_zhen": {
+        "class": TransformerHubInterface,
+        "args":{
+            "model_name_or_path": "hfhub://thu-coai/transformer_base_translation_zhen"
+        },
+        "examples": ["我是一个高速的机器翻译模型。"],
+        "expected_load_time": 10
+    }
+}
+
+compare_available_types = {
+    "Translation Zh-En: DA-Transformer v.s. Autoregressive Transformer": {
+        "models": ['dat_base_translation_zhen', 'transformer_base_translation_zhen'],
+        "examples": ["我是一个高速的机器翻译模型。", "非自回归模型可以用来加速自然语言生成。",
+                     "使用本服务前，用户必须同意以下条款：该服务是仅供非商业用途的研究预览。它不保证输出文本的正确性。本服务可能会收集用户数据以供将来研究。"],
+        "placeholder": "请输入一个中文句子。 (The model will translate the input into English.)"
+    },
+    "Question Generation: DA-Transformer v.s. MASS": {
+        "models": ['dat_uncased_squad', "mass_uncased_squad"],
+        "examples": ["Two [SEP] Two additional teams of 40 attendants each will accompany the flame on its mainland China route.", "DA-Transformer [SEP] Directed Acyclic Transformer (DA-Transformer) is a non-autoregressive sequence-to-sequence model designed for parallel text generation."],
+        "placeholder": "Answer [SEP] Your Passage Here (the answer should be appearred in the passage)."
+    },
+    "Translation En-De: DA-Transformer v.s. Autoregressive Transformer": {
+        "models": ['dat_base_translation_ende', 'transformer_base_translation_ende'],
+        "examples": ["I am a fast translation model.", "Non-autoregressive models are designed for fast natural language generation.",
+                     "By using this service, users are required to agree to the following terms: The service is a research preview intended for non-commercial use only."],
+        "placeholder": "Any English sentence here. (The model will translate the input into German.)"
+    },
+}
+
+detail_available_types = {
+    "Translation Zh-En": {
+        "model": 'dat_base_translation_zhen',
+        "examples": compare_available_types['Translation Zh-En: DA-Transformer v.s. Autoregressive Transformer']["examples"],
+        "placeholder": compare_available_types['Translation Zh-En: DA-Transformer v.s. Autoregressive Transformer']["placeholder"]
+    },
+    "Question Generation": {
+        "model": 'dat_uncased_squad',
+        "examples": compare_available_types['Question Generation: DA-Transformer v.s. MASS']["examples"],
+        "placeholder": compare_available_types['Question Generation: DA-Transformer v.s. MASS']["placeholder"]
+    },
+    "Translation En-De": {
+        "model": 'dat_base_translation_ende',
+        "examples": compare_available_types['Translation En-De: DA-Transformer v.s. Autoregressive Transformer']["examples"],
+        "placeholder": compare_available_types['Translation En-De: DA-Transformer v.s. Autoregressive Transformer']["placeholder"],
+    },
+}
+
+models = {}
+workers = None
+
+def softplus(x, beta=1):
+    return math.log1p(math.exp(-abs(x * beta))) / beta + max(x, 0)
+
+def get_fake_progress(min_progress, max_progress, used_time, expected_time):
+    percentage = max(1 - softplus(expected_time - used_time) / expected_time, 0)
+    return min_progress + (max_progress - min_progress) * percentage
+
+def generate(model, model_input):
+    return {"output": model.translate(model_input)}
+
+def generate_detail(model, model_input):
+    output, graph_info = model.generate_graph(model_input)
+    return {"output": output, "graph_info": graph_info}
+
+def load_model(model_name):
+    assert model_name in available_models
+    logger.info(f"start loading {model_name}")
+    model = available_models[model_name]['class'].from_pretrained(**available_models[model_name]['args'])
+    return model
+
+def warmup_model(model, model_name):
+    model.translate(available_models[model_name]['examples'][0])
+
+def submit(model_name, model_input, generate_fn, request: gr.Request, progress=gr.Progress()):
+    assert workers is not None, "No workers"
+    current_progress = 0
+
+    progress(0, desc="Downloading Checkpoints and Loading Models")
+    if model_name not in models:
+        load_start = time.time()
+        future = workers.submit(load_model, model_name)
+        while True:
+            try:
+                model = future.result(timeout=1)
+                break
+            except concurrent.futures._base.TimeoutError as _:
+                progress(get_fake_progress(min_progress=current_progress, max_progress=0.8, used_time=time.time() - load_start, expected_time=available_models[model_name]['expected_load_time']),
+                        desc="Downloading Checkpoints and Loading Models")
+        logger.info(f"Model Loaded: {model_name} Load Time: {time.time() - load_start}")
+        current_progress = 0.8
+        models[model_name] = model
+    else:
+        model = models[model_name]
+
+    # warmup for better inference time
+    progress(current_progress, desc="Downloading Checkpoints and Loading Models")
+    if current_progress == 0.8:
+        target_progress = 0.9
+    else:
+        target_progress = 0.5
+    warmup_start = time.time()
+    future = workers.submit(warmup_model, model, model_name)
+    while True:
+        try:
+            result = future.result(timeout=1)
+            break
+        except concurrent.futures._base.TimeoutError as _:
+            progress(get_fake_progress(min_progress=current_progress, max_progress=target_progress, used_time=time.time() - warmup_start, expected_time=1),
+                        desc="Downloading Checkpoints and Loading Models")
+    current_progress = target_progress
+
+    # running
+    progress(current_progress, desc="Running")
+    try:
+        generate_start = time.time()
+        future = workers.submit(generate_fn, model, model_input)
+        while True:
+            try:
+                result = future.result(timeout=1)
+                break
+            except concurrent.futures._base.TimeoutError as _:
+                progress(get_fake_progress(min_progress=current_progress, max_progress=1, used_time=time.time() - generate_start, expected_time=1),
+                            desc="Running")
+        inference_time = time.time() - generate_start
+
+        result_abbrev = {}
+        for key, value in result.items():
+            log_str = str(value)
+            if len(log_str) > 1024:
+                log_str = log_str[:1024] + "..."
+            result_abbrev[key] = log_str
+        logger.info(f"Input: [{model_input}] Output: [{result_abbrev}] Inference Time: {inference_time}")
+        return result, inference_time
+    except RuntimeError as err:
+        return f"Runtime Error: {str(err)}", 0
+
+
+def compare_init_state(model_selector):
+    model1 = compare_available_types[model_selector]['models'][0]
+    model2 = compare_available_types[model_selector]['models'][1]
+    state = [{"model_name": model1}, {"model_name": model2}]
+    return state
+
+def compare_refresh(model_selector, samples):
+    model1 = compare_available_types[model_selector]['models'][0]
+    model2 = compare_available_types[model_selector]['models'][1]
+    model_output1 = gr.Textbox.update(visible=True, label=model1)
+    model_output2 = gr.Textbox.update(visible=True, label=model2)
+    model_input = gr.Textbox.update(value="", placeholder=compare_available_types[model_selector]['placeholder'])
+    samples.clear()
+    samples += [[x]for x in compare_available_types[model_selector]['examples']]
+    examples = gr.Dataset.update(samples=samples)
+    model_speed = gr.Plot.update(visible=False)
+    return model_input, model_output1, model_output2, examples, samples, model_speed
+
+def compare_submit(model_input, idx, state, request: gr.Request, progress=gr.Progress()):
+    model_name = state[idx]['model_name']
+    model_output, inference_time = submit(model_name, model_input, generate, request, progress)
+    state[idx]['inference_time'] = inference_time
+    return model_output['output'], state
+
+def compare_dataset_click(examples, samples):
+    return samples[examples][0]
+
+def compare_show_plot(state):
+    x = [state[0]['model_name'], state[1]['model_name']]
+    y = [state[0]['inference_time'], state[1]['inference_time']]
+
+    fig = plt.figure(figsize=(12, 2.5))
+    ax = plt.subplot(111)
+    bars = ax.barh(x, y, 0.75)
+    ax.bar_label(bars, fmt="%.2f")
+    ax.set_yticks(np.arange(len(x)), labels=x)
+    ax.set_xlabel('Inference Time on CPU (s)')
+    plt.tight_layout()
+    # plt.subplots_adjust(left=0.1, bottom=0.1, right=0.9, top=0.9, wspace=0, hspace=0)
+
+    return gr.Row.update(visible=True), gr.Plot.update(value=fig, visible=True)
+
+def compare_clear():
+    return "", "", "", gr.Row.update(visible=False)
+
+example_list = []
+
+def build_tab_compare():
+    state = gr.State()
+    samples = gr.State(example_list)
+
+    available_type_names = list(compare_available_types.keys())
+    with gr.Row(elem_id="compare_model_selector_row"):
+        model_selector = gr.Dropdown(
+            choices=available_type_names,
+            value=available_type_names[0] if len(available_type_names) > 0 else "",
+            interactive=True,
+            show_label=False,
+            container=False)  # Recommended usage
+
+        # model_selector = gr.Dropdown(
+        #     choices=available_type_names,
+        #     value=available_type_names[0] if len(available_type_names) > 0 else "",
+        #     interactive=True,
+        #     show_label=False).style(container=False)
+
+    with gr.Row(elem_id="compare_model_input"):
+        model_input = gr.Textbox(lines=5, label="input")
+        # examples = gr.Dataset(examples=[], inputs=[model_input], elem_id="compare_examples")
+    examples = gr.Dataset(components=[model_input],
+        label="Examples",
+        type='index',
+        samples=example_list,
+        visible=True
+    )
+
+    # with gr.Row(elem_id="compare_examples"):
+
+    with gr.Row():
+        clear_btn = gr.Button(value="Clear")
+        submit_btn = gr.Button(value="Submit", variant="primary")
+
+    # with gr.Accordion("Parameters", open=False, visible=False) as parameter_row:
+    #     temperature = gr.Slider(minimum=0.0, maximum=1.0, value=0.7, step=0.1, interactive=True, label="Temperature",)
+    #     max_output_tokens = gr.Slider(minimum=0, maximum=1024, value=512, step=64, interactive=True, label="Max output tokens",)
+
+    with gr.Row(elem_id="compare_model_output"):
+        model_output1 = gr.Textbox(lines=5, label="output", visible=False)
+        model_output2 = gr.Textbox(lines=5, label="output", visible=False)
+
+    with gr.Row(elem_id="compare_model_speed", visible=False) as row:
+        with gr.Column():
+            model_speed = gr.Plot(value=None, label="Speed")
+            compare_hints = gr.Markdown("**Note the above time is measured on a free cloud server, which does not use GPU and is thus different from the setting in the papers.**")
+
+    model_selector.change(compare_refresh, [model_selector, samples], [model_input, model_output1, model_output2, examples, samples, model_speed])
+
+    clear_btn.click(compare_clear, None, [model_input, model_output1, model_output2, row])
+
+    submit_btn.click(compare_init_state, [model_selector], [state]).\
+                then(compare_submit, [model_input, gr.Number(value=0, visible=False, precision=0), state], [model_output1, state]).\
+                then(compare_submit, [model_input, gr.Number(value=1, visible=False, precision=0), state], [model_output2, state]).\
+                then(compare_show_plot, [state], [row, model_speed])
+    # submit_btn.click(compare_show_plot, [state], [model_speed])
+
+    examples.click(compare_dataset_click, [examples, samples], [model_input])
+
+    def load(fn):
+        fn(compare_refresh, [model_selector, samples], [model_input, model_output1, model_output2, examples, samples])
+
+    return load
+
+def detail_init_state(model_selector):
+    model = detail_available_types[model_selector]['model']
+    state = {"model_name": model, "cnt": 0}
+    return state
+
+def detail_refresh(model_selector, samples):
+    model = detail_available_types[model_selector]['model']
+    model_output = gr.Textbox.update(visible=True, label=model)
+    model_input = gr.Textbox.update(value="", placeholder=detail_available_types[model_selector]['placeholder'])
+    samples.clear()
+    samples += [[x]for x in detail_available_types[model_selector]['examples']]
+    examples = gr.Dataset.update(samples=samples)
+    model_speed = gr.Plot.update(visible=False)
+    return model_input, model_output, examples, samples, model_speed
+
+def detail_submit(model_input, state, request: gr.Request, progress=gr.Progress()):
+    model_name = state['model_name']
+    model_output, inference_time = submit(model_name, model_input, generate_detail, request, progress)
+    state['inference_time'] = inference_time
+    state["graph_info"] = model_output['graph_info']
+    # html_code = open("graph.html").read()
+
+    # state["cnt"] += 1
+    # if state["cnt"] > 2:
+    #     html_code += r"""<script type="text/javascript">addNode();</script>\n"""
+    # print(html_code)
+    
+    return model_output['output'], state, gr.Row.update(visible=True), json.dumps(state)
+
+def detail_dataset_click(examples, samples):
+    return samples[examples][0]
+
+def detail_clear():
+    return "", "", gr.Row.update(visible=False)
+
+def build_tab_detail():
+
+    state = gr.State()
+    samples = gr.State(example_list)
+
+    available_type_names = list(detail_available_types.keys())
+    with gr.Row(elem_id="detail_model_selector_row"):
+        model_selector = gr.Dropdown(
+            choices=available_type_names,
+            value=available_type_names[0] if len(available_type_names) > 0 else "",
+            interactive=True,
+            show_label=False,
+            container=False)
+
+    with gr.Row(elem_id="detail_model_input"):
+        model_input = gr.Textbox(lines=5, label="input")
+        # examples = gr.Dataset(examples=[], inputs=[model_input], elem_id="compare_examples")
+    examples = gr.Dataset(components=[model_input],
+        label="Examples",
+        type='index',
+        samples=example_list,
+        visible=True
+    )
+
+    # with gr.Row(elem_id="compare_examples"):
+
+    with gr.Row():
+        clear_btn = gr.Button(value="Clear")
+        submit_btn = gr.Button(value="Submit", variant="primary")
+
+    # with gr.Accordion("Parameters", open=False, visible=False) as parameter_row:
+    #     temperature = gr.Slider(minimum=0.0, maximum=1.0, value=0.7, step=0.1, interactive=True, label="Temperature",)
+    #     max_output_tokens = gr.Slider(minimum=0, maximum=1024, value=512, step=64, interactive=True, label="Max output tokens",)
+
+    with gr.Row(elem_id="detail_model_output"):
+        model_output = gr.Textbox(lines=5, label="output", visible=False)
+
+    with gr.Row(visible=False) as dag_graph:
+        with gr.Column(scale=1.8):
+            html = gr.HTML(open("graph.html").read())
+        with gr.Column(scale=1):
+            minimum_node_pass_prob = gr.Slider(0, 1, value=0.2, label="Show nodes with passing probability greater than", info="Nodes that predict the output sequence are always visible")
+            minimum_edge_prob = gr.Slider(0, 1, value=0.1, label="Show edges with transition probability greater than")
+            max_out_edge_num = gr.Slider(1, 10, value=5, step=1, label="Show top-k outgoing edges with k")
+            max_out_edge_prob = gr.Slider(0, 1, value=0.9, label="Show top-p outgoing edges with p")
+            force_in_edge = gr.Checkbox(True, label="Show at least one incoming edge for each node")
+            show_node_detail = gr.Checkbox(False, label="Show verbose node information")
+            show_edge_label = gr.Checkbox(False, label="Show transition probability")
+            network_refresh = gr.Button(value="Reinitialize DAG Visualization")
+        graph_parameters = [minimum_node_pass_prob, minimum_edge_prob, max_out_edge_num, max_out_edge_prob, force_in_edge, show_node_detail, show_edge_label]
+
+    js_state = gr.Textbox(visible=False)
+
+    model_selector.change(detail_refresh, [model_selector, samples], [model_input, model_output, examples, samples])
+
+    clear_btn.click(detail_clear, None, [model_input, model_output, dag_graph])
+
+    graph_create_js = """(state_str, minimum_node_pass_prob, minimum_edge_prob, max_out_edge_num, max_out_edge_prob, force_in_edge, show_node_detail, show_edge_label) => {
+                            var state = JSON.parse(state_str);
+                            var options = {
+                                minimum_node_pass_prob: minimum_node_pass_prob,
+                                minimum_edge_prob: minimum_edge_prob,
+                                max_out_edge_num: max_out_edge_num,
+                                max_out_edge_prob: max_out_edge_prob,
+                                force_in_edge: force_in_edge,
+                                show_node_detail: show_node_detail,
+                                show_edge_label: show_edge_label,
+                            }
+                            startNetwork(state.graph_info, options);
+                    }"""
+    graph_update_js = """(minimum_node_pass_prob, minimum_edge_prob, max_out_edge_num, max_out_edge_prob, force_in_edge, show_node_detail, show_edge_label) => {
+                            var options = {
+                                minimum_node_pass_prob: minimum_node_pass_prob,
+                                minimum_edge_prob: minimum_edge_prob,
+                                max_out_edge_num: max_out_edge_num,
+                                max_out_edge_prob: max_out_edge_prob,
+                                force_in_edge: force_in_edge,
+                                show_node_detail: show_node_detail,
+                                show_edge_label: show_edge_label,
+                            }
+                            updateNetwork(options);
+                    }"""
+    submit_btn.click(detail_init_state, [model_selector], [state]).\
+                then(detail_submit, [model_input, state], [model_output, state, dag_graph, js_state]).\
+                then(None, [js_state] + graph_parameters, None, _js=graph_create_js)
+    network_refresh.click(None, [js_state] + graph_parameters, None, _js=graph_create_js)
+    minimum_node_pass_prob.change(None, graph_parameters, None, _js=graph_update_js)
+    minimum_edge_prob.change(None, graph_parameters, None, _js=graph_update_js)
+    max_out_edge_num.change(None, graph_parameters, None, _js=graph_update_js)
+    max_out_edge_prob.change(None, graph_parameters, None, _js=graph_update_js)
+    force_in_edge.select(None, graph_parameters, None, _js=graph_update_js)
+    show_node_detail.select(None, graph_parameters, None, _js=graph_update_js)
+    show_edge_label.select(None, graph_parameters, None, _js=graph_update_js)
+
+    examples.click(detail_dataset_click, [examples, samples], [model_input])
+
+    def load(fn):
+        fn(detail_refresh, [model_selector, samples], [model_input, model_output, examples, samples])
+
+    return load
+
+def build_demo():
+    with gr.Blocks(title="DA-Transformer Demo", theme=gr.themes.Base(), css=css) as demo:
+        gr.Markdown(notice_markdown)
+
+        with gr.Tab("DA-Transformer Inspection") as detail_tab:
+            detail_load = build_tab_detail()
+        detail_load(detail_tab.select)
+
+        with gr.Tab("Speed Comparison") as compare_tab:
+            compare_load = build_tab_compare()
+        compare_load(compare_tab.select)
+
+        gr.Markdown(learn_more_markdown)
+
+        detail_load(demo.load)
+
+        demo.load(None,None,None,_js=open("global.js").read())
+    return demo
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--host", type=str, default="0.0.0.0")
+    parser.add_argument("--port", type=int)
+    parser.add_argument("--concurrency-count", type=int, default=1)
+    parser.add_argument("--share", action="store_true")
+    args = parser.parse_args()
+    logger.info(f"args: {args}")
+
+    workers = concurrent.futures.ThreadPoolExecutor(max_workers=1)
+    demo = build_demo()
+    demo.queue(concurrency_count=args.concurrency_count, status_update_rate=10,
+               api_open=False).launch(server_name=args.host, server_port=args.port,
+                                      share=args.share, max_threads=5)

global.js

@@ -0,0 +1,446 @@
+() => {
+    var script = document.createElement("script");
+    script.src = "https://visjs.github.io/vis-network/standalone/umd/vis-network.min.js";
+    document.head.appendChild(script);
+
+    // var current_options;
+    // var node_pass_prob, prelen, max_path, node_tokens, node_probs, links, tarlen;
+    // var visible_nodes, map_max_path;
+    // var visible_edges, sorted_links_out, sorted_links_in;
+    // var nodes, edges, network;
+    // var disable_edge_select;
+    showing_node = undefined;
+    showing_edge = undefined;
+    
+    function obtain_node_object(node_id) {
+        var title = "Node " + node_id.toString() + " (Pass Prob = " + node_pass_prob[node_id].toFixed(2) + ")\n";
+        title += "==========\n"
+        for(var j=0;j<5;j++){
+            if(node_probs[node_id][j]> 0.01){
+                title += node_tokens[node_id][j] + "\t" + node_probs[node_id][j].toFixed(2) + "\n";
+            }
+        }
+
+        var currentNode = {"id": node_id, "label": node_tokens[node_id][0], "value":node_pass_prob[node_id], "title": title};
+        if (current_options.show_node_detail){
+            currentNode.label = title;
+        }
+        if(map_max_path[node_id] !== undefined){
+            currentNode.x = map_max_path[node_id].x;
+            currentNode.y = map_max_path[node_id].y;
+            currentNode.fixed = true;
+            currentNode.mass = 5;
+            currentNode.color = {border: "red", background: "orange", highlight: {border: "red", background: "#ffcc66"}};
+        }
+        return currentNode;
+    }
+
+    function set_node_visibility(node_id, flag){
+        if(visible_nodes[node_id] == flag) return;
+        visible_nodes[node_id] = flag;
+        if(flag){
+            nodes.add(obtain_node_object(node_id));
+        }else{
+            nodes.remove(node_id);
+        }
+    }
+
+    function update_visible_nodes(clear_state=false){
+        if(typeof visible_nodes === "undefined" || clear_state){
+            visible_nodes = [...Array(prelen)].map((_, __) => false); // size: [prelen]
+            map_max_path = {}; // {node_id -> position on max_path}
+            var accumulated_x = 0;
+            for (var i=0;i<tarlen;i++){
+                accumulated_x += node_tokens[max_path[i]][0].length * 5
+                var y = Math.floor(Math.random() * 3) * 100 - 100;
+                map_max_path[max_path[i]] = {position: i, x: accumulated_x, y:y};
+                accumulated_x += node_tokens[max_path[i]][0].length * 5 + 100;
+            }
+            nodes = new vis.DataSet();
+        }
+
+        for (var i=0;i<prelen;i++){
+            if(node_pass_prob[i] >= current_options.minimum_node_pass_prob || map_max_path[i] !== undefined){
+                set_node_visibility(i, true);
+            }else{
+                set_node_visibility(i, false);
+            }
+        }
+    }
+
+    
+    function update_node_details(){
+        for(var i=0;i<prelen;i++) if (visible_nodes[i]){
+            currentNode = obtain_node_object(i);
+            nodes.updateOnly(currentNode);
+        }
+    }
+
+    function obtain_edge_object(i, j){
+        var edge_id = i.toString() + "-" + j.toString();
+        var label = links[i][j].toFixed(2);
+        var pass_label = (node_pass_prob[i] * links[i][j]).toFixed(2);
+        var title = "From Node " + i.toString() + " to Node " + j.toString() + "\n" + "Transition Probability:" + label + "\nPassing Probability:" + pass_label;
+        var currentEdge = {id: edge_id,
+            from: i, to: j, value: links[i][j] * node_pass_prob[i], title: title};
+        if (map_max_path[i] !== undefined && map_max_path[j] !== undefined && map_max_path[i].position + 1 == map_max_path[j].position){
+            currentEdge.color = "red";
+        }
+        if(current_options.show_edge_label){
+            currentEdge.label = label;
+        }else{
+            currentEdge.label = " ";
+        }
+
+        return currentEdge;
+    }
+
+    function set_edge_visibility(i, j, flag){
+        if(visible_edges[i][j] == flag) return;
+        visible_edges[i][j] = flag;
+        if(flag){
+            edges.add(obtain_edge_object(i, j));
+        }else{
+            var edge_id = i.toString() + "-" + j.toString();
+            edges.remove(edge_id);
+        }
+    }
+
+    function update_visible_edges(clear_state=false){
+        if(typeof visible_edges === "undefined" || clear_state){
+            visible_edges = [...Array(prelen)].map((_, __) => {return [...Array(prelen)].map((_, __) => false);}); // size: [prelen, prelen]
+
+            sorted_links_out = []; // size: [prelen, prelen]
+            for (var i=0;i<prelen;i++){
+                // sort the out edge according to their prob
+                sorted_links_out.push(links[i].map((val, idx) => {return {'idx': idx, 'val': val};}).
+                    sort((v1, v2) => v2.val - v1.val)
+                );
+            }
+
+            sorted_links_in = []; // size: [prelen, prelen], element {idx, val}, from big to small
+            for (var i=0;i<prelen;i++){
+                links_in = []
+                for(var j=0;j<prelen;j++) links_in.push({idx: j, val: links[j][i] * node_pass_prob[j]})
+                sorted_links_in.push(links_in.sort((v1, v2) => v2.val - v1.val));
+            }
+            edges = new vis.DataSet();
+        }
+
+        var next_visible_edges = [...Array(prelen)].map((_, __) => {return [...Array(prelen)].map((_, __) => false);}); // size: [prelen, prelen]
+        var links_in_num = [...Array(prelen)].map((_, __) => 0);
+
+        for (var i=0;i<prelen - 1;i++){
+            if(!visible_nodes[i]) continue;
+
+            // select visible out edge of node i
+            var left_visible_edge_num = current_options.max_out_edge_num;
+            var left_visible_edge_prob = current_options.max_out_edge_prob;
+            for(var j=0; j<prelen;j++){
+                var idx = sorted_links_out[i][j]['idx'];
+                if (sorted_links_out[i][j]['val'] < current_options.minimum_edge_prob) break;
+                if (visible_nodes[idx]){
+                    links_in_num[idx]++;
+                    next_visible_edges[i][idx] = true;
+                    left_visible_edge_num--;
+                    left_visible_edge_prob -= sorted_links_out[i][j]['val'];
+                    if (left_visible_edge_num==0 || left_visible_edge_prob < 0){
+                        break;
+                    }
+                }
+            }
+        }
+
+        if(current_options.force_in_edge){
+            // add at least one in edge per node
+            for (var i=0;i<prelen;i++){
+                if(i == 0 || !visible_nodes[i] || links_in_num[i] > 0) continue;
+                for(var j=0; j<prelen;j++){
+                    var idx = sorted_links_in[i][j]['idx'];
+                    if (visible_nodes[idx]){
+                        next_visible_edges[idx][i] = true;
+                        break;
+                    }
+                }
+            }
+        }
+        
+        for(var i=0;i<prelen;i++){
+            for(var j=i+1;j<prelen;j++){
+                set_edge_visibility(i, j, next_visible_edges[i][j]);
+            }
+        }
+    }
+
+    function update_edge_label(){
+        for(var i=0;i<prelen;i++){
+            for(var j=i+1;j<prelen;j++) if (visible_edges[i][j]){
+                currentEdge = obtain_edge_object(i, j);
+                edges.updateOnly(currentEdge);
+            }
+        }
+    }
+
+    function customScalingFunction(min,max,total,value) {
+        min = 0;
+        var scale = 1 / (max - min);
+        return Math.max(0,(value - min)*scale);
+    }
+
+    function escapeHtml(unsafe)
+    {
+        return unsafe
+            .replace(/&/g, "&amp;")
+            .replace(/</g, "&lt;")
+            .replace(/>/g, "&gt;")
+            .replace(/"/g, "&quot;")
+            .replace(/'/g, "&#039;");
+    }
+    function get_jumpable_node(idx){
+        if(visible_nodes[idx]){
+            return "<a href=\"javascript:network.selectNodes(["+ idx.toString() +"]);show_hint_node(" + idx.toString() + ");\">" + idx.toString() + "</a>";
+        }else{
+            return "<a class=\"invisible\" href=\"javascript:show_hint_node(["+ idx.toString() +"]);network.unselectAll();\">" + idx.toString() + "</a>";
+        }
+    }
+    function get_jumpable_edge(i, j, label){
+        var edge_id = i.toString() + "-" + j.toString();
+        if(visible_edges[i][j]){
+            return "<a href=\"javascript:network.selectEdges(['"+ edge_id +"']);show_hint_edge('" + edge_id + "');\">" + label + "</a>";
+        }else{
+            return "<a class=\"invisible\" href=\"javascript:show_hint_edge('" + edge_id + "');network.unselectAll();\">" + label + "</a>";
+        }
+    }
+
+    show_hint_node = function(node_id){
+        showing_node = node_id;
+        showing_edge = undefined;
+        var title = "<p>You selected <b>Node " + node_id.toString() + "</b> ";
+        if (visible_nodes[node_id]){
+            title += "(<a href=\"javascript:network.fit({nodes:[" + node_id.toString() + "],animation:true});\">Find Me!</a>). "
+        }else{
+            title += "(Not shown). "
+        }
+        title += "Passing Probability: <b>" + node_pass_prob[node_id].toFixed(2) + "</b>. You can click the links below to jump to other edges or nodes.</p>";
+        document.getElementById("hintsupper").innerHTML = title;
+
+        title = "<table><thead><tr><th>Rank</th><th>Candidate</th><th>Probability</th></tr></thead><tbody>";
+        for (var j=0;j<5;j++){
+            title += "<tr><td>#" + (j+1).toString() + "</td><td>" + escapeHtml(node_tokens[node_id][j]) + "</td><td>" + node_probs[node_id][j].toFixed(2) + "</td></tr>";
+        }
+        title += "</tbody>"
+        title += "<p>Top-5 Token Candidates: </p>";
+        document.getElementById("hintsleft").innerHTML = title;
+
+        title = "<table><thead><tr><th>Rank</th><th>To</th><th>Transition Prob.</th><th>Passing Prob.</th></tr></thead><tbody>";
+        for (var j=0;j<prelen;j++){
+            var idx = sorted_links_out[node_id][j].idx;
+            if(j < 5 || visible_edges[node_id][idx]){
+                title += "<tr><td>" + get_jumpable_edge(node_id, idx, "#" + (j+1).toString()) + "</td><td>" + get_jumpable_node(idx) + "</td><td>" + links[node_id][idx].toFixed(2) + "</td><td>" +
+                    (node_pass_prob[node_id] * links[node_id][idx]).toFixed(2) + "</td></tr>";
+            }
+        }
+        title += "</tbody>"
+        title += "<p>Top Outgoing Edges: </p>"
+        document.getElementById("hintscenter").innerHTML = title;
+
+        title = "<table><thead><tr><th>Rank</th><th>From</th><th>Transition Prob.</th><th>Passing Prob.</th></tr></thead><tbody>";
+        for (var j=0;j<prelen;j++){
+            var idx = sorted_links_in[node_id][j].idx;
+            if(j < 5 || visible_edges[idx][node_id]){
+                title += "<tr><td>" + get_jumpable_edge(idx, node_id, "#" + (j+1).toString()) + "</td><td>" + get_jumpable_node(idx) + "</td><td>" + links[idx][node_id].toFixed(2) + "</td><td>" +
+                    (node_pass_prob[idx] * links[idx][node_id]).toFixed(2) + "</td></tr>";
+            }
+        }
+        title += "</tbody>"
+        title += "<p>Top Incoming Edges: </p>"
+        document.getElementById("hintsright").innerHTML = title;
+
+        document.getElementById("hintsbottom").innerHTML =   
+          "<br>"+ 
+          "Passing probability of a node V represents how likely the node will be choosen in a random path, i.e., P(V \\in A). <br>" + 
+          "Passing probability of an edge from U to V represents how likely the node V follows the node U in a random path, i.e., P(a_i = U && a_{i+1} = V). <br>" + 
+          "Token probability represents how likely a token is predicted on the given node, i.e., P(y_i| v_{a_{i}}). <br>" + 
+          "Transition probability represents how likely a specific node is following the given node, i.e. P(a_{i+1} | a_{i}).<br>"
+    }
+
+    show_hint_edge = function(edge_id){
+        showing_edge = edge_id;
+        showing_node = undefined;
+        var i = parseInt(edge_id.split("-")[0]);
+        var j = parseInt(edge_id.split("-")[1]);
+        var label = links[i][j].toFixed(2);
+        var passing_label = (links[i][j] * node_pass_prob[i]).toFixed(2);
+        var title = "You selected an edge from <b>Node " + get_jumpable_node(i) + " to Node " + get_jumpable_node(j) + "</b>."
+        if (visible_edges[i][j]){
+            title += "(<a href=\"javascript:network.fit({nodes:[" + i.toString() + "," + j.toString() + "],animation:true});\">Find Me!</a>). "
+        }else{
+            title += "(Not shown). "
+        }
+        title += "<br> You can click the links above to jump to the nodes. <br><br>"
+        title += "Transition Probability:<b>" + label + "</b><br>";
+        title += "Passing Probability:<b>" + passing_label + "</b><br>";
+        title += "<br>" + 
+            "Transition probability represents how likely a specific node is following the given node, i.e. P(a_{i+1} | a_{i}).<br>" + 
+            "Passing probability of an edge from U to V represents how likely the node V follows the node U in a random path, i.e., P(a_i = U && a_{i+1} = V). <br>"
+        document.getElementById("hintsupper").innerHTML = title;
+        document.getElementById("hintsleft").innerHTML = "";
+        document.getElementById("hintsright").innerHTML = "";
+        document.getElementById("hintscenter").innerHTML = "";
+        document.getElementById("hintsbottom").innerHTML = "";
+    }
+
+    function clear_hint(){
+        showing_node = undefined;
+        showing_edge = undefined;
+        document.getElementById("hintsupper").innerHTML = "Use scroll to zoom in or out. Select or Hover over nodes and edges for more information ... (Try dragging nodes to replace them.)";
+        document.getElementById("hintsleft").innerHTML = "";
+        document.getElementById("hintsright").innerHTML = "";
+        document.getElementById("hintscenter").innerHTML = "";
+        document.getElementById("hintsbottom").innerHTML = "";
+    }
+
+    startNetwork = function(graph_info, options) {
+        current_options = options;
+
+        global_graph_info = graph_info;
+        node_pass_prob = graph_info['node_pass_prob'][0] // size: [prelen]
+        prelen = node_pass_prob.length
+        max_path = graph_info['max_paths'][0] // size: [tarlen]
+        tarlen = max_path.length
+        node_tokens = graph_info['node_tokens'][0] // size: [prelen, 5]
+        node_probs = graph_info['node_probs'][0] // size: [prelen, 5]
+        links = graph_info['links'][0] // size: [prelen, prelen]
+
+        update_visible_nodes(true);
+        update_visible_edges(true);
+
+        // create a network
+        var container = document.getElementById("daggraph");
+        var data = {
+            nodes: nodes,
+            edges: edges,
+        };
+        network_options = {
+            nodes: {
+                shape: "ellipse",
+                scaling: {
+                    label: {
+                      min: 8,
+                      max: 20,
+                    },
+                    customScalingFunction: customScalingFunction,
+                },
+            },
+            edges: {
+                arrowStrikethrough: false,
+                arrows: "to",
+                smooth: {
+                    type: "continuous"
+                },
+                color: "#2B7CE9",
+                font: { align: "bottom" },
+                length: 120,
+                scaling: {
+                    min: 0.5,
+                    max: 3,
+                    label: {
+                        min: 8,
+                        max: 15,
+                    },
+                    customScalingFunction: customScalingFunction,
+                }
+            }
+        };
+        network = new vis.Network(container, data, network_options);
+
+        network.off("dragStart");
+        network.on("dragStart", function (params) {
+            var idx = this.getNodeAt(params.pointer.DOM);
+            if (idx !== undefined) {
+                // console.log("dragstart " + idx.toString());
+                if (map_max_path[idx] !== undefined){
+                    data.nodes.update({id: idx, fixed: false});
+                }
+            }
+        });
+        network.off("dragEnd");
+        network.on("dragEnd", function (params) {
+            var idx = this.getNodeAt(params.pointer.DOM);
+            if (idx !== undefined){
+                // console.log("dragend " + idx.toString());
+                if (map_max_path[idx] !== undefined){
+                    data.nodes.update({id: idx, fixed: true});
+                    map_max_path[idx].x = params.pointer.canvas.x;
+                    map_max_path[idx].y = params.pointer.canvas.y;
+                }
+            }
+        });
+
+        disable_edge_select = false;
+        network.off("selectNode");
+        network.on("selectNode", function (params) {
+            var node_id = params.nodes[0];
+            show_hint_node(node_id);
+            disable_edge_select = true;
+            setTimeout(() => {disable_edge_select=false;}, 200);
+        });
+
+        network.off("selectEdge");
+        network.on("selectEdge", function (params) {
+            if(disable_edge_select) return;
+            var edge_id = params.edges[0];
+            show_hint_edge(edge_id);
+        });
+
+        network.off("deselectNode");
+        network.on("deselectNode", function (params) {
+            clear_hint();
+            showing_node = undefined;
+            showing_edge = undefined;
+        });
+        network.off("deselectEdge");
+        network.on("deselectEdge", function (params) {
+            clear_hint();
+        });
+    }
+
+    updateNetwork = function(options) {
+        if(typeof node_pass_prob === "undefined") return;
+        old_options = current_options;
+        current_options = options;
+        if(options.minimum_node_pass_prob != old_options.minimum_node_pass_prob){
+            update_visible_nodes();
+        }
+        if(options.minimum_node_pass_prob != old_options.minimum_node_pass_prob ||
+            options.minimum_edge_prob != old_options.minimum_edge_prob ||
+            options.max_out_edge_num != old_options.max_out_edge_num ||
+            options.max_out_edge_prob != old_options.max_out_edge_prob ||
+            options.force_in_edge != old_options.force_in_edge){
+            update_visible_edges();
+        }
+        if(options.show_node_detail != old_options.show_node_detail){
+            if(options.show_node_detail) {
+                network_options.nodes.shape = "dot";
+                network_options.nodes.scaling.label.min=10;
+                network_options.nodes.scaling.label.max=10;
+            }else{
+                network_options.nodes.shape = "ellipse";
+                network_options.nodes.scaling.label.min=8;
+                network_options.nodes.scaling.label.max=20;
+            }
+            network.setOptions(network_options);
+            update_node_details();
+        }
+        if(options.show_edge_label != old_options.show_edge_label){
+            update_edge_label();
+        }
+
+        if(showing_node != undefined){
+            show_hint_node(showing_node);
+        }
+        if(showing_edge != undefined){
+            show_hint_edge(showing_edge);
+        }
+    }
+}

graph.html

@@ -0,0 +1,49 @@
+<style type="text/css">
+    #contents {
+        height: 400px;
+        margin: 0;
+        padding: 0;
+        box-sizing: border-box;
+        position: relative;
+    }
+    #hints {
+        height: 400px;
+        margin: 0;
+        padding: 0;
+        box-sizing: border-box;
+        position: relative;
+    }
+
+    #daggraph {
+        width: 100%;
+        height: 100%;
+        border: 1px solid #d3d3d3;
+        box-sizing: border-box;
+    }
+
+    #hints .invisible:link, #hints .invisible:visited {
+        color: gray;
+    }
+
+    .container {right: 0; text-align: center;}
+    .container .left, .container .center, .container .right { display: inline-block; padding: 5px 10px; text-align: left;}
+    .container .left { float: left; }
+    .container .center { margin: 0 auto; }
+    .container .right { float: right; }
+    .clear { clear: both; }
+    #hintsleft td, #hintsleft th, #hintscenter td, #hintscenter th, #hintsright td, #hintsright th {padding: 3px 8px; }
+</style>
+
+<div id="contents">
+<div id="daggraph"></div>
+</div>
+<div id="hints">
+    <div id="hintsupper">Use scroll to zoom in or out. Select or Hover over nodes and edges for more information ... (Try dragging nodes to replace them.)</div>
+    <div class="container">
+        <div class="left" id="hintsleft"></div>
+        <div class="center" id="hintscenter"></div>
+        <div class="right"  id="hintsright"></div>
+        <div class="clear"></div>
+    </div>
+    <div id="hintsbottom"></div>
+</div>

packages.txt

@@ -0,0 +1,10 @@
+build-essential
+cmake
+libboost-system-dev
+libboost-thread-dev
+libboost-program-options-dev
+libboost-test-dev
+libeigen3-dev
+zlib1g-dev
+libbz2-dev
+liblzma-dev