added additional models

app.py

@@ -39,7 +39,9 @@ PROMPTS = [
     '{0} in a suit',
     '{0} in a dress'
 ]
-OPEN_CLIP_MODEL_NAMES = ['ViT-B-32', 'ViT-B-16', 'ViT-L-14']
+OPEN_CLIP_LAION400M_MODEL_NAMES = ['ViT-B-32', 'ViT-B-16', 'ViT-L-14']
+OPEN_CLIP_LAION2B_MODEL_NAMES = [('ViT-B-32', 'laion2b_s34b_b79k'), ('ViT-L-14', 'laion2b_s32b_b82k')]
+OPEN_AI_MODELS = ['ViT-B-32', 'ViT-B-16', 'ViT-L-14']
 NUM_TOTAL_NAMES = 1_000
 SEED = 42
 MIN_NUM_CORRECT_PROMPT_PREDS = 1
@@ -52,7 +54,7 @@ EXAMPLE_IMAGE_URLS = read_actor_files(EDAMPLE_IMAGE_DIR)
 save_images_to_folder(os.path.join(EDAMPLE_IMAGE_DIR, 'images'), EXAMPLE_IMAGE_URLS)
 
 MODELS = {}
-for model_name in OPEN_CLIP_MODEL_NAMES:
+for model_name in OPEN_CLIP_LAION400M_MODEL_NAMES:
     dataset = 'LAION400M'
     model, _, preprocess = open_clip.create_model_and_transforms(
         model_name,
@@ -63,24 +65,55 @@ for model_name in OPEN_CLIP_MODEL_NAMES:
         'model_instance': model,
         'preprocessing': preprocess,
         'model_name': model_name,
-        'prompt_text_embeddings': torch.load(f'./prompt_text_embeddings/{model_name}_prompt_text_embeddings.pt')
+        'tokenizer': open_clip.get_tokenizer(model_name),
+        'prompt_text_embeddings': torch.load(f'./prompt_text_embeddings/{model_name}_{dataset.lower()}_prompt_text_embeddings.pt')
+    }
+
+for model_name, dataset_name in OPEN_CLIP_LAION2B_MODEL_NAMES:
+    dataset = 'LAION2B'
+    model, _, preprocess = open_clip.create_model_and_transforms(
+        model_name,
+        pretrained=dataset_name
+    )
+    model = model.eval()
+    MODELS[f'OpenClip {model_name} trained on {dataset}'] = {
+        'model_instance': model,
+        'preprocessing': preprocess,
+        'model_name': model_name,
+        'tokenizer': open_clip.get_tokenizer(model_name),
+        'prompt_text_embeddings': torch.load(f'./prompt_text_embeddings/{model_name}_{dataset.lower()}_prompt_text_embeddings.pt')
+    }
+
+for model_name in OPEN_AI_MODELS:
+    dataset = 'OpenAI'
+    model, _, preprocess = open_clip.create_model_and_transforms(
+        model_name,
+        pretrained=dataset.lower()
+    )
+    model = model.eval()
+    MODELS[f'OpenClip {model_name} trained by {dataset}'] = {
+        'model_instance': model,
+        'preprocessing': preprocess,
+        'model_name': model_name,
+        'tokenizer': open_clip.get_tokenizer(model_name),
+        'prompt_text_embeddings': torch.load(f'./prompt_text_embeddings/{model_name}_{dataset.lower()}_prompt_text_embeddings.pt')
     }
 
 FULL_NAMES_DF = pd.read_csv('full_names.csv', index_col=0)
 LAION_MEMBERSHIP_OCCURENCE = pd.read_csv('laion_membership_occurence_count.csv', index_col=0)
 
 EXAMPLE_ACTORS_BY_MODEL = {
-    "ViT-B-32": ["T._J._Thyne"],
-    "ViT-B-16": ["Barbara_Schöneberger", "Carolin_Kebekus"],
-    "ViT-L-14": ["Max_Giermann", "Nicole_De_Boer"]
+    ("ViT-B-32", "laion400m"): ["T._J._Thyne"],
+    ("ViT-B-16", "laion400m"): ["Barbara_Schöneberger", "Carolin_Kebekus"],
+    ("ViT-L-14", "laion400m"): ["Max_Giermann", "Nicole_De_Boer"]
 }
 
 EXAMPLES = []
-for model_name, person_names in EXAMPLE_ACTORS_BY_MODEL.items():
+for (model_name, dataset_name), person_names in EXAMPLE_ACTORS_BY_MODEL.items():
     for name in person_names:
         image_folder = os.path.join("./example_images/images/", name)
         for dd_model_name in MODELS.keys():
