embedding_alignment.ipynb

{
    "cells": [
        {
            "cell_type": "code",
            "execution_count": 1,
            "source": [
                "import logging\n",
                "import time\n",
                "\n",
                "from munch import Munch\n",
                "import numpy as np\n",
                "import torch\n",
                "import yaml\n",
                "\n",
                "from lcn.embedding_alignment import (\n",
                "    convex_init,\n",
                "    align,\n",
                "    align_original,\n",
                ")\n",
                "from lcn.alignment_utils import (\n",
                "    compute_accuracy,\n",
                "    compute_csls,\n",
                "    compute_nn,\n",
                "    load_lexicon,\n",
                "    load_vectors,\n",
                "    pickle_cache,\n",
                "    refine,\n",
                ")"
            ],
            "outputs": [],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": 2,
            "source": [
                "# Set up logging\n",
                "logger = logging.getLogger()\n",
                "logger.handlers = []\n",
                "ch = logging.StreamHandler()\n",
                "formatter = logging.Formatter(\n",
                "        fmt='%(asctime)s (%(levelname)s): %(message)s',\n",
                "        datefmt='%Y-%m-%d %H:%M:%S')\n",
                "ch.setFormatter(formatter)\n",
                "logger.addHandler(ch)\n",
                "logger.setLevel('INFO')"
            ],
            "outputs": [],
            "metadata": {}
        },
        {
            "cell_type": "markdown",
            "source": [
                "# Configuration"
            ],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": 3,
            "source": [
                "with open('configs/embedding_alignment.yaml', 'r') as c:\n",
                "    config_seml = yaml.safe_load(c)"
            ],
            "outputs": [],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": 4,
            "source": [
                "config = Munch(config_seml['fixed'])\n",
                "\n",
                "language_src = 'en'\n",
                "language_tgt = 'es'\n",
                "data_dir = \"./data\"  # Download the data first, as described in the README\n",
                "\n",
                "method = 'lcn'  # Change this for other methods: original, full, nystrom, multiscale, sparse, lcn\n",
                "config.update(config_seml[\"from_en\"][method][\"fixed\"])\n",
                "config.nystrom = None if config.nystrom == \"None\" else config.nystrom\n",
                "config.sparse = None if config.sparse == \"None\" else config.sparse\n",
                "original = method == \"original\"\n",
                "\n",
                "seed = 1111\n",
                "test = False  # Change to run on the test set\n",
                "device = \"cuda\"  # Change to cpu if necessary"
            ],
            "outputs": [],
            "metadata": {}
        },
        {
            "cell_type": "markdown",
            "source": [
                "# Load data"
            ],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": 5,
            "source": [
                "model_src = f\"{data_dir}/wiki.{language_src}.vec\"\n",
                "model_tgt = f\"{data_dir}/wiki.{language_tgt}.vec\"\n",
                "if test:\n",
                "    val_numbers = \"5000-6500\"\n",
                "else:\n",
                "    val_numbers = \"0-5000\"\n",
                "lexicon = f\"{data_dir}/{language_src}-{language_tgt}.{val_numbers}.txt\"\n",
                "\n",
                "sinkhorn_reg = torch.tensor(config.sinkhorn_reg, device=device)\n",
                "\n",
                "logging.info(\"*** Wasserstein Procrustes ***\")\n",
                "\n",
                "np.random.seed(seed)\n",
                "torch.manual_seed(seed)\n",
                "\n",
                "maxload = 200_000\n",
                "w_src, x_src = pickle_cache(\n",
                "    f\"{model_src}.pkl\",\n",
                "    load_vectors,\n",
                "    [model_src, maxload],\n",
                "    dict(norm=True, center=True),\n",
                ")\n",
                "w_tgt, x_tgt = pickle_cache(\n",
                "    f\"{model_tgt}.pkl\",\n",
                "    load_vectors,\n",
                "    [model_tgt, maxload],\n",
                "    dict(norm=True, center=True),\n",
                ")\n",
                "src2tgt, _ = load_lexicon(lexicon, w_src, w_tgt)\n",
                "\n",
                "x_src = torch.tensor(x_src, dtype=torch.float).to(device)\n",
                "x_tgt = torch.tensor(x_tgt, dtype=torch.float).to(device)"
            ],
            "outputs": [
                {
                    "output_type": "stream",
                    "name": "stderr",
                    "text": [
                        "2021-07-12 01:24:51 (INFO): *** Wasserstein Procrustes ***\n"
                    ]
                }
            ],
            "metadata": {}
        },
        {
            "cell_type": "markdown",
            "source": [
                "# Convex initialization"
            ],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": 6,
            "source": [
                "torch.cuda.synchronize()\n",
                "t0 = time.time()\n",
                "R0 = convex_init(\n",
                "    x_src[:config.ninit], x_tgt[:config.ninit], sinkhorn_reg=sinkhorn_reg, apply_sqrt=True\n",
                ")\n",
                "torch.cuda.synchronize()\n",
                "logging.info(f\"Done [{time.time() - t0:.1f} sec]\")"
            ],
            "outputs": [
                {
                    "output_type": "display_data",
                    "data": {
                        "text/plain": [
                            "HBox(children=(FloatProgress(value=0.0), HTML(value='')))"
                        ],
                        "application/vnd.jupyter.widget-view+json": {
                            "version_major": 2,
                            "version_minor": 0,
                            "model_id": "c3df7acd9afa4d31b1429197e57e3900"
                        }
                    },
                    "metadata": {}
                }
            ],
            "metadata": {}
        },
        {
            "cell_type": "markdown",
            "source": [
                "# Main part: Wasserstein Procrustes"
            ],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "source": [
                "config.nystrom"
            ],
            "outputs": [],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "source": [
                "torch.cuda.synchronize()\n",
                "t0 = time.time()\n",
                "if original:\n",
                "    R = align_original(x_src, x_tgt, R0.clone(), sinkhorn_reg)\n",
                "else:\n",
                "    R = align(\n",
                "        x_src,\n",
                "        x_tgt,\n",
                "        R0.clone(),\n",
                "        sinkhorn_reg=sinkhorn_reg,\n",
                "        nystrom=config.nystrom,\n",
                "        sparse=config.sparse,\n",
                "        lr=config.lr,\n",
                "        niter=config.niter,\n",
                "        lr_half_niter=config.lr_half_niter,\n",
                "        ntrain=config.ntrain,\n",
                "        print_niter=config.print_niter,\n",
                "    )\n",
                "torch.cuda.synchronize()\n",
                "runtime = time.time() - t0\n",
                "logging.info(f\"Done [{runtime:.1f} sec]\")"
            ],
            "outputs": [
                {
                    "output_type": "display_data",
                    "data": {
                        "text/plain": [
                            "HBox(children=(FloatProgress(value=0.0, max=300.0), HTML(value='')))"
                        ],
                        "application/vnd.jupyter.widget-view+json": {
                            "version_major": 2,
                            "version_minor": 0,
                            "model_id": "9874081ef7634109a1da02b4fef220ac"
                        }
                    },
                    "metadata": {}
                },
                {
                    "output_type": "stream",
                    "name": "stderr",
                    "text": [
                        "2021-07-12 01:23:52 (INFO): Done [51.3 sec]\n"
                    ]
                },
                {
                    "output_type": "stream",
                    "name": "stdout",
                    "text": [
                        "\n"
                    ]
                }
            ],
            "metadata": {}
        },
        {
            "cell_type": "markdown",
            "source": [
                "# Compute accuracies"
            ],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "source": [
                "if test:\n",
                "    logging.info(\"Evaluation on test set\")\n",
                "else:\n",
                "    logging.info(\"Evaluation on validation set\")\n",
                "\n",
                "x_tgt_rot = x_tgt @ R.T\n",
                "\n",
                "acc_nn = compute_accuracy(x_src, x_tgt_rot, src2tgt, compute_nn)\n",
                "logging.info(f\"NN precision@1: {100 * acc_nn:.2f}%\")\n",
                "\n",
                "acc_csls = compute_accuracy(x_src, x_tgt_rot, src2tgt, compute_csls)\n",
                "logging.info(f\"CSLS precision@1: {100 * acc_csls:.2f}%\")"
            ],
            "outputs": [
                {
                    "output_type": "stream",
                    "name": "stderr",
                    "text": [
                        "2021-07-12 01:23:53 (INFO): NN precision@1: 78.96%\n",
                        "2021-07-12 01:24:01 (INFO): CSLS precision@1: 82.24%\n"
                    ]
                }
            ],
            "metadata": {}
        },
        {
            "cell_type": "markdown",
            "source": [
                "# Refine embeddings"
            ],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "source": [
                "x_tgt_rot = refine(x_src, x_tgt_rot, src2tgt)\n",
                "\n",
                "acc_refined = compute_accuracy(x_src, x_tgt_rot, src2tgt, compute_csls)\n",
                "logging.info(f\"Refined CSLS precision@1: {100 * acc_refined:.2f}%\")"
            ],
            "outputs": [
                {
                    "output_type": "display_data",
                    "data": {
                        "text/plain": [
                            "HBox(children=(FloatProgress(value=0.0, max=5.0), HTML(value='')))"
                        ],
                        "application/vnd.jupyter.widget-view+json": {
                            "version_major": 2,
                            "version_minor": 0,
                            "model_id": "6510efa5bd70413290935112b152f9a2"
                        }
                    },
                    "metadata": {}
                },
                {
                    "output_type": "error",
                    "ename": "KeyboardInterrupt",
                    "evalue": "",
                    "traceback": [
                        "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
                        "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
                        "\u001b[0;32m<ipython-input-21-cf874159fb52>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mx_tgt_rot\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mrefine\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx_src\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx_tgt_rot\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msrc2tgt\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      2\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0macc_refined\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcompute_accuracy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx_src\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx_tgt_rot\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msrc2tgt\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcompute_csls\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0mlogging\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minfo\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf\"Refined CSLS precision@1: {100 * acc_refined:.2f}%\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
                        "\u001b[0;32m~/graph-distance/lcn_public/lcn/alignment_utils.py\u001b[0m in \u001b[0;36mrefine\u001b[0;34m(x_src, y_tgt, lexicon, lexicon_size, nepochs, disable_tqdm)\u001b[0m\n\u001b[1;32m    170\u001b[0m     \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mtqdm\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnepochs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdisable\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdisable_tqdm\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    171\u001b[0m         \u001b[0mpred_src\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcompute_csls\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx_src\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my_tgt\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0midx_src\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 172\u001b[0;31m         \u001b[0mpred_tgt\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcompute_csls\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0my_tgt\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx_src\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0midx_tgt\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    173\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    174\u001b[0m         \u001b[0midx_mutual_src\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0midx_src\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpred_tgt\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
                        "\u001b[0;32m~/graph-distance/lcn_public/lcn/alignment_utils.py\u001b[0m in \u001b[0;36mcompute_csls\u001b[0;34m(x_src, y_tgt, idx_src, k, sc_batch_size, sim_batch_size)\u001b[0m\n\u001b[1;32m    132\u001b[0m         \u001b[0msc_batch\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0my_tgt\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0mj\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m:\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m@\u001b[0m \u001b[0mx_src\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mT\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    133\u001b[0m         \u001b[0mdotprod\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtopk\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msc_batch\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mk\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mk\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdim\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msorted\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 134\u001b[0;31m         \u001b[0msc2\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0mj\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmean\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdotprod\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdim\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    135\u001b[0m         \u001b[0;32mdel\u001b[0m \u001b[0msc_batch\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    136\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
                        "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
                    ]
                }
            ],
            "metadata": {}
        }
    ],
    "metadata": {
        "orig_nbformat": 4,
        "language_info": {
            "name": "python"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 2
}