{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import open_clip\n",
    "import torch\n",
    "from tqdm.notebook import tqdm\n",
    "import pandas as pd\n",
    "import os\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "PROMPTS = [\n",
    "    '{0}',\n",
    "    'an image of {0}',\n",
    "    'a photo of {0}',\n",
    "    '{0} on a photo',\n",
    "    'a photo of a person named {0}',\n",
    "    'a person named {0}',\n",
    "    'a man named {0}',\n",
    "    'a woman named {0}',\n",
    "    'the name of the person is {0}',\n",
    "    'a photo of a person with the name {0}',\n",
    "    '{0} at a gala',\n",
    "    'a photo of the celebrity {0}',\n",
    "    'actor {0}',\n",
    "    'actress {0}',\n",
    "    'a colored photo of {0}',\n",
    "    'a black and white photo of {0}',\n",
    "    'a cool photo of {0}',\n",
    "    'a cropped photo of {0}',\n",
    "    'a cropped image of {0}',\n",
    "    '{0} in a suit',\n",
    "    '{0} in a dress'\n",
    "]\n",
    "MODEL_NAMES = ['ViT-B-32', 'ViT-B-16', 'ViT-L-14']\n",
    "SEED = 42"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "outputs": [],
   "source": [
    "# init clip\n",
    "models = {}\n",
    "preprocessings = {}\n",
    "tokenizers = {}\n",
    "for model_name in MODEL_NAMES:\n",
    "    model, _, preprocess = open_clip.create_model_and_transforms(model_name, pretrained='laion400m_e32')\n",
    "    preprocessings[model_name] = preprocess\n",
    "    model = model.eval()\n",
    "    models[model_name] = model\n",
    "    tokenizers[model_name] = open_clip.get_tokenizer(model_name)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "outputs": [],
   "source": [
    "# define a function to get the predictions for an actor/actress\n",
    "@torch.no_grad()\n",
    "def get_text_embeddings(model, context, context_batchsize=1_000, use_tqdm=False):\n",
    "    context_batchsize = context_batchsize * torch.cuda.device_count()\n",
    "    # if there is not batches for the context unsqueeze it\n",
    "    if context.dim() < 3:\n",
    "        context = context.unsqueeze(0)\n",
    "\n",
    "    # get the batch size, the number of labels and the sequence length\n",
    "    seq_len = context.shape[-1]\n",
    "    viewed_context = context.view(-1, seq_len)\n",
    "\n",
    "    text_features = []\n",
    "    for context_batch_idx in tqdm(range(0, len(viewed_context), context_batchsize), desc=\"Calculating Text Embeddings\",\n",
    "                                  disable=not use_tqdm):\n",
    "        context_batch = viewed_context[context_batch_idx:context_batch_idx + context_batchsize]\n",
    "        batch_text_features = model.encode_text(context_batch, normalize=True).cpu()\n",
    "\n",
    "        text_features.append(batch_text_features)\n",
    "    text_features = torch.cat(text_features).view(list(context.shape[:-1]) + [-1])\n",
    "\n",
    "    return text_features"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "outputs": [
    {
     "data": {
      "text/plain": "     first_name sex  last_name\n0        Eliana   f   Cardenas\n1       Meghann   f    Daniels\n2           Ada   f  Stevenson\n3          Elsa   f    Leblanc\n4          Avah   f    Lambert\n...         ...  ..        ...\n9995      Kasen   m     Barker\n9996     Camryn   m    Roberts\n9997      Henry   m   Whitaker\n9998       Adin   m   Richards\n9999    Charley   m     Herman\n\n[10000 rows x 3 columns]",
      "text/html": "<div>\n<style scoped>\n    .dataframe tbody tr th:only-of-type {\n        vertical-align: middle;\n    }\n\n    .dataframe tbody tr th {\n        vertical-align: top;\n    }\n\n    .dataframe thead th {\n        text-align: right;\n    }\n</style>\n<table border=\"1\" class=\"dataframe\">\n  <thead>\n    <tr style=\"text-align: right;\">\n      <th></th>\n      <th>first_name</th>\n      <th>sex</th>\n      <th>last_name</th>\n    </tr>\n  </thead>\n  <tbody>\n    <tr>\n      <th>0</th>\n      <td>Eliana</td>\n      <td>f</td>\n      <td>Cardenas</td>\n    </tr>\n    <tr>\n      <th>1</th>\n      <td>Meghann</td>\n      <td>f</td>\n      <td>Daniels</td>\n    </tr>\n    <tr>\n      <th>2</th>\n      <td>Ada</td>\n      <td>f</td>\n      <td>Stevenson</td>\n    </tr>\n    <tr>\n      <th>3</th>\n      <td>Elsa</td>\n      <td>f</td>\n      <td>Leblanc</td>\n    </tr>\n    <tr>\n      <th>4</th>\n      <td>Avah</td>\n      <td>f</td>\n      <td>Lambert</td>\n    </tr>\n    <tr>\n      <th>...</th>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n    </tr>\n    <tr>\n      <th>9995</th>\n      <td>Kasen</td>\n      <td>m</td>\n      <td>Barker</td>\n    </tr>\n    <tr>\n      <th>9996</th>\n      <td>Camryn</td>\n      <td>m</td>\n      <td>Roberts</td>\n    </tr>\n    <tr>\n      <th>9997</th>\n      <td>Henry</td>\n      <td>m</td>\n      <td>Whitaker</td>\n    </tr>\n    <tr>\n      <th>9998</th>\n      <td>Adin</td>\n      <td>m</td>\n      <td>Richards</td>\n    </tr>\n    <tr>\n      <th>9999</th>\n      <td>Charley</td>\n      <td>m</td>\n      <td>Herman</td>\n    </tr>\n  </tbody>\n</table>\n<p>10000 rows × 3 columns</p>\n</div>"
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# load the possible names\n",
    "possible_names = pd.read_csv('./full_names.csv', index_col=0)\n",
    "possible_names\n",
    "# possible_names_list = (possible_names['first_name'] + ' ' + possible_names['last_name']).tolist()\n",
    "# possible_names_list[:5]"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "outputs": [
    {
     "data": {
      "text/plain": "     first_name sex  last_name         prompt_0                     prompt_1  \\\n0        Eliana   f   Cardenas  Eliana Cardenas  an image of Eliana Cardenas   \n1       Meghann   f    Daniels  Meghann Daniels  an image of Meghann Daniels   \n2           Ada   f  Stevenson    Ada Stevenson    an image of Ada Stevenson   \n3          Elsa   f    Leblanc     Elsa Leblanc     an image of Elsa Leblanc   \n4          Avah   f    Lambert     Avah Lambert     an image of Avah Lambert   \n...         ...  ..        ...              ...                          ...   \n9995      Kasen   m     Barker     Kasen Barker     an image of Kasen Barker   \n9996     Camryn   m    Roberts   Camryn Roberts   an image of Camryn Roberts   \n9997      Henry   m   Whitaker   Henry Whitaker   an image of Henry Whitaker   \n9998       Adin   m   Richards    Adin Richards    an image of Adin Richards   \n9999    Charley   m     Herman   Charley Herman   an image of Charley Herman   \n\n                        prompt_2                    prompt_3  \\\n0     a photo of Eliana Cardenas  Eliana Cardenas on a photo   \n1     a photo of Meghann Daniels  Meghann Daniels on a photo   \n2       a photo of Ada Stevenson    Ada Stevenson on a photo   \n3        a photo of Elsa Leblanc     Elsa Leblanc on a photo   \n4        a photo of Avah Lambert     Avah Lambert on a photo   \n...                          ...                         ...   \n9995     a photo of Kasen Barker     Kasen Barker on a photo   \n9996   a photo of Camryn Roberts   Camryn Roberts on a photo   \n9997   a photo of Henry Whitaker   Henry Whitaker on a photo   \n9998    a photo of Adin Richards    Adin Richards on a photo   \n9999   a photo of Charley Herman   Charley Herman on a photo   \n\n                                       prompt_4  \\\n0     a photo of a person named Eliana Cardenas   \n1     a photo of a person named Meghann Daniels   \n2       a photo of a person named Ada Stevenson   \n3        a photo of a person named Elsa Leblanc   \n4        a photo of a person named Avah Lambert   \n...                                         ...   \n9995     a photo of a person named Kasen Barker   \n9996   a photo of a person named Camryn Roberts   \n9997   a photo of a person named Henry Whitaker   \n9998    a photo of a person named Adin Richards   \n9999   a photo of a person named Charley Herman   \n\n                            prompt_5                     prompt_6  ...  \\\n0     a person named Eliana Cardenas  a man named Eliana Cardenas  ...   \n1     a person named Meghann Daniels  a man named Meghann Daniels  ...   \n2       a person named Ada Stevenson    a man named Ada Stevenson  ...   \n3        a person named Elsa Leblanc     a man named Elsa Leblanc  ...   \n4        a person named Avah Lambert     a man named Avah Lambert  ...   \n...                              ...                          ...  ...   \n9995     a person named Kasen Barker     a man named Kasen Barker  ...   \n9996   a person named Camryn Roberts   a man named Camryn Roberts  ...   \n9997   a person named Henry Whitaker   a man named Henry Whitaker  ...   \n9998    a person named Adin Richards    a man named Adin Richards  ...   \n9999   a person named Charley Herman   a man named Charley Herman  ...   \n\n                                     prompt_11              prompt_12  \\\n0     a photo of the celebrity Eliana Cardenas  actor Eliana Cardenas   \n1     a photo of the celebrity Meghann Daniels  actor Meghann Daniels   \n2       a photo of the celebrity Ada Stevenson    actor Ada Stevenson   \n3        a photo of the celebrity Elsa Leblanc     actor Elsa Leblanc   \n4        a photo of the celebrity Avah Lambert     actor Avah Lambert   \n...                                        ...                    ...   \n9995     a photo of the celebrity Kasen Barker     actor Kasen Barker   \n9996   a photo of the celebrity Camryn Roberts   actor Camryn Roberts   \n9997   a photo of the celebrity Henry Whitaker   actor Henry Whitaker   \n9998    a photo of the celebrity Adin Richards    actor Adin Richards   \n9999   a photo of the celebrity Charley Herman   actor Charley Herman   \n\n                    prompt_13                           prompt_14  \\\n0     actress Eliana Cardenas  a colored photo of Eliana Cardenas   \n1     actress Meghann Daniels  a colored photo of Meghann Daniels   \n2       actress Ada Stevenson    a colored photo of Ada Stevenson   \n3        actress Elsa Leblanc     a colored photo of Elsa Leblanc   \n4        actress Avah Lambert     a colored photo of Avah Lambert   \n...                       ...                                 ...   \n9995     actress Kasen Barker     a colored photo of Kasen Barker   \n9996   actress Camryn Roberts   a colored photo of Camryn Roberts   \n9997   actress Henry Whitaker   a colored photo of Henry Whitaker   \n9998    actress Adin Richards    a colored photo of Adin Richards   \n9999   actress Charley Herman   a colored photo of Charley Herman   \n\n                                       prompt_15  \\\n0     a black and white photo of Eliana Cardenas   \n1     a black and white photo of Meghann Daniels   \n2       a black and white photo of Ada Stevenson   \n3        a black and white photo of Elsa Leblanc   \n4        a black and white photo of Avah Lambert   \n...                                          ...   \n9995     a black and white photo of Kasen Barker   \n9996   a black and white photo of Camryn Roberts   \n9997   a black and white photo of Henry Whitaker   \n9998    a black and white photo of Adin Richards   \n9999   a black and white photo of Charley Herman   \n\n                            prompt_16                           prompt_17  \\\n0     a cool photo of Eliana Cardenas  a cropped photo of Eliana Cardenas   \n1     a cool photo of Meghann Daniels  a cropped photo of Meghann Daniels   \n2       a cool photo of Ada Stevenson    a cropped photo of Ada Stevenson   \n3        a cool photo of Elsa Leblanc     a cropped photo of Elsa Leblanc   \n4        a cool photo of Avah Lambert     a cropped photo of Avah Lambert   \n...                               ...                                 ...   \n9995     a cool photo of Kasen Barker     a cropped photo of Kasen Barker   \n9996   a cool photo of Camryn Roberts   a cropped photo of Camryn Roberts   \n9997   a cool photo of Henry Whitaker   a cropped photo of Henry Whitaker   \n9998    a cool photo of Adin Richards    a cropped photo of Adin Richards   \n9999   a cool photo of Charley Herman   a cropped photo of Charley Herman   \n\n                               prompt_18                  prompt_19  \\\n0     a cropped image of Eliana Cardenas  Eliana Cardenas in a suit   \n1     a cropped image of Meghann Daniels  Meghann Daniels in a suit   \n2       a cropped image of Ada Stevenson    Ada Stevenson in a suit   \n3        a cropped image of Elsa Leblanc     Elsa Leblanc in a suit   \n4        a cropped image of Avah Lambert     Avah Lambert in a suit   \n...                                  ...                        ...   \n9995     a cropped image of Kasen Barker     Kasen Barker in a suit   \n9996   a cropped image of Camryn Roberts   Camryn Roberts in a suit   \n9997   a cropped image of Henry Whitaker   Henry Whitaker in a suit   \n9998    a cropped image of Adin Richards    Adin Richards in a suit   \n9999   a cropped image of Charley Herman   Charley Herman in a suit   \n\n                       prompt_20  \n0     Eliana Cardenas in a dress  \n1     Meghann Daniels in a dress  \n2       Ada Stevenson in a dress  \n3        Elsa Leblanc in a dress  \n4        Avah Lambert in a dress  \n...                          ...  \n9995     Kasen Barker in a dress  \n9996   Camryn Roberts in a dress  \n9997   Henry Whitaker in a dress  \n9998    Adin Richards in a dress  \n9999   Charley Herman in a dress  \n\n[10000 rows x 24 columns]",
      "text/html": "<div>\n<style scoped>\n    .dataframe tbody tr th:only-of-type {\n        vertical-align: middle;\n    }\n\n    .dataframe tbody tr th {\n        vertical-align: top;\n    }\n\n    .dataframe thead th {\n        text-align: right;\n    }\n</style>\n<table border=\"1\" class=\"dataframe\">\n  <thead>\n    <tr style=\"text-align: right;\">\n      <th></th>\n      <th>first_name</th>\n      <th>sex</th>\n      <th>last_name</th>\n      <th>prompt_0</th>\n      <th>prompt_1</th>\n      <th>prompt_2</th>\n      <th>prompt_3</th>\n      <th>prompt_4</th>\n      <th>prompt_5</th>\n      <th>prompt_6</th>\n      <th>...</th>\n      <th>prompt_11</th>\n      <th>prompt_12</th>\n      <th>prompt_13</th>\n      <th>prompt_14</th>\n      <th>prompt_15</th>\n      <th>prompt_16</th>\n      <th>prompt_17</th>\n      <th>prompt_18</th>\n      <th>prompt_19</th>\n      <th>prompt_20</th>\n    </tr>\n  </thead>\n  <tbody>\n    <tr>\n      <th>0</th>\n      <td>Eliana</td>\n      <td>f</td>\n      <td>Cardenas</td>\n      <td>Eliana Cardenas</td>\n      <td>an image of Eliana Cardenas</td>\n      <td>a photo of Eliana Cardenas</td>\n      <td>Eliana Cardenas on a photo</td>\n      <td>a photo of a person named Eliana Cardenas</td>\n      <td>a person named Eliana Cardenas</td>\n      <td>a man named Eliana Cardenas</td>\n      <td>...</td>\n      <td>a photo of the celebrity Eliana Cardenas</td>\n      <td>actor Eliana Cardenas</td>\n      <td>actress Eliana Cardenas</td>\n      <td>a colored photo of Eliana Cardenas</td>\n      <td>a black and white photo of Eliana Cardenas</td>\n      <td>a cool photo of Eliana Cardenas</td>\n      <td>a cropped photo of Eliana Cardenas</td>\n      <td>a cropped image of Eliana Cardenas</td>\n      <td>Eliana Cardenas in a suit</td>\n      <td>Eliana Cardenas in a dress</td>\n    </tr>\n    <tr>\n      <th>1</th>\n      <td>Meghann</td>\n      <td>f</td>\n      <td>Daniels</td>\n      <td>Meghann Daniels</td>\n      <td>an image of Meghann Daniels</td>\n      <td>a photo of Meghann Daniels</td>\n      <td>Meghann Daniels on a photo</td>\n      <td>a photo of a person named Meghann Daniels</td>\n      <td>a person named Meghann Daniels</td>\n      <td>a man named Meghann Daniels</td>\n      <td>...</td>\n      <td>a photo of the celebrity Meghann Daniels</td>\n      <td>actor Meghann Daniels</td>\n      <td>actress Meghann Daniels</td>\n      <td>a colored photo of Meghann Daniels</td>\n      <td>a black and white photo of Meghann Daniels</td>\n      <td>a cool photo of Meghann Daniels</td>\n      <td>a cropped photo of Meghann Daniels</td>\n      <td>a cropped image of Meghann Daniels</td>\n      <td>Meghann Daniels in a suit</td>\n      <td>Meghann Daniels in a dress</td>\n    </tr>\n    <tr>\n      <th>2</th>\n      <td>Ada</td>\n      <td>f</td>\n      <td>Stevenson</td>\n      <td>Ada Stevenson</td>\n      <td>an image of Ada Stevenson</td>\n      <td>a photo of Ada Stevenson</td>\n      <td>Ada Stevenson on a photo</td>\n      <td>a photo of a person named Ada Stevenson</td>\n      <td>a person named Ada Stevenson</td>\n      <td>a man named Ada Stevenson</td>\n      <td>...</td>\n      <td>a photo of the celebrity Ada Stevenson</td>\n      <td>actor Ada Stevenson</td>\n      <td>actress Ada Stevenson</td>\n      <td>a colored photo of Ada Stevenson</td>\n      <td>a black and white photo of Ada Stevenson</td>\n      <td>a cool photo of Ada Stevenson</td>\n      <td>a cropped photo of Ada Stevenson</td>\n      <td>a cropped image of Ada Stevenson</td>\n      <td>Ada Stevenson in a suit</td>\n      <td>Ada Stevenson in a dress</td>\n    </tr>\n    <tr>\n      <th>3</th>\n      <td>Elsa</td>\n      <td>f</td>\n      <td>Leblanc</td>\n      <td>Elsa Leblanc</td>\n      <td>an image of Elsa Leblanc</td>\n      <td>a photo of Elsa Leblanc</td>\n      <td>Elsa Leblanc on a photo</td>\n      <td>a photo of a person named Elsa Leblanc</td>\n      <td>a person named Elsa Leblanc</td>\n      <td>a man named Elsa Leblanc</td>\n      <td>...</td>\n      <td>a photo of the celebrity Elsa Leblanc</td>\n      <td>actor Elsa Leblanc</td>\n      <td>actress Elsa Leblanc</td>\n      <td>a colored photo of Elsa Leblanc</td>\n      <td>a black and white photo of Elsa Leblanc</td>\n      <td>a cool photo of Elsa Leblanc</td>\n      <td>a cropped photo of Elsa Leblanc</td>\n      <td>a cropped image of Elsa Leblanc</td>\n      <td>Elsa Leblanc in a suit</td>\n      <td>Elsa Leblanc in a dress</td>\n    </tr>\n    <tr>\n      <th>4</th>\n      <td>Avah</td>\n      <td>f</td>\n      <td>Lambert</td>\n      <td>Avah Lambert</td>\n      <td>an image of Avah Lambert</td>\n      <td>a photo of Avah Lambert</td>\n      <td>Avah Lambert on a photo</td>\n      <td>a photo of a person named Avah Lambert</td>\n      <td>a person named Avah Lambert</td>\n      <td>a man named Avah Lambert</td>\n      <td>...</td>\n      <td>a photo of the celebrity Avah Lambert</td>\n      <td>actor Avah Lambert</td>\n      <td>actress Avah Lambert</td>\n      <td>a colored photo of Avah Lambert</td>\n      <td>a black and white photo of Avah Lambert</td>\n      <td>a cool photo of Avah Lambert</td>\n      <td>a cropped photo of Avah Lambert</td>\n      <td>a cropped image of Avah Lambert</td>\n      <td>Avah Lambert in a suit</td>\n      <td>Avah Lambert in a dress</td>\n    </tr>\n    <tr>\n      <th>...</th>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n    </tr>\n    <tr>\n      <th>9995</th>\n      <td>Kasen</td>\n      <td>m</td>\n      <td>Barker</td>\n      <td>Kasen Barker</td>\n      <td>an image of Kasen Barker</td>\n      <td>a photo of Kasen Barker</td>\n      <td>Kasen Barker on a photo</td>\n      <td>a photo of a person named Kasen Barker</td>\n      <td>a person named Kasen Barker</td>\n      <td>a man named Kasen Barker</td>\n      <td>...</td>\n      <td>a photo of the celebrity Kasen Barker</td>\n      <td>actor Kasen Barker</td>\n      <td>actress Kasen Barker</td>\n      <td>a colored photo of Kasen Barker</td>\n      <td>a black and white photo of Kasen Barker</td>\n      <td>a cool photo of Kasen Barker</td>\n      <td>a cropped photo of Kasen Barker</td>\n      <td>a cropped image of Kasen Barker</td>\n      <td>Kasen Barker in a suit</td>\n      <td>Kasen Barker in a dress</td>\n    </tr>\n    <tr>\n      <th>9996</th>\n      <td>Camryn</td>\n      <td>m</td>\n      <td>Roberts</td>\n      <td>Camryn Roberts</td>\n      <td>an image of Camryn Roberts</td>\n      <td>a photo of Camryn Roberts</td>\n      <td>Camryn Roberts on a photo</td>\n      <td>a photo of a person named Camryn Roberts</td>\n      <td>a person named Camryn Roberts</td>\n      <td>a man named Camryn Roberts</td>\n      <td>...</td>\n      <td>a photo of the celebrity Camryn Roberts</td>\n      <td>actor Camryn Roberts</td>\n      <td>actress Camryn Roberts</td>\n      <td>a colored photo of Camryn Roberts</td>\n      <td>a black and white photo of Camryn Roberts</td>\n      <td>a cool photo of Camryn Roberts</td>\n      <td>a cropped photo of Camryn Roberts</td>\n      <td>a cropped image of Camryn Roberts</td>\n      <td>Camryn Roberts in a suit</td>\n      <td>Camryn Roberts in a dress</td>\n    </tr>\n    <tr>\n      <th>9997</th>\n      <td>Henry</td>\n      <td>m</td>\n      <td>Whitaker</td>\n      <td>Henry Whitaker</td>\n      <td>an image of Henry Whitaker</td>\n      <td>a photo of Henry Whitaker</td>\n      <td>Henry Whitaker on a photo</td>\n      <td>a photo of a person named Henry Whitaker</td>\n      <td>a person named Henry Whitaker</td>\n      <td>a man named Henry Whitaker</td>\n      <td>...</td>\n      <td>a photo of the celebrity Henry Whitaker</td>\n      <td>actor Henry Whitaker</td>\n      <td>actress Henry Whitaker</td>\n      <td>a colored photo of Henry Whitaker</td>\n      <td>a black and white photo of Henry Whitaker</td>\n      <td>a cool photo of Henry Whitaker</td>\n      <td>a cropped photo of Henry Whitaker</td>\n      <td>a cropped image of Henry Whitaker</td>\n      <td>Henry Whitaker in a suit</td>\n      <td>Henry Whitaker in a dress</td>\n    </tr>\n    <tr>\n      <th>9998</th>\n      <td>Adin</td>\n      <td>m</td>\n      <td>Richards</td>\n      <td>Adin Richards</td>\n      <td>an image of Adin Richards</td>\n      <td>a photo of Adin Richards</td>\n      <td>Adin Richards on a photo</td>\n      <td>a photo of a person named Adin Richards</td>\n      <td>a person named Adin Richards</td>\n      <td>a man named Adin Richards</td>\n      <td>...</td>\n      <td>a photo of the celebrity Adin Richards</td>\n      <td>actor Adin Richards</td>\n      <td>actress Adin Richards</td>\n      <td>a colored photo of Adin Richards</td>\n      <td>a black and white photo of Adin Richards</td>\n      <td>a cool photo of Adin Richards</td>\n      <td>a cropped photo of Adin Richards</td>\n      <td>a cropped image of Adin Richards</td>\n      <td>Adin Richards in a suit</td>\n      <td>Adin Richards in a dress</td>\n    </tr>\n    <tr>\n      <th>9999</th>\n      <td>Charley</td>\n      <td>m</td>\n      <td>Herman</td>\n      <td>Charley Herman</td>\n      <td>an image of Charley Herman</td>\n      <td>a photo of Charley Herman</td>\n      <td>Charley Herman on a photo</td>\n      <td>a photo of a person named Charley Herman</td>\n      <td>a person named Charley Herman</td>\n      <td>a man named Charley Herman</td>\n      <td>...</td>\n      <td>a photo of the celebrity Charley Herman</td>\n      <td>actor Charley Herman</td>\n      <td>actress Charley Herman</td>\n      <td>a colored photo of Charley Herman</td>\n      <td>a black and white photo of Charley Herman</td>\n      <td>a cool photo of Charley Herman</td>\n      <td>a cropped photo of Charley Herman</td>\n      <td>a cropped image of Charley Herman</td>\n      <td>Charley Herman in a suit</td>\n      <td>Charley Herman in a dress</td>\n    </tr>\n  </tbody>\n</table>\n<p>10000 rows × 24 columns</p>\n</div>"
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# populate the prompts with the possible names\n",
    "prompts = []\n",
    "for idx, row in possible_names.iterrows():\n",
    "    df_dict = row.to_dict()\n",
    "    name = f'{row[\"first_name\"]} {row[\"last_name\"]}'\n",
    "    for prompt_idx, prompt in enumerate(PROMPTS):\n",
    "        df_dict[f'prompt_{prompt_idx}'] = prompt.format(name)\n",
    "    prompts.append(df_dict)\n",
    "prompts = pd.DataFrame(prompts)\n",
    "prompts"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "outputs": [],
   "source": [
    "label_context_vecs = []\n",
    "for i in range(len(PROMPTS)):\n",
    "    context = open_clip.tokenize(prompts[f'prompt_{i}'].to_numpy())\n",
    "    label_context_vecs.append(context)\n",
    "label_context_vecs = torch.stack(label_context_vecs)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "outputs": [
    {
     "data": {
      "text/plain": "Calculating Text Embeddings:   0%|          | 0/210 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "4267d43b498f481db5cbf7e709c9ace3"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "Calculating Text Embeddings:   0%|          | 0/210 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "34a21714ab4d42b2beaa3024bcdd8fdd"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "Calculating Text Embeddings:   0%|          | 0/210 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "3278ad478d7d455da8b03d954fbc4558"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "label_context_vecs = label_context_vecs.to(device)\n",
    "\n",
    "text_embeddings_per_model = {}\n",
    "for model_name, model in models.items():\n",
    "    model = model.to(device)\n",
    "    text_embeddings = get_text_embeddings(model, label_context_vecs, use_tqdm=True, context_batchsize=1_000)\n",
    "    text_embeddings_per_model[model_name] = text_embeddings\n",
    "    model = model.cpu()\n",
    "\n",
    "label_context_vecs = label_context_vecs.cpu()"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "outputs": [],
   "source": [
    "# save the calculated embeddings to a file\n",
    "if not os.path.exists('./prompt_text_embeddings'):\n",
    "    os.makedirs('./prompt_text_embeddings')"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "outputs": [],
   "source": [
    "for model_name, _ in models.items():\n",
    "    torch.save(\n",
    "        text_embeddings_per_model[model_name],\n",
    "        f'./prompt_text_embeddings/{model_name}_prompt_text_embeddings.pt'\n",
    "    )"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}