diff --git a/notebooks/03_eval.ipynb b/notebooks/03_eval.ipynb
index 1f50138..3284485 100644
--- a/notebooks/03_eval.ipynb
+++ b/notebooks/03_eval.ipynb
@@ -505,30 +505,6 @@
     "[captions_stack[idx] for idx in sorted_idxs_list[0]][:4]"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": 21,
-   "id": "1d150544",
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "['X-ray luminous galaxy clusters, eMACS clusters, Balmer Break Galaxies, Lyman-break galaxies, gravitational telescopes; identify most powerful gravitational telescopes for future studies, constrain mass distribution within extreme systems, improve understanding of galaxy-galaxy and galaxy-gas interactions, unveil Balmer Break Galaxies at z~2, detect Lyman-break galaxies at z>6',\n",
-       " 'gravitationally lensed galaxies, Lyman Break Galaxies, star-forming galaxies, UV-bright galaxies, lensing galaxies; construct detailed lensing models, probe mass and light profiles of lensing galaxies, constrain spectral energy distributions and star formation histories of lensed galaxies, study properties of high redshift galaxies, analyze morphologies of lensed galaxies',\n",
-       " 'gravitationally lensed galaxies, Lyman Break Galaxies, star-forming galaxies, UV-bright galaxies, lensing galaxies; construct detailed lensing models, probe mass and light profiles of lensing galaxies, constrain spectral energy distributions and star formation histories of lensed galaxies, study properties of high redshift galaxies, analyze morphologies of lensed galaxies',\n",
-       " 'galaxy clusters, lensed galaxies, dark matter, dark energy, supernovae; confront LambdaCDM model predictions for cluster mass profiles, create detailed maps of cluster internal structure, discover and analyze 1000 new multiply-imaged lensed galaxies to AB=26, test constancy of dark energy with time using supernovae at z>1.5, probe progenitor evolution with supernovae']"
-      ]
-     },
-     "execution_count": 21,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "[captions_stack[idx] for idx in sorted_idxs_list[-1]][:4]"
-   ]
-  },
   {
    "cell_type": "markdown",
    "id": "91b37ce4-8636-430c-b150-43a7b517c166",
@@ -539,7 +515,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 29,
+   "execution_count": 39,
    "id": "c7d0e94c-c4ba-4f3a-a919-7c0f922161ea",
    "metadata": {},
    "outputs": [],
@@ -549,7 +525,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 30,
+   "execution_count": 40,
    "id": "8b29c296-2a3c-4a08-8113-09618e5bee0e",
    "metadata": {},
    "outputs": [
@@ -564,7 +540,7 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "100%|██████████| 100/100 [03:51<00:00,  2.31s/it]\n"
+      " 31%|███       | 31/100 [01:30<03:17,  2.87s/it]"
      ]
     }
    ],
@@ -659,6 +635,67 @@
     "# plt.savefig(\"../paper/plots/tti_base.pdf\", dpi=500)"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8d052942",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "from matplotlib.transforms import blended_transform_factory\n",
