figure-Copy1.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "425b5b23",
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "from importlib import reload\n",
    "\n",
    "import torch\n",
    "from torch.nn.functional import softmax, log_softmax\n",
    "\n",
    "from datasets import load_dataset\n",
    "\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
    "\n",
    "from util import mt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "92a91a78",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = 'EleutherAI/gpt-neo-125M'\n",
    "model_name = 'facebook/xglm-7.5B'\n",
    "model_name = 'facebook/xglm-4.5B'\n",
    "model_name = 'facebook/xglm-2.9B'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5170b939",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "model = AutoModelForCausalLM.from_pretrained(model_name).eval().half().cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a8be7dcd",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/gscratch/zlab/ahai/miniconda3/lib/python3.9/site-packages/huggingface_hub/utils/_deprecation.py:97: FutureWarning: Deprecated argument(s) used in 'dataset_info': token. Will not be supported from version '0.12'.\n",
      "  warnings.warn(message, FutureWarning)\n",
      "Found cached dataset flores_101 (/gscratch/zlab/ahai/hf/datasets/gsarti___flores_101/mya/1.0.0/e663cc717b274f2cef5142df786973abbf1966a7115d99509c062936d77d4335)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "414c2ae4a60d409da9b797f269776213",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/gscratch/zlab/ahai/miniconda3/lib/python3.9/site-packages/huggingface_hub/utils/_deprecation.py:97: FutureWarning: Deprecated argument(s) used in 'dataset_info': token. Will not be supported from version '0.12'.\n",
      "  warnings.warn(message, FutureWarning)\n",
      "Found cached dataset flores_101 (/gscratch/zlab/ahai/hf/datasets/gsarti___flores_101/eng/1.0.0/e663cc717b274f2cef5142df786973abbf1966a7115d99509c062936d77d4335)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "db1931b1014e4518a8da1bda576f5d9c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Found cached dataset flores_101 (/gscratch/zlab/ahai/hf/datasets/gsarti___flores_101/eng/1.0.0/e663cc717b274f2cef5142df786973abbf1966a7115d99509c062936d77d4335)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4bf34c6ac1264c218bc3b88f0674d60f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "src = load_dataset(\"gsarti/flores_101\", 'mya')\n",
    "tgt = load_dataset(\"gsarti/flores_101\", 'eng')\n",
    "eng = load_dataset(\"gsarti/flores_101\", 'eng')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "5fa59a92",
   "metadata": {},
   "outputs": [],
   "source": [
    "src_lang = 'Burmese'\n",
    "tgt_lang = 'English'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e2d1befc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokenize(s):\n",
    "  ids = []\n",
    "  i = 0\n",
    "  t = 0\n",
    "  while t < len(s):\n",
    "    c = s[t]\n",
    "    if c == '\\n':\n",
    "      ids += s_ids(s[i:t]) + [2]\n",
    "      i = t+1\n",
    "      t = t+1\n",
    "    else:\n",
    "      t += 1\n",
    "  ids += s_ids(s[i:])\n",
    "  return ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a27eb68a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def s_ids(s):\n",
    "  return tokenizer(s).input_ids[1:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "0eaa9be4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_demo_prompt(src_endonym, tgt_endonym, src, tgt, k):\n",
    "  src_points = random.sample(list(src), k=k)\n",
    "  tgt_points = [ tgt[point['id']-1] for point in src_points ]\n",
    "  demos = []\n",
    "  for s, t in zip(src_points, tgt_points):\n",
    "    demo = f\"{src_endonym} : {s['sentence']} {tgt_endonym} : {t['sentence']}\"\n",
    "    demos.append(demo)\n",
    "  demo_prompt_ids = [2]\n",
    "  for demo in demos:\n",
    "    demo_prompt_ids.extend(s_ids(demo))\n",
    "    demo_prompt_ids.append(2)\n",
    "  return demo_prompt_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e6dc4e3",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62b40acd",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8c6c019",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec2dd5a2",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 1443,
   "id": "63953d94",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<module 'util.mt' from '/mmfs1/gscratch/zlab/ahai/repo/infogain/util/mt.py'>"
      ]
     },
     "execution_count": 1443,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "reload(mt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "48c418ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "src_points = random.sample(list(src['dev']), k=10)\n",
    "tgt_points = [ tgt['dev'][point['id']-1] for point in src_points ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d891d63c",
   "metadata": {},
   "outputs": [],
   "source": [
    "demo_prompt_ids = mt.make_demo_prompt(src_points, tgt_points, src_lang, tgt_lang, tokenizer)\n",
    "pmi_demo_prompt_ids = mt.make_demo_prompt(tgt_points, src_points, tgt_lang, src_lang, tokenizer)\n",
    "uncond_demo_prompt_ids = mt.make_uncond_demo_prompt(tgt_points, tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "b5bd48d6",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'idx' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m/tmp/ipykernel_2594/3208072265.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0minference_point\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0msrc\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'dev'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0midx\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      2\u001b[0m \u001b[0mcond_prompt_ids\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0muncond_demo_prompt_ids\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0ms_ids\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf\"{src_lang} : {inference_point['sentence']} {tgt_lang} :\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0mcond_prompt_t\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mLongTensor\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mcond_prompt_ids\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0muncond_prompt_ids\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdemo_prompt_ids\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0ms_ids\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf\"{src_lang} : {tgt_lang} :\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[0muncond_prompt_t\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mLongTensor\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'idx' is not defined"
     ]
    }
   ],
   "source": [
    "inference_point = src['dev'][idx] \n",
    "cond_prompt_ids = uncond_demo_prompt_ids + s_ids(f\"{src_lang} : {inference_point['sentence']} {tgt_lang} :\")\n",
    "cond_prompt_t = torch.LongTensor([cond_prompt_ids]).to(model.device)\n",
    "uncond_prompt_ids = demo_prompt_ids + s_ids(f\"{src_lang} : {tgt_lang} :\")\n",
    "uncond_prompt_t = torch.LongTensor([[2]]).to(model.device)\n",
    "print(idx)\n",
    "print()\n",
    "print(eng['dev'][idx]['sentence'])\n",
    "print()\n",
    "print(src['dev'][idx]['sentence'])\n",
    "print()\n",
    "print(tgt['dev'][idx]['sentence'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 392,
   "id": "a65d8bbb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "When the tenants started sharing what had occurred to them, most of the families involved suddenly realized that Carolyn Wilson of the OHA had stolen their security deposits, and skipped out of town.\n",
      "\n",
      "Kun vuokralaiset alkoivat kertoa, mitä heille oli tapahtunut, useimmat perheet, joita asia koski, käsittivät pian, että Oaklandin asuntoviraston edustaja Carolyn Wilson varastanut heidän takuutalletuksensa ja paennut kaupungista.\n",
      "\n",
      "When the tenants started sharing what had occurred to them, most of the families involved suddenly realized that Carolyn Wilson of the OHA had stolen their security deposits, and skipped out of town.\n"
     ]
    }
   ],
   "source": [
    "inference_point = src['dev'][idx] \n",
    "cond_prompt_ids = demo_prompt_ids + s_ids(f\"{src_lang} : {inference_point['sentence']} {tgt_lang} :\")\n",
    "cond_prompt_t = torch.LongTensor([cond_prompt_ids]).to(model.device)\n",
    "uncond_prompt_ids = demo_prompt_ids + s_ids(f\"{src_lang} : {tgt_lang} :\")\n",
    "uncond_prompt_t = torch.LongTensor([uncond_demo_prompt_ids]).to(model.device)\n",
    "print(eng['dev'][idx]['sentence'])\n",
    "print()\n",
    "print(src['dev'][idx]['sentence'])\n",
    "print()\n",
    "print(tgt['dev'][idx]['sentence'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d83f249",
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "  gen_ids = mt.cfg_p_greedy(model, 0.2, cond_prompt_t, uncond_prompt_t, 50, tokenizer=tokenizer, p=0.5)\n",
    "print()\n",
    "print(tokenizer.decode(gen_ids[0].tolist()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "4bca9d0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "idx = -1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "ba8a828f",
   "metadata": {},
   "outputs": [],
   "source": [
    "idx += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "f586006d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "On Monday, scientists from the Stanford University School of Medicine announced the invention of a new diagnostic tool that can sort cells by type: a tiny printable chip that can be manufactured using standard inkjet printers for possibly about one U.S. cent each.\n",
      "\n",
      "တနင်္လာနေ့တွင် စတန်းဖို့ဒ်တက္ကသိုလ် ဆေးကျောင်းမှ သိပ္ပံပညာရှင်များသည် ဆဲလ်များကို အမျိုးအစားအလိုက် စီစဉ်နိုင်သော ရောဂါခွဲခြားစစ်ဆေးမှု ကိရိယာအသစ် တီထွင်မှုအကြောင်းကို ကြေညာခဲ့သည်- ၎င်းသည် စံနှုန်းမီ မင်စက်ကလေးများဖြင့် ပုံဖော်သည့် ပရင်တာများကို သုံးကာ ထုတ်လုပ်နိုင်သော အလွန်သေးငယ်သည့် တစ်ခုလျှင် U.S. ဆင့်တစ်ပြားသာသာရှိသော ပရင့်လုပ်၍ရသော အပြားလေးဖြစ်ပါသည်။\n",
      "\n",
      "On Monday, scientists from the Stanford University School of Medicine announced the invention of a new diagnostic tool that can sort cells by type: a tiny printable chip that can be manufactured using standard inkjet printers for possibly about one U.S. cent each.\n"
     ]
    }
   ],
   "source": [
    "inference_point = src['dev'][idx] \n",
    "cond_prompt_ids = demo_prompt_ids + mt.xglm_ids(f\"{src_lang} : {inference_point['sentence']} {tgt_lang} :\", tokenizer)\n",
    "cond_prompt_t = torch.LongTensor([cond_prompt_ids]).to(model.device)\n",
    "uncond_prompt_ids = demo_prompt_ids + mt.xglm_ids(f\"{tgt_lang} :\", tokenizer)\n",
    "uncond_prompt_t = torch.LongTensor([uncond_prompt_ids]).to(model.device)\n",
    "print(eng['dev'][idx]['sentence'])\n",
    "print()\n",
    "print(src['dev'][idx]['sentence'])\n",
    "print() \n",
    "print(tgt['dev'][idx]['sentence'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 723,
   "id": "0f7cd54d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<module 'util.mt' from '/mmfs1/gscratch/zlab/ahai/repo/infogain/util/mt.py'>"
      ]
     },
     "execution_count": 723,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "reload(mt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 724,
   "id": "8b6c8284",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Mr. Ackley was born in Hong Kong and attended Harvard University and Harvard Law School.</s> Burmese : ၎င်းသည် ၎င်း၏ ဇာတ်ကောင်ကို ပုံဖော်ရာတွင် ၎င်း၏ မိသားစုနှင့် ၎\n"
     ]
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "  gen_ids = mt.tp_greedy(model, 1, cond_prompt_t, uncond_prompt_t, 50, tokenizer=tokenizer, p=0.8)\n",
    "print()\n",
    "print(tokenizer.decode(gen_ids[0].tolist()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 725,
   "id": "237ad239",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mr. Aung was born in Hong Kong and attended the University of New York and Harvard Law School.</s> Burmese : ၎င်းသည် ၎င်း၏ ဇာတ်ကောင်ကို ပုံဖော်ရာတွင် ၎င်း၏ ဇာတ်\n"
     ]
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "  greedy_gen_ids = mt.greedy(model, cond_prompt_t, 50)\n",
    "print(tokenizer.decode(greedy_gen_ids[0].tolist()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "1930e238",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tuesday evening, the University of Toronto's School of Biological Sciences announced the development of a new test that can classify cells by their type.</s> Burmese : ရွှေ့ပြောင်းအလုပ်သမားများသည် ၎င်းတို့၏ မိသားစုဝင်များကို \n"
     ]
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "  greedy_gen_ids = mt.p_sample(model, cond_prompt_t, 50, 0.5, temp=0.5)\n",
    "print(tokenizer.decode(greedy_gen_ids[0].tolist()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 737,
   "id": "cd0b394b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<module 'util.mt' from '/mmfs1/gscratch/zlab/ahai/repo/infogain/util/mt.py'>"
      ]
     },
     "execution_count": 737,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "reload(mt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "7a5dd66a",
   "metadata": {},
   "outputs": [],
   "source": [
    "inference_point = src['dev'][idx]\n",
    "pmi_prompt_ids = pmi_demo_prompt_ids + mt.xglm_ids(f\"{tgt_lang} : {tokenizer.decode(greedy_gen_ids[0])} {src_lang} :\", tokenizer)\n",
    "pmi_prompt_t = torch.LongTensor([pmi_prompt_ids]).to(model.device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "9e921b6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "gen_ids = torch.LongTensor([mt.cleanup(greedy_gen_ids[0].tolist())]).to(model.device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "17ba837b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 31])"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gen_ids.size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "727b7d95",
   "metadata": {},
   "outputs": [],
   "source": [
    "pmi_check_t = torch.cat([pmi_prompt_t, gen_ids], dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "956affb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "res = model(pmi_check_t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "3a85bd96",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/conda-bld/pytorch_1659484806139/work/aten/src/ATen/native/cuda/ScatterGatherKernel.cu:144: operator(): block: [0,0,0], thread: [0,0,0] Assertion `idx_dim >= 0 && idx_dim < index_size && \"index out of bounds\"` failed.\n",
      "/opt/conda/conda-bld/pytorch_1659484806139/work/aten/src/ATen/native/cuda/ScatterGatherKernel.cu:144: operator(): block: [0,0,0], thread: [1,0,0] Assertion `idx_dim >= 0 && idx_dim < index_size && \"index out of bounds\"` failed.\n",
      "/opt/conda/conda-bld/pytorch_1659484806139/work/aten/src/ATen/native/cuda/ScatterGatherKernel.cu:144: operator(): block: [0,0,0], thread: [3,0,0] Assertion `idx_dim >= 0 && idx_dim < index_size && \"index out of bounds\"` failed.\n",
      "/opt/conda/conda-bld/pytorch_1659484806139/work/aten/src/ATen/native/cuda/ScatterGatherKernel.cu:144: operator(): block: [0,0,0], thread: [4,0,0] Assertion `idx_dim >= 0 && idx_dim < index_size && \"index out of bounds\"` failed.\n",
      "/opt/conda/conda-bld/pytorch_1659484806139/work/aten/src/ATen/native/cuda/ScatterGatherKernel.cu:144: operator(): block: [0,0,0], thread: [5,0,0] Assertion `idx_dim >= 0 && idx_dim < index_size && \"index out of bounds\"` failed.\n",
      "/opt/conda/conda-bld/pytorch_1659484806139/work/aten/src/ATen/native/cuda/ScatterGatherKernel.cu:144: operator(): block: [0,0,0], thread: [6,0,0] Assertion `idx_dim >= 0 && idx_dim < index_size && \"index out of bounds\"` failed.\n",
      "/opt/conda/conda-bld/pytorch_1659484806139/work/aten/src/ATen/native/cuda/ScatterGatherKernel.cu:144: operator(): block: [0,0,0], thread: [9,0,0] Assertion `idx_dim >= 0 && idx_dim < index_size && \"index out of bounds\"` failed.\n",
      "/opt/conda/conda-bld/pytorch_1659484806139/work/aten/src/ATen/native/cuda/ScatterGatherKernel.cu:144: operator(): block: [0,0,0], thread: [10,0,0] Assertion `idx_dim >= 0 && idx_dim < index_size && \"index out of bounds\"` failed.\n",
      "/opt/conda/conda-bld/pytorch_1659484806139/work/aten/src/ATen/native/cuda/ScatterGatherKernel.cu:144: operator(): block: [0,0,0], thread: [11,0,0] Assertion `idx_dim >= 0 && idx_dim < index_size && \"index out of bounds\"` failed.\n",
      "/opt/conda/conda-bld/pytorch_1659484806139/work/aten/src/ATen/native/cuda/ScatterGatherKernel.cu:144: operator(): block: [0,0,0], thread: [12,0,0] Assertion `idx_dim >= 0 && idx_dim < index_size && \"index out of bounds\"` failed.\n",
      "/opt/conda/conda-bld/pytorch_1659484806139/work/aten/src/ATen/native/cuda/ScatterGatherKernel.cu:144: operator(): block: [0,0,0], thread: [13,0,0] Assertion `idx_dim >= 0 && idx_dim < index_size && \"index out of bounds\"` failed.\n",
      "/opt/conda/conda-bld/pytorch_1659484806139/work/aten/src/ATen/native/cuda/ScatterGatherKernel.cu:144: operator(): block: [0,0,0], thread: [14,0,0] Assertion `idx_dim >= 0 && idx_dim < index_size && \"index out of bounds\"` failed.\n",
      "/opt/conda/conda-bld/pytorch_1659484806139/work/aten/src/ATen/native/cuda/ScatterGatherKernel.cu:144: operator(): block: [0,0,0], thread: [15,0,0] Assertion `idx_dim >= 0 && idx_dim < index_size && \"index out of bounds\"` failed.\n",
      "/opt/conda/conda-bld/pytorch_1659484806139/work/aten/src/ATen/native/cuda/ScatterGatherKernel.cu:144: operator(): block: [0,0,0], thread: [16,0,0] Assertion `idx_dim >= 0 && idx_dim < index_size && \"index out of bounds\"` failed.\n",
      "/opt/conda/conda-bld/pytorch_1659484806139/work/aten/src/ATen/native/cuda/ScatterGatherKernel.cu:144: operator(): block: [0,0,0], thread: [17,0,0] Assertion `idx_dim >= 0 && idx_dim < index_size && \"index out of bounds\"` failed.\n",
      "/opt/conda/conda-bld/pytorch_1659484806139/work/aten/src/ATen/native/cuda/ScatterGatherKernel.cu:144: operator(): block: [0,0,0], thread: [19,0,0] Assertion `idx_dim >= 0 && idx_dim < index_size && \"index out of bounds\"` failed.\n",
      "/opt/conda/conda-bld/pytorch_1659484806139/work/aten/src/ATen/native/cuda/ScatterGatherKernel.cu:144: operator(): block: [0,0,0], thread: [20,0,0] Assertion `idx_dim >= 0 && idx_dim < index_size && \"index out of bounds\"` failed.\n",
      "/opt/conda/conda-bld/pytorch_1659484806139/work/aten/src/ATen/native/cuda/ScatterGatherKernel.cu:144: operator(): block: [0,0,0], thread: [21,0,0] Assertion `idx_dim >= 0 && idx_dim < index_size && \"index out of bounds\"` failed.\n",
      "/opt/conda/conda-bld/pytorch_1659484806139/work/aten/src/ATen/native/cuda/ScatterGatherKernel.cu:144: operator(): block: [0,0,0], thread: [22,0,0] Assertion `idx_dim >= 0 && idx_dim < index_size && \"index out of bounds\"` failed.\n",
      "/opt/conda/conda-bld/pytorch_1659484806139/work/aten/src/ATen/native/cuda/ScatterGatherKernel.cu:144: operator(): block: [0,0,0], thread: [23,0,0] Assertion `idx_dim >= 0 && idx_dim < index_size && \"index out of bounds\"` failed.\n",
      "/opt/conda/conda-bld/pytorch_1659484806139/work/aten/src/ATen/native/cuda/ScatterGatherKernel.cu:144: operator(): block: [0,0,0], thread: [24,0,0] Assertion `idx_dim >= 0 && idx_dim < index_size && \"index out of bounds\"` failed.\n",
      "/opt/conda/conda-bld/pytorch_1659484806139/work/aten/src/ATen/native/cuda/ScatterGatherKernel.cu:144: operator(): block: [0,0,0], thread: [25,0,0] Assertion `idx_dim >= 0 && idx_dim < index_size && \"index out of bounds\"` failed.\n",
      "/opt/conda/conda-bld/pytorch_1659484806139/work/aten/src/ATen/native/cuda/ScatterGatherKernel.cu:144: operator(): block: [0,0,0], thread: [26,0,0] Assertion `idx_dim >= 0 && idx_dim < index_size && \"index out of bounds\"` failed.\n",
      "/opt/conda/conda-bld/pytorch_1659484806139/work/aten/src/ATen/native/cuda/ScatterGatherKernel.cu:144: operator(): block: [0,0,0], thread: [27,0,0] Assertion `idx_dim >= 0 && idx_dim < index_size && \"index out of bounds\"` failed.\n",
      "/opt/conda/conda-bld/pytorch_1659484806139/work/aten/src/ATen/native/cuda/ScatterGatherKernel.cu:144: operator(): block: [0,0,0], thread: [28,0,0] Assertion `idx_dim >= 0 && idx_dim < index_size && \"index out of bounds\"` failed.\n"
     ]
    },
    {
     "ename": "RuntimeError",
     "evalue": "CUDA error: device-side assert triggered\nCUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.\nFor debugging consider passing CUDA_LAUNCH_BLOCKING=1.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;32m/gscratch/zlab/ahai/miniconda3/lib/python3.9/site-packages/IPython/core/formatters.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, obj)\u001b[0m\n\u001b[1;32m    700\u001b[0m                 \u001b[0mtype_pprinters\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtype_printers\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    701\u001b[0m                 deferred_pprinters=self.deferred_printers)\n\u001b[0;32m--> 702\u001b[0;31m             \u001b[0mprinter\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpretty\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mobj\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    703\u001b[0m             \u001b[0mprinter\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mflush\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    704\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mstream\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgetvalue\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/gscratch/zlab/ahai/miniconda3/lib/python3.9/site-packages/IPython/lib/pretty.py\u001b[0m in \u001b[0;36mpretty\u001b[0;34m(self, obj)\u001b[0m\n\u001b[1;32m    392\u001b[0m                         \u001b[0;32mif\u001b[0m \u001b[0mcls\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mobject\u001b[0m\u001b[0;31m \u001b[0m\u001b[0;31m\\\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    393\u001b[0m                                 \u001b[0;32mand\u001b[0m \u001b[0mcallable\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcls\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__dict__\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'__repr__'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 394\u001b[0;31m                             \u001b[0;32mreturn\u001b[0m \u001b[0m_repr_pprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mobj\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcycle\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    395\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    396\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0m_default_pprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mobj\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcycle\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/gscratch/zlab/ahai/miniconda3/lib/python3.9/site-packages/IPython/lib/pretty.py\u001b[0m in \u001b[0;36m_repr_pprint\u001b[0;34m(obj, p, cycle)\u001b[0m\n\u001b[1;32m    698\u001b[0m     \u001b[0;34m\"\"\"A pprint that just redirects to the normal repr function.\"\"\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    699\u001b[0m     \u001b[0;31m# Find newlines and replace them with p.break_()\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 700\u001b[0;31m     \u001b[0moutput\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mrepr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mobj\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    701\u001b[0m     \u001b[0mlines\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0moutput\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msplitlines\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    702\u001b[0m     \u001b[0;32mwith\u001b[0m \u001b[0mp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgroup\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/gscratch/zlab/ahai/miniconda3/lib/python3.9/site-packages/torch/_tensor.py\u001b[0m in \u001b[0;36m__repr__\u001b[0;34m(self, tensor_contents)\u001b[0m\n\u001b[1;32m    336\u001b[0m                                          tensor_contents=tensor_contents)\n\u001b[1;32m    337\u001b[0m         \u001b[0;31m# All strings are unicode in Python 3.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 338\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_tensor_str\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_str\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtensor_contents\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtensor_contents\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    339\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    340\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgradient\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minputs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/gscratch/zlab/ahai/miniconda3/lib/python3.9/site-packages/torch/_tensor_str.py\u001b[0m in \u001b[0;36m_str\u001b[0;34m(self, tensor_contents)\u001b[0m\n\u001b[1;32m    479\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0m_str\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtensor_contents\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    480\u001b[0m     \u001b[0;32mwith\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mno_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 481\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0m_str_intern\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtensor_contents\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtensor_contents\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m/gscratch/zlab/ahai/miniconda3/lib/python3.9/site-packages/torch/_tensor_str.py\u001b[0m in \u001b[0;36m_str_intern\u001b[0;34m(inp, tensor_contents)\u001b[0m\n\u001b[1;32m    445\u001b[0m                         \u001b[0mtensor_str\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_tensor_str\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto_dense\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mindent\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    446\u001b[0m                     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 447\u001b[0;31m                         \u001b[0mtensor_str\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_tensor_str\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mindent\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    448\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    449\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlayout\u001b[0m \u001b[0;34m!=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstrided\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/gscratch/zlab/ahai/miniconda3/lib/python3.9/site-packages/torch/_tensor_str.py\u001b[0m in \u001b[0;36m_tensor_str\u001b[0;34m(self, indent)\u001b[0m\n\u001b[1;32m    268\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0m_tensor_str_with_formatter\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mindent\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msummarize\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mreal_formatter\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mimag_formatter\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    269\u001b[0m     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 270\u001b[0;31m         \u001b[0mformatter\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_Formatter\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mget_summarized_data\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0msummarize\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    271\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0m_tensor_str_with_formatter\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mindent\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msummarize\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mformatter\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    272\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/gscratch/zlab/ahai/miniconda3/lib/python3.9/site-packages/torch/_tensor_str.py\u001b[0m in \u001b[0;36m__init__\u001b[0;34m(self, tensor)\u001b[0m\n\u001b[1;32m    101\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    102\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 103\u001b[0;31m             \u001b[0mnonzero_finite_vals\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmasked_select\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtensor_view\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0misfinite\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtensor_view\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m&\u001b[0m \u001b[0mtensor_view\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mne\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    104\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    105\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mnonzero_finite_vals\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnumel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: CUDA error: device-side assert triggered\nCUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.\nFor debugging consider passing CUDA_LAUNCH_BLOCKING=1."
     ]
    }
   ],
   "source": [
    "res.logits[0, -1-gen_ids.size(1):-1].gather(0, gen_ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "ca9f3663",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 31])"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gen_ids.size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "3d91c4c3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 981, 256008])"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res.logits.size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34124290",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "830a1cc0",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "id": "8227260e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.return_types.topk(\n",
       "values=tensor([-31.4399, -31.4399, -14.8043, -14.8043, -14.8043, -14.8043, -14.3989,\n",
       "        -14.3989, -14.3989, -14.3989, -14.3989, -14.3989, -14.3989, -14.1112,\n",
       "        -14.1112, -14.1112, -14.1112, -14.1112, -14.1112, -14.1112, -14.1112,\n",
       "        -14.1112, -14.1112, -14.1112, -14.1112, -13.8880, -13.8880, -13.8880,\n",
       "        -13.8880, -13.8880, -13.8880, -13.8880, -13.8880, -13.8880, -13.8880,\n",
       "        -13.8880, -13.8880, -13.8880, -13.8880, -13.8880, -13.8880, -13.8880,\n",
       "        -13.8880, -13.8880, -13.8880, -13.8880, -13.8880, -13.8880, -13.8880,\n",
       "        -13.8880, -13.8880, -13.8880, -13.8880, -13.8880, -13.8880, -13.7057,\n",
       "        -13.7057, -13.7057, -13.7057, -13.7057, -13.7057, -13.7057, -13.7057,\n",
       "        -13.7057, -13.7057, -13.7057, -13.7057, -13.7057, -13.7057, -13.7057,\n",
       "        -13.7057, -13.7057, -13.7057, -13.7057, -13.7057, -13.7057, -13.7057,\n",
       "        -13.7057, -13.7057, -13.7057, -13.7057, -13.7057, -13.7057, -13.7057,\n",
       "        -13.7057, -13.7057, -13.7057, -13.7057, -13.7057, -13.7057, -13.7057,\n",
       "        -13.7057, -13.7057, -13.7057, -13.7057, -13.7057, -13.7057, -13.7057,\n",
       "        -13.7057, -13.7057], device='cuda:0'),\n",
       "indices=tensor([202800, 246898, 162646, 201312,  35749,   6092, 222231, 224336,   6794,\n",
       "        209299, 230486, 218803,  60691, 244982, 236946, 190425,  76594,  68842,\n",
       "         73763,  32844, 126552, 146480,  18665, 172867, 169850, 211013, 246405,\n",
       "        237009, 254379, 219347, 232540, 235445,  91723,  94186,  58618,  68274,\n",
       "         33661,  42271,  37159, 197516, 114620, 210815, 213183, 203655, 216841,\n",
       "        219282, 228462, 161697, 154000, 102926, 100613, 100984, 182369, 119325,\n",
       "        185575, 198657, 105342, 191362, 202803, 205307, 130045, 123898, 130431,\n",
       "        133283,  25006,  24821,  15562,  18134,  61337,  48797,  61484,   8219,\n",
       "         68399,  96949,  91114,  95528, 109271, 106704, 111824, 115251, 174701,\n",
       "        171254, 165133, 165793, 159518, 151082, 162956, 163262, 185326, 184985,\n",
       "        180071, 181178, 187327, 187136, 187866, 188906, 150727, 142741, 134512,\n",
       "        141213], device='cuda:0'))"
      ]
     },
     "execution_count": 85,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "temp = 1\n",
    "alpha = 1\n",
    "eps = torch.finfo(torch.float32).eps\n",
    "with torch.no_grad():\n",
    "  c_lls = log_softmax(model(cond_prompt_t).logits[0][-1] / temp, dim=0)\n",
    "  u_lls = log_softmax(model(uncond_prompt_t).logits[0][-1] / temp, dim=0)\n",
    "  pmi = c_lls - u_lls\n",
    "  pmi = pmi.long()\n",
    "  pmi_dist = pmi - pmi.min() + eps\n",
    "  pmi_dist = pmi_dist / pmi_dist.sum()\n",
    "  pmi_lls = pmi_dist.log()\n",
    "#   pmi.div_(pmi.sum())\n",
    "#   apmi = alpha*pmi + (1-alpha)*torch.ones_like(pmi).div_(pmi.size(0))\n",
    "# tokenizer.decode(torch.topk(apmi, 10).indices)\n",
    "torch.topk(pmi_lls, 100, largest=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "id": "3633ce8c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-13.7058])"
      ]
     },
     "execution_count": 78,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.Tensor([1.1160e-06]).log()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1989,
   "id": "88858c56",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.return_types.topk(\n",
      "values=tensor([-0.6235, -0.8735, -5.3750, -5.5000, -6.0000, -6.2500, -6.2500, -6.6250,\n",
      "        -7.0000, -7.2500], device='cuda:0', dtype=torch.float16),\n",
      "indices=tensor([   268, 139910,     46,    984,    345,  21992, 102595,     32,   2686,\n",
      "           826], device='cuda:0'))\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_52376/1101773222.py:2: UserWarning: Implicit dimension choice for log_softmax has been deprecated. Change the call to include dim=X as an argument.\n",
      "  print(torch.topk(log_softmax(model(cond_prompt_t).logits[0][-1] / temp), 10))\n"
     ]
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "  print(torch.topk(log_softmax(model(cond_prompt_t).logits[0][-1] / temp), 10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 775,
   "id": "d1b7ca3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_ids = torch.LongTensor([tokenizer('translate into english, 中文： 谢谢你！').input_ids[:-1]]).to(model.device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 232,
   "id": "93d7c56f",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_ids = torch.LongTensor([tokenize('中文： 谢谢你! english:')]).to(model.device)                          "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 233,
   "id": "0602b35a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[6, 19773, 13, 6, 13305, 13305, 710, 35, 83763, 13]]"
      ]
     },
     "execution_count": 233,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prompt_ids.tolist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 351,
   "id": "fe912cb2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "XGLMConfig {\n",
       "  \"_name_or_path\": \"facebook/xglm-2.9B\",\n",
       "  \"activation_dropout\": 0.0,\n",
       "  \"activation_function\": \"gelu\",\n",
       "  \"architectures\": [\n",
       "    \"XGLMForCausalLM\"\n",
       "  ],\n",
       "  \"attention_dropout\": 0.1,\n",
       "  \"attention_heads\": 16,\n",
       "  \"bos_token_id\": 0,\n",
       "  \"d_model\": 2048,\n",
       "  \"decoder_start_token_id\": 2,\n",
       "  \"dropout\": 0.1,\n",
       "  \"eos_token_id\": 2,\n",
       "  \"ffn_dim\": 8192,\n",
       "  \"init_std\": 0.02,\n",
       "  \"layerdrop\": 0.0,\n",
       "  \"max_position_embeddings\": 2048,\n",
       "  \"model_type\": \"xglm\",\n",
       "  \"num_layers\": 48,\n",
       "  \"pad_token_id\": 1,\n",
       "  \"scale_embedding\": true,\n",
       "  \"transformers_version\": \"4.22.2\",\n",
       "  \"use_cache\": true,\n",
       "  \"vocab_size\": 256008\n",
       "}"
      ]
     },
     "execution_count": 351,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 275,
   "id": "cdd568ae",
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "expected Tensor as element 0 in argument 0, but got list",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m/tmp/ipykernel_52376/3632386642.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;32mwith\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mno_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m   \u001b[0mgen_ids\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmt\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msample\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mprompt_ids\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m200\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m/mmfs1/gscratch/zlab/ahai/repo/infogain/util/mt.py\u001b[0m in \u001b[0;36msample\u001b[0;34m(model, prompt_ids, length, temp)\u001b[0m\n\u001b[1;32m     14\u001b[0m     \u001b[0mgen_ids\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mLongTensor\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     15\u001b[0m     \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlength\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 16\u001b[0;31m         \u001b[0mcur_ids\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mprompt_ids\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgen_ids\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdim\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     17\u001b[0m         \u001b[0mlogits\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcur_ids\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlogits\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     18\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mtemp\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mTypeError\u001b[0m: expected Tensor as element 0 in argument 0, but got list"
     ]
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "  gen_ids = mt.sample(model, prompt_ids, 200)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 309,
   "id": "39940ea3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Leading researchers say that the test can detect cancer, tuberculosis, HIV and malaria in patients who live in low-income countries, where for example, the survival rate of breast cancer is half that of the richer countries.</s> Finnish.</s> Finnish : Hän sanoi, että hän sanoi, että hän sanoi, että, että, että, että, että, että, että, että, että, että, että, että, että, että, että, että,'"
      ]
     },
     "execution_count": 309,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.decode(gen_ids[0].tolist())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 231,
   "id": "25d4f61a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "中文: 谢谢你! | english:\n"
     ]
    }
   ],
   "source": [
    "print(tokenizer.decode(prompt_ids[0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "5b70b758",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Stanfordin yliopiston lääketieteen laitoksen tutkijat ilmoittivat maanantaina uuden diagnostiikkatyökalun keksimisestä: solut tyypin mukaan lajitteleva pienenpieni tulostettava siru, joka voidaan valmistaa normaaleilla mustesuihkutulostimilla mahdollisesti noin yhden Yhdysvaltain sentin kappalehintaan.'"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "inference_point['sentence']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 209,
   "id": "9a74a4e1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[]"
      ]
     },
     "execution_count": 209,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "s_ids('')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 219,
   "id": "e9f6fcad",
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokenize(s):\n",
    "  ids = []\n",
    "  i = 0\n",
    "  t = 0\n",
    "  while t < len(s):\n",
    "    c = s[t]\n",
    "    if c == '\\n':\n",
    "      ids += s_ids(s[i:t]) + [2]\n",
    "      i = t+1\n",
    "      t = t+1\n",
    "    else:\n",
    "      t += 1\n",
    "  ids += s_ids(s[i:])\n",
    "  return ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c8169866",
   "metadata": {},
   "outputs": [],
   "source": [
    "vocab = tokenizer.get_vocab()\n",
    "tokens = list(sorted(tokenizer.get_vocab(), key=lambda k: vocab[k]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 121,
   "id": "c9b91f8e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['<s>',\n",
       " '<pad>',\n",
       " '</s>',\n",
       " '<unk>',\n",
       " ',',\n",
       " '.',\n",
       " '▁',\n",
       " 's',\n",
       " '-',\n",
       " 'a',\n",
       " '▁de',\n",
       " '▁a',\n",
       " 'e',\n",
       " ':',\n",
       " 'i']"
      ]
     },
     "execution_count": 121,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokens[:15]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 305,
   "id": "c9b43d40",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['SPECIAL_TOKENS_ATTRIBUTES',\n",
       " '__annotations__',\n",
       " '__call__',\n",
       " '__class__',\n",
       " '__delattr__',\n",
       " '__dict__',\n",
       " '__dir__',\n",
       " '__doc__',\n",
       " '__eq__',\n",
       " '__format__',\n",
       " '__ge__',\n",
       " '__getattribute__',\n",
       " '__gt__',\n",
       " '__hash__',\n",
       " '__init__',\n",
       " '__init_subclass__',\n",
       " '__le__',\n",
       " '__len__',\n",
       " '__lt__',\n",
       " '__module__',\n",
       " '__ne__',\n",
       " '__new__',\n",
       " '__reduce__',\n",
       " '__reduce_ex__',\n",
       " '__repr__',\n",
       " '__setattr__',\n",
       " '__sizeof__',\n",
       " '__str__',\n",
       " '__subclasshook__',\n",
       " '__weakref__',\n",
       " '_add_tokens',\n",
       " '_additional_special_tokens',\n",
       " '_auto_class',\n",
       " '_batch_encode_plus',\n",
       " '_bos_token',\n",
       " '_call_one',\n",
       " '_cls_token',\n",
       " '_convert_encoding',\n",
       " '_convert_id_to_token',\n",
       " '_convert_token_to_id_with_added_voc',\n",
       " '_create_repo',\n",
       " '_decode',\n",
       " '_decode_use_source_tokenizer',\n",
       " '_encode_plus',\n",
       " '_eos_token',\n",
       " '_eventual_warn_about_too_long_sequence',\n",
       " '_eventually_correct_t5_max_length',\n",
       " '_from_pretrained',\n",
       " '_get_files_timestamps',\n",
       " '_get_padding_truncation_strategies',\n",
       " '_in_target_context_manager',\n",
       " '_mask_token',\n",
       " '_pad',\n",
       " '_pad_token',\n",
       " '_pad_token_type_id',\n",
       " '_processor_class',\n",
       " '_save_pretrained',\n",
       " '_sep_token',\n",
       " '_set_processor_class',\n",
       " '_switch_to_input_mode',\n",
       " '_switch_to_target_mode',\n",
       " '_tokenizer',\n",
       " '_unk_token',\n",
       " '_upload_modified_files',\n",
       " 'add_special_tokens',\n",
       " 'add_tokens',\n",
       " 'additional_special_tokens',\n",
       " 'additional_special_tokens_ids',\n",
       " 'all_special_ids',\n",
       " 'all_special_tokens',\n",
       " 'all_special_tokens_extended',\n",
       " 'as_target_tokenizer',\n",
       " 'backend_tokenizer',\n",
       " 'batch_decode',\n",
       " 'batch_encode_plus',\n",
       " 'bos_token',\n",
       " 'bos_token_id',\n",
       " 'build_inputs_with_special_tokens',\n",
       " 'can_save_slow_tokenizer',\n",
       " 'clean_up_tokenization',\n",
       " 'cls_token',\n",
       " 'cls_token_id',\n",
       " 'convert_ids_to_tokens',\n",
       " 'convert_tokens_to_ids',\n",
       " 'convert_tokens_to_string',\n",
       " 'create_token_type_ids_from_sequences',\n",
       " 'decode',\n",
       " 'decoder',\n",
       " 'deprecation_warnings',\n",
       " 'encode',\n",
       " 'encode_plus',\n",
       " 'eos_token',\n",
       " 'eos_token_id',\n",
       " 'from_pretrained',\n",
       " 'get_added_vocab',\n",
       " 'get_special_tokens_mask',\n",
       " 'get_vocab',\n",
       " 'init_inputs',\n",
       " 'init_kwargs',\n",
       " 'is_fast',\n",
       " 'mask_token',\n",
       " 'mask_token_id',\n",
       " 'max_len_sentences_pair',\n",
       " 'max_len_single_sentence',\n",
       " 'max_model_input_sizes',\n",
       " 'model_input_names',\n",
       " 'model_max_length',\n",
       " 'name_or_path',\n",
       " 'num_madeup_words',\n",
       " 'num_special_tokens_to_add',\n",
       " 'pad',\n",
       " 'pad_token',\n",
       " 'pad_token_id',\n",
       " 'pad_token_type_id',\n",
       " 'padding_side',\n",
       " 'prepare_for_model',\n",
       " 'prepare_seq2seq_batch',\n",
       " 'pretrained_init_configuration',\n",
       " 'pretrained_vocab_files_map',\n",
       " 'push_to_hub',\n",
       " 'register_for_auto_class',\n",
       " 'sanitize_special_tokens',\n",
       " 'save_pretrained',\n",
       " 'save_vocabulary',\n",
       " 'sep_token',\n",
       " 'sep_token_id',\n",
       " 'set_truncation_and_padding',\n",
       " 'slow_tokenizer_class',\n",
       " 'special_tokens_map',\n",
       " 'special_tokens_map_extended',\n",
       " 'tokenize',\n",
       " 'train_new_from_iterator',\n",
       " 'truncate_sequences',\n",
       " 'truncation_side',\n",
       " 'unk_token',\n",
       " 'unk_token_id',\n",
       " 'verbose',\n",
       " 'vocab',\n",
       " 'vocab_file',\n",
       " 'vocab_files_names',\n",
       " 'vocab_size']"
      ]
     },
     "execution_count": 305,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dir(tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 306,
   "id": "253e0e5e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1000000000000000019884624838656"
      ]
     },
     "execution_count": 306,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.model_max_length"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 319,
   "id": "5f05feed",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_demo_prompt(src_points, tgt_points, src_name, tgt_name, src, tgt, k, tokenizer):\n",
    "    demos = []\n",
    "    for s, t in zip(src_points, tgt_points):\n",
    "        assert(s['id'] == t['id'])\n",
    "        demo = f\"{src_name} : {s['sentence']} {tgt_name} : {t['sentence']}\"\n",
    "        demos.append(demo)\n",
    "    demo_prompt_ids = [2]\n",
    "    for demo in demos:\n",
    "      demo_prompt_ids.extend(mt.xglm_tokenize(demo, tokenizer))\n",
    "      demo_prompt_ids.append(2)\n",
    "    return demo_prompt_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 323,
   "id": "c629e8cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "src_points = random.sample(list(src['dev']), k=3)\n",
    "tgt_points = [ tgt['dev'][point['id']-1] for point in src_points ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 334,
   "id": "4737f0b0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<module 'util.mt' from '/mmfs1/gscratch/zlab/ahai/repo/infogain/util/mt.py'>"
      ]
     },
     "execution_count": 334,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "reload(mt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 330,
   "id": "563a6b24",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " finnish : Kiertoajelut ovat halvempia suurille ryhmille, joten jos liikut yksin tai vain yhden ystävän kanssa, koeta tutustua muihin ihmisiin ja muodostaa 4–6 hengen ryhmä parempaa yksilöhintaa varten. english : Tours are cheaper for larger groups, so if you're by yourself or with just one friend, try to meet other people and form a group of four to six for a better per-person rate.\n",
      " finnish : Vaatimukset on tarkoitettu tekemään maiden välisestä muuttovirrasta järjestelmällisen. english : These requirements are designed to provide an organized migratory flow between both countries.\n",
      " finnish : Varo: pikkukaupunkien baarit eivät täälläpäin aina ole parhaita ajanviettopaikkoja vieraspaikkakuntalaisille. english : Beware: small-town bars here are not always good places for the out-of-state visitor to hang out.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(mt.post_proc(tokenizer.decode(mt.make_demo_prompt(src_points, tgt_points, 'finnish', 'english', tokenizer))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 331,
   "id": "2ccd9bbf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_uncond_demo_prompt(tgt_points, tokenizer):\n",
    "    uncond_demo_prompt_ids = [2]\n",
    "    for t in tgt_points:\n",
    "      uncond_demo_prompt_ids.extend(mt.xglm_tokenize(t['sentence'], tokenizer))\n",
    "      uncond_demo_prompt_ids.append(2)\n",
    "    return uncond_demo_prompt_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 333,
   "id": "09b18b80",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " Tours are cheaper for larger groups, so if you're by yourself or with just one friend, try to meet other people and form a group of four to six for a better per-person rate.\n",
      " These requirements are designed to provide an organized migratory flow between both countries.\n",
      " Beware: small-town bars here are not always good places for the out-of-state visitor to hang out.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(mt.post_proc(tokenizer.decode(make_uncond_demo_prompt(tgt_points, tokenizer))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f168d750",
   "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.9.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}