tokenizer(stable) .ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "_cell_guid": "dffec90f-2a93-4648-af8f-8d726ec124fd",
    "_uuid": "7d0bea29-d695-46e1-b876-35fc9aa00eb6",
    "papermill": {
     "duration": 0.004161,
     "end_time": "2023-02-14T10:03:26.838664",
     "exception": false,
     "start_time": "2023-02-14T10:03:26.834503",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Related module  docs\n",
    "* https://github.com/huggingface/tokenizers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Note\n",
    "* add \".to(device)\" when move to transformer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# hyparameter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "bos_id = 0\n",
    "eos_id = 1\n",
    "pad_id = 2\n",
    "gpu_num = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "_cell_guid": "916f1c81-fb6d-4303-84ab-6bca1826e449",
    "_uuid": "2950befb-81d9-4ed6-ad52-3dcd21171bd2",
    "papermill": {
     "duration": 0.002951,
     "end_time": "2023-02-14T10:03:26.844969",
     "exception": false,
     "start_time": "2023-02-14T10:03:26.842018",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Download dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "_cell_guid": "8862fbd2-1edf-4ad7-97b7-4640da2948d2",
    "_uuid": "14555048-265e-4033-a43c-34f7bfe40b7f",
    "execution": {
     "iopub.execute_input": "2023-02-26T06:29:42.007925Z",
     "iopub.status.busy": "2023-02-26T06:29:42.007327Z",
     "iopub.status.idle": "2023-02-26T06:33:53.117371Z",
     "shell.execute_reply": "2023-02-26T06:33:53.116201Z",
     "shell.execute_reply.started": "2023-02-26T06:29:42.007785Z"
    },
    "jupyter": {
     "outputs_hidden": false
    },
    "papermill": {
     "duration": 1789.034246,
     "end_time": "2023-02-14T10:33:15.882295",
     "exception": false,
     "start_time": "2023-02-14T10:03:26.848049",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# from datasets import load_dataset\n",
    "# dataset = load_dataset(\"wmt14\", 'de-en', split='train')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/data2/zrs/.conda/envs/pytorch/lib/python3.8/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "from datasets import load_from_disk\n",
    "dataset = load_from_disk('dataset')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "_cell_guid": "ab77634e-9556-4525-aeac-95aab979f9eb",
    "_uuid": "26e03e5f-3d6c-4315-95fa-6daae0278211",
    "papermill": {
     "duration": 0.006511,
     "end_time": "2023-02-14T10:33:15.894836",
     "exception": false,
     "start_time": "2023-02-14T10:33:15.888325",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# write raw sectences into txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "_cell_guid": "b9db514a-f0a2-440a-b490-a96738fb9944",
    "_uuid": "6b4ee72d-dfa3-49db-8f31-a915f85cb2b5",
    "execution": {
     "iopub.execute_input": "2023-02-26T06:33:53.120002Z",
     "iopub.status.busy": "2023-02-26T06:33:53.119587Z",
     "iopub.status.idle": "2023-02-26T06:42:01.497803Z",
     "shell.execute_reply": "2023-02-26T06:42:01.496696Z",
     "shell.execute_reply.started": "2023-02-26T06:33:53.119969Z"
    },
    "jupyter": {
     "outputs_hidden": false
    },
    "papermill": {
     "duration": 5090.464084,
     "end_time": "2023-02-14T11:58:06.365706",
     "exception": false,
     "start_time": "2023-02-14T10:33:15.901622",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# with open(\"en.txt\",'w') as f:\n",
    "#     for i in range(len(dataset)):\n",
    "#         f.write(dataset[i]['translation']['en']+'\\n')\n",
    "        \n",
    "# with open(\"de.txt\",'w') as f:\n",
    "#     for i in range(len(dataset)):\n",
    "#         f.write(dataset[i]['translation']['de']+'\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "_cell_guid": "22bd7295-5d49-44f9-b373-facad72b0ffa",
    "_uuid": "784cf351-ac3b-4905-887e-a1db69cd5d38",
    "execution": {
     "iopub.execute_input": "2023-02-26T06:42:01.499425Z",
     "iopub.status.busy": "2023-02-26T06:42:01.499127Z",
     "iopub.status.idle": "2023-02-26T06:42:01.504184Z",
     "shell.execute_reply": "2023-02-26T06:42:01.503344Z",
     "shell.execute_reply.started": "2023-02-26T06:42:01.499401Z"
    },
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# dataset['translation']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "_cell_guid": "fa128c03-bb2f-44bf-8487-b68551594eff",
    "_uuid": "9feb6e8f-45fe-41cf-9fc1-997ae3933b32",
    "execution": {
     "iopub.execute_input": "2023-02-26T06:42:01.506424Z",
     "iopub.status.busy": "2023-02-26T06:42:01.505438Z",
     "iopub.status.idle": "2023-02-26T06:42:01.521754Z",
     "shell.execute_reply": "2023-02-26T06:42:01.520526Z",
     "shell.execute_reply.started": "2023-02-26T06:42:01.506379Z"
    },
    "jupyter": {
     "outputs_hidden": false
    },
    "papermill": {
     "duration": 0.030896,
     "end_time": "2023-02-14T11:58:06.409866",
     "exception": false,
     "start_time": "2023-02-14T11:58:06.378970",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Somehow Zuma must find a way to honor his own generation’s commitment to racial justice and national liberation, while empowering the masses who daily suffer the sting of class differences and yearn for material gain.'"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset[4508784]['translation']['en']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "_cell_guid": "b37daf84-cddb-4205-b47a-8fa700c66099",
    "_uuid": "00b63607-0efd-4db6-90fe-fac6a6fbc369",
    "execution": {
     "iopub.execute_input": "2023-02-26T06:42:01.525471Z",
     "iopub.status.busy": "2023-02-26T06:42:01.525140Z",
     "iopub.status.idle": "2023-02-26T06:42:01.533533Z",
     "shell.execute_reply": "2023-02-26T06:42:01.532424Z",
     "shell.execute_reply.started": "2023-02-26T06:42:01.525442Z"
    },
    "jupyter": {
     "outputs_hidden": false
    },
    "papermill": {
     "duration": 0.015873,
     "end_time": "2023-02-14T11:58:06.432177",
     "exception": false,
     "start_time": "2023-02-14T11:58:06.416304",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4508785"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_cell_guid": "b9ac2d98-08c1-4bfe-af60-e296e04dcd43",
    "_uuid": "41f562c7-bf59-429a-84f7-f8fd0a26abff",
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "_cell_guid": "0e0664d2-05ca-43f7-85ed-775b1322b8a5",
    "_uuid": "64c9c492-49a3-4f3e-b8be-1eeafe6c5b3f",
    "papermill": {
     "duration": 0.005895,
     "end_time": "2023-02-14T11:58:06.447012",
     "exception": false,
     "start_time": "2023-02-14T11:58:06.441117",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# train a tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ! pip install tokenizers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "_cell_guid": "b0cbeacd-5399-4b38-a776-a9185a2fc233",
    "_uuid": "a6a99de4-9a0d-4384-860f-296fe8c09820"
   },
   "source": [
    "### model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "_cell_guid": "5cbf015b-1c91-470b-9825-c8ed4813ef0c",
    "_uuid": "1bc47256-0397-482c-9041-f7ad5a29e09b",
    "execution": {
     "iopub.execute_input": "2023-02-26T06:42:01.535252Z",
     "iopub.status.busy": "2023-02-26T06:42:01.534971Z",
     "iopub.status.idle": "2023-02-26T06:42:01.570225Z",
     "shell.execute_reply": "2023-02-26T06:42:01.569201Z",
     "shell.execute_reply.started": "2023-02-26T06:42:01.535227Z"
    },
    "jupyter": {
     "outputs_hidden": false
    },
    "papermill": {
     "duration": null,
     "end_time": null,
     "exception": false,
     "start_time": "2023-02-14T11:58:06.452913",
     "status": "running"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from tokenizers import Tokenizer\n",
    "from tokenizers.models import BPE\n",
    "\n",
    "tokenizer = Tokenizer(BPE())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "_cell_guid": "dc886065-fde6-4944-a860-137e51a8be0d",
    "_uuid": "b2e15ffd-15c7-4b71-b1e3-f23973545c7b"
   },
   "source": [
    "### Normalization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "_cell_guid": "3d8d7942-ea39-424e-9adc-63029ec2f785",
    "_uuid": "11a9cc10-a66e-4642-99fd-cdc825e0b9ba"
   },
   "source": [
    "### pre_tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "_cell_guid": "c570b8c8-045f-4e50-8932-fc835bde9d78",
    "_uuid": "8a1c6f2e-6b08-4f9b-8c7b-39b1ba0f61f2",
    "execution": {
     "iopub.execute_input": "2023-02-26T06:42:01.571705Z",
     "iopub.status.busy": "2023-02-26T06:42:01.571410Z",
     "iopub.status.idle": "2023-02-26T06:42:01.577484Z",
     "shell.execute_reply": "2023-02-26T06:42:01.576466Z",
     "shell.execute_reply.started": "2023-02-26T06:42:01.571660Z"
    },
    "jupyter": {
     "outputs_hidden": false
    },
    "papermill": {
     "duration": null,
     "end_time": null,
     "exception": null,
     "start_time": null,
     "status": "pending"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from tokenizers.pre_tokenizers import ByteLevel\n",
    "\n",
    "tokenizer.pre_tokenizer = ByteLevel()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "_cell_guid": "8135caef-545d-4dad-a616-ffa91a550927",
    "_uuid": "4fbc05b9-a640-423c-b7d4-71f35747625f"
   },
   "source": [
    "### post_processor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "_cell_guid": "335e3777-6a34-4b51-a25b-1a838ac5ea32",
    "_uuid": "faf37196-94a9-4e50-b597-d4f18b9c0da0",
    "execution": {
     "iopub.execute_input": "2023-02-26T06:42:01.579371Z",
     "iopub.status.busy": "2023-02-26T06:42:01.579025Z",
     "iopub.status.idle": "2023-02-26T06:42:01.592575Z",
     "shell.execute_reply": "2023-02-26T06:42:01.591474Z",
     "shell.execute_reply.started": "2023-02-26T06:42:01.579337Z"
    },
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "from tokenizers.processors import TemplateProcessing\n",
    "tokenizer.post_processor = TemplateProcessing(\n",
    "    single=\"<bos> $A <eos>\",\n",
    "    special_tokens=[(\"<bos>\", 0), (\"<eos>\", 1)],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "_cell_guid": "e446e83d-e57b-448b-8c01-1681cbe43bb1",
    "_uuid": "a46af1aa-6de0-493d-81e4-cfdd3a6f907f",
    "execution": {
     "iopub.execute_input": "2023-02-26T02:49:17.090295Z",
     "iopub.status.busy": "2023-02-26T02:49:17.089482Z",
     "iopub.status.idle": "2023-02-26T02:49:17.158323Z",
     "shell.execute_reply": "2023-02-26T02:49:17.157147Z",
     "shell.execute_reply.started": "2023-02-26T02:49:17.090209Z"
    }
   },
   "source": [
    "### decoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "_cell_guid": "40fb4655-cd4d-4da1-809b-2a003a0690f4",
    "_uuid": "e6e3c531-dbb1-435b-93a6-77b08a2e3468",
    "execution": {
     "iopub.execute_input": "2023-02-26T06:42:01.594590Z",
     "iopub.status.busy": "2023-02-26T06:42:01.593896Z",
     "iopub.status.idle": "2023-02-26T06:42:01.603245Z",
     "shell.execute_reply": "2023-02-26T06:42:01.602333Z",
     "shell.execute_reply.started": "2023-02-26T06:42:01.594560Z"
    },
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "from tokenizers.decoders import ByteLevel\n",
    "tokenizer.decoder = ByteLevel()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "_cell_guid": "0de9bc79-74cc-4133-8d42-e35ba6af687c",
    "_uuid": "2f323083-b992-4139-b2fe-f40871422f71"
   },
   "source": [
    "# other setting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "_cell_guid": "c9e70136-8e13-4a9e-b7b0-16839473acb9",
    "_uuid": "5b1622df-9b00-4b73-aa9d-a80266987b2a",
    "execution": {
     "iopub.execute_input": "2023-02-26T06:42:01.607450Z",
     "iopub.status.busy": "2023-02-26T06:42:01.606618Z",
     "iopub.status.idle": "2023-02-26T06:42:01.614100Z",
     "shell.execute_reply": "2023-02-26T06:42:01.613294Z",
     "shell.execute_reply.started": "2023-02-26T06:42:01.607411Z"
    },
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "tokenizer.enable_padding(pad_id=2,pad_token='<pad>')\n",
    "# tokenizer.enable_truncation(maxlen=513)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "_cell_guid": "efb89169-be3c-4528-8b9d-bcbe360b1c38",
    "_uuid": "d1bf3247-046a-4963-b209-ae47e09a7cda"
   },
   "source": [
    "# Train by raw file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "_cell_guid": "529bb0dc-53da-4983-865a-724092a06c39",
    "_uuid": "d5de2487-d45b-4b02-bb76-f94ffc3ad48c",
    "execution": {
     "iopub.execute_input": "2023-02-26T06:42:01.616300Z",
     "iopub.status.busy": "2023-02-26T06:42:01.615276Z"
    },
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from tokenizers.trainers import BpeTrainer\n",
    "\n",
    "trainer = BpeTrainer(vocab_size=37000, show_progress=True, special_tokens=[\"<bos>\", \"<eos>\", \"<pad>\"])\n",
    "tokenizer.train(files=[\"de.txt\", \"en.txt\"], trainer=trainer)\n",
    "\n",
    "tokenizer.save(\"tokenizer.json\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "_cell_guid": "3b1606f3-ac1f-48e0-af01-5992527fd2f1",
    "_uuid": "874192b1-a261-4b1b-a8ae-7dc7b0084dd4"
   },
   "source": [
    "# train from memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "_cell_guid": "4ee7c700-5365-4dc5-949f-8468dbc9be8b",
    "_uuid": "cf8639ad-857c-4984-8ec9-2b423d4fc849",
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# def batch_iterator(batch_size=65536):\n",
    "#     for i in range(0, len(dataset), batch_size):\n",
    "#         yield dataset[i : i + batch_size]['translation']['en']\n",
    "#         yield dataset[i : i + batch_size]['translation']['de']\n",
    "        \n",
    "        \n",
    "# from tokenizers.trainers import BpeTrainer\n",
    "\n",
    "# trainer = BpeTrainer(vocab_size=37000, show_progress=True, special_tokens=[\"<bos>\", \"<eos>\", \"<pad>\"])\n",
    "# tokenizer.train_from_iterator(batch_iterator(),trainer=trainer,length=len(dataset))\n",
    "\n",
    "# tokenizer.save(\"tokenizer.json\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "_cell_guid": "369c4fce-6668-4bbf-85da-a8384582465f",
    "_uuid": "ebd2febf-cc11-4923-9930-ba7390c43585",
    "papermill": {
     "duration": null,
     "end_time": null,
     "exception": null,
     "start_time": null,
     "status": "pending"
    },
    "tags": []
   },
   "source": [
    "# use tokenizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### add data to memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "de_en_pairs = []\n",
    "for i in range(len(dataset)):\n",
    "    de_en_pairs.append((dataset[i]['translation']['de'],dataset[i]['translation']['en']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "de_en_pairs = sorted(de_en_pairs,key=lambda x:len(x[0])+len(x[1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "en_sents=[]\n",
    "de_sents=[]\n",
    "for pairs in de_en_pairs:\n",
    "    en_sents.append(pairs[1])\n",
    "    de_sents.append(pairs[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# en_sents[-1],de_sents[-1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train DataLoader"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### in paper there are 25000 source tokens and 25000 target tokens,here 4096 sourece tokens and 4096 target tokens for 1 gpu\n",
    "* for safe, 3125 for 1gpu and use gradient accumulation to update for 8 batch\n",
    "* for safe, 1024 for 1gpu for 24 batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "from tokenizers import Tokenizer\n",
    "tokenizer = Tokenizer.from_file(\"tokenizer.json\")\n",
    "def batch_generator(dataset,gpu_num=1,max_len=3125):\n",
    "    en_cnt = 0\n",
    "    de_cnt = 0\n",
    "    en_batch = []\n",
    "    de_batch = []\n",
    "    batch_size = 0\n",
    "    for pairs in dataset:\n",
    "        \n",
    "        en_batch.append(pairs[1])\n",
    "        de_batch.append(pairs[0])\n",
    "        en_cnt += len(pairs[1])\n",
    "        de_cnt += len(pairs[0])\n",
    "        batch_size += 1\n",
    "        \n",
    "        if batch_size%gpu_num == 0:          \n",
    "            if en_cnt + de_cnt > max_len*gpu_num:\n",
    "\n",
    "                en_output = tokenizer.encode_batch(en_batch[:-gpu_num])\n",
    "                de_output = tokenizer.encode_batch(de_batch[:-gpu_num])\n",
    "                \n",
    "                \n",
    "                en_ids = [] \n",
    "                de_ids = []\n",
    "                target_en_ids = []\n",
    "                target_de_ids = []\n",
    "                en_padding_mask = []\n",
    "                de_padding_mask = []\n",
    "\n",
    "                for en in en_output:\n",
    "                    en_ids.append(en.ids)\n",
    "                    target_en_ids.append(en.ids[1:]+[pad_id])\n",
    "                    en_padding_mask.append(en.attention_mask)\n",
    "                    \n",
    "                for de in de_output:\n",
    "                    de_ids.append(de.ids)\n",
    "                    target_de_ids.append(de.ids[1:]+[pad_id])\n",
    "                    de_padding_mask.append(de.attention_mask)              \n",
    "\n",
    "                yield torch.LongTensor(en_ids).t().contiguous(),\\\n",
    "                        torch.LongTensor(de_ids).t().contiguous(),\\\n",
    "                        torch.LongTensor(target_en_ids).t().contiguous(),\\\n",
    "                        torch.LongTensor(target_de_ids).t().contiguous(),\\\n",
    "                        torch.BoolTensor(1-np.array(en_padding_mask)),\\\n",
    "                        torch.BoolTensor(1-np.array(de_padding_mask))\n",
    "            \n",
    "\n",
    "                en_cnt = 0\n",
    "                de_cnt = 0            \n",
    "                en_batch = en_batch[-gpu_num:]\n",
    "                de_batch = de_batch[-gpu_num:]            \n",
    "\n",
    "    if en_ids:\n",
    "        yield torch.LongTensor(en_ids).t().contiguous(),\\\n",
    "                torch.LongTensor(de_ids).t().contiguous(),\\\n",
    "                torch.LongTensor(target_en_ids).t().contiguous(),\\\n",
    "                torch.LongTensor(target_de_ids).t().contiguous(),\\\n",
    "                torch.BoolTensor(1-np.array(en_padding_mask)),\\\n",
    "                torch.BoolTensor(1-np.array(de_padding_mask))\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "#"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### check (padding_mask in pytorch require pos of \\<pad> is True)\n",
    "* the shape of mask is (B,S), yes, it's batch first ,different from the input (S,B,E)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([5, 1115])\n",
      "tensor([[   0,    0,    0,  ...,    0,    0,    0],\n",
      "        [ 828,  828,  828,  ..., 2045,  418,  933],\n",
      "        [   1,    1,    1,  ...,   16,   16,   16],\n",
      "        [   2,    2,    2,  ...,   19,   23,   19],\n",
      "        [   2,    2,    2,  ...,    1,    1,    1]])\n",
      "torch.Size([1115, 5])\n",
      "tensor([[False, False, False,  True,  True],\n",
      "        [False, False, False,  True,  True],\n",
      "        [False, False, False,  True,  True],\n",
      "        ...,\n",
      "        [False, False, False, False, False],\n",
      "        [False, False, False, False, False],\n",
      "        [False, False, False, False, False]])\n",
      "torch.Size([5, 1115])\n",
      "tensor([[ 828,  828,  828,  ..., 2045,  418,  933],\n",
      "        [   1,    1,    1,  ...,   16,   16,   16],\n",
      "        [   2,    2,    2,  ...,   19,   23,   19],\n",
      "        [   2,    2,    2,  ...,    1,    1,    1],\n",
      "        [   2,    2,    2,  ...,    2,    2,    2]])\n",
      "******************************\n",
      "torch.Size([5, 1115])\n",
      "tensor([[   0,    0,    0,  ...,    0,    0,    0],\n",
      "        [ 828,  828,  828,  ..., 2045,  418,  933],\n",
      "        [   1,    1,    1,  ...,   14,   14,   14],\n",
      "        [   2,    2,    2,  ...,   19,   23,   19],\n",
      "        [   2,    2,    2,  ...,    1,    1,    1]])\n",
      "torch.Size([1115, 5])\n",
      "tensor([[False, False, False,  True,  True],\n",
      "        [False, False, False,  True,  True],\n",
      "        [False, False, False,  True,  True],\n",
      "        ...,\n",
      "        [False, False, False, False, False],\n",
      "        [False, False, False, False, False],\n",
      "        [False, False, False, False, False]])\n",
      "torch.Size([5, 1115])\n",
      "tensor([[ 828,  828,  828,  ..., 2045,  418,  933],\n",
      "        [   1,    1,    1,  ...,   14,   14,   14],\n",
      "        [   2,    2,    2,  ...,   19,   23,   19],\n",
      "        [   2,    2,    2,  ...,    1,    1,    1],\n",
      "        [   2,    2,    2,  ...,    2,    2,    2]])\n",
      " 5.1\n",
      " 5,1\n",
      "<bos> 5.1<eos>\n",
      "<bos> 5,1<eos>\n",
      "torch.Size([6, 412])\n",
      "tensor([[   0,    0,    0,  ...,    0,    0,    0],\n",
      "        [ 732,  418,  418,  ..., 1773, 1773, 1773],\n",
      "        [  16,   16,   16,  ...,  828,  828,  828],\n",
      "        [  26,   26,   27,  ...,    1,    1,    1],\n",
      "        [   1,    1,    1,  ...,    2,    2,    2],\n",
      "        [   2,    2,    2,  ...,    2,    2,    2]])\n",
      "torch.Size([412, 6])\n",
      "tensor([[False, False, False, False, False,  True],\n",
      "        [False, False, False, False, False,  True],\n",
      "        [False, False, False, False, False,  True],\n",
      "        ...,\n",
      "        [False, False, False, False,  True,  True],\n",
      "        [False, False, False, False,  True,  True],\n",
      "        [False, False, False, False,  True,  True]])\n",
      "torch.Size([6, 412])\n",
      "tensor([[ 732,  418,  418,  ..., 1773, 1773, 1773],\n",
      "        [  16,   16,   16,  ...,  828,  828,  828],\n",
      "        [  26,   26,   27,  ...,    1,    1,    1],\n",
      "        [   1,    1,    1,  ...,    2,    2,    2],\n",
      "        [   2,    2,    2,  ...,    2,    2,    2],\n",
      "        [   2,    2,    2,  ...,    2,    2,    2]])\n",
      "******************************\n",
      "torch.Size([6, 412])\n",
      "tensor([[   0,    0,    0,  ...,    0,    0,    0],\n",
      "        [ 732,  418,  418,  ..., 1773, 1773, 1773],\n",
      "        [  14,   14,   14,  ...,  828,  828,  828],\n",
      "        [  26,   26,   27,  ...,    1,    1,    1],\n",
      "        [   1,    1,    1,  ...,    2,    2,    2],\n",
      "        [   2,    2,    2,  ...,    2,    2,    2]])\n",
      "torch.Size([412, 6])\n",
      "tensor([[False, False, False, False, False,  True],\n",
      "        [False, False, False, False, False,  True],\n",
      "        [False, False, False, False, False,  True],\n",
      "        ...,\n",
      "        [False, False, False, False,  True,  True],\n",
      "        [False, False, False, False,  True,  True],\n",
      "        [False, False, False, False,  True,  True]])\n",
      "torch.Size([6, 412])\n",
      "tensor([[ 732,  418,  418,  ..., 1773, 1773, 1773],\n",
      "        [  14,   14,   14,  ...,  828,  828,  828],\n",
      "        [  26,   26,   27,  ...,    1,    1,    1],\n",
      "        [   1,    1,    1,  ...,    2,    2,    2],\n",
      "        [   2,    2,    2,  ...,    2,    2,    2],\n",
      "        [   2,    2,    2,  ...,    2,    2,    2]])\n",
      "    .\n",
      "    .\n",
      "<bos>    .<eos><pad><pad>\n",
      "<bos>    .<eos><pad><pad>\n",
      "torch.Size([4, 391])\n",
      "tensor([[   0,    0,    0,  ...,    0,    0,    0],\n",
      "        [1773, 1773, 1773,  ..., 1773, 1773, 1773],\n",
      "        [ 828,  828,  828,  ...,  828,  828,  828],\n",
      "        [   1,    1,    1,  ...,    1,    1,    1]])\n",
      "torch.Size([391, 4])\n",
      "tensor([[False, False, False, False],\n",
      "        [False, False, False, False],\n",
      "        [False, False, False, False],\n",
      "        ...,\n",
      "        [False, False, False, False],\n",
      "        [False, False, False, False],\n",
      "        [False, False, False, False]])\n",
      "torch.Size([4, 391])\n",
      "tensor([[1773, 1773, 1773,  ..., 1773, 1773, 1773],\n",
      "        [ 828,  828,  828,  ...,  828,  828,  828],\n",
      "        [   1,    1,    1,  ...,    1,    1,    1],\n",
      "        [   2,    2,    2,  ...,    2,    2,    2]])\n",
      "******************************\n",
      "torch.Size([4, 391])\n",
      "tensor([[   0,    0,    0,  ...,    0,    0,    0],\n",
      "        [1773, 1773, 1773,  ..., 1773, 1773, 1773],\n",
      "        [ 828,  828,  828,  ...,  828,  828,  828],\n",
      "        [   1,    1,    1,  ...,    1,    1,    1]])\n",
      "torch.Size([391, 4])\n",
      "tensor([[False, False, False, False],\n",
      "        [False, False, False, False],\n",
      "        [False, False, False, False],\n",
      "        ...,\n",
      "        [False, False, False, False],\n",
      "        [False, False, False, False],\n",
      "        [False, False, False, False]])\n",
      "torch.Size([4, 391])\n",
      "tensor([[1773, 1773, 1773,  ..., 1773, 1773, 1773],\n",
      "        [ 828,  828,  828,  ...,  828,  828,  828],\n",
      "        [   1,    1,    1,  ...,    1,    1,    1],\n",
      "        [   2,    2,    2,  ...,    2,    2,    2]])\n",
      "    .\n",
      "    .\n",
      "<bos>    .<eos>\n",
      "<bos>    .<eos>\n",
      "torch.Size([4, 391])\n",
      "tensor([[   0,    0,    0,  ...,    0,    0,    0],\n",
      "        [1773, 1773, 1773,  ..., 1773, 1773, 1773],\n",
      "        [ 828,  828,  828,  ...,  828,  828,  828],\n",
      "        [   1,    1,    1,  ...,    1,    1,    1]])\n",
      "torch.Size([391, 4])\n",
      "tensor([[False, False, False, False],\n",
      "        [False, False, False, False],\n",
      "        [False, False, False, False],\n",
      "        ...,\n",
      "        [False, False, False, False],\n",
      "        [False, False, False, False],\n",
      "        [False, False, False, False]])\n",
      "torch.Size([4, 391])\n",
      "tensor([[1773, 1773, 1773,  ..., 1773, 1773, 1773],\n",
      "        [ 828,  828,  828,  ...,  828,  828,  828],\n",
      "        [   1,    1,    1,  ...,    1,    1,    1],\n",
      "        [   2,    2,    2,  ...,    2,    2,    2]])\n",
      "******************************\n",
      "torch.Size([4, 391])\n",
      "tensor([[   0,    0,    0,  ...,    0,    0,    0],\n",
      "        [1773, 1773, 1773,  ..., 1773, 1773, 1773],\n",
      "        [ 828,  828,  828,  ...,  828,  828,  828],\n",
      "        [   1,    1,    1,  ...,    1,    1,    1]])\n",
      "torch.Size([391, 4])\n",
      "tensor([[False, False, False, False],\n",
      "        [False, False, False, False],\n",
      "        [False, False, False, False],\n",
      "        ...,\n",
      "        [False, False, False, False],\n",
      "        [False, False, False, False],\n",
      "        [False, False, False, False]])\n",
      "torch.Size([4, 391])\n",
      "tensor([[1773, 1773, 1773,  ..., 1773, 1773, 1773],\n",
      "        [ 828,  828,  828,  ...,  828,  828,  828],\n",
      "        [   1,    1,    1,  ...,    1,    1,    1],\n",
      "        [   2,    2,    2,  ...,    2,    2,    2]])\n",
      "    .\n",
      "    .\n",
      "<bos>    .<eos>\n",
      "<bos>    .<eos>\n",
      "torch.Size([4, 391])\n",
      "tensor([[   0,    0,    0,  ...,    0,    0,    0],\n",
      "        [1773, 1773, 1773,  ..., 1773, 1773, 1773],\n",
      "        [ 828,  828,  828,  ...,  828,  828,  828],\n",
      "        [   1,    1,    1,  ...,    1,    1,    1]])\n",
      "torch.Size([391, 4])\n",
      "tensor([[False, False, False, False],\n",
      "        [False, False, False, False],\n",
      "        [False, False, False, False],\n",
      "        ...,\n",
      "        [False, False, False, False],\n",
      "        [False, False, False, False],\n",
      "        [False, False, False, False]])\n",
      "torch.Size([4, 391])\n",
      "tensor([[1773, 1773, 1773,  ..., 1773, 1773, 1773],\n",
      "        [ 828,  828,  828,  ...,  828,  828,  828],\n",
      "        [   1,    1,    1,  ...,    1,    1,    1],\n",
      "        [   2,    2,    2,  ...,    2,    2,    2]])\n",
      "******************************\n",
      "torch.Size([4, 391])\n",
      "tensor([[   0,    0,    0,  ...,    0,    0,    0],\n",
      "        [1773, 1773, 1773,  ..., 1773, 1773, 1773],\n",
      "        [ 828,  828,  828,  ...,  828,  828,  828],\n",
      "        [   1,    1,    1,  ...,    1,    1,    1]])\n",
      "torch.Size([391, 4])\n",
      "tensor([[False, False, False, False],\n",
      "        [False, False, False, False],\n",
      "        [False, False, False, False],\n",
      "        ...,\n",
      "        [False, False, False, False],\n",
      "        [False, False, False, False],\n",
      "        [False, False, False, False]])\n",
      "torch.Size([4, 391])\n",
      "tensor([[1773, 1773, 1773,  ..., 1773, 1773, 1773],\n",
      "        [ 828,  828,  828,  ...,  828,  828,  828],\n",
      "        [   1,    1,    1,  ...,    1,    1,    1],\n",
      "        [   2,    2,    2,  ...,    2,    2,    2]])\n",
      "    .\n",
      "    .\n",
      "<bos>    .<eos>\n",
      "<bos>    .<eos>\n",
      "torch.Size([5, 391])\n",
      "tensor([[   0,    0,    0,  ...,    0,    0,    0],\n",
      "        [1773, 1773, 1773,  ..., 2930, 2930, 2930],\n",
      "        [ 828,  828,  828,  ...,   16,   16,   16],\n",
      "        [   1,    1,    1,  ...,    1,    1,    1],\n",
      "        [   2,    2,    2,  ...,    2,    2,    2]])\n",
      "torch.Size([391, 5])\n",
      "tensor([[False, False, False, False,  True],\n",
      "        [False, False, False, False,  True],\n",
      "        [False, False, False, False,  True],\n",
      "        ...,\n",
      "        [False, False, False, False,  True],\n",
      "        [False, False, False, False,  True],\n",
      "        [False, False, False, False,  True]])\n",
      "torch.Size([5, 391])\n",
      "tensor([[1773, 1773, 1773,  ..., 2930, 2930, 2930],\n",
      "        [ 828,  828,  828,  ...,   16,   16,   16],\n",
      "        [   1,    1,    1,  ...,    1,    1,    1],\n",
      "        [   2,    2,    2,  ...,    2,    2,    2],\n",
      "        [   2,    2,    2,  ...,    2,    2,    2]])\n",
      "******************************\n",
      "torch.Size([5, 391])\n",
      "tensor([[    0,     0,     0,  ...,     0,     0,     0],\n",
      "        [ 1773,  1773,  1773,  ..., 14499, 14499, 14499],\n",
      "        [  828,   828,   828,  ...,    16,    16,    16],\n",
      "        [    1,     1,     1,  ...,     1,     1,     1],\n",
      "        [    2,     2,     2,  ...,     2,     2,     2]])\n",
      "torch.Size([391, 5])\n",
      "tensor([[False, False, False, False,  True],\n",
      "        [False, False, False, False,  True],\n",
      "        [False, False, False, False,  True],\n",
      "        ...,\n",
      "        [False, False, False, False,  True],\n",
      "        [False, False, False, False,  True],\n",
      "        [False, False, False, False,  True]])\n",
      "torch.Size([5, 391])\n",
      "tensor([[ 1773,  1773,  1773,  ..., 14499, 14499, 14499],\n",
      "        [  828,   828,   828,  ...,    16,    16,    16],\n",
      "        [    1,     1,     1,  ...,     1,     1,     1],\n",
      "        [    2,     2,     2,  ...,     2,     2,     2],\n",
      "        [    2,     2,     2,  ...,     2,     2,     2]])\n",
      " No.\n",
      " Nein.\n",
      "<bos> No.<eos><pad>\n",
      "<bos> Nein.<eos><pad>\n",
      "torch.Size([7, 360])\n",
      "tensor([[   0,    0,    0,  ...,    0,    0,    0],\n",
      "        [2930, 2930, 2930,  ..., 8419, 1773, 8419],\n",
      "        [  16,   16,   16,  ...,   33,  835,   33],\n",
      "        ...,\n",
      "        [   2,    2,    2,  ...,    2,    1,    2],\n",
      "        [   2,    2,    2,  ...,    2,    2,    2],\n",
      "        [   2,    2,    2,  ...,    2,    2,    2]])\n",
      "torch.Size([360, 7])\n",
      "tensor([[False, False, False,  ...,  True,  True,  True],\n",
      "        [False, False, False,  ...,  True,  True,  True],\n",
      "        [False, False, False,  ...,  True,  True,  True],\n",
      "        ...,\n",
      "        [False, False, False,  ...,  True,  True,  True],\n",
      "        [False, False, False,  ..., False,  True,  True],\n",
      "        [False, False, False,  ...,  True,  True,  True]])\n",
      "torch.Size([7, 360])\n",
      "tensor([[2930, 2930, 2930,  ..., 8419, 1773, 8419],\n",
      "        [  16,   16,   16,  ...,   33,  835,   33],\n",
      "        [   1,    1,    1,  ...,    1,   16,    1],\n",
      "        ...,\n",
      "        [   2,    2,    2,  ...,    2,    2,    2],\n",
      "        [   2,    2,    2,  ...,    2,    2,    2],\n",
      "        [   2,    2,    2,  ...,    2,    2,    2]])\n",
      "******************************\n",
      "torch.Size([6, 360])\n",
      "tensor([[    0,     0,     0,  ...,     0,     0,     0],\n",
      "        [14499, 14499, 14499,  ...,  9031,  1773,  9031],\n",
      "        [   16,    16,    16,  ...,    33,   835,    33],\n",
      "        [    1,     1,     1,  ...,     1,    16,     1],\n",
      "        [    2,     2,     2,  ...,     2,     1,     2],\n",
      "        [    2,     2,     2,  ...,     2,     2,     2]])\n",
      "torch.Size([360, 6])\n",
      "tensor([[False, False, False, False,  True,  True],\n",
      "        [False, False, False, False,  True,  True],\n",
      "        [False, False, False, False,  True,  True],\n",
      "        ...,\n",
      "        [False, False, False, False,  True,  True],\n",
      "        [False, False, False, False, False,  True],\n",
      "        [False, False, False, False,  True,  True]])\n",
      "torch.Size([6, 360])\n",
      "tensor([[14499, 14499, 14499,  ...,  9031,  1773,  9031],\n",
      "        [   16,    16,    16,  ...,    33,   835,    33],\n",
      "        [    1,     1,     1,  ...,     1,    16,     1],\n",
      "        [    2,     2,     2,  ...,     2,     1,     2],\n",
      "        [    2,     2,     2,  ...,     2,     2,     2],\n",
      "        [    2,     2,     2,  ...,     2,     2,     2]])\n",
      " Why?\n",
      " Warum?\n",
      "<bos> Why?<eos><pad><pad><pad>\n",
      "<bos> Warum?<eos><pad><pad>\n"
     ]
    }
   ],
   "source": [
    "\n",
    "for en_ids,de_ids,target_en_ids,target_de_ids,\\\n",
    "    en_padding_mask,de_padding_mask in batch_generator(dataset=de_en_pairs,gpu_num=1):\n",
    "    print(en_ids.shape)\n",
    "    print(en_ids)\n",
    "    \n",
    "    print(en_padding_mask.shape)\n",
    "    print(en_padding_mask)    \n",
    "    \n",
    "    print(target_en_ids.shape)\n",
    "    print(target_en_ids)\n",
    "\n",
    "    print(\"*\"*30)\n",
    "    print(de_ids.shape)\n",
    "    print(de_ids)\n",
    "    \n",
    "    print(de_padding_mask.shape)\n",
    "    print(de_padding_mask)  \n",
    "    \n",
    "    print(target_de_ids.shape)\n",
    "    print(target_de_ids)   \n",
    "    \n",
    "    en_sent = en_ids.t().tolist()\n",
    "    de_sent = de_ids.t().tolist()\n",
    "    print(tokenizer.decode(en_sent[-1]))\n",
    "    print(tokenizer.decode(de_sent[-1]))\n",
    "    print(tokenizer.decode(en_sent[-1],skip_special_tokens = False))\n",
    "    print(tokenizer.decode(de_sent[-1],skip_special_tokens = False))\n",
    "    if (en_ids.shape[0]>=7):\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Vaild dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Found cached dataset wmt14 (/data2/zrs/.cache/huggingface/datasets/wmt14/de-en/1.0.0/2de185b074515e97618524d69f5e27ee7545dcbed4aa9bc1a4235710ffca33f4)\n"
     ]
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "valid_dataset = load_dataset(\"wmt14\", 'de-en', split='validation')\n",
    "\n",
    "valid_de_en_pairs = []\n",
    "for i in range(len(valid_dataset)):\n",
    "    valid_de_en_pairs.append((valid_dataset[i]['translation']['de'],valid_dataset[i]['translation']['en']))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### check"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([38, 12])\n",
      "tensor([[    0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "             0,     0],\n",
      "        [  264,  5864,  2218,  7993,  3804,   368, 23837,   711,   791,  1089,\n",
      "           401,  6729],\n",
      "        [ 5864,   228,    14,    14,  4776,   438,   235,   278, 16021,   213,\n",
      "           836,    14],\n",
      "        [  228,  5395,   229,  5864,   278,  5409, 12283,   235, 14078,  2351,\n",
      "         10165,   981],\n",
      "        [ 3979, 14585, 12961,   228,  2011,    14,    14,   438, 28485,    14,\n",
      "          7805,  7805],\n",
      "        [  262,   762,   228, 23392,    28,   229,   229,  6831,  3246, 17776,\n",
      "          1870,   617],\n",
      "        [ 7203,  1413,  6718, 13655,   981,  2322,  4830,   344,   229, 36027,\n",
      "         20431,  5889],\n",
      "        [  229,   471, 16279,   914,   836,   491,  1111,   213,  4229, 34801,\n",
      "           262,  4838],\n",
      "        [  305,   229,   438,  5975,  6393, 23112,   389,  4810,  2686,   229,\n",
      "          2902,  7216],\n",
      "        [   15,   898,   213,  4490,   491, 20479,  4213,   254,  6706,  9670,\n",
      "           213, 14132],\n",
      "        [ 4367,   262, 19536,   254,   473,   229,   326,  4830,    14,   254,\n",
      "          7686,    14],\n",
      "        [  625,  8728,    14, 18517,   213,  4830,   229,  5532,   507,   229,\n",
      "          9940, 32072],\n",
      "        [  254, 18517, 24541, 10485,  7580,  4642,  8761,   473, 10009,  1430,\n",
      "          5066,   411],\n",
      "        [10898, 10485,   344,   235,  4859,  1371,   254,  9880,  1233,   262,\n",
      "           261,   229],\n",
      "        [    1,    16, 18517,   229,   335,    16, 14463,   836,  9408,  2619,\n",
      "         12799,  1430],\n",
      "        [    2,     1, 10485,  2641,  2442,     1,  6706,  7805,   836,   235,\n",
      "           254,   262],\n",
      "        [    2,     2,   278,  1111,   291,     2,   235,  2198,  5864, 17254,\n",
      "          1317,  6060],\n",
      "        [    2,     2,   359,   235,  3113,     2,   229,  2158,   228,  1111,\n",
      "         19444,   396],\n",
      "        [    2,     2,   234,   213,    15,     2,  2641,  3221,  7527,  1074,\n",
      "            16,   213],\n",
      "        [    2,     2,   211, 15660,   640,     2,  1111,   229,  5730,  8024,\n",
      "             1,  2442],\n",
      "        [    2,     2,   235,    16,    16,     2,    16, 10607,   235,   235,\n",
      "             2,   291],\n",
      "        [    2,     2,   229,     1,     1,     2,     1,   448,  5877,  9107,\n",
      "             2,   335],\n",
      "        [    2,     2,  2641,     2,     2,     2,     2,  7216,  1111,  7700,\n",
      "             2, 10165],\n",
      "        [    2,     2,  1111,     2,     2,     2,     2,  1494,    16,    16,\n",
      "             2,  2134],\n",
      "        [    2,     2,  1138,     2,     2,     2,     2,   755,     1,     1,\n",
      "             2,   261],\n",
      "        [    2,     2,   229,     2,     2,     2,     2,  2799,     2,     2,\n",
      "             2, 10731],\n",
      "        [    2,     2,  1633,     2,     2,     2,     2,    16,     2,     2,\n",
      "             2,   229],\n",
      "        [    2,     2,   254,     2,     2,     2,     2,     1,     2,     2,\n",
      "             2,  1430],\n",
      "        [    2,     2,  1073,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,   262],\n",
      "        [    2,     2, 14476,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,  2619],\n",
      "        [    2,     2,   471,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,   254],\n",
      "        [    2,     2,  3710,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,  2284],\n",
      "        [    2,     2,  4045,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,   378],\n",
      "        [    2,     2,    16,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,   213],\n",
      "        [    2,     2,     1,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,  8747],\n",
      "        [    2,     2,     2,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,  4961],\n",
      "        [    2,     2,     2,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,    16],\n",
      "        [    2,     2,     2,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,     1]])\n",
      "torch.Size([12, 38])\n",
      "tensor([[False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True,  True,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True,  True,  True,  True,  True],\n",
      "        [False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False, False,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True,  True,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True,  True,  True,  True,  True],\n",
      "        [False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False,  True,  True,  True],\n",
      "        [False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False,  True,  True,  True,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True,  True,  True,  True,  True],\n",
      "        [False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False,  True,  True,  True,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True,  True,  True,  True,  True],\n",
      "        [False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False, False,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True,  True,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True,  True,  True,  True,  True],\n",
      "        [False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False,  True,  True,  True,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True,  True,  True,  True,  True],\n",
      "        [False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False, False, False, False,  True,  True,\n",
      "          True,  True,  True,  True,  True,  True,  True,  True],\n",
      "        [False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True,  True,  True,  True,  True],\n",
      "        [False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True,  True,  True,  True,  True],\n",
      "        [False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False, False, False, False, False, False,\n",
      "          True,  True,  True,  True,  True,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True,  True,  True,  True,  True],\n",
      "        [False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False, False, False, False]])\n",
      "torch.Size([38, 12])\n",
      "tensor([[  264,  5864,  2218,  7993,  3804,   368, 23837,   711,   791,  1089,\n",
      "           401,  6729],\n",
      "        [ 5864,   228,    14,    14,  4776,   438,   235,   278, 16021,   213,\n",
      "           836,    14],\n",
      "        [  228,  5395,   229,  5864,   278,  5409, 12283,   235, 14078,  2351,\n",
      "         10165,   981],\n",
      "        [ 3979, 14585, 12961,   228,  2011,    14,    14,   438, 28485,    14,\n",
      "          7805,  7805],\n",
      "        [  262,   762,   228, 23392,    28,   229,   229,  6831,  3246, 17776,\n",
      "          1870,   617],\n",
      "        [ 7203,  1413,  6718, 13655,   981,  2322,  4830,   344,   229, 36027,\n",
      "         20431,  5889],\n",
      "        [  229,   471, 16279,   914,   836,   491,  1111,   213,  4229, 34801,\n",
      "           262,  4838],\n",
      "        [  305,   229,   438,  5975,  6393, 23112,   389,  4810,  2686,   229,\n",
      "          2902,  7216],\n",
      "        [   15,   898,   213,  4490,   491, 20479,  4213,   254,  6706,  9670,\n",
      "           213, 14132],\n",
      "        [ 4367,   262, 19536,   254,   473,   229,   326,  4830,    14,   254,\n",
      "          7686,    14],\n",
      "        [  625,  8728,    14, 18517,   213,  4830,   229,  5532,   507,   229,\n",
      "          9940, 32072],\n",
      "        [  254, 18517, 24541, 10485,  7580,  4642,  8761,   473, 10009,  1430,\n",
      "          5066,   411],\n",
      "        [10898, 10485,   344,   235,  4859,  1371,   254,  9880,  1233,   262,\n",
      "           261,   229],\n",
      "        [    1,    16, 18517,   229,   335,    16, 14463,   836,  9408,  2619,\n",
      "         12799,  1430],\n",
      "        [    2,     1, 10485,  2641,  2442,     1,  6706,  7805,   836,   235,\n",
      "           254,   262],\n",
      "        [    2,     2,   278,  1111,   291,     2,   235,  2198,  5864, 17254,\n",
      "          1317,  6060],\n",
      "        [    2,     2,   359,   235,  3113,     2,   229,  2158,   228,  1111,\n",
      "         19444,   396],\n",
      "        [    2,     2,   234,   213,    15,     2,  2641,  3221,  7527,  1074,\n",
      "            16,   213],\n",
      "        [    2,     2,   211, 15660,   640,     2,  1111,   229,  5730,  8024,\n",
      "             1,  2442],\n",
      "        [    2,     2,   235,    16,    16,     2,    16, 10607,   235,   235,\n",
      "             2,   291],\n",
      "        [    2,     2,   229,     1,     1,     2,     1,   448,  5877,  9107,\n",
      "             2,   335],\n",
      "        [    2,     2,  2641,     2,     2,     2,     2,  7216,  1111,  7700,\n",
      "             2, 10165],\n",
      "        [    2,     2,  1111,     2,     2,     2,     2,  1494,    16,    16,\n",
      "             2,  2134],\n",
      "        [    2,     2,  1138,     2,     2,     2,     2,   755,     1,     1,\n",
      "             2,   261],\n",
      "        [    2,     2,   229,     2,     2,     2,     2,  2799,     2,     2,\n",
      "             2, 10731],\n",
      "        [    2,     2,  1633,     2,     2,     2,     2,    16,     2,     2,\n",
      "             2,   229],\n",
      "        [    2,     2,   254,     2,     2,     2,     2,     1,     2,     2,\n",
      "             2,  1430],\n",
      "        [    2,     2,  1073,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,   262],\n",
      "        [    2,     2, 14476,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,  2619],\n",
      "        [    2,     2,   471,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,   254],\n",
      "        [    2,     2,  3710,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,  2284],\n",
      "        [    2,     2,  4045,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,   378],\n",
      "        [    2,     2,    16,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,   213],\n",
      "        [    2,     2,     1,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,  8747],\n",
      "        [    2,     2,     2,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,  4961],\n",
      "        [    2,     2,     2,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,    16],\n",
      "        [    2,     2,     2,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,     1],\n",
      "        [    2,     2,     2,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,     2]])\n",
      "******************************\n",
      "torch.Size([54, 12])\n",
      "tensor([[    0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
      "             0,     0],\n",
      "        [ 2234,   487,  7428,   487, 23645,   368,  1124,   368,  3497, 36672,\n",
      "           487,  5526],\n",
      "        [ 2764, 36781,  9004, 30925,   374,  1011, 10083,  1011, 19293,  1481,\n",
      "          1730,  3224],\n",
      "        [14811,   269,   336, 18670,  2220,  6572,   337,  6572,   596, 17776,\n",
      "          2965,   468],\n",
      "        [ 8624, 30020, 12961,   269,    28, 22877, 15349,   596,   669,  4130,\n",
      "         28115,   697],\n",
      "        [ 4725, 12777,   228, 30020,   797,   797,   530,   267,   340,  1160,\n",
      "         11411,   797],\n",
      "        [   14, 20761,  7167, 12777,  1730,  2020,   267,  5809,  7893, 27505,\n",
      "            14, 11228],\n",
      "        [  526,  1040, 13431,   620,  6171,  8677,  1317,   269,  1141,  6153,\n",
      "           474,  3729],\n",
      "        [  269,  2288,   217,   235,   468,   336,    15,  9207,  4229,   362,\n",
      "           267,   267],\n",
      "        [ 4438,   388,   381,  1432,   462, 10915,    36,  5800,  2686,   857,\n",
      "         14805,  1065],\n",
      "        [ 3753,   269,   673,  1755, 20364,  2620,  5378,  2038,   242,  9107,\n",
      "           673,  9195],\n",
      "        [  358,  5551, 24502,   235,   379,   269,  1246,  2158,  2762,  9555,\n",
      "          7913,   381],\n",
      "        [10898,    14,   345,   340,   267,  2763,   381,  1617, 12262,    14,\n",
      "          2816,   267],\n",
      "        [ 7676,   340,    14,  2763,  2965,    16,   267, 11228,    14,   267,\n",
      "           500,   545],\n",
      "        [26939,  2965,  3558, 13191, 29147,     1,  7909, 36490,   643,   267,\n",
      "          3408,  1145],\n",
      "        [    1,  3394,   480,   804, 18386,     2,   269,    14,  1807, 19523,\n",
      "           270,   593],\n",
      "        [    2, 28640, 21318,  5975,    16,     2,  7893,   267,  1233,   385,\n",
      "           673, 18045],\n",
      "        [    2,   337,    14, 11071,     1,     2,   235,   336,  9408,  2965,\n",
      "         26548, 30718],\n",
      "        [    2, 12259,   474,   358,     2,     2,   340,  4458,  1617,  4681,\n",
      "           269, 10619],\n",
      "        [    2,    16,   269,  2965,     2,     2,  4791,   381,  2764,   235,\n",
      "          1317, 15339],\n",
      "        [    2,     1,  2965,  3394,     2,     2,  3060,   267, 14811, 17254,\n",
      "            15,    14],\n",
      "        [    2,     2,  3394, 28640,     2,     2,  7307, 15661,  8624,  3060,\n",
      "          2697,   336],\n",
      "        [    2,     2, 28640, 34322,     2,     2,    16,   603,  6859, 29769,\n",
      "            67,  2559],\n",
      "        [    2,     2,   235,    16,     2,     2,     1,   340,   235,    16,\n",
      "          1197,   381],\n",
      "        [    2,     2,   340,     1,     2,     2,     2,  2306,  5877,     1,\n",
      "         10434,   248],\n",
      "        [    2,     2,  2763,     2,     2,     2,     2,   424,  3060,     2,\n",
      "           671, 25962],\n",
      "        [    2,     2, 10494,     2,     2,     2,     2,  2437, 34322,     2,\n",
      "         13054,  7122],\n",
      "        [    2,     2,   211,     2,     2,     2,     2, 33425,   468,     2,\n",
      "            16,    14],\n",
      "        [    2,     2,   374,     2,     2,     2,     2,    16,  6457,     2,\n",
      "             1,   462],\n",
      "        [    2,     2,   428,     2,     2,     2,     2,     1,    16,     2,\n",
      "             2,   590],\n",
      "        [    2,     2,   267,     2,     2,     2,     2,     2,     1,     2,\n",
      "             2,  2965],\n",
      "        [    2,     2,  5576,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,  3788],\n",
      "        [    2,     2,   269,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,   428],\n",
      "        [    2,     2,  1141,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2, 14805],\n",
      "        [    2,     2, 34334,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,   337],\n",
      "        [    2,     2,  6942,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2, 21607],\n",
      "        [    2,     2,  1377,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,    14],\n",
      "        [    2,     2,   320,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,   270],\n",
      "        [    2,     2,  9801,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,  4934],\n",
      "        [    2,     2,  1177,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2, 35544],\n",
      "        [    2,     2,    16,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,   336],\n",
      "        [    2,     2,     1,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,  2965],\n",
      "        [    2,     2,     2,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,  3662],\n",
      "        [    2,     2,     2,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2, 17527],\n",
      "        [    2,     2,     2,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,  3639],\n",
      "        [    2,     2,     2,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,    14],\n",
      "        [    2,     2,     2,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,   381],\n",
      "        [    2,     2,     2,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,   267],\n",
      "        [    2,     2,     2,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,   430],\n",
      "        [    2,     2,     2,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,  7028],\n",
      "        [    2,     2,     2,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,  7399],\n",
      "        [    2,     2,     2,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2, 18357],\n",
      "        [    2,     2,     2,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,    16],\n",
      "        [    2,     2,     2,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,     1]])\n",
      "torch.Size([12, 54])\n",
      "tensor([[False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False, False,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True,  True,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True,  True,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True,  True,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True],\n",
      "        [False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False, False, False, False, False, False,\n",
      "         False,  True,  True,  True,  True,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True,  True,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True,  True,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True],\n",
      "        [False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False,  True,  True,  True,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True],\n",
      "        [False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True,  True,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True,  True,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True],\n",
      "        [False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False, False, False, False,  True,  True,\n",
      "          True,  True,  True,  True,  True,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True,  True,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True,  True,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True],\n",
      "        [False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True,  True,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True,  True,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True,  True,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True],\n",
      "        [False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False,  True,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True,  True,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True,  True,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True],\n",
      "        [False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False, False, False, False, False, False,\n",
      "          True,  True,  True,  True,  True,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True,  True,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True],\n",
      "        [False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False, False, False, False, False, False,\n",
      "         False,  True,  True,  True,  True,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True,  True,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True],\n",
      "        [False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True,  True,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True,  True,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True],\n",
      "        [False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False, False, False, False, False,  True,\n",
      "          True,  True,  True,  True,  True,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True,  True,  True,  True,  True,  True,  True,\n",
      "          True,  True,  True,  True],\n",
      "        [False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False, False, False, False, False, False, False,\n",
      "         False, False, False, False]])\n",
      "torch.Size([54, 12])\n",
      "tensor([[ 2234,   487,  7428,   487, 23645,   368,  1124,   368,  3497, 36672,\n",
      "           487,  5526],\n",
      "        [ 2764, 36781,  9004, 30925,   374,  1011, 10083,  1011, 19293,  1481,\n",
      "          1730,  3224],\n",
      "        [14811,   269,   336, 18670,  2220,  6572,   337,  6572,   596, 17776,\n",
      "          2965,   468],\n",
      "        [ 8624, 30020, 12961,   269,    28, 22877, 15349,   596,   669,  4130,\n",
      "         28115,   697],\n",
      "        [ 4725, 12777,   228, 30020,   797,   797,   530,   267,   340,  1160,\n",
      "         11411,   797],\n",
      "        [   14, 20761,  7167, 12777,  1730,  2020,   267,  5809,  7893, 27505,\n",
      "            14, 11228],\n",
      "        [  526,  1040, 13431,   620,  6171,  8677,  1317,   269,  1141,  6153,\n",
      "           474,  3729],\n",
      "        [  269,  2288,   217,   235,   468,   336,    15,  9207,  4229,   362,\n",
      "           267,   267],\n",
      "        [ 4438,   388,   381,  1432,   462, 10915,    36,  5800,  2686,   857,\n",
      "         14805,  1065],\n",
      "        [ 3753,   269,   673,  1755, 20364,  2620,  5378,  2038,   242,  9107,\n",
      "           673,  9195],\n",
      "        [  358,  5551, 24502,   235,   379,   269,  1246,  2158,  2762,  9555,\n",
      "          7913,   381],\n",
      "        [10898,    14,   345,   340,   267,  2763,   381,  1617, 12262,    14,\n",
      "          2816,   267],\n",
      "        [ 7676,   340,    14,  2763,  2965,    16,   267, 11228,    14,   267,\n",
      "           500,   545],\n",
      "        [26939,  2965,  3558, 13191, 29147,     1,  7909, 36490,   643,   267,\n",
      "          3408,  1145],\n",
      "        [    1,  3394,   480,   804, 18386,     2,   269,    14,  1807, 19523,\n",
      "           270,   593],\n",
      "        [    2, 28640, 21318,  5975,    16,     2,  7893,   267,  1233,   385,\n",
      "           673, 18045],\n",
      "        [    2,   337,    14, 11071,     1,     2,   235,   336,  9408,  2965,\n",
      "         26548, 30718],\n",
      "        [    2, 12259,   474,   358,     2,     2,   340,  4458,  1617,  4681,\n",
      "           269, 10619],\n",
      "        [    2,    16,   269,  2965,     2,     2,  4791,   381,  2764,   235,\n",
      "          1317, 15339],\n",
      "        [    2,     1,  2965,  3394,     2,     2,  3060,   267, 14811, 17254,\n",
      "            15,    14],\n",
      "        [    2,     2,  3394, 28640,     2,     2,  7307, 15661,  8624,  3060,\n",
      "          2697,   336],\n",
      "        [    2,     2, 28640, 34322,     2,     2,    16,   603,  6859, 29769,\n",
      "            67,  2559],\n",
      "        [    2,     2,   235,    16,     2,     2,     1,   340,   235,    16,\n",
      "          1197,   381],\n",
      "        [    2,     2,   340,     1,     2,     2,     2,  2306,  5877,     1,\n",
      "         10434,   248],\n",
      "        [    2,     2,  2763,     2,     2,     2,     2,   424,  3060,     2,\n",
      "           671, 25962],\n",
      "        [    2,     2, 10494,     2,     2,     2,     2,  2437, 34322,     2,\n",
      "         13054,  7122],\n",
      "        [    2,     2,   211,     2,     2,     2,     2, 33425,   468,     2,\n",
      "            16,    14],\n",
      "        [    2,     2,   374,     2,     2,     2,     2,    16,  6457,     2,\n",
      "             1,   462],\n",
      "        [    2,     2,   428,     2,     2,     2,     2,     1,    16,     2,\n",
      "             2,   590],\n",
      "        [    2,     2,   267,     2,     2,     2,     2,     2,     1,     2,\n",
      "             2,  2965],\n",
      "        [    2,     2,  5576,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,  3788],\n",
      "        [    2,     2,   269,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,   428],\n",
      "        [    2,     2,  1141,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2, 14805],\n",
      "        [    2,     2, 34334,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,   337],\n",
      "        [    2,     2,  6942,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2, 21607],\n",
      "        [    2,     2,  1377,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,    14],\n",
      "        [    2,     2,   320,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,   270],\n",
      "        [    2,     2,  9801,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,  4934],\n",
      "        [    2,     2,  1177,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2, 35544],\n",
      "        [    2,     2,    16,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,   336],\n",
      "        [    2,     2,     1,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,  2965],\n",
      "        [    2,     2,     2,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,  3662],\n",
      "        [    2,     2,     2,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2, 17527],\n",
      "        [    2,     2,     2,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,  3639],\n",
      "        [    2,     2,     2,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,    14],\n",
      "        [    2,     2,     2,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,   381],\n",
      "        [    2,     2,     2,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,   267],\n",
      "        [    2,     2,     2,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,   430],\n",
      "        [    2,     2,     2,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,  7028],\n",
      "        [    2,     2,     2,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,  7399],\n",
      "        [    2,     2,     2,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2, 18357],\n",
      "        [    2,     2,     2,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,    16],\n",
      "        [    2,     2,     2,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,     1],\n",
      "        [    2,     2,     2,     2,     2,     2,     2,     2,     2,     2,\n",
      "             2,     2]])\n",
      " Furthermore, these laws also reduce early voting periods, invalidate the right to register as a voter on election day and withdraw the right to vote of citizens with a criminal record.\n",
      " Darüber hinaus werden durch diese Gesetze ebenfalls die Zeiträume für die vorzeitige Stimmabgabe verkürzt, das Recht für ungültig erklärt, sich am Wahltag als Wähler zu registrieren, und Staatsbürgern das Wahlrecht abgesprochen, für die eine Gerichtsakte vorliegt.\n",
      "<bos> Furthermore, these laws also reduce early voting periods, invalidate the right to register as a voter on election day and withdraw the right to vote of citizens with a criminal record.<eos>\n",
      "<bos> Darüber hinaus werden durch diese Gesetze ebenfalls die Zeiträume für die vorzeitige Stimmabgabe verkürzt, das Recht für ungültig erklärt, sich am Wahltag als Wähler zu registrieren, und Staatsbürgern das Wahlrecht abgesprochen, für die eine Gerichtsakte vorliegt.<eos>\n"
     ]
    }
   ],
   "source": [
    "\n",
    "for en_ids,de_ids,target_en_ids,target_de_ids,\\\n",
    "    en_padding_mask,de_padding_mask in batch_generator(dataset=valid_de_en_pairs,gpu_num=1):\n",
    "    print(en_ids.shape)\n",
    "    print(en_ids)\n",
    "    \n",
    "    print(en_padding_mask.shape)\n",
    "    print(en_padding_mask)    \n",
    "    \n",
    "    print(target_en_ids.shape)\n",
    "    print(target_en_ids)\n",
    "\n",
    "    print(\"*\"*30)\n",
    "    print(de_ids.shape)\n",
    "    print(de_ids)\n",
    "    \n",
    "    print(de_padding_mask.shape)\n",
    "    print(de_padding_mask)  \n",
    "    \n",
    "    print(target_de_ids.shape)\n",
    "    print(target_de_ids)   \n",
    "    \n",
    "    en_sent = en_ids.t().tolist()\n",
    "    de_sent = de_ids.t().tolist()\n",
    "    print(tokenizer.decode(en_sent[-1]))\n",
    "    print(tokenizer.decode(de_sent[-1]))\n",
    "    print(tokenizer.decode(en_sent[-1],skip_special_tokens = False))\n",
    "    print(tokenizer.decode(de_sent[-1],skip_special_tokens = False))\n",
    "    if (en_ids.shape[0]>=7):\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "_cell_guid": "d9287aa2-5140-476b-89d3-2838a8936af7",
    "_uuid": "7234e2f7-561e-403b-8587-7ca6b96eb43a"
   },
   "source": [
    "# other info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "_cell_guid": "29699e46-b572-4b46-b220-139b2095deb2",
    "_uuid": "e84b64b6-9fac-4bc8-8551-f36bf4b3697b",
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37000\n",
      "0\n",
      "1\n",
      "2\n"
     ]
    }
   ],
   "source": [
    "print(tokenizer.get_vocab_size())\n",
    "print(tokenizer.token_to_id('<bos>'))\n",
    "print(tokenizer.token_to_id('<eos>'))\n",
    "print(tokenizer.token_to_id('<pad>'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}