+    "from matplotlib import rc\n",
+    "import matplotlib\n",
+    "\n",
+    "rc(\"text\", usetex=True)\n",
+    "rc(\"text.latex\", preamble=r\"\\usepackage{color}\")\n",
+    "rc(\"text.latex\", preamble=r\"\\usepackage{xcolor}\")\n",
+    "\n",
+    "# Assuming 'tokenizer', 'model', 'image_feat_stack', 'config', and 'cosine_similarity' are already defined\n",
+    "\n",
+    "def get_image_from_stack(images_stack, idx):\n",
+    "    return images_stack[idx // config.training.batch_size][idx % config.training.batch_size]\n",
+    "\n",
+    "# List of text queries\n",
+    "# queries = [\"globular clusters\", \"dwarf galaxy\", \"SN1987A\", \"cluster lensing\"] \n",
+    "queries = [\"Jupiter\", \"dwarf galaxy\", \"SN1987A\", \"strong lensing\"] \n",
+    "\n",
+    "# Number of columns for the images\n",
+    "m_cols = 4  # Change this as needed\n",
+    "\n",
+    "sorted_idxs_list = []\n",
+    "\n",
+    "save_dir = \"../paper/plots/data/\"\n",
+    "\n",
+    "# Process each query\n",
+    "for row, query in enumerate(queries):\n",
+    "    query_tokenized = tokenizer([query], padding=\"max_length\", truncation=True, max_length=77, return_tensors=\"np\")\n",
+    "    text_feat = model.get_text_features(**query_tokenized, params=restored_state.params)\n",
+    "    sorted_idxs = np.argsort(jax.vmap(cosine_similarity, in_axes=(0, None))(np.vstack(image_base_feat_stack), text_feat[0]))[::-1]\n",
+    "\n",
+    "    sorted_idxs_list.append(sorted_idxs)\n",
+    "\n",
+    "    with open(save_dir + f\"/query_tti_base_{row}.txt\", 'w') as file:\n",
+    "        file.write(query)\n",
+    "\n",
+    "    # Loop through the images for this query and plot them\n",
+    "    for col in range(m_cols):\n",
+    "\n",
+    "        idx = sorted_idxs[col]\n",
+    "\n",
+    "        plt.figure(figsize=(8,8))\n",
+    "        plt.imshow(get_image_from_stack(images_base_stack, idx))\n",
+    "        plt.axis('off')\n",
+    "        plt.tight_layout()\n",
+    "        plt.savefig(save_dir + f\"/img_tti_base_{row}_{col}.pdf\")\n",
+    "        plt.close()\n",
+    "\n",
+    "        with open(save_dir + f\"/propid_tti_base_{row}_{col}.txt\", 'w') as file:\n",
+    "            idy = str(get_propid_phenomena(captions_base_stack[idx]))\n",
+    "            url = f\"\\\\href{{https://archive.stsci.edu/proposal_search.php?id={idy}&mission=hst}}{{{idy}}}\"\n",
+    "            file.write(url)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "f9b5cea6-c2b2-4684-b29f-69877fd290eb",
diff --git a/paper/plots/data/img_tti_base_0_0.pdf b/paper/plots/data/img_tti_base_0_0.pdf
new file mode 100644
index 0000000..c6e045c
Binary files /dev/null and b/paper/plots/data/img_tti_base_0_0.pdf differ
diff --git a/paper/plots/data/img_tti_base_0_1.pdf b/paper/plots/data/img_tti_base_0_1.pdf
new file mode 100644
index 0000000..f986178
Binary files /dev/null and b/paper/plots/data/img_tti_base_0_1.pdf differ
diff --git a/paper/plots/data/img_tti_base_0_2.pdf b/paper/plots/data/img_tti_base_0_2.pdf
new file mode 100644
index 0000000..ebc14e0
Binary files /dev/null and b/paper/plots/data/img_tti_base_0_2.pdf differ
diff --git a/paper/plots/data/img_tti_base_0_3.pdf b/paper/plots/data/img_tti_base_0_3.pdf
new file mode 100644
index 0000000..17e1143
Binary files /dev/null and b/paper/plots/data/img_tti_base_0_3.pdf differ
diff --git a/paper/plots/data/img_tti_base_1_0.pdf b/paper/plots/data/img_tti_base_1_0.pdf
new file mode 100644
index 0000000..082ee61
Binary files /dev/null and b/paper/plots/data/img_tti_base_1_0.pdf differ
diff --git a/paper/plots/data/img_tti_base_1_1.pdf b/paper/plots/data/img_tti_base_1_1.pdf
new file mode 100644
index 0000000..ce81544
Binary files /dev/null and b/paper/plots/data/img_tti_base_1_1.pdf differ
diff --git a/paper/plots/data/img_tti_base_1_2.pdf b/paper/plots/data/img_tti_base_1_2.pdf
new file mode 100644
index 0000000..5a38cd8
Binary files /dev/null and b/paper/plots/data/img_tti_base_1_2.pdf differ
diff --git a/paper/plots/data/img_tti_base_1_3.pdf b/paper/plots/data/img_tti_base_1_3.pdf
new file mode 100644
index 0000000..8c56939
Binary files /dev/null and b/paper/plots/data/img_tti_base_1_3.pdf differ
diff --git a/paper/plots/data/img_tti_base_2_0.pdf b/paper/plots/data/img_tti_base_2_0.pdf
new file mode 100644
index 0000000..dbf8315
Binary files /dev/null and b/paper/plots/data/img_tti_base_2_0.pdf differ
diff --git a/paper/plots/data/img_tti_base_2_1.pdf b/paper/plots/data/img_tti_base_2_1.pdf
new file mode 100644
index 0000000..16f606b
Binary files /dev/null and b/paper/plots/data/img_tti_base_2_1.pdf differ
diff --git a/paper/plots/data/img_tti_base_2_2.pdf b/paper/plots/data/img_tti_base_2_2.pdf
new file mode 100644
index 0000000..56c3880
Binary files /dev/null and b/paper/plots/data/img_tti_base_2_2.pdf differ
diff --git a/paper/plots/data/img_tti_base_2_3.pdf b/paper/plots/data/img_tti_base_2_3.pdf
new file mode 100644
index 0000000..0bc4a55
Binary files /dev/null and b/paper/plots/data/img_tti_base_2_3.pdf differ
diff --git a/paper/plots/data/img_tti_base_3_0.pdf b/paper/plots/data/img_tti_base_3_0.pdf
new file mode 100644
index 0000000..b84170d
Binary files /dev/null and b/paper/plots/data/img_tti_base_3_0.pdf differ
diff --git a/paper/plots/data/img_tti_base_3_1.pdf b/paper/plots/data/img_tti_base_3_1.pdf
new file mode 100644
index 0000000..542b03b
Binary files /dev/null and b/paper/plots/data/img_tti_base_3_1.pdf differ
diff --git a/paper/plots/data/img_tti_base_3_2.pdf b/paper/plots/data/img_tti_base_3_2.pdf
new file mode 100644
index 0000000..33d6962
Binary files /dev/null and b/paper/plots/data/img_tti_base_3_2.pdf differ
diff --git a/paper/plots/data/img_tti_base_3_3.pdf b/paper/plots/data/img_tti_base_3_3.pdf
new file mode 100644
index 0000000..cca485f
Binary files /dev/null and b/paper/plots/data/img_tti_base_3_3.pdf differ
diff --git a/paper/plots/data/propid_tti_base_0_0.txt b/paper/plots/data/propid_tti_base_0_0.txt
new file mode 100644
index 0000000..4ba4d5f
--- /dev/null
+++ b/paper/plots/data/propid_tti_base_0_0.txt
@@ -0,0 +1 @@
+\href{https://archive.stsci.edu/proposal_search.php?id=16293&mission=hst}{16293}
\ No newline at end of file
diff --git a/paper/plots/data/propid_tti_base_0_1.txt b/paper/plots/data/propid_tti_base_0_1.txt
new file mode 100644
index 0000000..cfda18e
--- /dev/null
+++ b/paper/plots/data/propid_tti_base_0_1.txt
@@ -0,0 +1 @@
+\href{https://archive.stsci.edu/proposal_search.php?id=6753&mission=hst}{6753}
\ No newline at end of file
diff --git a/paper/plots/data/propid_tti_base_0_2.txt b/paper/plots/data/propid_tti_base_0_2.txt
new file mode 100644
index 0000000..09443d4
--- /dev/null
+++ b/paper/plots/data/propid_tti_base_0_2.txt
@@ -0,0 +1 @@
+\href{https://archive.stsci.edu/proposal_search.php?id=13472&mission=hst}{13472}
\ No newline at end of file
diff --git a/paper/plots/data/propid_tti_base_0_3.txt b/paper/plots/data/propid_tti_base_0_3.txt
new file mode 100644
index 0000000..b317190
--- /dev/null
+++ b/paper/plots/data/propid_tti_base_0_3.txt
@@ -0,0 +1 @@
+\href{https://archive.stsci.edu/proposal_search.php?id=16264&mission=hst}{16264}
\ No newline at end of file
diff --git a/paper/plots/data/propid_tti_base_1_0.txt b/paper/plots/data/propid_tti_base_1_0.txt
new file mode 100644
index 0000000..a9b84ed
--- /dev/null
+++ b/paper/plots/data/propid_tti_base_1_0.txt
@@ -0,0 +1 @@
+\href{https://archive.stsci.edu/proposal_search.php?id=16843&mission=hst}{16843}
\ No newline at end of file
diff --git a/paper/plots/data/propid_tti_base_1_1.txt b/paper/plots/data/propid_tti_base_1_1.txt
new file mode 100644
index 0000000..bcb6952
--- /dev/null
+++ b/paper/plots/data/propid_tti_base_1_1.txt
@@ -0,0 +1 @@
+\href{https://archive.stsci.edu/proposal_search.php?id=17128&mission=hst}{17128}
\ No newline at end of file
diff --git a/paper/plots/data/propid_tti_base_1_2.txt b/paper/plots/data/propid_tti_base_1_2.txt
new file mode 100644
index 0000000..d0bc4a6
--- /dev/null
+++ b/paper/plots/data/propid_tti_base_1_2.txt
@@ -0,0 +1 @@
+\href{https://archive.stsci.edu/proposal_search.php?id=12196&mission=hst}{12196}
\ No newline at end of file
diff --git a/paper/plots/data/propid_tti_base_1_3.txt b/paper/plots/data/propid_tti_base_1_3.txt
new file mode 100644
index 0000000..b9c9126
--- /dev/null
+++ b/paper/plots/data/propid_tti_base_1_3.txt
@@ -0,0 +1 @@
+\href{https://archive.stsci.edu/proposal_search.php?id=14160&mission=hst}{14160}
\ No newline at end of file
diff --git a/paper/plots/data/propid_tti_base_2_0.txt b/paper/plots/data/propid_tti_base_2_0.txt
new file mode 100644
index 0000000..a9b84ed
--- /dev/null
+++ b/paper/plots/data/propid_tti_base_2_0.txt
@@ -0,0 +1 @@
+\href{https://archive.stsci.edu/proposal_search.php?id=16843&mission=hst}{16843}
\ No newline at end of file
diff --git a/paper/plots/data/propid_tti_base_2_1.txt b/paper/plots/data/propid_tti_base_2_1.txt
new file mode 100644
index 0000000..f89e500
--- /dev/null
+++ b/paper/plots/data/propid_tti_base_2_1.txt
@@ -0,0 +1 @@
+\href{https://archive.stsci.edu/proposal_search.php?id=5963&mission=hst}{5963}
\ No newline at end of file
diff --git a/paper/plots/data/propid_tti_base_2_2.txt b/paper/plots/data/propid_tti_base_2_2.txt
new file mode 100644
index 0000000..26ffcdd
--- /dev/null
+++ b/paper/plots/data/propid_tti_base_2_2.txt
@@ -0,0 +1 @@
+\href{https://archive.stsci.edu/proposal_search.php?id=9709&mission=hst}{9709}
\ No newline at end of file
diff --git a/paper/plots/data/propid_tti_base_2_3.txt b/paper/plots/data/propid_tti_base_2_3.txt
new file mode 100644
index 0000000..bcb6952
--- /dev/null
+++ b/paper/plots/data/propid_tti_base_2_3.txt
@@ -0,0 +1 @@
+\href{https://archive.stsci.edu/proposal_search.php?id=17128&mission=hst}{17128}
\ No newline at end of file
diff --git a/paper/plots/data/propid_tti_base_3_0.txt b/paper/plots/data/propid_tti_base_3_0.txt
new file mode 100644
index 0000000..51b2083
--- /dev/null
+++ b/paper/plots/data/propid_tti_base_3_0.txt
@@ -0,0 +1 @@
+\href{https://archive.stsci.edu/proposal_search.php?id=12945&mission=hst}{12945}
\ No newline at end of file
diff --git a/paper/plots/data/propid_tti_base_3_1.txt b/paper/plots/data/propid_tti_base_3_1.txt
new file mode 100644
index 0000000..21e64e2
--- /dev/null
+++ b/paper/plots/data/propid_tti_base_3_1.txt
@@ -0,0 +1 @@
+\href{https://archive.stsci.edu/proposal_search.php?id=11096&mission=hst}{11096}
\ No newline at end of file
diff --git a/paper/plots/data/propid_tti_base_3_2.txt b/paper/plots/data/propid_tti_base_3_2.txt
new file mode 100644
index 0000000..2418a29
--- /dev/null
+++ b/paper/plots/data/propid_tti_base_3_2.txt
@@ -0,0 +1 @@
+\href{https://archive.stsci.edu/proposal_search.php?id=16684&mission=hst}{16684}
\ No newline at end of file
diff --git a/paper/plots/data/propid_tti_base_3_3.txt b/paper/plots/data/propid_tti_base_3_3.txt
new file mode 100644
index 0000000..51b2083
--- /dev/null
+++ b/paper/plots/data/propid_tti_base_3_3.txt
@@ -0,0 +1 @@
+\href{https://archive.stsci.edu/proposal_search.php?id=12945&mission=hst}{12945}
\ No newline at end of file
diff --git a/paper/plots/data/query_tti_base_0.txt b/paper/plots/data/query_tti_base_0.txt
new file mode 100644
index 0000000..ce1d052
--- /dev/null
+++ b/paper/plots/data/query_tti_base_0.txt
@@ -0,0 +1 @@
+Jupiter
\ No newline at end of file
diff --git a/paper/plots/data/query_tti_base_1.txt b/paper/plots/data/query_tti_base_1.txt
new file mode 100644
index 0000000..cdd8a9a
--- /dev/null
+++ b/paper/plots/data/query_tti_base_1.txt
@@ -0,0 +1 @@
+dwarf galaxy
\ No newline at end of file
diff --git a/paper/plots/data/query_tti_base_2.txt b/paper/plots/data/query_tti_base_2.txt
new file mode 100644
index 0000000..f71fa1a
--- /dev/null
+++ b/paper/plots/data/query_tti_base_2.txt
@@ -0,0 +1 @@
+SN1987A
\ No newline at end of file
diff --git a/paper/plots/data/query_tti_base_3.txt b/paper/plots/data/query_tti_base_3.txt
new file mode 100644
index 0000000..0812d9a
--- /dev/null
+++ b/paper/plots/data/query_tti_base_3.txt
@@ -0,0 +1 @@
+strong lensing
\ No newline at end of file