-            if model_name not in dd_model_name:
+            if not (model_name.lower() in dd_model_name.lower() and dataset_name.lower() in dd_model_name.lower()):
                 continue
 
             EXAMPLES.append([
@@ -139,7 +172,7 @@ CSS = """
             transform: translateY(10px);
             background: white;
         }
-        
+
         .dark .footer {
             border-color: #303030;
         }
@@ -221,8 +254,8 @@ gr.Files.preprocess = preprocess
 
 @torch.no_grad()
 def calculate_text_embeddings(model_name, prompts):
-    tokenizer = open_clip.get_tokenizer(MODELS[model_name]['model_name'])
-    context_vecs = open_clip.tokenize(prompts)
+    tokenizer = MODELS[model_name]['tokenizer']
+    context_vecs = tokenizer(prompts)
 
     model_instance = MODELS[model_name]['model_instance']
 
@@ -509,7 +542,8 @@ with block as demo:
             with gr.Column():
                 model_dd = gr.Dropdown(label="CLIP Model", choices=list(MODELS.keys()),
                                        value=list(MODELS.keys())[0])
-                true_name = gr.Textbox(label='Name of Person:', lines=1, value=DEFAULT_INITIAL_NAME)
+                true_name = gr.Textbox(label='Name of Person (make sure it matches the prompts):', lines=1, value=DEFAULT_INITIAL_NAME,
+                                       every=5)
                 prompts = gr.Dataframe(
                     value=[[x.format(DEFAULT_INITIAL_NAME) for x in PROMPTS]],
                     label='Prompts Used (hold shift to scroll sideways):',

calculate_text_embeddings.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": null,
    "metadata": {
     "collapsed": true
    },
@@ -39,33 +39,70 @@
     "    '{0} in a suit',\n",
     "    '{0} in a dress'\n",
     "]\n",
-    "MODEL_NAMES = ['ViT-B-32', 'ViT-B-16', 'ViT-L-14']\n",
+    "OPEN_CLIP_LAION400M_MODEL_NAMES = ['ViT-B-32', 'ViT-B-16', 'ViT-L-14']\n",
+    "OPEN_CLIP_LAION2B_MODEL_NAMES = [('ViT-B-32', 'laion2b_s34b_b79k') , ('ViT-L-14', 'laion2b_s32b_b82k')]\n",
+    "OPEN_AI_MODELS = ['ViT-B-32', 'ViT-B-16', 'ViT-L-14']\n",
     "SEED = 42"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
+   "metadata": {
+    "collapsed": false
+   },
    "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",
+    "MODELS = {}\n",
+    "for model_name in OPEN_CLIP_LAION400M_MODEL_NAMES:\n",
+    "    dataset = 'LAION400M'\n",
+    "    model, _, preprocess = open_clip.create_model_and_transforms(\n",
+    "        model_name,\n",
+    "        pretrained=f'{dataset.lower()}_e32'\n",
+    "    )\n",
     "    model = model.eval()\n",
-    "    models[model_name] = model\n",
-    "    tokenizers[model_name] = open_clip.get_tokenizer(model_name)"
-   ],
-   "metadata": {
-    "collapsed": false
-   }
+    "    MODELS[(model_name, dataset.lower())] = {\n",
+    "        'model_instance': model,\n",
+    "        'preprocessing': preprocess,\n",
+    "        'model_name': model_name,\n",
+    "        'tokenizer': open_clip.get_tokenizer(model_name),\n",
+    "    }\n",
+    "\n",
+    "for model_name, dataset_name in OPEN_CLIP_LAION2B_MODEL_NAMES:\n",
+    "    dataset = 'LAION2B'\n",
+    "    model, _, preprocess = open_clip.create_model_and_transforms(\n",
+    "        model_name,\n",
+    "        pretrained = dataset_name\n",
+    "    )\n",
+    "    model = model.eval()\n",
+    "    MODELS[(model_name, dataset.lower())] = {\n",
+    "        'model_instance': model,\n",
+    "        'preprocessing': preprocess,\n",
+    "        'model_name': model_name,\n",
+    "        'tokenizer': open_clip.get_tokenizer(model_name)\n",
+    "    }\n",
+    "\n",
+    "for model_name in OPEN_AI_MODELS:\n",
+    "    dataset = 'OpenAI'\n",
+    "    model, _, preprocess = open_clip.create_model_and_transforms(\n",
+    "        model_name,\n",
+    "        pretrained=dataset.lower()\n",
+    "    )\n",
+    "    model = model.eval()\n",
+    "    MODELS[(model_name, dataset.lower())] = {\n",
+    "        'model_instance': model,\n",
+    "        'preprocessing': preprocess,\n",
+    "        'model_name': model_name,\n",
+    "        'tokenizer': open_clip.get_tokenizer(model_name)\n",
+    "    }"
+   ]
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
+   "metadata": {
+    "collapsed": false
+   },
    "outputs": [],
    "source": [
     "# define a function to get the predictions for an actor/actress\n",
@@ -90,50 +127,30 @@
     "    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"
-    }
-   ],
+   "execution_count": null,
+   "metadata": {
+    "collapsed": false
+   },
+   "outputs": [],
    "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"
-    }
-   ],
+   "execution_count": null,
+   "metadata": {
+    "collapsed": false
+   },
+   "outputs": [],
    "source": [
     "# populate the prompts with the possible names\n",
     "prompts = []\n",
@@ -145,119 +162,83 @@
     "    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)"
-   ],
+   "execution_count": null,
    "metadata": {
     "collapsed": false
-   }
+   },
+   "outputs": [],
+   "source": [
+    "label_context_vecs_per_model = {}\n",
+    "for dict_key, model_dict in MODELS.items():\n",
+    "    label_context_vecs = []\n",
+    "    for i in range(len(PROMPTS)):\n",
+    "        context = model_dict['tokenizer'](prompts[f'prompt_{i}'].to_numpy())\n",
+    "        label_context_vecs.append(context)\n",
+    "    label_context_vecs = torch.stack(label_context_vecs)\n",
+    "    label_context_vecs_per_model[dict_key] = label_context_vecs"
+   ]
   },
   {
    "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"
-    }
-   ],
+   "execution_count": null,
+   "metadata": {
+    "collapsed": false
+   },
+   "outputs": [],
    "source": [
-    "label_context_vecs = label_context_vecs.to(device)\n",
-    "\n",
     "text_embeddings_per_model = {}\n",
-    "for model_name, model in models.items():\n",
+    "for dict_key, model_dict in MODELS.items():\n",
+    "    label_context_vecs = label_context_vecs_per_model[dict_key].to(device)\n",
+    "    model = model_dict['model_instance']\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",
+    "    text_embeddings = get_text_embeddings(model, label_context_vecs, use_tqdm=True, context_batchsize=5_000)\n",
+    "    text_embeddings_per_model[dict_key] = text_embeddings\n",
     "    model = model.cpu()\n",
+    "    label_context_vecs = label_context_vecs.cpu()\n",
     "\n",
     "label_context_vecs = label_context_vecs.cpu()"
-   ],
-   "metadata": {
-    "collapsed": false
-   }
+   ]
   },
   {
    "cell_type": "code",
-   "execution_count": 18,
+   "execution_count": null,
+   "metadata": {
+    "collapsed": false
+   },
    "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,
+   "execution_count": null,
+   "metadata": {
+    "collapsed": false
+   },
    "outputs": [],
    "source": [
-    "for model_name, _ in models.items():\n",
+    "for (model_name, dataset_name), model_dict 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",
+    "        text_embeddings_per_model[(model_name, dataset_name)],\n",
+    "        f'./prompt_text_embeddings/{model_name}_{dataset_name}_prompt_text_embeddings.pt'\n",
     "    )"
-   ],
-   "metadata": {
-    "collapsed": false
-   }
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "outputs": [],
-   "source": [],
    "metadata": {
     "collapsed": false
-   }
+   },
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {
@@ -269,14 +250,14 @@
   "language_info": {
    "codemirror_mode": {
     "name": "ipython",
-    "version": 2
+    "version": 3
    },
    "file_extension": ".py",
    "mimetype": "text/x-python",
    "name": "python",
    "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython2",
-   "version": "2.7.6"
+   "pygments_lexer": "ipython3",
+   "version": "3.8.13"
   }
  },
  "nbformat": 4,

prompt_text_embeddings/{ViT-B-16_prompt_text_embeddings.pt → ViT-B-16_laion400m_prompt_text_embeddings.pt}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:21a5aaf59ac44f37f2903c4b9bfc2cc1025757420de6aa500aafacbb9797f9ab
-size 430080883
+oid sha256:2eba829d60be9ec74485ad0ccdc6cd93c599bb8c0ed3036c099a19ab71fa251a
+size 430080977

prompt_text_embeddings/{ViT-B-32_prompt_text_embeddings.pt → ViT-B-16_openai_prompt_text_embeddings.pt}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:5c842c10a0982de17892a86dd9f00f5858eda62795d1fdf37904046ebc565685
-size 430080883
+oid sha256:c86948737c065233154162deb78c14bfd827eb731df087da082a64d2540f88b6
+size 430080968

prompt_text_embeddings/{ViT-L-14_prompt_text_embeddings.pt → ViT-B-32_laion2b_prompt_text_embeddings.pt}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:4fd82898cfd5d233aec8b3c644af5ddd71265eef1d6c63fdc4d1e6167d5450a9
-size 645120883
+oid sha256:fe5c5abd8f02ae34eb97ed192ec67e6345cb44df2c60c00bf71d1fe86d06f9d4
+size 430080971

prompt_text_embeddings/ViT-B-32_laion400m_prompt_text_embeddings.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:da2108f9b80d59975ad30ee72405a31f1d722a00cf22d54fe3523784e6706151
+size 430080977

prompt_text_embeddings/ViT-B-32_openai_prompt_text_embeddings.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7a576a46ca84d9794d4f4e82eed2146f104411462f364efce61274c585c5546c
+size 430080968

prompt_text_embeddings/ViT-L-14_laion2b_prompt_text_embeddings.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:62b2313b873c93cf18c08faa46a1a7088cf7c832abdd28ddaedd0e46624c693d
+size 645120971

prompt_text_embeddings/ViT-L-14_laion400m_prompt_text_embeddings.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ae5074916120756e59f8a38516f11a3a9c2c962843cff75de7947247a74c3ee6
+size 645120977

prompt_text_embeddings/ViT-L-14_openai_prompt_text_embeddings.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:66f9c693c897997cc38160eacae8ff6547312f026ae680431833c7b0898a9a44
+size 645120968