# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     name: python3
# ---

# + [markdown] id="B03QGHBzbPOG" colab_type="text"
# To start this Jupyter Dash app, please run all the cells below. Then, click on the **temporary** URL at the end of the last cell to open the app.

# + id="JvrnXoYRbRYj" colab_type="code" colab={}
# !pip install -q jupyter-dash==0.3.0rc1 dash-bootstrap-components transformers

# + id="QPZchYwk5cml" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 51} outputId="1bfafdbe-4213-4bf1-9057-03448d9351d9"
import time

import dash
import dash_html_components as html
import dash_core_components as dcc
import dash_bootstrap_components as dbc
from dash.dependencies import Input, Output, State
from jupyter_dash import JupyterDash
import torch
from transformers import MarianMTModel, MarianTokenizer
import nltk

nltk.download("punkt")
from nltk.tokenize import sent_tokenize


# + id="0ypKd69vaGV3" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 328, "referenced_widgets": ["fb0d3301bbcb45ea826f6448398230d2", "478a6aee30f94fb9ab5496bd3bca68f5", "3da2e50c0daf4f97bbd70fc069a96c3f", "823f1106e3d845fd91412aa93a0060a3", "<KEY>", "<KEY>", "<KEY>", "<KEY>", "d5f6a096f7f748559a8e439d88a3f751", "ebe0824e016640e99ca053baafd85c38", "<KEY>", "688c69ea3fd54ad28545e324c7c5ba1b", "88b61848f5fb44b9ad1ce3b1483e92c7", "3abca65c686449d1a1ff58d60380201a", "<KEY>", "d4dcfef4a35f4560afc05883242d85a8", "<KEY>", "<KEY>", "<KEY>", "323209780dcb4ef1ae7d88c2aa68ffe4", "791591fba544410f8a58fa10fb2a4732", "<KEY>", "583aa6f8761a41d9b3becd58af30f224", "<KEY>", "e3ea38d93f5c40b2826117ca4f19464f", "<KEY>", "<KEY>", "<KEY>", "<KEY>", "<KEY>", "691dee45106d41f8ae9d85f7507c6cc2", "<KEY>", "f1eb10b41ed3491b8e84ee06524cb627", "<KEY>", "<KEY>", "c2216ffa110245038d7d1e156e7f9e12", "842e75a82cf040119a8ce046bcc93516", "<KEY>", "441a0db485a8456f89be2e46fd56db2b", "<KEY>", "f2c5a2cee11a489aa8f8d53008c336c0", "b3320e1bc34d48c688c336e3bce40804", "<KEY>", "0e92ca047f45424691fc0c47a7981640", "<KEY>", "1d7ab36303e8471cafc41adbb89a8919", "deaf862f72c04e9f98ee9a084d8cd981", "ef1ee2c295db4200941acf6e4eca0be2"]} outputId="60b3b2f6-e7f1-4be3-d934-17fd3a7363ca"
# Choose device and load model
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Device: {device}")
model_name = "Helsinki-NLP/opus-mt-en-ROMANCE"
tokenizer = MarianTokenizer.from_pretrained(model_name)
model = MarianMTModel.from_pretrained(model_name).to(device)
model.eval();

# + id="cea8pGXMaHww" colab_type="code" colab={}
# Create app
app = JupyterDash(__name__, external_stylesheets=[dbc.themes.BOOTSTRAP])

# Define Layout
app.layout = dbc.Container(
    fluid=True,
    children=[
        html.H1("Dash Neural Machine Translation"),
        html.Hr(),
        dbc.Spinner(
            dbc.Row(
                [
                    dbc.Col(dbc.Button("Translate", id="button-translate"), width=2),
                    dbc.Col(
                        html.Div(id="time-output", style={"margin-top": "8px"}),
                        width=10,
                    ),
                ],
                style={"margin-bottom": "15px"},
            )
        ),
        dbc.Row(
            [
                dbc.Col(
                    [
                        dbc.InputGroup(
                            [
                                dbc.InputGroupAddon(
                                    "Source Language", addon_type="prepend"
                                ),
                                dbc.Select(
                                    id="source-language",
                                    options=[{"label": "English", "value": "en"}],
                                    value="en",
                                ),
                            ]
                        ),
                        dbc.Textarea(
                            id="source-text",
                            style={"margin-top": "15px", "height": "65vh"},
                        ),
                    ]
                ),
                dbc.Col(
                    [
                        dbc.InputGroup(
                            [
                                dbc.InputGroupAddon(
                                    "Target Language", addon_type="prepend"
                                ),
                                dbc.Select(
                                    id="target-language",
                                    options=[
                                        {
                                            "label": v.replace(">>", "").replace(
                                                "<<", ""
                                            ),
                                            "value": v,
                                        }
                                        for v in tokenizer.supported_language_codes
                                    ],
                                    value=">>fr<<",
                                ),
                            ]
                        ),
                        dbc.Textarea(
                            id="target-text",
                            style={"margin-top": "15px", "height": "65vh"},
                        ),
                    ]
                ),
            ]
        ),
    ],
)


# + id="Xc_Oq_9baCLc" colab_type="code" colab={}
@app.callback(
    [Output("target-text", "value"), Output("time-output", "children")],
    [
        Input("button-translate", "n_clicks"),
        Input("source-language", "value"),
        Input("target-language", "value"),
    ],
    [State("source-text", "value")],
)
def translate(n_clicks, src_lang, tgt_lang, src_text):
    if src_text is None or src_text == "":
        return "", "Did not run."

    t0 = time.time()

    # Create the batched input
    text_in = [
        tgt_lang + " " + sent.replace("\n", " ") for sent in sent_tokenize(src_text)
    ]
    batch = tokenizer.prepare_translation_batch(text_in)
    for k in batch:
        batch[k] = batch[k].to(device)

    # Run the model and decode the output
    translated = model.generate(**batch)
    tgt_text = [tokenizer.decode(t, skip_special_tokens=True) for t in translated]

    t1 = time.time()
    time_output = f"Translated on {device} in {t1-t0:.2f}s"

    return " ".join(tgt_text), time_output


# + id="WOYpVyJWaBdV" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 51} outputId="1e8340c0-0230-40a0-d69e-1019cac01250"
app.run_server(mode='external')

StarcoderdataJupyterScriptsDedupFiltered