Prefix_detection_SFR_Mistral.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3102f184-ce19-4a85-a719-a9285fa8c12d",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ef0aec88-592c-4b5c-bf37-0e367c449d28",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5bbf3bce35b143749377c4ab49c919dd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import torch\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "\n",
    "MODEL_NAME = \"mistralai\"\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(\"mistralai/Mistral-7B-v0.1\", cache_dir=\"/mistralai/\", torch_dtype=torch.float16).to(\n",
    "        \"cuda\")\n",
    "tok = AutoTokenizer.from_pretrained(\"mistralai/Mistral-7B-v0.1\", cache_dir=\"/mistralai/\", torch_dtype=torch.float16)\n",
    "\n",
    "for param in model.parameters():\n",
    "    param.requires_grad = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5e371e1b-15f6-442c-a22d-963bb96c18c6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "MistralForCausalLM(\n",
       "  (model): MistralModel(\n",
       "    (embed_tokens): Embedding(32000, 4096)\n",
       "    (layers): ModuleList(\n",
       "      (0-31): 32 x MistralDecoderLayer(\n",
       "        (self_attn): MistralAttention(\n",
       "          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (k_proj): Linear(in_features=4096, out_features=1024, bias=False)\n",
       "          (v_proj): Linear(in_features=4096, out_features=1024, bias=False)\n",
       "          (o_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (rotary_emb): MistralRotaryEmbedding()\n",
       "        )\n",
       "        (mlp): MistralMLP(\n",
       "          (gate_proj): Linear(in_features=4096, out_features=14336, bias=False)\n",
       "          (up_proj): Linear(in_features=4096, out_features=14336, bias=False)\n",
       "          (down_proj): Linear(in_features=14336, out_features=4096, bias=False)\n",
       "          (act_fn): SiLU()\n",
       "        )\n",
       "        (input_layernorm): MistralRMSNorm()\n",
       "        (post_attention_layernorm): MistralRMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (norm): MistralRMSNorm()\n",
       "  )\n",
       "  (lm_head): Linear(in_features=4096, out_features=32000, bias=False)\n",
       ")"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "901ad08e-ab8d-4490-81ae-c714d0a0aefd",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = \"cuda\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c1d83d3e-97bb-44b1-885b-4d646d7948bb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d6d8e070cf37497eb9cbffbb8a5a28de",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch import Tensor\n",
    "from transformers import AutoTokenizer, AutoModel\n",
    "\n",
    "# load model and tokenizer\n",
    "model = AutoModel.from_pretrained('Salesforce/SFR-Embedding-Mistral', cache_dir=\"/SFR_Mistral/\", torch_dtype=torch.float16).to(device)\n",
    "\n",
    "# no special token\n",
    "tok = AutoTokenizer.from_pretrained('Salesforce/SFR-Embedding-Mistral', cache_dir=\"/SFR_Mistral/\", torch_dtype=torch.float16, add_special_tokens=False) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3f689c89-8859-4d5e-84c7-86bd51e04be9",
   "metadata": {},
   "outputs": [],
   "source": [
    "for param in model.parameters():\n",
    "    param.requires_grad = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "63f91121-0dd5-4ddd-9f1a-2cd4715c808d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tok.add_special_tokens({'pad_token': tok.eos_token})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "d7ad010f-6b17-49ab-b592-7578e62d01db",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_model_layer_attentions(sentence, token_ids, model, tokenizer, layer):\n",
    "    input_ids = tokenizer(sentence, padding='max_length', max_length=100, return_tensors=\"pt\")[\"input_ids\"].to(device)\n",
    "    lst = token_ids + input_ids.tolist()[0]\n",
    "    input_ids = torch.tensor([lst], dtype=torch.int).to(device)\n",
    "    with torch.no_grad():\n",
    "        output = model(input_ids, output_attentions=True)\n",
    "    \n",
    "    flattened_attentions = [attn_layer.flatten() for attn_layer in output.attentions]\n",
    "    flattened_vector = torch.cat(flattened_attentions)  # This concatenates them into a single vector\n",
    "    return flattened_vector\n",
    "\n",
    "def get_model_layer_hidden_states(sentence, token_ids, model, tokenizer, layer):\n",
    "    input_ids = tokenizer(sentence, padding='max_length', max_length=100, return_tensors=\"pt\")[\"input_ids\"].to(device)\n",
    "    lst = token_ids + input_ids.tolist()[0]\n",
    "    input_ids = torch.tensor([lst], dtype=torch.int).to(device)\n",
    "    with torch.no_grad():\n",
    "        output = model(input_ids, output_hidden_states=True, output_attentions=True)\n",
    "    flattened_hidden_states = [attn_layer.flatten() for attn_layer in output.hidden_states[-1:]]\n",
    "    flattened_vector = torch.cat(flattened_hidden_states)  # This concatenates them into a single vector\n",
    "    return flattened_vector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "e24c67e1-6293-433e-b8d5-4fcef6f4aac7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch import Tensor\n",
    "from transformers import AutoTokenizer, AutoModel\n",
    "\n",
    "def last_token_pool(last_hidden_states: Tensor,\n",
    "                 attention_mask: Tensor) -> Tensor:\n",
    "    left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])\n",
    "    if left_padding:\n",
    "        return last_hidden_states[:, -1]\n",
    "    else:\n",
    "        sequence_lengths = attention_mask.sum(dim=1) - 1\n",
    "        batch_size = last_hidden_states.shape[0]\n",
    "        return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]\n",
    "\n",
    "def get_detailed_instruct(task_description: str, query: str) -> str:\n",
    "    return f'Instruct: {task_description}\\nQuery: {query}'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32e18981-3f9c-4bdc-b164-e64fa6e3a355",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "id": "5272ad05-5091-4310-9284-e30ebf6a9577",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys([507, 508, 509, 510, 511, 512, 513, 514, 515, 516, 517, 518, 519, 520, 521, 522, 523, 524, 525, 526, 527, 528, 529, 530, 531, 532, 533, 534, 535, 536, 537, 538, 539, 540, 541, 542, 543, 544, 545, 546, 547, 548, 549, 550, 551, 552, 553, 554, 555, 556, 557, 558, 559, 560, 561, 562, 563, 564, 565, 566, 567, 568, 569, 570, 571, 572, 573, 574, 575, 576, 577, 578, 579, 580, 581, 582, 583, 584, 585, 586, 587, 588, 589, 590, 591, 592, 593, 594, 595, 596, 597, 598, 599, 600, 601, 602, 603, 604, 605, 606, 607, 608, 609, 610, 611, 612, 613, 614, 615])"
      ]
     },
     "execution_count": 89,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pickle\n",
    "with open('SFR_Mistral_IMDB_1.pkl', 'rb') as f:\n",
    "    json_str_1 = pickle.load(f)\n",
    "\n",
    "json_str_1.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "id": "4363d0d9-af51-48b0-9b6e-23b6820c9b56",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys([400, 401, 402, 403, 404, 405, 406, 407, 408, 409, 410, 411, 412, 413, 414, 415, 416, 417, 418, 419, 420, 421, 422, 423, 424, 425, 426, 427, 428, 429, 430, 431, 432, 433, 434, 435, 436, 437, 438, 439, 440, 441, 442, 443, 444, 445, 446, 447, 448, 449, 450, 451, 452, 453, 454, 455, 456, 457, 458, 459, 460, 461, 462, 463, 464, 465, 466, 467, 468, 469, 470, 471, 472, 473, 474, 475, 476, 477, 478, 479, 480, 481, 482, 483, 484, 485, 486, 487, 488, 489, 490, 491, 492, 493, 494, 495, 496, 497, 498, 499, 500, 501, 502, 503, 504, 505, 506])"
      ]
     },
     "execution_count": 90,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pickle\n",
    "with open('SFR_Mistral_IMDB_2.pkl', 'rb') as f:\n",
    "    json_str_2 = pickle.load(f)\n",
    "\n",
    "json_str_2.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "id": "dd676122-6e81-4f18-bfba-8f7cc85251b3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys([267, 268, 269, 270, 271, 272, 273, 274, 275, 276, 277, 278, 279, 280, 281, 282, 283, 284, 285, 286, 287, 288, 289, 290, 291, 292, 293, 294, 295, 296, 297, 298, 299, 300, 301, 302, 303, 304, 305, 306, 307, 308, 309, 310, 311, 312, 313, 314, 315, 316, 317, 318, 319, 320, 321, 322, 323, 324, 325, 326, 327, 328, 329, 330, 331, 332, 333, 334, 335, 336, 337, 338, 339, 340, 341, 342, 343, 344, 345, 346, 347, 348, 349, 350, 351, 352, 353, 354, 355, 356, 357, 358, 359, 360, 361, 362, 363, 364, 365, 366, 367, 368, 369, 370, 371, 372, 373])"
      ]
     },
     "execution_count": 91,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pickle\n",
    "with open('SFR_Mistral_IMDB_3.pkl', 'rb') as f:\n",
    "    json_str_3 = pickle.load(f)\n",
    "\n",
    "json_str_3.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "id": "0e7acbbf-16d3-4b7c-b1c7-8a9e4ee3882a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys([208, 209, 210, 211, 212, 213, 214, 215, 216, 217, 218, 219, 220, 221, 222, 223, 224, 225, 226, 227, 228, 229, 230, 231, 232, 233, 234, 235, 236, 237, 238, 239, 240, 241, 242, 243, 244, 245, 246, 247, 248, 249, 250, 251, 252, 253, 254, 255, 256, 257, 258, 259, 260, 261, 262, 263, 264, 265, 266])"
      ]
     },
     "execution_count": 92,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pickle\n",
    "with open('SFR_Mistral_IMDB_4.pkl', 'rb') as f:\n",
    "    json_str_4 = pickle.load(f)\n",
    "\n",
    "json_str_4.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "id": "b7573555-41bd-4bcf-8d40-cc0bc7714438",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys([180, 181, 182, 183, 184, 185, 186, 187, 188, 189, 190, 191, 192, 193, 194, 195, 196, 197, 198, 199, 200, 201, 202, 203, 204, 205, 206, 207])"
      ]
     },
     "execution_count": 93,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pickle\n",
    "with open('SFR_Mistral_IMDB_5.pkl', 'rb') as f:\n",
    "    json_str_5 = pickle.load(f)\n",
    "\n",
    "json_str_5.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "id": "33b23ed5-e87f-4452-b17d-10ffac2990e6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys([110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 164, 165])"
      ]
     },
     "execution_count": 94,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pickle\n",
    "with open('SFR_Mistral_IMDB_6.pkl', 'rb') as f:\n",
    "    json_str_6 = pickle.load(f)\n",
    "\n",
    "json_str_6.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "id": "2bcc542c-5a1e-4113-acb9-9add9c7b417b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys([54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110])"
      ]
     },
     "execution_count": 95,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pickle\n",
    "with open('SFR_Mistral_IMDB_7.pkl', 'rb') as f:\n",
    "    json_str_7 = pickle.load(f)\n",
    "\n",
    "json_str_7.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "id": "f18361f2-a7be-4632-a15c-37f72db2f1bc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys([27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39])"
      ]
     },
     "execution_count": 96,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pickle\n",
    "with open('SFR_Mistral_IMDB_8.pkl', 'rb') as f:\n",
    "    json_str_8 = pickle.load(f)\n",
    "\n",
    "json_str_8.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "id": "52cdea34-7f51-410e-8712-d15cd5b526f8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17])"
      ]
     },
     "execution_count": 97,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pickle\n",
    "with open('SFR_Mistral_IMDB_9.pkl', 'rb') as f:\n",
    "    json_str_9 = pickle.load(f)\n",
    "\n",
    "json_str_9.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "id": "6f95b69b-5c4a-4956-9571-0011acb72765",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys([18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53])"
      ]
     },
     "execution_count": 98,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pickle\n",
    "with open('SFR_Mistral_IMDB_10.pkl', 'rb') as f:\n",
    "    json_str_10 = pickle.load(f)\n",
    "\n",
    "json_str_10.keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8bd487e9",
   "metadata": {},
   "source": [
    "# Count perturbation performance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "id": "b9afd7a4-01d7-4464-8fd7-9e497d67cea5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0, 0.2247191011235955, 0.2247191011235955)"
      ]
     },
     "execution_count": 111,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hidden_states_pos = []\n",
    "layer = -1\n",
    "counter_success = 0\n",
    "counter = 0\n",
    "total = 0\n",
    "for json_str in [json_str_4, json_str_5, json_str_6, json_str_7, json_str_8, json_str_9, json_str_10]:\n",
    "    for index in json_str:\n",
    "        if 'is_success' not in json_str[index].keys():\n",
    "            continue\n",
    "        if json_str[index]['is_success'] > 0:\n",
    "            counter += 1\n",
    "            if json_str[index]['is_success'] > 1:\n",
    "                counter_success += 1\n",
    "            prompt_text = json_str[index][\"prompt_text\"]\n",
    "            token_ids = json_str[index][\"token_ids\"]\n",
    "            \n",
    "        total += 1\n",
    "len(hidden_states_pos), counter / total, counter_success / total"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dde02e8c",
   "metadata": {},
   "source": [
    "# Detect the GGPP prefix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "id": "45bbf2ee-b0ab-4197-9bc9-5f5e8f8dde47",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "100"
      ]
     },
     "execution_count": 99,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "json_str_filter = {}\n",
    "j = 0\n",
    "for json_str in [json_str_1, json_str_2, json_str_3, json_str_4, json_str_5, json_str_6, json_str_7, json_str_8, json_str_9, json_str_10]:\n",
    "    for i in json_str:\n",
    "        if j >= 100:\n",
    "            break\n",
    "        if json_str[i]['is_success'] > 1:\n",
    "            json_str_filter[i] = json_str[i]\n",
    "            j += 1\n",
    "j"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 134,
   "id": "d21c9673-d7c6-4113-bc49-dbb9790f9095",
   "metadata": {},
   "outputs": [],
   "source": [
    "import csv\n",
    "from struct import pack, unpack\n",
    "import lzma\n",
    "# import blosc\n",
    "import numpy as np\n",
    "import pickle\n",
    "import torch\n",
    "\n",
    "import csv\n",
    "\n",
    "tconst_title = {}\n",
    "count = 0\n",
    "with open(\"/title.basics.tsv\") as fd:\n",
    "    rd = csv.reader(fd, delimiter=\"\\t\", quotechar='\"')\n",
    "    for row in rd:\n",
    "        if count > 0:\n",
    "            \n",
    "            tconst_title[row[0]] = row[2]\n",
    "            \n",
    "        count += 1\n",
    "\n",
    "device = torch.device(\"cuda\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 135,
   "id": "9b697c47-f2f0-49e9-9edc-e417a8359a05",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompts = []\n",
    "name_tt = {}\n",
    "nconst_name = []\n",
    "nconst_profession = []\n",
    "nconst_embedding = []\n",
    "\n",
    "with open(\"/datasets/name.basics.tsv\") as fd:\n",
    "    rd = csv.reader(fd, delimiter=\"\\t\", quotechar='\"')\n",
    "    counter = 0\n",
    "    \n",
    "    for row in rd:\n",
    "        if counter > 1000:           \n",
    "            break\n",
    "        if counter > 0:            \n",
    "            prompt = row[1] +\" was born in \"+row[2]+\", \"+\"and died in \"+row[3]+\". He/She's primary professions are \"+', '.join(map(str, row[4].split(\",\")))+\".\"\n",
    "            \n",
    "            tts = row[5].split(\",\")\n",
    "            prompt += \" He/She is known for movies:\"\n",
    "            for t in tts:\n",
    "                if t == tts[-1]:\n",
    "                    prompt += \" '\" + tconst_title[t]+\"'.\"\n",
    "                else:\n",
    "                    prompt += \" '\" + tconst_title[t]+\"',\"\n",
    "\n",
    "            max_length = 526\n",
    "            input_texts = prompt\n",
    "            batch_dict = tok(input_texts, max_length=max_length, padding=True, truncation=True, return_tensors=\"pt\").to(device)\n",
    "            with torch.no_grad():\n",
    "                outputs = model(**batch_dict)\n",
    "                embedding = last_token_pool(outputs.last_hidden_state, batch_dict['attention_mask']).to(\"cpu\")\n",
    "        \n",
    "            nconst_embedding.append(embedding[0])\n",
    "            prompts.append(prompt)\n",
    "\n",
    "        counter += 1\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 136,
   "id": "499954c1-4dc1-4478-8211-3d608929b422",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "task_query:  According to IMDB dataset and your knowledges, what movies Terrence Malick has worked on and what were her/his roles?\n",
      "top_k:  [516 653 413 765 415  75 264 646 179 973]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Michael Mann has worked on and what were her/his roles?\n",
      "top_k:  [519 340 435 879 839 322 139 608 775 298]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Colm Meaney has worked on and what were her/his roles?\n",
      "top_k:  [537 558 341 637 392 538 982 295 545 535]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Roger Moore has worked on and what were her/his roles?\n",
      "top_k:  [548 210 124 725 616 960 323  56 192 934]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Judd Nelson has worked on and what were her/his roles?\n",
      "top_k:  [554 686 545 207  22 130 882 638 664 362]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Leonard Nimoy has worked on and what were her/his roles?\n",
      "top_k:  [558 637 852 652 994 493 478 909 982 827]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Alan Parker has worked on and what were her/his roles?\n",
      "top_k:  [569 646  82 613 867 264 204 272 735 942]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Jason Patric has worked on and what were her/his roles?\n",
      "top_k:  [573 890 663 918 463 686 545 615  98 865]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Ron Perlman has worked on and what were her/his roles?\n",
      "top_k:  [578 164 464 441 608 545 579 605 416 613]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Joseph L. has worked on and what were her/his roles?\n",
      "top_k:  [580 724 427 983 389 252 790  24 181 712]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Joe Pesci has worked on and what were her/his roles?\n",
      "top_k:  [581 500 731 724 648 133 831 608 216 985]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Robert Picardo has worked on and what were her/his roles?\n",
      "top_k:  [584 535 295 532 493 558 615 835 637 852]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Dennis Quaid has worked on and what were her/his roles?\n",
      "top_k:  [597 453 664 134 598 343 362 611 199 437]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Robert Redford has worked on and what were her/his roles?\n",
      "top_k:  [601 379  52 133 374 264 997 615 835 735]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies John C. has worked on and what were her/his roles?\n",
      "top_k:  [130 421 603 456 625 289 690 454 517 134]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Burt Reynolds has worked on and what were her/his roles?\n",
      "top_k:  [607  43 819   8  35  11 307 944 985 979]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Giovanni Ribisi has worked on and what were her/his roles?\n",
      "top_k:  [609 545 362  27 153 542 664 903 202 613]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Peter Greenaway has worked on and what were her/his roles?\n",
      "top_k:  [424 951 633 563 179 263 942   4 973 560]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Steve Guttenberg has worked on and what were her/his roles?\n",
      "top_k:  [429 711 536 187 113 413 664  70 835 210]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Daryl Hannah has worked on and what were her/his roles?\n",
      "top_k:  [434 355 706 105 664 896 538 343 674 356]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Dan Hedaya has worked on and what were her/his roles?\n",
      "top_k:  [444 731 361 836 545 762 357 686 718 664]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies William Hurt has worked on and what were her/his roles?\n",
      "top_k:  [457 456  33 512 286 244 199 635 637 902]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Raul Julia has worked on and what were her/his roles?\n",
      "top_k:  [470  27 209 490 743 966  64 895 872  78]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Michael Keaton has worked on and what were her/his roles?\n",
      "top_k:  [473 322 139 149 519 545 698 244 460 298]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Fritz Lang has worked on and what were her/his roles?\n",
      "top_k:  [484  18  16   4  76  75 973 427 252 539]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Brandon Lee has worked on and what were her/his roles?\n",
      "top_k:  [487  44   7 333 968 882 498 118 451 387]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Christopher Lloyd has worked on and what were her/his roles?\n",
      "top_k:  [501 488 976 685 362 393 416 149 558 356]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Julia Louis-Dreyfus has worked on and what were her/his roles?\n",
      "top_k:  [505 209 193 922 687 356 724 688 210 105]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Stephen Baldwin has worked on and what were her/his roles?\n",
      "top_k:  [285 836 286 283 961 284 664 835 545 409]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Michael Biehn has worked on and what were her/his roles?\n",
      "top_k:  [298 460 933 545 664 199 538 283 996 291]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Ernest Borgnine has worked on and what were her/his roles?\n",
      "top_k:  [307  43 902 615 387 863 840 985 944   7]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Charles Bronson has worked on and what were her/his roles?\n",
      "top_k:  [313   7 335  44 987  31 487 121 985 536]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Barbara Carrera has worked on and what were her/his roles?\n",
      "top_k:  [323 818 827 960 118 893 934 936 265 210]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Jackie Chan has worked on and what were her/his roles?\n",
      "top_k:  [328 333  44 196 240 301 705 498 118  83]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Francis Ford has worked on and what were her/his roles?\n",
      "top_k:  [337 530 405  75 147 183 767 413 484 401]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Roger Corman has worked on and what were her/his roles?\n",
      "top_k:  [338 247 210 342 811 790 340 735 539 548]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Olivier De has worked on and what were her/his roles?\n",
      "top_k:  [358 846 800  13  58 662 321 960 732 564]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Brian De has worked on and what were her/his roles?\n",
      "top_k:  [360 734 730 901 973 706 305 943 826 282]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Gérard Depardieu has worked on and what were her/his roles?\n",
      "top_k:  [366  90  85 605 732 899 808 803 321 271]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Walt Disney has worked on and what were her/his roles?\n",
      "top_k:  [369 427 864 252 873 121 972 833  82 704]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Martin Scorsese has worked on and what were her/his roles?\n",
      "top_k:  [216 974 639 133 264 630 840 137  79 839]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Charlie Sheen has worked on and what were her/his roles?\n",
      "top_k:  [220 639 121  31 545 664 105 792 706 192]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Will Smith has worked on and what were her/his roles?\n",
      "top_k:  [225 244 585 245 644 457 199  92 353 635]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Sylvester Stallone has worked on and what were her/his roles?\n",
      "top_k:  [229 355 615 664 663 481 985 633 184 231]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Quentin Tarantino has worked on and what were her/his roles?\n",
      "top_k:  [232 811 171 608 618  40 662 234 513 491]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Jean-Claude Van has worked on and what were her/his roles?\n",
      "top_k:  [240 605 850 899 766 794 615 248 892 418]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Mark Wahlberg has worked on and what were her/his roles?\n",
      "top_k:  [241 353 433 686 545 189 608 500 188 244]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Bruce Willis has worked on and what were her/his roles?\n",
      "top_k:  [245  44 309 615  92 244 608 985 225 199]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies John Woo has worked on and what were her/his roles?\n",
      "top_k:  [246 823 289 625 690 956 454 333  77 405]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Edward D. has worked on and what were her/his roles?\n",
      "top_k:  [247 437  63 589 798 252 427 891 380 714]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Robert Altman has worked on and what were her/his roles?\n",
      "top_k:  [264 735 973 252 835 912 708 251  52 796]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Heather Locklear has worked on and what were her/his roles?\n",
      "top_k:  [180 161 503 524 616 428 896 210 356 343]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies George Lucas has worked on and what were her/his roles?\n",
      "top_k:  [183 107 433 354 566  40  26 122 401 785]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Ewan McGregor has worked on and what were her/his roles?\n",
      "top_k:  [190 703 189 159  26 536 433 124 664 635]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Julianne Moore has worked on and what were her/his roles?\n",
      "top_k:  [193 209 522 495 192 299 752 123 655 475]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Jack Nicholson has worked on and what were her/his roles?\n",
      "top_k:  [196 492 559 615 545 301 114 686 883 244]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Ryan Phillippe has worked on and what were her/his roles?\n",
      "top_k:  [201 211 449 202 706 545 575 359 608 441]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Keanu Reeves has worked on and what were her/his roles?\n",
      "top_k:  [205  35 663 125 545 333 118 706 441 189]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Pierce Brosnan has worked on and what were her/his roles?\n",
      "top_k:  [111 124 548 847 979 301  43 764 936 867]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Wes Craven has worked on and what were her/his roles?\n",
      "top_k:  [126 435 646 247 342 116 790 248 273 340]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Robert De has worked on and what were her/his roles?\n",
      "top_k:  [133 374 379 615  52 264 493 708 252 359]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Johnny Depp has worked on and what were her/his roles?\n",
      "top_k:  [135 361 690 603 196 456 545 466 133 517]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Cameron Diaz has worked on and what were her/his roles?\n",
      "top_k:  [138 105 297 118 210 538 177 115 472 476]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Michael Douglas has worked on and what were her/his roles?\n",
      "top_k:  [139  17 473 244 519 460 322 340 210 608]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Clint Eastwood has worked on and what were her/his roles?\n",
      "top_k:  [141  21 316  14  10 985  70 379 335   7]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Erika Eleniak has worked on and what were her/his roles?\n",
      "top_k:  [142 941 574 108 922 222 282 616 343 118]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Morgan Freeman has worked on and what were her/his roles?\n",
      "top_k:  [150 608 244   7 558 468 545 391 417 985]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Ron Howard has worked on and what were her/his roles?\n",
      "top_k:  [164 464 435 711 815 208 708 796 942 749]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Alfred Newman has worked on and what were her/his roles?\n",
      "top_k:  [ 54 928  24  76  81  69  34  64 709  66]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Laurence Olivier has worked on and what were her/his roles?\n",
      "top_k:  [ 58  13 563   1 179  45 846  50  23  79]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Edward G. has worked on and what were her/his roles?\n",
      "top_k:  [247  63 589 437 252 682 798 380 427 891]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies John Wayne has worked on and what were her/his roles?\n",
      "top_k:  [ 77 405 953 972 729 134  10 289  67 421]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Natalie Wood has worked on and what were her/his roles?\n",
      "top_k:  [ 80 203  37 363  29 247  27 764 248   5]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Alan Miller has worked on and what were her/his roles?\n",
      "top_k:  [ 82 646 569 256 613 867 272  41 264 739]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Henner Hofmann has worked on and what were her/his roles?\n",
      "top_k:  [ 84 886 727  16 355 450 441 321 449 706]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Louis de has worked on and what were her/his roles?\n",
      "top_k:  [ 85 314 732 319 321 365 662 362 570 960]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Aleksey Korenev has worked on and what were her/his roles?\n",
      "top_k:  [ 87  86  26  40 646 262 284 480 476 706]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies John Cleese has worked on and what were her/his roles?\n",
      "top_k:  [ 91   3 134  23 456 454 633 421 130 690]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies J. Reifel has worked on and what were her/his roles?\n",
      "top_k:  [ 93 359 701 210 226 622  37 252 491  16]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Pamela Anderson has worked on and what were her/his roles?\n",
      "top_k:  [ 96 898 426 574 644 210 941 108 757 616]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Dan Aykroyd has worked on and what were her/his roles?\n",
      "top_k:  [100   3 900 731 444 119 393 557 244 195]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Kim Basinger has worked on and what were her/his roles?\n",
      "top_k:  [106   1 498  11 933 896 868 674 210 677]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Matthew Broderick has worked on and what were her/his roles?\n",
      "top_k:  [110 545 189 353 221 497 368 981 244  29]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Brigitte Bardot has worked on and what were her/his roles?\n",
      "top_k:  [  2 973 943 321 732 481 301  53   1 760]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Bette Davis has worked on and what were her/his roles?\n",
      "top_k:  [ 11 877 540   1 841  71 868  16 933  43]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Doris Day has worked on and what were her/his roles?\n",
      "top_k:  [ 12  37  16 363  29  22  38 915  27  70]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Olivia de has worked on and what were her/his roles?\n",
      "top_k:  [ 13 846  58 555 662 356 564 358 363  38]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies James Dean has worked on and what were her/his roles?\n",
      "top_k:  [ 14  70 248 542 362 363 985  50 789 999]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Kirk Douglas has worked on and what were her/his roles?\n",
      "top_k:  [ 17  74 475 139  38  43  70  37  21  62]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Federico Fellini has worked on and what were her/his roles?\n",
      "top_k:  [ 18  51 484 932   4  64 970 319 732 876]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies John Gielgud has worked on and what were her/his roles?\n",
      "top_k:  [ 23 456 856  58 682  26 867  50 405   8]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Jerry Goldsmith has worked on and what were her/his roles?\n",
      "top_k:  [ 24 928 289  34  54  15  69  76  64  81]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Rita Hayworth has worked on and what were her/his roles?\n",
      "top_k:  [ 27   5 543  37  21 877 363  42  16  53]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Margaux Hemingway has worked on and what were her/his roles?\n",
      "top_k:  [ 28 445  29 481  16  53  37 210 428  27]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Alfred Hitchcock has worked on and what were her/his roles?\n",
      "top_k:  [ 32  37  70 864   5  29 646 735   6  54]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Deborah Kerr has worked on and what were her/his roles?\n",
      "top_k:  [ 38  37  30  29 678 217 179 363  70  71]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Sophia Loren has worked on and what were her/his roles?\n",
      "top_k:  [ 46  51 772 732  18  53 422 764   1  16]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Henry Mancini has worked on and what were her/his roles?\n",
      "top_k:  [ 48 928  64  54  81 289 709  24  34  76]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies Groucho Marx has worked on and what were her/his roles?\n",
      "top_k:  [ 49  53  27  21  16 121 543 682  51 210]\n",
      "task_query:  According to IMDB dataset and your knowledges, what movies James Mason has worked on and what were her/his roles?\n",
      "top_k:  [ 50  70  14 335 248 341   8 867 468 985]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(100, 0)"
      ]
     },
     "execution_count": 136,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from DGD.help_functions import get_embedding, perturb_sentence, rank_tokens_by_importance, get_initial_ids, user_prompt_generation, check_success_origin, check_MSEloss, get_first_output_token, get_neighbor_ids, check_success_target_origin, get_relevant_documents\n",
    "from DGD.DGD import get_optimized_prefix_embedding\n",
    "import random\n",
    "\n",
    "# Set dataset\n",
    "dataset = \"imdb\"\n",
    "query_mode = \"for_objects\"\n",
    "\n",
    "counter = 0\n",
    "\n",
    "hidden_states_positive = []\n",
    "layer = -1\n",
    "origin_token_list = {}\n",
    "for index in json_str_filter:\n",
    "\n",
    "    user_prompt = json_str_filter[index][\"user_prompt\"]\n",
    "    initial_firstk = 10\n",
    "    initial_ids = [random.randint(0, 32000) for _ in range(initial_firstk)]\n",
    "    user_prompt_id = index\n",
    "    target_prompt_id = index\n",
    "    \n",
    "    topk_result = 10\n",
    "    top_k = get_relevant_documents(torch.stack(nconst_embedding), user_prompt, topk_result, model, tok)\n",
    "    is_success = check_success_origin(target_prompt_id, user_prompt_id, top_k)\n",
    "    if is_success <= 0:\n",
    "        counter += 1\n",
    "        \n",
    "    origin_token_list[index] = initial_ids\n",
    "len(origin_token_list.keys()), counter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 137,
   "id": "29de570d-5bf5-4b61-a917-f41e1460b74e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 137,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "counter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 142,
   "id": "a0a083cd-a39e-4635-ac15-9ec9ee6b0e25",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "According to IMDB dataset and your knowledges, what movies Terrence Malick has worked on and what were her/his roles?\n",
      "[21511, 393, 1375, 403, 5381, 1401, 25346, 14800, 13127, 3999]\n",
      "According to IMDB dataset and your knowledges, what movies Michael Mann has worked on and what were her/his roles?\n",
      "[17332, 24158, 403, 7043, 3616, 12567, 28740, 21309, 24978, 28784]\n",
      "According to IMDB dataset and your knowledges, what movies Colm Meaney has worked on and what were her/his roles?\n",
      "[20991, 27542, 263, 4285, 17300, 4764, 6283, 18526, 31739, 5202]\n",
      "According to IMDB dataset and your knowledges, what movies Roger Moore has worked on and what were her/his roles?\n",
      "[21511, 1739, 973, 403, 5381, 1259, 7826, 6926, 20215, 23262]\n",
      "According to IMDB dataset and your knowledges, what movies Judd Nelson has worked on and what were her/his roles?\n",
      "[393, 15536, 4262, 1896, 16889, 403, 2236, 27191, 24694, 22193]\n",
      "According to IMDB dataset and your knowledges, what movies Leonard Nimoy has worked on and what were her/his roles?\n",
      "[3217, 1286, 11395, 20550, 28710, 403, 4583, 3931, 3406, 18081]\n",
      "According to IMDB dataset and your knowledges, what movies Alan Parker has worked on and what were her/his roles?\n",
      "[1337, 28708, 19560, 12820, 5381, 456, 27574, 11187, 27486, 27671]\n",
      "According to IMDB dataset and your knowledges, what movies Jason Patric has worked on and what were her/his roles?\n",
      "[18486, 14603, 15418, 349, 21664, 22013, 27932, 12871, 19014, 3169]\n",
      "According to IMDB dataset and your knowledges, what movies Ron Perlman has worked on and what were her/his roles?\n",
      "[1295, 375, 2238, 7393, 3589, 31958, 4449, 8070, 5985, 18970]\n",
      "According to IMDB dataset and your knowledges, what movies Joseph L. has worked on and what were her/his roles?\n",
      "[19512, 20008, 403, 5381, 297, 28705, 28740, 28774, 28734, 28770]\n",
      "According to IMDB dataset and your knowledges, what movies Joe Pesci has worked on and what were her/his roles?\n",
      "[17821, 12998, 22923, 3156, 22913, 5230, 19551, 28705, 30793, 21922]\n",
      "According to IMDB dataset and your knowledges, what movies Robert Picardo has worked on and what were her/his roles?\n",
      "[15444, 19673, 403, 11997, 13918, 26729, 1279, 28598, 4060, 23432]\n",
      "According to IMDB dataset and your knowledges, what movies Dennis Quaid has worked on and what were her/his roles?\n",
      "[6081, 13002, 943, 282, 2047, 22917, 892, 1406, 812, 11030]\n",
      "According to IMDB dataset and your knowledges, what movies Robert Redford has worked on and what were her/his roles?\n",
      "[524, 777, 597, 11578, 303, 3213, 5381, 13682, 27253, 28723]\n",
      "According to IMDB dataset and your knowledges, what movies John C. has worked on and what were her/his roles?\n",
      "[6886, 330, 646, 403, 5381, 13434, 28705, 19176, 28560, 10292]\n",
      "According to IMDB dataset and your knowledges, what movies Burt Reynolds has worked on and what were her/his roles?\n",
      "[16745, 301, 3011, 12677, 403, 20395, 59, 14800, 5488, 9268]\n",
      "According to IMDB dataset and your knowledges, what movies Giovanni Ribisi has worked on and what were her/his roles?\n",
      "[7647, 13946, 3690, 820, 458, 403, 10314, 20126, 9077, 4307]\n",
      "According to IMDB dataset and your knowledges, what movies Peter Greenaway has worked on and what were her/his roles?\n",
      "[6359, 7564, 28716, 444, 403, 3690, 24908, 24006, 27093, 30195]\n",
      "According to IMDB dataset and your knowledges, what movies Steve Guttenberg has worked on and what were her/his roles?\n",
      "[13517, 28194, 15713, 403, 26182, 24694, 4700, 13482, 8958, 22281]\n",
      "According to IMDB dataset and your knowledges, what movies Daryl Hannah has worked on and what were her/his roles?\n",
      "[23373, 11395, 27313, 403, 18486, 5399, 3406, 22893, 23633, 2600]\n",
      "According to IMDB dataset and your knowledges, what movies Dan Hedaya has worked on and what were her/his roles?\n",
      "[19848, 361, 14813, 1294, 403, 16897, 31197, 985, 17420, 22266]\n",
      "According to IMDB dataset and your knowledges, what movies William Hurt has worked on and what were her/his roles?\n",
      "[12177, 9465, 403, 11170, 297, 12938, 3859, 22414, 1099, 30664]\n",
      "According to IMDB dataset and your knowledges, what movies Raul Julia has worked on and what were her/his roles?\n",
      "[976, 2350, 1010, 2162, 465, 403, 29813, 17156, 29678, 19010]\n",
      "According to IMDB dataset and your knowledges, what movies Michael Keaton has worked on and what were her/his roles?\n",
      "[10562, 271, 912, 524, 694, 349, 26364, 20728, 6005, 22279]\n",
      "According to IMDB dataset and your knowledges, what movies Fritz Lang has worked on and what were her/his roles?\n",
      "[6522, 2025, 834, 1096, 20564, 20541, 13975, 28173, 28536, 25999]\n",
      "According to IMDB dataset and your knowledges, what movies Brandon Lee has worked on and what were her/his roles?\n",
      "[28396, 418, 5067, 403, 28400, 23301, 3850, 11173, 6976, 8890]\n",
      "According to IMDB dataset and your knowledges, what movies Christopher Lloyd has worked on and what were her/his roles?\n",
      "[5376, 1296, 3247, 13131, 403, 5381, 3010, 4902, 17491, 28932]\n",
      "According to IMDB dataset and your knowledges, what movies Julia Louis-Dreyfus has worked on and what were her/his roles?\n",
      "[4573, 2047, 7303, 19068, 17623, 25858, 28808, 13958, 24953, 27975]\n",
      "According to IMDB dataset and your knowledges, what movies Stephen Baldwin has worked on and what were her/his roles?\n",
      "[12042, 365, 834, 311, 632, 403, 6216, 6622, 3051, 15704]\n",
      "According to IMDB dataset and your knowledges, what movies Michael Biehn has worked on and what were her/his roles?\n",
      "[21511, 393, 1845, 28764, 403, 25987, 297, 11631, 21491, 18081]\n",
      "According to IMDB dataset and your knowledges, what movies Ernest Borgnine has worked on and what were her/his roles?\n",
      "[401, 8354, 13550, 403, 9134, 3217, 8545, 20416, 27657, 31685]\n",
      "According to IMDB dataset and your knowledges, what movies Charles Bronson has worked on and what were her/his roles?\n",
      "[13194, 17083, 508, 386, 403, 21259, 24025, 28705, 28740, 28774]\n",
      "According to IMDB dataset and your knowledges, what movies Barbara Carrera has worked on and what were her/his roles?\n",
      "[14115, 6505, 628, 403, 26926, 3406, 17864, 31739, 14115, 13842]\n",
      "According to IMDB dataset and your knowledges, what movies Jackie Chan has worked on and what were her/his roles?\n",
      "[18486, 14603, 15418, 403, 19683, 297, 14814, 3128, 13976, 21002]\n",
      "According to IMDB dataset and your knowledges, what movies Francis Ford has worked on and what were her/his roles?\n",
      "[16668, 2875, 23735, 403, 5381, 7179, 13934, 31484, 415, 23359]\n",
      "According to IMDB dataset and your knowledges, what movies Roger Corman has worked on and what were her/his roles?\n",
      "[10860, 17161, 411, 403, 10355, 297, 6180, 25581, 23672, 30195]\n",
      "According to IMDB dataset and your knowledges, what movies Olivier De has worked on and what were her/his roles?\n",
      "[5368, 4401, 1896, 321, 2047, 403, 5381, 3999, 28705, 13649]\n",
      "According to IMDB dataset and your knowledges, what movies Brian De has worked on and what were her/his roles?\n",
      "[365, 593, 28710, 24930, 1265, 403, 5381, 297, 12340, 12246]\n",
      "According to IMDB dataset and your knowledges, what movies Gérard Depardieu has worked on and what were her/his roles?\n",
      "[1744, 794, 7562, 6559, 4677, 28519, 2128, 11502, 7346, 28774]\n",
      "According to IMDB dataset and your knowledges, what movies Walt Disney has worked on and what were her/his roles?\n",
      "[13131, 1010, 364, 3186, 403, 7826, 19079, 31739, 19244, 4527]\n",
      "According to IMDB dataset and your knowledges, what movies Martin Scorsese has worked on and what were her/his roles?\n",
      "[7149, 384, 16849, 349, 25363, 22232, 10265, 5485, 23565, 13915]\n",
      "According to IMDB dataset and your knowledges, what movies Charlie Sheen has worked on and what were her/his roles?\n",
      "[318, 11849, 3405, 1742, 403, 27136, 21209, 28705, 16531, 30707]\n",
      "According to IMDB dataset and your knowledges, what movies Will Smith has worked on and what were her/his roles?\n",
      "[9101, 15878, 28710, 403, 10088, 28878, 6586, 13852, 3164, 4602]\n",
      "According to IMDB dataset and your knowledges, what movies Sylvester Stallone has worked on and what were her/his roles?\n",
      "[6715, 12807, 28726, 308, 403, 519, 23883, 19010, 9351, 4205]\n",
      "According to IMDB dataset and your knowledges, what movies Quentin Tarantino has worked on and what were her/his roles?\n",
      "[5480, 9734, 3520, 863, 21259, 297, 13918, 28740, 28774, 1047]\n",
      "According to IMDB dataset and your knowledges, what movies Jean-Claude Van has worked on and what were her/his roles?\n",
      "[18709, 1343, 28790, 3068, 403, 4660, 297, 21780, 16945, 3169]\n",
      "According to IMDB dataset and your knowledges, what movies Mark Wahlberg has worked on and what were her/his roles?\n",
      "[1471, 309, 28718, 4930, 1079, 403, 15975, 15797, 1439, 1715]\n",
      "According to IMDB dataset and your knowledges, what movies Bruce Willis has worked on and what were her/his roles?\n",
      "[3189, 16885, 384, 28742, 1534, 2224, 403, 5381, 3628, 9329]\n",
      "According to IMDB dataset and your knowledges, what movies John Woo has worked on and what were her/his roles?\n",
      "[4660, 2463, 1689, 10806, 28173, 4222, 10522, 26307, 26307, 22379]\n",
      "According to IMDB dataset and your knowledges, what movies Edward D. has worked on and what were her/his roles?\n",
      "[1295, 5214, 11775, 3396, 5381, 5158, 23345, 28740, 28774, 28782]\n",
      "According to IMDB dataset and your knowledges, what movies Robert Altman has worked on and what were her/his roles?\n",
      "[24704, 1375, 15661, 1528, 5381, 26463, 10983, 21936, 100, 14671]\n",
      "According to IMDB dataset and your knowledges, what movies Heather Locklear has worked on and what were her/his roles?\n",
      "[4797, 28092, 453, 403, 10749, 11340, 5488, 18297, 21041, 28750]\n",
      "According to IMDB dataset and your knowledges, what movies George Lucas has worked on and what were her/his roles?\n",
      "[17332, 365, 426, 403, 10749, 13044, 28705, 28740, 28774, 2675]\n",
      "According to IMDB dataset and your knowledges, what movies Ewan McGregor has worked on and what were her/his roles?\n",
      "[420, 566, 11469, 3049, 349, 25136, 23768, 22568, 11946, 8922]\n",
      "According to IMDB dataset and your knowledges, what movies Julianne Moore has worked on and what were her/his roles?\n",
      "[1366, 5403, 8364, 28971, 28092, 29918, 17309, 21969, 21364, 21149]\n",
      "According to IMDB dataset and your knowledges, what movies Jack Nicholson has worked on and what were her/his roles?\n",
      "[17821, 293, 26424, 3047, 17332, 10521, 28906, 403, 5381, 28161]\n",
      "According to IMDB dataset and your knowledges, what movies Ryan Phillippe has worked on and what were her/his roles?\n",
      "[6522, 1174, 15661, 28971, 21259, 24312, 28705, 28740, 19433, 10647]\n",
      "According to IMDB dataset and your knowledges, what movies Keanu Reeves has worked on and what were her/his roles?\n",
      "[5459, 475, 28723, 12407, 403, 21107, 3406, 31739, 15583, 22932]\n",
      "According to IMDB dataset and your knowledges, what movies Pierce Brosnan has worked on and what were her/his roles?\n",
      "[22706, 25533, 403, 4157, 6906, 24726, 24240, 30464, 8827, 20017]\n",
      "According to IMDB dataset and your knowledges, what movies Wes Craven has worked on and what were her/his roles?\n",
      "[4797, 20096, 9383, 19785, 28723, 422, 31811, 23236, 15161, 0]\n",
      "According to IMDB dataset and your knowledges, what movies Robert De has worked on and what were her/his roles?\n",
      "[8103, 24794, 5162, 18934, 5381, 26045, 25154, 28495, 19420, 15704]\n",
      "According to IMDB dataset and your knowledges, what movies Johnny Depp has worked on and what were her/his roles?\n",
      "[7240, 5776, 28726, 1126, 6304, 5381, 21072, 20541, 13695, 24068]\n",
      "According to IMDB dataset and your knowledges, what movies Cameron Diaz has worked on and what were her/his roles?\n",
      "[408, 469, 20452, 26171, 1294, 403, 5381, 11721, 28723, 22095]\n",
      "According to IMDB dataset and your knowledges, what movies Michael Douglas has worked on and what were her/his roles?\n",
      "[3474, 10171, 27085, 2164, 349, 29157, 30600, 30831, 25554, 28495]\n",
      "According to IMDB dataset and your knowledges, what movies Clint Eastwood has worked on and what were her/his roles?\n",
      "[5625, 299, 1739, 13160, 19215, 29053, 18023, 19969, 18540, 11607]\n",
      "According to IMDB dataset and your knowledges, what movies Erika Eleniak has worked on and what were her/his roles?\n",
      "[30620, 420, 755, 28721, 554, 403, 1573, 3861, 21259, 18529]\n",
      "According to IMDB dataset and your knowledges, what movies Morgan Freeman has worked on and what were her/his roles?\n",
      "[22387, 723, 6359, 1265, 18881, 21996, 15464, 10596, 10522, 12642]\n",
      "According to IMDB dataset and your knowledges, what movies Ron Howard has worked on and what were her/his roles?\n",
      "[28161, 5526, 263, 441, 5976, 25645, 22625, 28517, 24628, 31685]\n",
      "According to IMDB dataset and your knowledges, what movies Alfred Newman has worked on and what were her/his roles?\n",
      "[22706, 25533, 28763, 5381, 17976, 20893, 12781, 27266, 4112, 8908]\n",
      "According to IMDB dataset and your knowledges, what movies Laurence Olivier has worked on and what were her/his roles?\n",
      "[6808, 351, 834, 3494, 562, 4241, 297, 1478, 23998, 30664]\n",
      "According to IMDB dataset and your knowledges, what movies Edward G. has worked on and what were her/his roles?\n",
      "[2740, 28710, 15052, 7322, 403, 5381, 297, 28705, 28740, 28774]\n",
      "According to IMDB dataset and your knowledges, what movies John Wayne has worked on and what were her/his roles?\n",
      "[5376, 1296, 16235, 13131, 403, 5381, 4394, 21936, 20541, 92]\n",
      "According to IMDB dataset and your knowledges, what movies Natalie Wood has worked on and what were her/his roles?\n",
      "[1366, 495, 27542, 263, 403, 30664, 21419, 93, 20541, 9351]\n",
      "According to IMDB dataset and your knowledges, what movies Alan Miller has worked on and what were her/his roles?\n",
      "[10805, 9734, 1523, 14949, 5381, 297, 28894, 24196, 28774, 28782]\n",
      "According to IMDB dataset and your knowledges, what movies Henner Hofmann has worked on and what were her/his roles?\n",
      "[7274, 16589, 403, 5381, 3448, 28705, 16461, 4117, 28784, 28740]\n",
      "According to IMDB dataset and your knowledges, what movies Louis de has worked on and what were her/his roles?\n",
      "[6502, 19973, 403, 5381, 657, 22112, 11521, 24741, 7633, 7445]\n",
      "According to IMDB dataset and your knowledges, what movies Aleksey Korenev has worked on and what were her/his roles?\n",
      "[21388, 5112, 5655, 29924, 4985, 570, 403, 14648, 17127, 5488]\n",
      "According to IMDB dataset and your knowledges, what movies John Cleese has worked on and what were her/his roles?\n",
      "[4299, 24396, 403, 6606, 13026, 21259, 4747, 1434, 25721, 10522]\n",
      "According to IMDB dataset and your knowledges, what movies J. Reifel has worked on and what were her/his roles?\n",
      "[2215, 26419, 31925, 5381, 19567, 17512, 28740, 28774, 28782, 28734]\n",
      "According to IMDB dataset and your knowledges, what movies Pamela Anderson has worked on and what were her/his roles?\n",
      "[28396, 3838, 498, 992, 28711, 11909, 19010, 23142, 15108, 23607]\n",
      "According to IMDB dataset and your knowledges, what movies Dan Aykroyd has worked on and what were her/his roles?\n",
      "[7163, 3552, 420, 2654, 3554, 2590, 26858, 6407, 27193, 25533]\n",
      "According to IMDB dataset and your knowledges, what movies Kim Basinger has worked on and what were her/his roles?\n",
      "[12332, 401, 643, 403, 29417, 29854, 11377, 15376, 17426, 20541]\n",
      "According to IMDB dataset and your knowledges, what movies Matthew Broderick has worked on and what were her/his roles?\n",
      "[8999, 515, 18213, 403, 13958, 19252, 10685, 27616, 1302, 1302]\n",
      "According to IMDB dataset and your knowledges, what movies Brigitte Bardot has worked on and what were her/his roles?\n",
      "[7772, 1091, 30814, 0, 26171, 1294, 403, 27325, 3797, 29905]\n",
      "According to IMDB dataset and your knowledges, what movies Bette Davis has worked on and what were her/his roles?\n",
      "[26586, 9236, 2915, 595, 5057, 30273, 29655, 9952, 19010, 22675]\n",
      "According to IMDB dataset and your knowledges, what movies Doris Day has worked on and what were her/his roles?\n",
      "[14115, 6505, 628, 403, 21304, 29905, 15376, 12870, 3937, 27768]\n",
      "According to IMDB dataset and your knowledges, what movies Olivia de has worked on and what were her/his roles?\n",
      "[13195, 985, 269, 403, 5381, 8449, 20541, 22977, 20541, 9268]\n",
      "According to IMDB dataset and your knowledges, what movies James Dean has worked on and what were her/his roles?\n",
      "[10785, 1136, 1079, 6304, 15451, 25443, 18881, 8149, 347, 31387]\n",
      "According to IMDB dataset and your knowledges, what movies Kirk Douglas has worked on and what were her/his roles?\n",
      "[10117, 5311, 6448, 26984, 5381, 23998, 24033, 22799, 28750, 15778]\n",
      "According to IMDB dataset and your knowledges, what movies Federico Fellini has worked on and what were her/his roles?\n",
      "[10117, 5311, 6448, 8166, 5381, 17569, 27616, 0, 19613, 17054]\n",
      "According to IMDB dataset and your knowledges, what movies John Gielgud has worked on and what were her/his roles?\n",
      "[627, 7804, 31150, 365, 291, 761, 28971, 3628, 19314, 29905]\n",
      "According to IMDB dataset and your knowledges, what movies Jerry Goldsmith has worked on and what were her/his roles?\n",
      "[351, 5675, 17079, 403, 5381, 3219, 19314, 13334, 10200, 26203]\n",
      "According to IMDB dataset and your knowledges, what movies Rita Hayworth has worked on and what were her/his roles?\n",
      "[2215, 334, 381, 468, 349, 24971, 11834, 9952, 20541, 24746]\n",
      "According to IMDB dataset and your knowledges, what movies Margaux Hemingway has worked on and what were her/his roles?\n",
      "[1744, 794, 7562, 6559, 8263, 17632, 20761, 8307, 10040, 1393]\n",
      "According to IMDB dataset and your knowledges, what movies Alfred Hitchcock has worked on and what were her/his roles?\n",
      "[4797, 8221, 1079, 3867, 2352, 11721, 24472, 24825, 26703, 31739]\n",
      "According to IMDB dataset and your knowledges, what movies Deborah Kerr has worked on and what were her/his roles?\n",
      "[23540, 6815, 28712, 7593, 14877, 29220, 13958, 30195, 9392, 9646]\n",
      "According to IMDB dataset and your knowledges, what movies Sophia Loren has worked on and what were her/his roles?\n",
      "[23373, 23515, 811, 28710, 10744, 340, 22650, 10522, 21312, 6122]\n",
      "According to IMDB dataset and your knowledges, what movies Henry Mancini has worked on and what were her/his roles?\n",
      "[15425, 24150, 1464, 28574, 22370, 61, 30443, 28675, 21507, 13009]\n",
      "According to IMDB dataset and your knowledges, what movies Groucho Marx has worked on and what were her/his roles?\n",
      "[5480, 30195, 24618, 864, 19673, 24062, 5381, 4624, 17328, 4714]\n",
      "According to IMDB dataset and your knowledges, what movies James Mason has worked on and what were her/his roles?\n",
      "[7240, 14195, 21208, 20541, 30255, 28705, 28782, 12919, 28787, 7445]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "100"
      ]
     },
     "execution_count": 142,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from DGD.help_functions import get_embedding, perturb_sentence, rank_tokens_by_importance, get_initial_ids, user_prompt_generation, check_success, check_MSEloss, get_first_output_token, get_neighbor_ids\n",
    "from DGD.DGD import get_optimized_prefix_embedding\n",
    "\n",
    "# Set dataset\n",
    "dataset = \"imdb\"\n",
    "query_mode = \"for_objects\"\n",
    "\n",
    "hidden_states_negative = []\n",
    "layer = -1\n",
    "for index in json_str_filter:\n",
    "    prompt_text = json_str_filter[index][\"prompt_text\"]\n",
    "    token_ids = json_str_filter[index][\"token_ids\"]\n",
    "    object = \"\"\n",
    "    sentence = user_prompt_generation(prompt_text, object, dataset, query_mode)\n",
    "    print(sentence)\n",
    "    print(token_ids)\n",
    "    hidden_state = get_model_layer_hidden_states(sentence, token_ids, model, tok, layer)\n",
    "    hidden_states_negative.append(hidden_state.to(\"cpu\"))\n",
    "len(hidden_states_negative)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 143,
   "id": "2f3c1ac3-1fd5-44ea-899f-8d1d129ed54d",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "According to IMDB dataset and your knowledges, what movies Terrence Malick has worked on and what were her/his roles?\n",
      "[9376, 4889, 26013, 28444, 8040, 26770, 24555, 19725, 8141, 20266]\n",
      "According to IMDB dataset and your knowledges, what movies Michael Mann has worked on and what were her/his roles?\n",
      "[28164, 12506, 2075, 9593, 21554, 24638, 1185, 15282, 20487, 31520]\n",
      "According to IMDB dataset and your knowledges, what movies Colm Meaney has worked on and what were her/his roles?\n",
      "[10151, 4283, 6794, 17052, 15596, 8603, 2362, 28164, 21659, 24509]\n",
      "According to IMDB dataset and your knowledges, what movies Roger Moore has worked on and what were her/his roles?\n",
      "[8751, 12707, 25467, 8239, 268, 6266, 5262, 2995, 3180, 19574]\n",
      "According to IMDB dataset and your knowledges, what movies Judd Nelson has worked on and what were her/his roles?\n",
      "[18029, 31205, 2527, 7836, 8680, 17662, 24113, 9950, 27361, 22216]\n",
      "According to IMDB dataset and your knowledges, what movies Leonard Nimoy has worked on and what were her/his roles?\n",
      "[17568, 23701, 13546, 28814, 12857, 18510, 28447, 30308, 23434, 244]\n",
      "According to IMDB dataset and your knowledges, what movies Alan Parker has worked on and what were her/his roles?\n",
      "[6281, 25560, 14900, 2585, 27652, 15231, 10377, 2264, 13384, 22072]\n",
      "According to IMDB dataset and your knowledges, what movies Jason Patric has worked on and what were her/his roles?\n",
      "[11077, 10330, 9297, 13103, 10492, 30562, 29226, 14299, 1378, 13008]\n",
      "According to IMDB dataset and your knowledges, what movies Ron Perlman has worked on and what were her/his roles?\n",
      "[27947, 9146, 7002, 1998, 14051, 3897, 1246, 23709, 1247, 16229]\n",
      "According to IMDB dataset and your knowledges, what movies Joseph L. has worked on and what were her/his roles?\n",
      "[6963, 422, 10009, 28753, 28516, 5163, 7430, 7461, 20532, 14766]\n",
      "According to IMDB dataset and your knowledges, what movies Joe Pesci has worked on and what were her/his roles?\n",
      "[3574, 27839, 13374, 21369, 16830, 6825, 27395, 9600, 26721, 12169]\n",
      "According to IMDB dataset and your knowledges, what movies Robert Picardo has worked on and what were her/his roles?\n",
      "[14272, 6398, 27465, 25347, 11597, 8093, 8893, 16018, 5608, 1779]\n",
      "According to IMDB dataset and your knowledges, what movies Dennis Quaid has worked on and what were her/his roles?\n",
      "[21156, 8848, 12680, 8746, 2673, 23601, 18297, 29066, 27356, 22766]\n",
      "According to IMDB dataset and your knowledges, what movies Robert Redford has worked on and what were her/his roles?\n",
      "[23290, 18052, 5249, 24622, 14192, 20966, 17616, 1138, 20736, 18330]\n",
      "According to IMDB dataset and your knowledges, what movies John C. has worked on and what were her/his roles?\n",
      "[26010, 2208, 15717, 5805, 15745, 23350, 4490, 20558, 30607, 278]\n",
      "According to IMDB dataset and your knowledges, what movies Burt Reynolds has worked on and what were her/his roles?\n",
      "[781, 3783, 18919, 1224, 25280, 22843, 25596, 226, 4270, 31287]\n",
      "According to IMDB dataset and your knowledges, what movies Giovanni Ribisi has worked on and what were her/his roles?\n",
      "[20771, 31862, 27698, 31689, 23438, 19861, 19873, 3654, 25269, 26168]\n",
      "According to IMDB dataset and your knowledges, what movies Peter Greenaway has worked on and what were her/his roles?\n",
      "[3395, 2776, 30610, 18077, 19308, 8312, 26003, 31117, 26673, 8277]\n",
      "According to IMDB dataset and your knowledges, what movies Steve Guttenberg has worked on and what were her/his roles?\n",
      "[12296, 20916, 27953, 18533, 5791, 12436, 28374, 672, 1726, 31745]\n",
      "According to IMDB dataset and your knowledges, what movies Daryl Hannah has worked on and what were her/his roles?\n",
      "[15846, 3201, 1078, 11555, 12717, 18324, 30549, 28689, 21585, 815]\n",
      "According to IMDB dataset and your knowledges, what movies Dan Hedaya has worked on and what were her/his roles?\n",
      "[7174, 5355, 26907, 8791, 19790, 11442, 10601, 14769, 3631, 17209]\n",
      "According to IMDB dataset and your knowledges, what movies William Hurt has worked on and what were her/his roles?\n",
      "[5101, 22693, 14881, 15029, 10648, 6330, 8299, 16204, 7465, 16063]\n",
      "According to IMDB dataset and your knowledges, what movies Raul Julia has worked on and what were her/his roles?\n",
      "[8171, 19867, 23709, 26463, 9281, 30223, 29559, 14882, 21991, 13729]\n",
      "According to IMDB dataset and your knowledges, what movies Michael Keaton has worked on and what were her/his roles?\n",
      "[29266, 10290, 27342, 28884, 6185, 6001, 15435, 9027, 20919, 16551]\n",
      "According to IMDB dataset and your knowledges, what movies Fritz Lang has worked on and what were her/his roles?\n",
      "[27518, 18213, 11632, 8531, 29316, 23589, 8895, 11711, 4936, 22667]\n",
      "According to IMDB dataset and your knowledges, what movies Brandon Lee has worked on and what were her/his roles?\n",
      "[5827, 5841, 19298, 16049, 8771, 27067, 20041, 13656, 14426, 25964]\n",
      "According to IMDB dataset and your knowledges, what movies Christopher Lloyd has worked on and what were her/his roles?\n",
      "[3017, 14131, 3035, 9888, 20464, 24398, 3113, 26362, 24850, 28403]\n",
      "According to IMDB dataset and your knowledges, what movies Julia Louis-Dreyfus has worked on and what were her/his roles?\n",
      "[6531, 2334, 1454, 20505, 18321, 27503, 17767, 25885, 20066, 13352]\n",
      "According to IMDB dataset and your knowledges, what movies Stephen Baldwin has worked on and what were her/his roles?\n",
      "[2390, 29349, 10354, 5060, 21826, 16355, 1980, 11138, 29242, 11124]\n",
      "According to IMDB dataset and your knowledges, what movies Michael Biehn has worked on and what were her/his roles?\n",
      "[19724, 16365, 21927, 31657, 22620, 21341, 2417, 16465, 6524, 4829]\n",
      "According to IMDB dataset and your knowledges, what movies Ernest Borgnine has worked on and what were her/his roles?\n",
      "[12343, 15533, 16545, 4636, 7483, 26983, 13676, 6315, 9426, 13531]\n",
      "According to IMDB dataset and your knowledges, what movies Charles Bronson has worked on and what were her/his roles?\n",
      "[16268, 10663, 24796, 5942, 19968, 20234, 12247, 20538, 12782, 7471]\n",
      "According to IMDB dataset and your knowledges, what movies Barbara Carrera has worked on and what were her/his roles?\n",
      "[817, 11376, 5502, 28606, 27160, 4301, 30419, 30759, 4542, 14168]\n",
      "According to IMDB dataset and your knowledges, what movies Jackie Chan has worked on and what were her/his roles?\n",
      "[23871, 27299, 7625, 6000, 5781, 30538, 10782, 11383, 26913, 10744]\n",
      "According to IMDB dataset and your knowledges, what movies Francis Ford has worked on and what were her/his roles?\n",
      "[15537, 6133, 31828, 20635, 1978, 26204, 13830, 30751, 11531, 21097]\n",
      "According to IMDB dataset and your knowledges, what movies Roger Corman has worked on and what were her/his roles?\n",
      "[18847, 6518, 27230, 10934, 24418, 9483, 28446, 30775, 21463, 29005]\n",
      "According to IMDB dataset and your knowledges, what movies Olivier De has worked on and what were her/his roles?\n",
      "[1551, 31206, 25903, 18655, 13490, 26322, 1379, 2463, 21306, 2628]\n",
      "According to IMDB dataset and your knowledges, what movies Brian De has worked on and what were her/his roles?\n",
      "[8918, 22155, 10728, 27270, 29325, 26299, 16474, 3799, 1887, 14746]\n",
      "According to IMDB dataset and your knowledges, what movies Gérard Depardieu has worked on and what were her/his roles?\n",
      "[14383, 1632, 23307, 23560, 1721, 20471, 18752, 14054, 31660, 18176]\n",
      "According to IMDB dataset and your knowledges, what movies Walt Disney has worked on and what were her/his roles?\n",
      "[23352, 15143, 24323, 1498, 16085, 8944, 4767, 22583, 31705, 12538]\n",
      "According to IMDB dataset and your knowledges, what movies Martin Scorsese has worked on and what were her/his roles?\n",
      "[9071, 31615, 10473, 25026, 20800, 9486, 16944, 6609, 22471, 26405]\n",
      "According to IMDB dataset and your knowledges, what movies Charlie Sheen has worked on and what were her/his roles?\n",
      "[8897, 24153, 9026, 2081, 26096, 8915, 10060, 25818, 18846, 5580]\n",
      "According to IMDB dataset and your knowledges, what movies Will Smith has worked on and what were her/his roles?\n",
      "[18649, 15078, 14722, 19788, 21984, 9792, 22859, 19520, 21746, 24940]\n",
      "According to IMDB dataset and your knowledges, what movies Sylvester Stallone has worked on and what were her/his roles?\n",
      "[13715, 28274, 23486, 2329, 21363, 25138, 8490, 21116, 26884, 30763]\n",
      "According to IMDB dataset and your knowledges, what movies Quentin Tarantino has worked on and what were her/his roles?\n",
      "[30188, 13491, 12215, 24578, 24368, 15598, 7979, 31120, 17393, 3429]\n",
      "According to IMDB dataset and your knowledges, what movies Jean-Claude Van has worked on and what were her/his roles?\n",
      "[10184, 24785, 19259, 23410, 28687, 21262, 15286, 20359, 26757, 12359]\n",
      "According to IMDB dataset and your knowledges, what movies Mark Wahlberg has worked on and what were her/his roles?\n",
      "[25527, 7099, 31151, 15353, 16096, 28304, 26907, 21137, 25322, 2405]\n",
      "According to IMDB dataset and your knowledges, what movies Bruce Willis has worked on and what were her/his roles?\n",
      "[1873, 14364, 27106, 5709, 20296, 31681, 29338, 4188, 15678, 23087]\n",
      "According to IMDB dataset and your knowledges, what movies John Woo has worked on and what were her/his roles?\n",
      "[28376, 1250, 22754, 31725, 9940, 20060, 22763, 1477, 30625, 23793]\n",
      "According to IMDB dataset and your knowledges, what movies Edward D. has worked on and what were her/his roles?\n",
      "[12523, 29435, 30061, 31638, 28166, 17069, 20958, 8501, 23090, 8723]\n",
      "According to IMDB dataset and your knowledges, what movies Robert Altman has worked on and what were her/his roles?\n",
      "[9122, 28402, 1330, 3682, 15861, 13955, 7915, 31624, 10245, 20498]\n",
      "According to IMDB dataset and your knowledges, what movies Heather Locklear has worked on and what were her/his roles?\n",
      "[468, 4746, 27683, 22609, 23798, 18652, 31245, 27158, 9872, 7256]\n",
      "According to IMDB dataset and your knowledges, what movies George Lucas has worked on and what were her/his roles?\n",
      "[11102, 30112, 31114, 23718, 14364, 2274, 15368, 30333, 29191, 27437]\n",
      "According to IMDB dataset and your knowledges, what movies Ewan McGregor has worked on and what were her/his roles?\n",
      "[4757, 3929, 16052, 1426, 17121, 1418, 5113, 22065, 1229, 7533]\n",
      "According to IMDB dataset and your knowledges, what movies Julianne Moore has worked on and what were her/his roles?\n",
      "[6280, 10779, 2151, 22026, 30723, 7621, 25068, 15652, 20410, 12684]\n",
      "According to IMDB dataset and your knowledges, what movies Jack Nicholson has worked on and what were her/his roles?\n",
      "[5501, 26807, 28480, 843, 10053, 3881, 30404, 1455, 9858, 13728]\n",
      "According to IMDB dataset and your knowledges, what movies Ryan Phillippe has worked on and what were her/his roles?\n",
      "[12615, 6174, 1404, 26592, 20674, 8245, 10737, 1417, 6875, 29501]\n",
      "According to IMDB dataset and your knowledges, what movies Keanu Reeves has worked on and what were her/his roles?\n",
      "[27921, 30237, 5435, 9182, 11495, 29137, 7783, 31393, 27356, 4241]\n",
      "According to IMDB dataset and your knowledges, what movies Pierce Brosnan has worked on and what were her/his roles?\n",
      "[17055, 26008, 1876, 25441, 12329, 16854, 15821, 27925, 54, 17552]\n",
      "According to IMDB dataset and your knowledges, what movies Wes Craven has worked on and what were her/his roles?\n",
      "[24891, 999, 3153, 20936, 8805, 8704, 24673, 17919, 3798, 5610]\n",
      "According to IMDB dataset and your knowledges, what movies Robert De has worked on and what were her/his roles?\n",
      "[26377, 31561, 22153, 25206, 9504, 24843, 31088, 6838, 1309, 29580]\n",
      "According to IMDB dataset and your knowledges, what movies Johnny Depp has worked on and what were her/his roles?\n",
      "[5048, 27629, 1725, 28061, 12555, 15160, 5632, 15715, 5812, 17070]\n",
      "According to IMDB dataset and your knowledges, what movies Cameron Diaz has worked on and what were her/his roles?\n",
      "[3111, 14032, 6152, 30737, 28879, 19104, 3308, 4846, 10470, 6216]\n",
      "According to IMDB dataset and your knowledges, what movies Michael Douglas has worked on and what were her/his roles?\n",
      "[10265, 14769, 7160, 6498, 2356, 8477, 1049, 8326, 7083, 2542]\n",
      "According to IMDB dataset and your knowledges, what movies Clint Eastwood has worked on and what were her/his roles?\n",
      "[2109, 24706, 13990, 1545, 84, 14012, 4357, 23336, 17979, 28587]\n",
      "According to IMDB dataset and your knowledges, what movies Erika Eleniak has worked on and what were her/his roles?\n",
      "[30592, 17301, 8225, 5905, 7561, 11175, 11832, 4612, 10738, 20784]\n",
      "According to IMDB dataset and your knowledges, what movies Morgan Freeman has worked on and what were her/his roles?\n",
      "[28135, 20986, 1755, 10481, 8170, 17816, 17140, 2628, 12681, 28038]\n",
      "According to IMDB dataset and your knowledges, what movies Ron Howard has worked on and what were her/his roles?\n",
      "[30755, 19027, 25702, 30004, 11368, 31435, 11141, 24752, 5683, 285]\n",
      "According to IMDB dataset and your knowledges, what movies Alfred Newman has worked on and what were her/his roles?\n",
      "[18488, 15779, 30217, 23736, 21301, 8158, 27486, 22811, 8971, 12613]\n",
      "According to IMDB dataset and your knowledges, what movies Laurence Olivier has worked on and what were her/his roles?\n",
      "[16615, 19938, 13196, 18076, 16441, 21368, 11156, 4912, 13060, 18904]\n",
      "According to IMDB dataset and your knowledges, what movies Edward G. has worked on and what were her/his roles?\n",
      "[18735, 15546, 6494, 11253, 24770, 6708, 8657, 5370, 20022, 8114]\n",
      "According to IMDB dataset and your knowledges, what movies John Wayne has worked on and what were her/his roles?\n",
      "[23861, 10971, 10884, 18123, 26161, 27871, 30013, 8437, 9467, 27987]\n",
      "According to IMDB dataset and your knowledges, what movies Natalie Wood has worked on and what were her/his roles?\n",
      "[20579, 31399, 12351, 12089, 28516, 26536, 29955, 9314, 14605, 16746]\n",
      "According to IMDB dataset and your knowledges, what movies Alan Miller has worked on and what were her/his roles?\n",
      "[24316, 16075, 23176, 1602, 9999, 18272, 6966, 22714, 14055, 4066]\n",
      "According to IMDB dataset and your knowledges, what movies Henner Hofmann has worked on and what were her/his roles?\n",
      "[7565, 8747, 28550, 19604, 6419, 6939, 12499, 18589, 5398, 30502]\n",
      "According to IMDB dataset and your knowledges, what movies Louis de has worked on and what were her/his roles?\n",
      "[16079, 20816, 22680, 17930, 29112, 7462, 21143, 19198, 12609, 9022]\n",
      "According to IMDB dataset and your knowledges, what movies Aleksey Korenev has worked on and what were her/his roles?\n",
      "[23511, 10842, 30586, 11406, 26443, 31843, 16565, 17076, 9974, 1388]\n",
      "According to IMDB dataset and your knowledges, what movies John Cleese has worked on and what were her/his roles?\n",
      "[19146, 16191, 9467, 18000, 2539, 14487, 16840, 6250, 4263, 3166]\n",
      "According to IMDB dataset and your knowledges, what movies J. Reifel has worked on and what were her/his roles?\n",
      "[16525, 25715, 24415, 22549, 22032, 17386, 29600, 2519, 8791, 658]\n",
      "According to IMDB dataset and your knowledges, what movies Pamela Anderson has worked on and what were her/his roles?\n",
      "[28084, 9261, 25863, 21170, 2918, 1180, 15517, 23214, 29967, 20589]\n",
      "According to IMDB dataset and your knowledges, what movies Dan Aykroyd has worked on and what were her/his roles?\n",
      "[15168, 7916, 2251, 19709, 16535, 28883, 27286, 237, 5042, 8783]\n",
      "According to IMDB dataset and your knowledges, what movies Kim Basinger has worked on and what were her/his roles?\n",
      "[25368, 31162, 9559, 17412, 23437, 15636, 31011, 15660, 5394, 15716]\n",
      "According to IMDB dataset and your knowledges, what movies Matthew Broderick has worked on and what were her/his roles?\n",
      "[31726, 8835, 8424, 1883, 30414, 5576, 11292, 12400, 17466, 6214]\n",
      "According to IMDB dataset and your knowledges, what movies Brigitte Bardot has worked on and what were her/his roles?\n",
      "[3527, 22703, 22927, 15123, 26638, 116, 2660, 13324, 14254, 7704]\n",
      "According to IMDB dataset and your knowledges, what movies Bette Davis has worked on and what were her/his roles?\n",
      "[17791, 17105, 21044, 16700, 7722, 12592, 13106, 25579, 23550, 16387]\n",
      "According to IMDB dataset and your knowledges, what movies Doris Day has worked on and what were her/his roles?\n",
      "[16813, 29445, 18845, 12903, 24447, 20707, 13450, 24084, 903, 1089]\n",
      "According to IMDB dataset and your knowledges, what movies Olivia de has worked on and what were her/his roles?\n",
      "[18872, 11039, 22985, 29076, 23719, 15005, 5752, 29081, 23904, 3719]\n",
      "According to IMDB dataset and your knowledges, what movies James Dean has worked on and what were her/his roles?\n",
      "[29175, 9139, 27443, 9380, 29376, 11686, 2256, 6808, 25948, 31940]\n",
      "According to IMDB dataset and your knowledges, what movies Kirk Douglas has worked on and what were her/his roles?\n",
      "[28866, 11006, 1513, 29630, 28639, 22514, 31664, 17603, 28155, 25166]\n",
      "According to IMDB dataset and your knowledges, what movies Federico Fellini has worked on and what were her/his roles?\n",
      "[16553, 4215, 19192, 31270, 13241, 8014, 15584, 1001, 25940, 14922]\n",
      "According to IMDB dataset and your knowledges, what movies John Gielgud has worked on and what were her/his roles?\n",
      "[21681, 11993, 2226, 25525, 25058, 22049, 15081, 3606, 2859, 16437]\n",
      "According to IMDB dataset and your knowledges, what movies Jerry Goldsmith has worked on and what were her/his roles?\n",
      "[22000, 21203, 2948, 19597, 26295, 7124, 28156, 19528, 24782, 3829]\n",
      "According to IMDB dataset and your knowledges, what movies Rita Hayworth has worked on and what were her/his roles?\n",
      "[5835, 12548, 18447, 30029, 26603, 29935, 11472, 9580, 23726, 8149]\n",
      "According to IMDB dataset and your knowledges, what movies Margaux Hemingway has worked on and what were her/his roles?\n",
      "[22198, 4108, 19345, 7778, 30661, 18906, 1343, 14837, 14121, 10527]\n",
      "According to IMDB dataset and your knowledges, what movies Alfred Hitchcock has worked on and what were her/his roles?\n",
      "[15266, 88, 9861, 18624, 15771, 14800, 6157, 6569, 949, 10256]\n",
      "According to IMDB dataset and your knowledges, what movies Deborah Kerr has worked on and what were her/his roles?\n",
      "[19642, 7973, 930, 1746, 5603, 8432, 17646, 22104, 24994, 9565]\n",
      "According to IMDB dataset and your knowledges, what movies Sophia Loren has worked on and what were her/his roles?\n",
      "[1353, 30308, 14339, 13973, 22520, 30157, 14635, 6454, 19567, 30931]\n",
      "According to IMDB dataset and your knowledges, what movies Henry Mancini has worked on and what were her/his roles?\n",
      "[6757, 8404, 4782, 23158, 7049, 25278, 8593, 2024, 27809, 20586]\n",
      "According to IMDB dataset and your knowledges, what movies Groucho Marx has worked on and what were her/his roles?\n",
      "[17834, 23104, 14475, 14330, 6282, 27641, 3240, 12896, 25709, 19307]\n",
      "According to IMDB dataset and your knowledges, what movies James Mason has worked on and what were her/his roles?\n",
      "[7513, 26358, 4139, 25022, 19102, 11757, 23988, 2654, 8541, 11004]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "100"
      ]
     },
     "execution_count": 143,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from DGD.help_functions import get_embedding, perturb_sentence, rank_tokens_by_importance, get_initial_ids, user_prompt_generation, check_success, check_MSEloss, get_first_output_token, get_neighbor_ids\n",
    "from DGD.DGD import get_optimized_prefix_embedding\n",
    "\n",
    "# Set dataset\n",
    "dataset = \"imdb\"\n",
    "query_mode = \"for_objects\"\n",
    "\n",
    "hidden_states_positive = []\n",
    "layer = -1\n",
    "counter = 0\n",
    "for index in json_str_filter:\n",
    "    # if counter == 100:\n",
    "    #     break\n",
    "    prompt_text = json_str_filter[index][\"prompt_text\"]\n",
    "    initial_ids = origin_token_list[index]\n",
    "    object = \"\"\n",
    "    sentence = user_prompt_generation(prompt_text, object, dataset, query_mode)\n",
    "    print(sentence)\n",
    "    print(initial_ids)\n",
    "    hidden_state = get_model_layer_hidden_states(sentence, initial_ids, model, tok, layer)\n",
    "    hidden_states_positive.append(hidden_state.to(\"cpu\"))\n",
    "    counter += 1\n",
    "len(hidden_states_positive)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 144,
   "id": "e86cad5f-8c4b-49da-8c94-a267be62a350",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n",
      "torch.Size([450560])\n"
     ]
    }
   ],
   "source": [
    "X_1 = []\n",
    "y_1 = []\n",
    "X_2 = []\n",
    "y_2 = []\n",
    "counter = 0\n",
    "size = hidden_states_negative[0].shape[0]\n",
    "for matrix in hidden_states_positive:\n",
    "    if counter == 100:\n",
    "        break\n",
    "    if matrix.shape[0] != size:\n",
    "        continue\n",
    "    print(matrix.shape)\n",
    "    X_1.append(matrix.to(\"cpu\"))\n",
    "    y_1.append(0)\n",
    "    \n",
    "    counter += 1\n",
    "    \n",
    "counter = 0\n",
    "for matrix in hidden_states_negative:\n",
    "    if counter == 100:\n",
    "        break\n",
    "    if matrix.shape[0] != size:\n",
    "        continue\n",
    "    print(matrix.shape)\n",
    "    X_2.append(matrix.to(\"cpu\"))\n",
    "    y_2.append(1)\n",
    "\n",
    "    counter += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 145,
   "id": "4679e9d9-a717-4109-aba8-0bb09f112761",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(100, 100)"
      ]
     },
     "execution_count": 145,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(X_1), len(X_2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 146,
   "id": "dcb4ef2f-6772-4ba4-af0d-f3df0cff8e1a",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "120 120 80 80\n",
      "training_time:  4.46355907805264\n",
      "response_time:  0.20115510933101177\n",
      "0 :\n",
      "1.0\n",
      " \n",
      "120 120 80 80\n",
      "training_time:  4.4981503849849105\n",
      "response_time:  0.20134748984128237\n",
      "1 :\n",
      "1.0\n",
      " \n",
      "120 120 80 80\n",
      "training_time:  5.7243487681262195\n",
      "response_time:  0.20093853818252683\n",
      "2 :\n",
      "1.0\n",
      " \n",
      "120 120 80 80\n",
      "training_time:  6.029270411003381\n",
      "response_time:  0.20127836661413312\n",
      "3 :\n",
      "1.0\n",
      " \n",
      "120 120 80 80\n",
      "training_time:  3.6956862648949027\n",
      "response_time:  0.20144648803398013\n",
      "4 :\n",
      "1.0\n",
      " \n",
      "120 120 80 80\n",
      "training_time:  5.4898805767297745\n",
      "response_time:  0.22186102718114853\n",
      "5 :\n",
      "1.0\n",
      " \n",
      "120 120 80 80\n",
      "training_time:  4.394320995081216\n",
      "response_time:  0.2013642732053995\n",
      "6 :\n",
      "1.0\n",
      " \n",
      "120 120 80 80\n",
      "training_time:  4.403729849029332\n",
      "response_time:  0.2185397632420063\n",
      "7 :\n",
      "1.0\n",
      " \n",
      "120 120 80 80\n",
      "training_time:  3.665802648290992\n",
      "response_time:  0.20120144123211503\n",
      "8 :\n",
      "1.0\n",
      " \n",
      "120 120 80 80\n",
      "training_time:  4.561057776212692\n",
      "response_time:  0.2010862729512155\n",
      "9 :\n",
      "1.0\n",
      " \n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(1.0, 0.9936708860759493, 0.9875, 1.0, 0.99375)"
      ]
     },
     "execution_count": 146,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Train:test = 6:4, hidden_states, imdb\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.linear_model import LogisticRegression\n",
    "from sklearn.metrics import roc_auc_score, f1_score, recall_score, precision_score, accuracy_score\n",
    "from sklearn import metrics\n",
    "from sklearn.neural_network import MLPClassifier\n",
    "import time\n",
    "\n",
    "auroc_lst = []\n",
    "f1_lst = []\n",
    "recall_lst = []\n",
    "precision_lst = []\n",
    "accuracy_lst = []\n",
    "training_times = []\n",
    "response_times = []\n",
    "for seed in range(0, 10):\n",
    "    shuffled_X_1, shuffled_X_2, shuffled_y_1, shuffled_y_2 = split_four_lists_randomly(X_1, X_2, y_1, y_2, seed)\n",
    "    X_train = shuffled_X_1[:60] + shuffled_X_2[:60]\n",
    "    X_test = shuffled_X_1[60:] + shuffled_X_2[60:]\n",
    "    y_train = shuffled_y_1[:60] + shuffled_y_2[:60]\n",
    "    y_test = shuffled_y_1[60:] + shuffled_y_2[60:]\n",
    "\n",
    "    X_train, y_train = split_lists_randomly(X_train, y_train, seed)\n",
    "    X_test, y_test = split_lists_randomly(X_test, y_test, seed)\n",
    "    print(len(X_train), len(y_train), len(X_test), len(y_test))\n",
    "    clf = LogisticRegression()\n",
    "    \n",
    "    # Start time\n",
    "    start_time = time.perf_counter()\n",
    "    \n",
    "    clf.fit(X_train, y_train)\n",
    "    \n",
    "    # End time\n",
    "    end_time = time.perf_counter()\n",
    "\n",
    "    # Calculate response time\n",
    "    training_time = end_time - start_time\n",
    "    training_times.append(training_time)\n",
    "\n",
    "    print(\"training_time: \", training_time)\n",
    "\n",
    "    # Start time\n",
    "    start_time = time.perf_counter()\n",
    "    \n",
    "    y_pred = clf.predict(X_test)\n",
    "    \n",
    "    # End time\n",
    "    end_time = time.perf_counter()\n",
    "    \n",
    "    # Calculate response time\n",
    "    response_time = end_time - start_time\n",
    "    response_times.append(response_time)\n",
    "\n",
    "    print(\"response_time: \", response_time)\n",
    "    \n",
    "    y_pred_proba = clf.predict_proba(X_test)[:, 1]\n",
    "    auroc = roc_auc_score(y_test, y_pred_proba)\n",
    "    f1 = f1_score(y_test, y_pred)\n",
    "    recall = recall_score(y_test, y_pred)\n",
    "    precision = precision_score(y_test, y_pred)\n",
    "    accuracy = accuracy_score(y_test, y_pred)\n",
    "\n",
    "        \n",
    "    auroc_lst.append(auroc)\n",
    "    f1_lst.append(f1)\n",
    "    recall_lst.append(recall)\n",
    "    precision_lst.append(precision)\n",
    "    accuracy_lst.append(accuracy)\n",
    "    print(str(seed) + \" :\")\n",
    "    print(auroc)\n",
    "    print(\" \")\n",
    "sum(auroc_lst) / len(auroc_lst), sum(f1_lst) / len(f1_lst), sum(recall_lst) / len(recall_lst), sum(precision_lst) / len(precision_lst), sum(accuracy_lst) / len(accuracy_lst)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47bd59a2-bf6a-45b6-ab5f-148bd38f481b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Train:test = 6:4, attentions, imdb\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.linear_model import LogisticRegression\n",
    "from sklearn.metrics import roc_auc_score, f1_score, recall_score, precision_score, accuracy_score\n",
    "from sklearn import metrics\n",
    "from sklearn.neural_network import MLPClassifier\n",
    "import time\n",
    "\n",
    "auroc_lst = []\n",
    "f1_lst = []\n",
    "recall_lst = []\n",
    "precision_lst = []\n",
    "accuracy_lst = []\n",
    "training_times = []\n",
    "response_times = []\n",
    "for seed in range(0, 10):\n",
    "    # X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.4, random_state=seed)\n",
    "    shuffled_X_1, shuffled_X_2, shuffled_y_1, shuffled_y_2 = split_four_lists_randomly(X_1, X_2, y_1, y_2, seed)\n",
    "    X_train = shuffled_X_1[:60] + shuffled_X_2[:60]\n",
    "    X_test = shuffled_X_1[60:] + shuffled_X_2[60:]\n",
    "    y_train = shuffled_y_1[:60] + shuffled_y_2[:60]\n",
    "    y_test = shuffled_y_1[60:] + shuffled_y_2[60:]\n",
    "\n",
    "    X_train, y_train = split_lists_randomly(X_train, y_train, seed)\n",
    "    X_test, y_test = split_lists_randomly(X_test, y_test, seed)\n",
    "    # clf = MLPClassifier(solver='lbfgs', alpha=1e-5, hidden_layer_sizes=(500, 10), random_state=1)\n",
    "    print(len(X_train), len(y_train), len(X_test), len(y_test))\n",
    "    clf = LogisticRegression()\n",
    "    \n",
    "    # Start time\n",
    "    start_time = time.perf_counter()\n",
    "    \n",
    "    clf.fit(X_train, y_train)\n",
    "    \n",
    "    # End time\n",
    "    end_time = time.perf_counter()\n",
    "\n",
    "    # Calculate response time\n",
    "    training_time = end_time - start_time\n",
    "    training_times.append(training_time)\n",
    "\n",
    "    print(\"training_time: \", training_time)\n",
    "\n",
    "    # Start time\n",
    "    start_time = time.perf_counter()\n",
    "    \n",
    "    y_pred = clf.predict(X_test)\n",
    "    \n",
    "    # End time\n",
    "    end_time = time.perf_counter()\n",
    "    \n",
    "    # Calculate response time\n",
    "    response_time = end_time - start_time\n",
    "    response_times.append(response_time)\n",
    "\n",
    "    print(\"response_time: \", response_time)\n",
    "    \n",
    "    y_pred_proba = clf.predict_proba(X_test)[:, 1]\n",
    "    auroc = roc_auc_score(y_test, y_pred_proba)\n",
    "    f1 = f1_score(y_test, y_pred)\n",
    "    recall = recall_score(y_test, y_pred)\n",
    "    precision = precision_score(y_test, y_pred)\n",
    "    accuracy = accuracy_score(y_test, y_pred)\n",
    "\n",
    "        \n",
    "    auroc_lst.append(auroc)\n",
    "    f1_lst.append(f1)\n",
    "    recall_lst.append(recall)\n",
    "    precision_lst.append(precision)\n",
    "    accuracy_lst.append(accuracy)\n",
    "    print(str(seed) + \" :\")\n",
    "    print(auroc)\n",
    "    print(\" \")\n",
    "sum(auroc_lst) / len(auroc_lst), sum(f1_lst) / len(f1_lst), sum(recall_lst) / len(recall_lst), sum(precision_lst) / len(precision_lst), sum(accuracy_lst) / len(accuracy_lst)"
   ]
  }
 ],
 "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.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}