faiss_aligner.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "faiss_aligner.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "9aa180a58a474394be35123dacca9ed5": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_b01f943acd3c4e02b2fceb8b6553c3d5",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_6fa6abb3ad0e4fc4ba1acf0330d4a0b4",
              "IPY_MODEL_d5222e263cc44c37b064bd756e388938"
            ]
          }
        },
        "b01f943acd3c4e02b2fceb8b6553c3d5": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "6fa6abb3ad0e4fc4ba1acf0330d4a0b4": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_03189a2f44e641b488c32c449186b3c9",
            "_dom_classes": [],
            "description": "100%",
            "_model_name": "FloatProgressModel",
            "bar_style": "success",
            "max": 1754318854,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 1754318854,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_92f902c85ec74fbc8837719096d145f1"
          }
        },
        "d5222e263cc44c37b064bd756e388938": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_0f65530843304ae1b3d4386f95b44e98",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 1.75G/1.75G [02:00&lt;00:00, 14.6MB/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_4c608041331644528f06b2b7f935bd59"
          }
        },
        "03189a2f44e641b488c32c449186b3c9": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "92f902c85ec74fbc8837719096d145f1": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "0f65530843304ae1b3d4386f95b44e98": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "4c608041331644528f06b2b7f935bd59": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "c79d329a97344f2bb4692ddbf6aeed8a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_ce0f45d557284abc9c773e8a55feebe8",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_7daebe5da0d445549c5c336933493ff7",
              "IPY_MODEL_d78f6be8e9c8483da61d25c3c10b0a85"
            ]
          }
        },
        "ce0f45d557284abc9c773e8a55feebe8": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "7daebe5da0d445549c5c336933493ff7": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_c5b2f29e930d48369810335ec60dc37a",
            "_dom_classes": [],
            "description": "Batches: 100%",
            "_model_name": "FloatProgressModel",
            "bar_style": "success",
            "max": 2194,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 2194,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_a622d6c500b141e8bf817d5661ced978"
          }
        },
        "d78f6be8e9c8483da61d25c3c10b0a85": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_9cc1274be69048babc7c1a2435ea33b1",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 2194/2194 [02:34&lt;00:00, 14.20it/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_761bb43205fb4ddb8d79898094497543"
          }
        },
        "c5b2f29e930d48369810335ec60dc37a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "a622d6c500b141e8bf817d5661ced978": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "9cc1274be69048babc7c1a2435ea33b1": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "761bb43205fb4ddb8d79898094497543": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "47f4ebd0caef4c76a89b8f7a4ca71629": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_819de80217704cf8b4389092a4eadeff",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_ce1fdc49ad2c43ca97ce566fc78df0e4",
              "IPY_MODEL_1eaf9a4bab22495588a786a5d844ba25"
            ]
          }
        },
        "819de80217704cf8b4389092a4eadeff": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "ce1fdc49ad2c43ca97ce566fc78df0e4": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_42b3b5054be54ae8a1db50d8aae504bf",
            "_dom_classes": [],
            "description": "Batches: 100%",
            "_model_name": "FloatProgressModel",
            "bar_style": "success",
            "max": 79,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 79,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_256fb839fba848a79c03850b0474424c"
          }
        },
        "1eaf9a4bab22495588a786a5d844ba25": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_39bbc7e27dfa44dca70956d145b82807",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 79/79 [00:05&lt;00:00, 13.24it/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_f46e90ec73274c23a0405ef338ec08d3"
          }
        },
        "42b3b5054be54ae8a1db50d8aae504bf": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "256fb839fba848a79c03850b0474424c": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "39bbc7e27dfa44dca70956d145b82807": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "f46e90ec73274c23a0405ef338ec08d3": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "B7UVlJNKpWzf"
      },
      "source": [
        "# For CPU\n",
        "!apt install libomp-dev\n",
        "# !pip install faiss\n",
        "!pip install faiss-gpu\n",
        "!pip install sentence_transformers"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "hgEucTeGphrh",
        "outputId": "c5921c83-9b1d-4d00-8581-a330538a4fc0"
      },
      "source": [
        "source_file = \"/content/drive/MyDrive/aligner-test-files/india_2020_en-ta/total_en_2020_January_sen.txt\"\n",
        "target_file = \"/content/drive/MyDrive/aligner-test-files/india_2020_en-ta/total_ta_2020_January_sen.txt\"\n",
        "\n",
        "with open(source_file, encoding=\"utf-16\") as f:\n",
        "    source_sentences = f.readlines()\n",
        "with open(target_file, encoding=\"utf-16\") as f:\n",
        "    target_sentences = f.readlines()\n",
        "\n",
        "source_sentences = list(map(lambda s: s.strip(), source_sentences))\n",
        "target_sentences = list(map(lambda s: s.strip(), target_sentences))\n",
        "print(len(source_sentences))\n",
        "print(len(target_sentences))"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "70199\n",
            "2497\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "L4qvjdhWpnOr"
      },
      "source": [
        "import faiss\n",
        "import numpy as np\n",
        "import time\n",
        "import gzip\n",
        "import lzma\n",
        "\n",
        "########  Functions to find and score candidates\n",
        "def score(x, y, fwd_mean, bwd_mean, margin):\n",
        "    return margin(x.dot(y), (fwd_mean + bwd_mean) / 2)\n",
        "\n",
        "\n",
        "def score_candidates(x, y, candidate_inds, fwd_mean, bwd_mean, margin):\n",
        "    scores = np.zeros(candidate_inds.shape)\n",
        "    for i in range(scores.shape[0]):\n",
        "        for j in range(scores.shape[1]):\n",
        "            k = candidate_inds[i, j]\n",
        "            scores[i, j] = score(x[i], y[k], fwd_mean[i], bwd_mean[k], margin)\n",
        "    return scores\n",
        "\n",
        "\n",
        "def kNN(x, y, k):\n",
        "    start_time = time.time()\n",
        "    print(\"Perform exact search\")\n",
        "    d = 64  \n",
        "    cpu_index = faiss.IndexFlatL2(y.shape[1])\n",
        "    print(\"index done\")\n",
        "    gpu_index = faiss.index_cpu_to_all_gpus(cpu_index)\n",
        "    print(\"added to index\")\n",
        "    gpu_index.add(np.ascontiguousarray(y) ) \n",
        "    print(\"add successfull\")\n",
        "    sim, ind = gpu_index.search(np.ascontiguousarray(x), k)    \n",
        "    print(\"Done: {:.2f} sec\".format(time.time()-start_time))\n",
        "    return sim, ind"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 102,
          "referenced_widgets": [
            "9aa180a58a474394be35123dacca9ed5",
            "b01f943acd3c4e02b2fceb8b6553c3d5",
            "6fa6abb3ad0e4fc4ba1acf0330d4a0b4",
            "d5222e263cc44c37b064bd756e388938",
            "03189a2f44e641b488c32c449186b3c9",
            "92f902c85ec74fbc8837719096d145f1",
            "0f65530843304ae1b3d4386f95b44e98",
            "4c608041331644528f06b2b7f935bd59"
          ]
        },
        "id": "Q84jm74UprRx",
        "outputId": "1b22d752-65c0-44c2-e1cc-af407feedb3b"
      },
      "source": [
        "%%time\n",
        "from sentence_transformers import SentenceTransformer, models\n",
        "import numpy as np\n",
        "import gzip\n",
        "import tqdm\n",
        "import torch\n",
        "import pandas as pd\n",
        "\n",
        "#Model we want to use for bitext mining. LaBSE achieves state-of-the-art performance\n",
        "model_name = 'LaBSE'\n",
        "model = SentenceTransformer(model_name, device='cuda')"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "9aa180a58a474394be35123dacca9ed5",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, max=1754318854.0), HTML(value='')))"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\n",
            "CPU times: user 51.5 s, sys: 10.7 s, total: 1min 2s\n",
            "Wall time: 2min 8s\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 120,
          "referenced_widgets": [
            "c79d329a97344f2bb4692ddbf6aeed8a",
            "ce0f45d557284abc9c773e8a55feebe8",
            "7daebe5da0d445549c5c336933493ff7",
            "d78f6be8e9c8483da61d25c3c10b0a85",
            "c5b2f29e930d48369810335ec60dc37a",
            "a622d6c500b141e8bf817d5661ced978",
            "9cc1274be69048babc7c1a2435ea33b1",
            "761bb43205fb4ddb8d79898094497543"
          ]
        },
        "id": "ec7hOHGSp-xZ",
        "outputId": "7e14d862-70bb-41b6-8458-fff9a0d4c5f7"
      },
      "source": [
        "%%time\n",
        "source_sentences = list(source_sentences)\n",
        "print(\"Encode source sentences\", len(source_sentences))\n",
        "source_embeddings = model.encode(source_sentences, show_progress_bar=True, convert_to_numpy=True)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Encode source sentences 70199\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "c79d329a97344f2bb4692ddbf6aeed8a",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, description='Batches', max=2194.0, style=ProgressStyle(description_wid…"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\n",
            "CPU times: user 2min 20s, sys: 15 s, total: 2min 35s\n",
            "Wall time: 2min 35s\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 120,
          "referenced_widgets": [
            "47f4ebd0caef4c76a89b8f7a4ca71629",
            "819de80217704cf8b4389092a4eadeff",
            "ce1fdc49ad2c43ca97ce566fc78df0e4",
            "1eaf9a4bab22495588a786a5d844ba25",
            "42b3b5054be54ae8a1db50d8aae504bf",
            "256fb839fba848a79c03850b0474424c",
            "39bbc7e27dfa44dca70956d145b82807",
            "f46e90ec73274c23a0405ef338ec08d3"
          ]
        },
        "id": "rQ7yuReGqBlS",
        "outputId": "4c5445a5-05ce-4843-9a4b-742470c16bc4"
      },
      "source": [
        "%%time\n",
        "target_sentences = list(target_sentences)\n",
        "print(\"Encode target sentences\", len(target_sentences))\n",
        "target_embeddings = model.encode(target_sentences, show_progress_bar=True, convert_to_numpy=True)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Encode target sentences 2497\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "47f4ebd0caef4c76a89b8f7a4ca71629",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, description='Batches', max=79.0, style=ProgressStyle(description_width…"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\n",
            "CPU times: user 6.02 s, sys: 46.1 ms, total: 6.06 s\n",
            "Wall time: 6.01 s\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "qgiOQ_UuqGCs",
        "outputId": "3f1d9833-6aad-4580-a55d-c87cafb3a180"
      },
      "source": [
        "%%time\n",
        "\n",
        "# We base the scoring on k nearest neighbors for each element\n",
        "knn_neighbors = 2\n",
        "\n",
        "# Min score for text pairs. Note, score can be larger than 1\n",
        "min_threshold = 1\n",
        "\n",
        "\n",
        "x = source_embeddings\n",
        "y = target_embeddings\n",
        "\n",
        "\n",
        "# Perform kNN in both directions\n",
        "x2y_sim, x2y_ind = kNN(x, y, knn_neighbors)\n",
        "x2y_mean = x2y_sim.mean(axis=1)\n",
        "print(x2y_mean[0])\n",
        "\n",
        "y2x_sim, y2x_ind = kNN(y, x, knn_neighbors)\n",
        "y2x_mean = y2x_sim.mean(axis=1)\n",
        "print(y2x_mean[0])\n",
        "\n",
        "# Compute forward and backward scores\n",
        "margin = lambda a, b: a / b\n",
        "fwd_scores = score_candidates(x, y, x2y_ind, x2y_mean, y2x_mean, margin)\n",
        "bwd_scores = score_candidates(y, x, y2x_ind, y2x_mean, x2y_mean, margin)\n",
        "fwd_best = x2y_ind[np.arange(x.shape[0]), fwd_scores.argmax(axis=1)]\n",
        "bwd_best = y2x_ind[np.arange(y.shape[0]), bwd_scores.argmax(axis=1)]\n",
        "\n",
        "indices = np.stack([np.concatenate([np.arange(x.shape[0]), bwd_best]), np.concatenate([fwd_best, np.arange(y.shape[0])])], axis=1)\n",
        "scores = np.concatenate([fwd_scores.max(axis=1), bwd_scores.max(axis=1)])\n",
        "seen_src, seen_trg = set(), set()\n",
        "\n",
        "#Extact list of parallel sentences\n",
        "src_out = []\n",
        "tgt_out = []\n",
        "score_out = []\n",
        "sentences_written = 0\n",
        "for i in np.argsort(-scores):\n",
        "    src_ind, trg_ind = indices[i]\n",
        "    src_ind = int(src_ind)\n",
        "    trg_ind = int(trg_ind)\n",
        "\n",
        "    if scores[i] < min_threshold:\n",
        "        break\n",
        "\n",
        "    if src_ind not in seen_src and trg_ind not in seen_trg:\n",
        "        seen_src.add(src_ind)\n",
        "        seen_trg.add(trg_ind)\n",
        "        # fOut.write(\"{:.4f}\\t{}\\t{}\\n\".format(scores[i], source_sentences[src_ind].replace(\"\\t\", \" \"), target_sentences[trg_ind].replace(\"\\t\", \" \")))\n",
        "        src_out.append(source_sentences[src_ind].replace(\"\\t\", \" \"))\n",
        "        tgt_out.append(target_sentences[trg_ind].replace(\"\\t\", \" \"))\n",
        "        score_out.append(scores[i])\n",
        "        sentences_written += 1\n"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Perform exact search\n",
            "index done\n",
            "added to index\n",
            "add successfull\n",
            "Done: 0.27 sec\n",
            "1.1421453\n",
            "Perform exact search\n",
            "index done\n",
            "added to index\n",
            "add successfull\n",
            "Done: 0.23 sec\n",
            "0.96488225\n",
            "CPU times: user 1.68 s, sys: 437 ms, total: 2.12 s\n",
            "Wall time: 1.99 s\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 414
        },
        "id": "s8i2a92HqJ-Z",
        "outputId": "97ccdfd6-7400-4d4a-c34b-2d081e6706da"
      },
      "source": [
        "df = pd.DataFrame(list(zip(src_out, tgt_out, score_out)),columns = ['src_out', 'tgt_out','score_out'])\n",
        "df"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/html": [
              "<div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>src_out</th>\n",
              "      <th>tgt_out</th>\n",
              "      <th>score_out</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>— Kunal Kamra (@kunalkamra88) January 28, 2020</td>\n",
              "      <td>— Kunal Kamra (@kunalkamra88) January 28, 2020</td>\n",
              "      <td>11.512191</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>— Ben Smith (@BuzzFeedBen) January 13, 2020</td>\n",
              "      <td>— Ben Smith (@BuzzFeedBen) January 13, 2020</td>\n",
              "      <td>6.531665</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>2</th>\n",
              "      <td>— Hardeep Singh Puri (@HardeepSPuri) January 2...</td>\n",
              "      <td>— Hardeep Singh Puri (@HardeepSPuri) January 2...</td>\n",
              "      <td>6.116069</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3</th>\n",
              "      <td>Imam, who is pursuing PhD in Modern Indian His...</td>\n",
              "      <td>நவீன இந்திய வரலாற்றில் பிஎச்டி படித்து வரும் இ...</td>\n",
              "      <td>5.515931</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>4</th>\n",
              "      <td>Imam, who is pursuing PhD in Modern Indian His...</td>\n",
              "      <td>நவீன இந்திய வரலாற்றில் பிஎச்டி படித்து வரும் இ...</td>\n",
              "      <td>5.515927</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>...</th>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>694</th>\n",
              "      <td>In September 2013, a fast-track court held the...</td>\n",
              "      <td>2013ம் ஆண்டு விரைவு நீதிமன்றங்கள் மூலமாக கூட்ட...</td>\n",
              "      <td>1.004952</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>695</th>\n",
              "      <td>Earlier this month, I had the opportunity to i...</td>\n",
              "      <td>இந்த மாதத்தின் தொடக்கத்தில், கருமமே கண்ணாக இரு...</td>\n",
              "      <td>1.004757</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>696</th>\n",
              "      <td>At Delhi’s protest march, called by the JNUSU ...</td>\n",
              "      <td>ஆர்ப்பாட்டத்தில், முன்னாள் ஜே. என். யு. மாணவர்...</td>\n",
              "      <td>1.004210</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>697</th>\n",
              "      <td>BJP is ruling in India.</td>\n",
              "      <td>டெல்லியை ஆளும் ஆம் ஆத்மிக்கும், மத்தியில் ஆளும...</td>\n",
              "      <td>1.002981</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>698</th>\n",
              "      <td>Sabarimala review pleas: SC bench led by CJI t...</td>\n",
              "      <td>சபரிமலை வழக்கு: மறுஆய்வு மனுக்களை விசாரிக்கவில...</td>\n",
              "      <td>1.001869</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "<p>699 rows × 3 columns</p>\n",
              "</div>"
            ],
            "text/plain": [
              "                                               src_out  ...  score_out\n",
              "0       — Kunal Kamra (@kunalkamra88) January 28, 2020  ...  11.512191\n",
              "1          — Ben Smith (@BuzzFeedBen) January 13, 2020  ...   6.531665\n",
              "2    — Hardeep Singh Puri (@HardeepSPuri) January 2...  ...   6.116069\n",
              "3    Imam, who is pursuing PhD in Modern Indian His...  ...   5.515931\n",
              "4    Imam, who is pursuing PhD in Modern Indian His...  ...   5.515927\n",
              "..                                                 ...  ...        ...\n",
              "694  In September 2013, a fast-track court held the...  ...   1.004952\n",
              "695  Earlier this month, I had the opportunity to i...  ...   1.004757\n",
              "696  At Delhi’s protest march, called by the JNUSU ...  ...   1.004210\n",
              "697                            BJP is ruling in India.  ...   1.002981\n",
              "698  Sabarimala review pleas: SC bench led by CJI t...  ...   1.001869\n",
              "\n",
              "[699 rows x 3 columns]"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 20
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3SW7zgdJtZSN"
      },
      "source": [
        "df.to_json(\"output_file.json\", force_ascii=False, indent=2, orient='records')"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}