Created
June 25, 2023 03:37
-
-
Save rabernat/40f53bba3a81aeb420e14872388c6fc1 to your computer and use it in GitHub Desktop.
Xarray as an embedded vector database
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
{ | |
"cells": [ | |
{ | |
"attachments": {}, | |
"cell_type": "markdown", | |
"metadata": {}, | |
"source": [ | |
"# Xarray as a Vector Database\n", | |
"\n", | |
"This little experiment shows that it is straightforward to use Xarray as an embedded vector database.\n", | |
"\n", | |
"We will create a custom [Xarray Accessor](https://docs.xarray.dev/en/stable/internals/extending-xarray.html) extension called `vector` which performs nearest-neighbor search across a multidimensional vector space. This is meant to represent the embedding space of an LLM.\n" | |
] | |
}, | |
{ | |
"cell_type": "code", | |
"execution_count": 1, | |
"metadata": {}, | |
"outputs": [], | |
"source": [ | |
"import xarray as xr\n", | |
"import numpy as np\n", | |
"from scipy.spatial import KDTree\n", | |
"\n", | |
"@xr.register_dataset_accessor(\"vector\")\n", | |
"class VectorAccessor:\n", | |
" def __init__(self, xarray_obj):\n", | |
" self._obj = xarray_obj\n", | |
"\n", | |
" def set_index(self, vector_variable=\"vectors\", embedding_dim=\"embedding\"):\n", | |
" vector_var = self._obj[vector_variable]\n", | |
" assert vector_var.ndim == 2\n", | |
" assert embedding_dim in vector_var.dims\n", | |
" self._sample_dim = [dim for dim in vector_var.dims if dim != embedding_dim][0]\n", | |
" self._embedding_dim = embedding_dim\n", | |
" self._vectors = vector_var.transpose(self._sample_dim, self._embedding_dim)\n", | |
" self._tree = KDTree(self._vectors.values)\n", | |
"\n", | |
" def query(\n", | |
" self,\n", | |
" vectors,\n", | |
" k=1,\n", | |
" distance_var_name=\"distance\",\n", | |
" k_var_name=\"k\",\n", | |
" query_vector_var_name=\"query_vectors\",\n", | |
" query_dim_name=\"query\",\n", | |
" **kwargs\n", | |
" ):\n", | |
" vectors = np.asarray(vectors)\n", | |
" assert vectors.ndim == 2\n", | |
" nqueries, nd = vectors.shape\n", | |
" assert nd == self._obj.dims[self._embedding_dim]\n", | |
"\n", | |
" # the expensive part\n", | |
" dist, idx = self._tree.query(vectors, k=k, **kwargs)\n", | |
"\n", | |
" # figure out the indexing of the result\n", | |
" if isinstance(k, int):\n", | |
" kcoords = np.arange(1, k+1)\n", | |
" nk = k\n", | |
" if k==1:\n", | |
" idx = idx[:, None]\n", | |
" dist = dist[:, None]\n", | |
" else:\n", | |
" kcoords = k # assume a list-like object\n", | |
" nk = len(k)\n", | |
"\n", | |
" # make sure we understand the result shape\n", | |
" assert dist.shape == idx.shape\n", | |
" assert idx.shape == (nqueries, nk)\n", | |
"\n", | |
" # select the samples\n", | |
" result_dims = [query_dim_name, self._sample_dim]\n", | |
" indexer = xr.DataArray(idx, dims=result_dims)\n", | |
" result = self._obj.isel({self._sample_dim: indexer})\n", | |
"\n", | |
" # augment the result with the distance and k coordinates\n", | |
" result.coords[distance_var_name] = xr.DataArray(dist, dims=result_dims)\n", | |
" result.coords[k_var_name] = xr.DataArray(kcoords, dims=[self._sample_dim])\n", | |
" result.coords[query_vector_var_name] = xr.DataArray(vectors, dims=[query_dim_name, self._embedding_dim])\n", | |
" return result\n" | |
] | |
}, | |
{ | |
"attachments": {}, | |
"cell_type": "markdown", | |
"metadata": {}, | |
"source": [ | |
"\n", | |
"To use this, first we create some vectors as a numpy array.\n", | |
"These are three vectors in a 3D space.\n" | |
] | |
}, | |
{ | |
"cell_type": "code", | |
"execution_count": 2, | |
"metadata": {}, | |
"outputs": [], | |
"source": [ | |
"vectors = np.array(\n", | |
" [\n", | |
" [1, 0, 0],\n", | |
" [0, 1, 0],\n", | |
" [0, 0, 1],\n", | |
" ]\n", | |
")\n" | |
] | |
}, | |
{ | |
"attachments": {}, | |
"cell_type": "markdown", | |
"metadata": {}, | |
"source": [ | |
"Now we wrap these with an Xarray dataset. We define two dimensions:\n", | |
"\n", | |
"- `sample`: representing the distinct sample documents that have been embedded using the LLM. There is one vector for each.\n", | |
"- `embedding` the vector space dimension.\n", | |
"\n", | |
"We make `vectors` a coordinate variable.\n", | |
"We also define a data variable `documents` representing the actual data from which the embeddings were generated." | |
] | |
}, | |
{ | |
"cell_type": "code", | |
"execution_count": 3, | |
"metadata": {}, | |
"outputs": [ | |
{ | |
"data": { | |
"text/html": [ | |
"<div><svg style=\"position: absolute; width: 0; height: 0; overflow: hidden\">\n", | |
"<defs>\n", | |
"<symbol id=\"icon-database\" viewBox=\"0 0 32 32\">\n", | |
"<path d=\"M16 0c-8.837 0-16 2.239-16 5v4c0 2.761 7.163 5 16 5s16-2.239 16-5v-4c0-2.761-7.163-5-16-5z\"></path>\n", | |
"<path d=\"M16 17c-8.837 0-16-2.239-16-5v6c0 2.761 7.163 5 16 5s16-2.239 16-5v-6c0 2.761-7.163 5-16 5z\"></path>\n", | |
"<path d=\"M16 26c-8.837 0-16-2.239-16-5v6c0 2.761 7.163 5 16 5s16-2.239 16-5v-6c0 2.761-7.163 5-16 5z\"></path>\n", | |
"</symbol>\n", | |
"<symbol id=\"icon-file-text2\" viewBox=\"0 0 32 32\">\n", | |
"<path d=\"M28.681 7.159c-0.694-0.947-1.662-2.053-2.724-3.116s-2.169-2.030-3.116-2.724c-1.612-1.182-2.393-1.319-2.841-1.319h-15.5c-1.378 0-2.5 1.121-2.5 2.5v27c0 1.378 1.122 2.5 2.5 2.5h23c1.378 0 2.5-1.122 2.5-2.5v-19.5c0-0.448-0.137-1.23-1.319-2.841zM24.543 5.457c0.959 0.959 1.712 1.825 2.268 2.543h-4.811v-4.811c0.718 0.556 1.584 1.309 2.543 2.268zM28 29.5c0 0.271-0.229 0.5-0.5 0.5h-23c-0.271 0-0.5-0.229-0.5-0.5v-27c0-0.271 0.229-0.5 0.5-0.5 0 0 15.499-0 15.5 0v7c0 0.552 0.448 1 1 1h7v19.5z\"></path>\n", | |
"<path d=\"M23 26h-14c-0.552 0-1-0.448-1-1s0.448-1 1-1h14c0.552 0 1 0.448 1 1s-0.448 1-1 1z\"></path>\n", | |
"<path d=\"M23 22h-14c-0.552 0-1-0.448-1-1s0.448-1 1-1h14c0.552 0 1 0.448 1 1s-0.448 1-1 1z\"></path>\n", | |
"<path d=\"M23 18h-14c-0.552 0-1-0.448-1-1s0.448-1 1-1h14c0.552 0 1 0.448 1 1s-0.448 1-1 1z\"></path>\n", | |
"</symbol>\n", | |
"</defs>\n", | |
"</svg>\n", | |
"<style>/* CSS stylesheet for displaying xarray objects in jupyterlab.\n", | |
" *\n", | |
" */\n", | |
"\n", | |
":root {\n", | |
" --xr-font-color0: var(--jp-content-font-color0, rgba(0, 0, 0, 1));\n", | |
" --xr-font-color2: var(--jp-content-font-color2, rgba(0, 0, 0, 0.54));\n", | |
" --xr-font-color3: var(--jp-content-font-color3, rgba(0, 0, 0, 0.38));\n", | |
" --xr-border-color: var(--jp-border-color2, #e0e0e0);\n", | |
" --xr-disabled-color: var(--jp-layout-color3, #bdbdbd);\n", | |
" --xr-background-color: var(--jp-layout-color0, white);\n", | |
" --xr-background-color-row-even: var(--jp-layout-color1, white);\n", | |
" --xr-background-color-row-odd: var(--jp-layout-color2, #eeeeee);\n", | |
"}\n", | |
"\n", | |
"html[theme=dark],\n", | |
"body[data-theme=dark],\n", | |
"body.vscode-dark {\n", | |
" --xr-font-color0: rgba(255, 255, 255, 1);\n", | |
" --xr-font-color2: rgba(255, 255, 255, 0.54);\n", | |
" --xr-font-color3: rgba(255, 255, 255, 0.38);\n", | |
" --xr-border-color: #1F1F1F;\n", | |
" --xr-disabled-color: #515151;\n", | |
" --xr-background-color: #111111;\n", | |
" --xr-background-color-row-even: #111111;\n", | |
" --xr-background-color-row-odd: #313131;\n", | |
"}\n", | |
"\n", | |
".xr-wrap {\n", | |
" display: block !important;\n", | |
" min-width: 300px;\n", | |
" max-width: 700px;\n", | |
"}\n", | |
"\n", | |
".xr-text-repr-fallback {\n", | |
" /* fallback to plain text repr when CSS is not injected (untrusted notebook) */\n", | |
" display: none;\n", | |
"}\n", | |
"\n", | |
".xr-header {\n", | |
" padding-top: 6px;\n", | |
" padding-bottom: 6px;\n", | |
" margin-bottom: 4px;\n", | |
" border-bottom: solid 1px var(--xr-border-color);\n", | |
"}\n", | |
"\n", | |
".xr-header > div,\n", | |
".xr-header > ul {\n", | |
" display: inline;\n", | |
" margin-top: 0;\n", | |
" margin-bottom: 0;\n", | |
"}\n", | |
"\n", | |
".xr-obj-type,\n", | |
".xr-array-name {\n", | |
" margin-left: 2px;\n", | |
" margin-right: 10px;\n", | |
"}\n", | |
"\n", | |
".xr-obj-type {\n", | |
" color: var(--xr-font-color2);\n", | |
"}\n", | |
"\n", | |
".xr-sections {\n", | |
" padding-left: 0 !important;\n", | |
" display: grid;\n", | |
" grid-template-columns: 150px auto auto 1fr 20px 20px;\n", | |
"}\n", | |
"\n", | |
".xr-section-item {\n", | |
" display: contents;\n", | |
"}\n", | |
"\n", | |
".xr-section-item input {\n", | |
" display: none;\n", | |
"}\n", | |
"\n", | |
".xr-section-item input + label {\n", | |
" color: var(--xr-disabled-color);\n", | |
"}\n", | |
"\n", | |
".xr-section-item input:enabled + label {\n", | |
" cursor: pointer;\n", | |
" color: var(--xr-font-color2);\n", | |
"}\n", | |
"\n", | |
".xr-section-item input:enabled + label:hover {\n", | |
" color: var(--xr-font-color0);\n", | |
"}\n", | |
"\n", | |
".xr-section-summary {\n", | |
" grid-column: 1;\n", | |
" color: var(--xr-font-color2);\n", | |
" font-weight: 500;\n", | |
"}\n", | |
"\n", | |
".xr-section-summary > span {\n", | |
" display: inline-block;\n", | |
" padding-left: 0.5em;\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in:disabled + label {\n", | |
" color: var(--xr-font-color2);\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in + label:before {\n", | |
" display: inline-block;\n", | |
" content: '►';\n", | |
" font-size: 11px;\n", | |
" width: 15px;\n", | |
" text-align: center;\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in:disabled + label:before {\n", | |
" color: var(--xr-disabled-color);\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in:checked + label:before {\n", | |
" content: '▼';\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in:checked + label > span {\n", | |
" display: none;\n", | |
"}\n", | |
"\n", | |
".xr-section-summary,\n", | |
".xr-section-inline-details {\n", | |
" padding-top: 4px;\n", | |
" padding-bottom: 4px;\n", | |
"}\n", | |
"\n", | |
".xr-section-inline-details {\n", | |
" grid-column: 2 / -1;\n", | |
"}\n", | |
"\n", | |
".xr-section-details {\n", | |
" display: none;\n", | |
" grid-column: 1 / -1;\n", | |
" margin-bottom: 5px;\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in:checked ~ .xr-section-details {\n", | |
" display: contents;\n", | |
"}\n", | |
"\n", | |
".xr-array-wrap {\n", | |
" grid-column: 1 / -1;\n", | |
" display: grid;\n", | |
" grid-template-columns: 20px auto;\n", | |
"}\n", | |
"\n", | |
".xr-array-wrap > label {\n", | |
" grid-column: 1;\n", | |
" vertical-align: top;\n", | |
"}\n", | |
"\n", | |
".xr-preview {\n", | |
" color: var(--xr-font-color3);\n", | |
"}\n", | |
"\n", | |
".xr-array-preview,\n", | |
".xr-array-data {\n", | |
" padding: 0 5px !important;\n", | |
" grid-column: 2;\n", | |
"}\n", | |
"\n", | |
".xr-array-data,\n", | |
".xr-array-in:checked ~ .xr-array-preview {\n", | |
" display: none;\n", | |
"}\n", | |
"\n", | |
".xr-array-in:checked ~ .xr-array-data,\n", | |
".xr-array-preview {\n", | |
" display: inline-block;\n", | |
"}\n", | |
"\n", | |
".xr-dim-list {\n", | |
" display: inline-block !important;\n", | |
" list-style: none;\n", | |
" padding: 0 !important;\n", | |
" margin: 0;\n", | |
"}\n", | |
"\n", | |
".xr-dim-list li {\n", | |
" display: inline-block;\n", | |
" padding: 0;\n", | |
" margin: 0;\n", | |
"}\n", | |
"\n", | |
".xr-dim-list:before {\n", | |
" content: '(';\n", | |
"}\n", | |
"\n", | |
".xr-dim-list:after {\n", | |
" content: ')';\n", | |
"}\n", | |
"\n", | |
".xr-dim-list li:not(:last-child):after {\n", | |
" content: ',';\n", | |
" padding-right: 5px;\n", | |
"}\n", | |
"\n", | |
".xr-has-index {\n", | |
" font-weight: bold;\n", | |
"}\n", | |
"\n", | |
".xr-var-list,\n", | |
".xr-var-item {\n", | |
" display: contents;\n", | |
"}\n", | |
"\n", | |
".xr-var-item > div,\n", | |
".xr-var-item label,\n", | |
".xr-var-item > .xr-var-name span {\n", | |
" background-color: var(--xr-background-color-row-even);\n", | |
" margin-bottom: 0;\n", | |
"}\n", | |
"\n", | |
".xr-var-item > .xr-var-name:hover span {\n", | |
" padding-right: 5px;\n", | |
"}\n", | |
"\n", | |
".xr-var-list > li:nth-child(odd) > div,\n", | |
".xr-var-list > li:nth-child(odd) > label,\n", | |
".xr-var-list > li:nth-child(odd) > .xr-var-name span {\n", | |
" background-color: var(--xr-background-color-row-odd);\n", | |
"}\n", | |
"\n", | |
".xr-var-name {\n", | |
" grid-column: 1;\n", | |
"}\n", | |
"\n", | |
".xr-var-dims {\n", | |
" grid-column: 2;\n", | |
"}\n", | |
"\n", | |
".xr-var-dtype {\n", | |
" grid-column: 3;\n", | |
" text-align: right;\n", | |
" color: var(--xr-font-color2);\n", | |
"}\n", | |
"\n", | |
".xr-var-preview {\n", | |
" grid-column: 4;\n", | |
"}\n", | |
"\n", | |
".xr-index-preview {\n", | |
" grid-column: 2 / 5;\n", | |
" color: var(--xr-font-color2);\n", | |
"}\n", | |
"\n", | |
".xr-var-name,\n", | |
".xr-var-dims,\n", | |
".xr-var-dtype,\n", | |
".xr-preview,\n", | |
".xr-attrs dt {\n", | |
" white-space: nowrap;\n", | |
" overflow: hidden;\n", | |
" text-overflow: ellipsis;\n", | |
" padding-right: 10px;\n", | |
"}\n", | |
"\n", | |
".xr-var-name:hover,\n", | |
".xr-var-dims:hover,\n", | |
".xr-var-dtype:hover,\n", | |
".xr-attrs dt:hover {\n", | |
" overflow: visible;\n", | |
" width: auto;\n", | |
" z-index: 1;\n", | |
"}\n", | |
"\n", | |
".xr-var-attrs,\n", | |
".xr-var-data,\n", | |
".xr-index-data {\n", | |
" display: none;\n", | |
" background-color: var(--xr-background-color) !important;\n", | |
" padding-bottom: 5px !important;\n", | |
"}\n", | |
"\n", | |
".xr-var-attrs-in:checked ~ .xr-var-attrs,\n", | |
".xr-var-data-in:checked ~ .xr-var-data,\n", | |
".xr-index-data-in:checked ~ .xr-index-data {\n", | |
" display: block;\n", | |
"}\n", | |
"\n", | |
".xr-var-data > table {\n", | |
" float: right;\n", | |
"}\n", | |
"\n", | |
".xr-var-name span,\n", | |
".xr-var-data,\n", | |
".xr-index-name div,\n", | |
".xr-index-data,\n", | |
".xr-attrs {\n", | |
" padding-left: 25px !important;\n", | |
"}\n", | |
"\n", | |
".xr-attrs,\n", | |
".xr-var-attrs,\n", | |
".xr-var-data,\n", | |
".xr-index-data {\n", | |
" grid-column: 1 / -1;\n", | |
"}\n", | |
"\n", | |
"dl.xr-attrs {\n", | |
" padding: 0;\n", | |
" margin: 0;\n", | |
" display: grid;\n", | |
" grid-template-columns: 125px auto;\n", | |
"}\n", | |
"\n", | |
".xr-attrs dt,\n", | |
".xr-attrs dd {\n", | |
" padding: 0;\n", | |
" margin: 0;\n", | |
" float: left;\n", | |
" padding-right: 10px;\n", | |
" width: auto;\n", | |
"}\n", | |
"\n", | |
".xr-attrs dt {\n", | |
" font-weight: normal;\n", | |
" grid-column: 1;\n", | |
"}\n", | |
"\n", | |
".xr-attrs dt:hover span {\n", | |
" display: inline-block;\n", | |
" background: var(--xr-background-color);\n", | |
" padding-right: 10px;\n", | |
"}\n", | |
"\n", | |
".xr-attrs dd {\n", | |
" grid-column: 2;\n", | |
" white-space: pre-wrap;\n", | |
" word-break: break-all;\n", | |
"}\n", | |
"\n", | |
".xr-icon-database,\n", | |
".xr-icon-file-text2,\n", | |
".xr-no-icon {\n", | |
" display: inline-block;\n", | |
" vertical-align: middle;\n", | |
" width: 1em;\n", | |
" height: 1.5em !important;\n", | |
" stroke-width: 0;\n", | |
" stroke: currentColor;\n", | |
" fill: currentColor;\n", | |
"}\n", | |
"</style><pre class='xr-text-repr-fallback'><xarray.Dataset>\n", | |
"Dimensions: (sample: 3, embedding: 3)\n", | |
"Coordinates:\n", | |
" vectors (sample, embedding) int64 1 0 0 0 1 0 0 0 1\n", | |
"Dimensions without coordinates: sample, embedding\n", | |
"Data variables:\n", | |
" documents (sample) <U1 'a' 'b' 'c'</pre><div class='xr-wrap' style='display:none'><div class='xr-header'><div class='xr-obj-type'>xarray.Dataset</div></div><ul class='xr-sections'><li class='xr-section-item'><input id='section-0cebc028-5e9d-461f-8729-230413376406' class='xr-section-summary-in' type='checkbox' disabled ><label for='section-0cebc028-5e9d-461f-8729-230413376406' class='xr-section-summary' title='Expand/collapse section'>Dimensions:</label><div class='xr-section-inline-details'><ul class='xr-dim-list'><li><span>sample</span>: 3</li><li><span>embedding</span>: 3</li></ul></div><div class='xr-section-details'></div></li><li class='xr-section-item'><input id='section-75d32f0e-c96e-4aae-8030-8d76ec98cc9d' class='xr-section-summary-in' type='checkbox' checked><label for='section-75d32f0e-c96e-4aae-8030-8d76ec98cc9d' class='xr-section-summary' >Coordinates: <span>(1)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><ul class='xr-var-list'><li class='xr-var-item'><div class='xr-var-name'><span>vectors</span></div><div class='xr-var-dims'>(sample, embedding)</div><div class='xr-var-dtype'>int64</div><div class='xr-var-preview xr-preview'>1 0 0 0 1 0 0 0 1</div><input id='attrs-2d3d0d5e-9067-48cf-be27-0a17761023a4' class='xr-var-attrs-in' type='checkbox' disabled><label for='attrs-2d3d0d5e-9067-48cf-be27-0a17761023a4' title='Show/Hide attributes'><svg class='icon xr-icon-file-text2'><use xlink:href='#icon-file-text2'></use></svg></label><input id='data-52a6b0d0-04bc-49a9-a41f-b976eb427b6e' class='xr-var-data-in' type='checkbox'><label for='data-52a6b0d0-04bc-49a9-a41f-b976eb427b6e' title='Show/Hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-var-attrs'><dl class='xr-attrs'></dl></div><div class='xr-var-data'><pre>array([[1, 0, 0],\n", | |
" [0, 1, 0],\n", | |
" [0, 0, 1]])</pre></div></li></ul></div></li><li class='xr-section-item'><input id='section-f40df0b8-340a-4767-986c-b675b6e77c3c' class='xr-section-summary-in' type='checkbox' checked><label for='section-f40df0b8-340a-4767-986c-b675b6e77c3c' class='xr-section-summary' >Data variables: <span>(1)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><ul class='xr-var-list'><li class='xr-var-item'><div class='xr-var-name'><span>documents</span></div><div class='xr-var-dims'>(sample)</div><div class='xr-var-dtype'><U1</div><div class='xr-var-preview xr-preview'>'a' 'b' 'c'</div><input id='attrs-9298dd23-a47c-4efa-87ad-2587b3a85af2' class='xr-var-attrs-in' type='checkbox' disabled><label for='attrs-9298dd23-a47c-4efa-87ad-2587b3a85af2' title='Show/Hide attributes'><svg class='icon xr-icon-file-text2'><use xlink:href='#icon-file-text2'></use></svg></label><input id='data-d94cb59b-2ff3-409d-88fb-9a843b4e51bd' class='xr-var-data-in' type='checkbox'><label for='data-d94cb59b-2ff3-409d-88fb-9a843b4e51bd' title='Show/Hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-var-attrs'><dl class='xr-attrs'></dl></div><div class='xr-var-data'><pre>array(['a', 'b', 'c'], dtype='<U1')</pre></div></li></ul></div></li><li class='xr-section-item'><input id='section-57a1ea50-1afe-4787-be0a-8d8a36fae76e' class='xr-section-summary-in' type='checkbox' disabled ><label for='section-57a1ea50-1afe-4787-be0a-8d8a36fae76e' class='xr-section-summary' title='Expand/collapse section'>Indexes: <span>(0)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><ul class='xr-var-list'></ul></div></li><li class='xr-section-item'><input id='section-df93befb-b8e7-46ed-8679-1a2ce6f167e8' class='xr-section-summary-in' type='checkbox' disabled ><label for='section-df93befb-b8e7-46ed-8679-1a2ce6f167e8' class='xr-section-summary' title='Expand/collapse section'>Attributes: <span>(0)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><dl class='xr-attrs'></dl></div></li></ul></div></div>" | |
], | |
"text/plain": [ | |
"<xarray.Dataset>\n", | |
"Dimensions: (sample: 3, embedding: 3)\n", | |
"Coordinates:\n", | |
" vectors (sample, embedding) int64 1 0 0 0 1 0 0 0 1\n", | |
"Dimensions without coordinates: sample, embedding\n", | |
"Data variables:\n", | |
" documents (sample) <U1 'a' 'b' 'c'" | |
] | |
}, | |
"execution_count": 3, | |
"metadata": {}, | |
"output_type": "execute_result" | |
} | |
], | |
"source": [ | |
"ds = xr.DataArray(\n", | |
" vectors,\n", | |
" dims=[\"sample\", \"embedding\"],\n", | |
" name=\"vectors\"\n", | |
").to_dataset().set_coords(\"vectors\")\n", | |
"ds[\"documents\"] = (\"sample\", [\"a\", \"b\", \"c\"])\n", | |
"ds.vector.set_index()\n", | |
"ds" | |
] | |
}, | |
{ | |
"attachments": {}, | |
"cell_type": "markdown", | |
"metadata": {}, | |
"source": [ | |
"Now we can query this vector space to find the nearest vectors.\n", | |
"\n", | |
"Single point query, single nearest neighbor:" | |
] | |
}, | |
{ | |
"cell_type": "code", | |
"execution_count": 4, | |
"metadata": {}, | |
"outputs": [ | |
{ | |
"data": { | |
"text/html": [ | |
"<div><svg style=\"position: absolute; width: 0; height: 0; overflow: hidden\">\n", | |
"<defs>\n", | |
"<symbol id=\"icon-database\" viewBox=\"0 0 32 32\">\n", | |
"<path d=\"M16 0c-8.837 0-16 2.239-16 5v4c0 2.761 7.163 5 16 5s16-2.239 16-5v-4c0-2.761-7.163-5-16-5z\"></path>\n", | |
"<path d=\"M16 17c-8.837 0-16-2.239-16-5v6c0 2.761 7.163 5 16 5s16-2.239 16-5v-6c0 2.761-7.163 5-16 5z\"></path>\n", | |
"<path d=\"M16 26c-8.837 0-16-2.239-16-5v6c0 2.761 7.163 5 16 5s16-2.239 16-5v-6c0 2.761-7.163 5-16 5z\"></path>\n", | |
"</symbol>\n", | |
"<symbol id=\"icon-file-text2\" viewBox=\"0 0 32 32\">\n", | |
"<path d=\"M28.681 7.159c-0.694-0.947-1.662-2.053-2.724-3.116s-2.169-2.030-3.116-2.724c-1.612-1.182-2.393-1.319-2.841-1.319h-15.5c-1.378 0-2.5 1.121-2.5 2.5v27c0 1.378 1.122 2.5 2.5 2.5h23c1.378 0 2.5-1.122 2.5-2.5v-19.5c0-0.448-0.137-1.23-1.319-2.841zM24.543 5.457c0.959 0.959 1.712 1.825 2.268 2.543h-4.811v-4.811c0.718 0.556 1.584 1.309 2.543 2.268zM28 29.5c0 0.271-0.229 0.5-0.5 0.5h-23c-0.271 0-0.5-0.229-0.5-0.5v-27c0-0.271 0.229-0.5 0.5-0.5 0 0 15.499-0 15.5 0v7c0 0.552 0.448 1 1 1h7v19.5z\"></path>\n", | |
"<path d=\"M23 26h-14c-0.552 0-1-0.448-1-1s0.448-1 1-1h14c0.552 0 1 0.448 1 1s-0.448 1-1 1z\"></path>\n", | |
"<path d=\"M23 22h-14c-0.552 0-1-0.448-1-1s0.448-1 1-1h14c0.552 0 1 0.448 1 1s-0.448 1-1 1z\"></path>\n", | |
"<path d=\"M23 18h-14c-0.552 0-1-0.448-1-1s0.448-1 1-1h14c0.552 0 1 0.448 1 1s-0.448 1-1 1z\"></path>\n", | |
"</symbol>\n", | |
"</defs>\n", | |
"</svg>\n", | |
"<style>/* CSS stylesheet for displaying xarray objects in jupyterlab.\n", | |
" *\n", | |
" */\n", | |
"\n", | |
":root {\n", | |
" --xr-font-color0: var(--jp-content-font-color0, rgba(0, 0, 0, 1));\n", | |
" --xr-font-color2: var(--jp-content-font-color2, rgba(0, 0, 0, 0.54));\n", | |
" --xr-font-color3: var(--jp-content-font-color3, rgba(0, 0, 0, 0.38));\n", | |
" --xr-border-color: var(--jp-border-color2, #e0e0e0);\n", | |
" --xr-disabled-color: var(--jp-layout-color3, #bdbdbd);\n", | |
" --xr-background-color: var(--jp-layout-color0, white);\n", | |
" --xr-background-color-row-even: var(--jp-layout-color1, white);\n", | |
" --xr-background-color-row-odd: var(--jp-layout-color2, #eeeeee);\n", | |
"}\n", | |
"\n", | |
"html[theme=dark],\n", | |
"body[data-theme=dark],\n", | |
"body.vscode-dark {\n", | |
" --xr-font-color0: rgba(255, 255, 255, 1);\n", | |
" --xr-font-color2: rgba(255, 255, 255, 0.54);\n", | |
" --xr-font-color3: rgba(255, 255, 255, 0.38);\n", | |
" --xr-border-color: #1F1F1F;\n", | |
" --xr-disabled-color: #515151;\n", | |
" --xr-background-color: #111111;\n", | |
" --xr-background-color-row-even: #111111;\n", | |
" --xr-background-color-row-odd: #313131;\n", | |
"}\n", | |
"\n", | |
".xr-wrap {\n", | |
" display: block !important;\n", | |
" min-width: 300px;\n", | |
" max-width: 700px;\n", | |
"}\n", | |
"\n", | |
".xr-text-repr-fallback {\n", | |
" /* fallback to plain text repr when CSS is not injected (untrusted notebook) */\n", | |
" display: none;\n", | |
"}\n", | |
"\n", | |
".xr-header {\n", | |
" padding-top: 6px;\n", | |
" padding-bottom: 6px;\n", | |
" margin-bottom: 4px;\n", | |
" border-bottom: solid 1px var(--xr-border-color);\n", | |
"}\n", | |
"\n", | |
".xr-header > div,\n", | |
".xr-header > ul {\n", | |
" display: inline;\n", | |
" margin-top: 0;\n", | |
" margin-bottom: 0;\n", | |
"}\n", | |
"\n", | |
".xr-obj-type,\n", | |
".xr-array-name {\n", | |
" margin-left: 2px;\n", | |
" margin-right: 10px;\n", | |
"}\n", | |
"\n", | |
".xr-obj-type {\n", | |
" color: var(--xr-font-color2);\n", | |
"}\n", | |
"\n", | |
".xr-sections {\n", | |
" padding-left: 0 !important;\n", | |
" display: grid;\n", | |
" grid-template-columns: 150px auto auto 1fr 20px 20px;\n", | |
"}\n", | |
"\n", | |
".xr-section-item {\n", | |
" display: contents;\n", | |
"}\n", | |
"\n", | |
".xr-section-item input {\n", | |
" display: none;\n", | |
"}\n", | |
"\n", | |
".xr-section-item input + label {\n", | |
" color: var(--xr-disabled-color);\n", | |
"}\n", | |
"\n", | |
".xr-section-item input:enabled + label {\n", | |
" cursor: pointer;\n", | |
" color: var(--xr-font-color2);\n", | |
"}\n", | |
"\n", | |
".xr-section-item input:enabled + label:hover {\n", | |
" color: var(--xr-font-color0);\n", | |
"}\n", | |
"\n", | |
".xr-section-summary {\n", | |
" grid-column: 1;\n", | |
" color: var(--xr-font-color2);\n", | |
" font-weight: 500;\n", | |
"}\n", | |
"\n", | |
".xr-section-summary > span {\n", | |
" display: inline-block;\n", | |
" padding-left: 0.5em;\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in:disabled + label {\n", | |
" color: var(--xr-font-color2);\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in + label:before {\n", | |
" display: inline-block;\n", | |
" content: '►';\n", | |
" font-size: 11px;\n", | |
" width: 15px;\n", | |
" text-align: center;\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in:disabled + label:before {\n", | |
" color: var(--xr-disabled-color);\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in:checked + label:before {\n", | |
" content: '▼';\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in:checked + label > span {\n", | |
" display: none;\n", | |
"}\n", | |
"\n", | |
".xr-section-summary,\n", | |
".xr-section-inline-details {\n", | |
" padding-top: 4px;\n", | |
" padding-bottom: 4px;\n", | |
"}\n", | |
"\n", | |
".xr-section-inline-details {\n", | |
" grid-column: 2 / -1;\n", | |
"}\n", | |
"\n", | |
".xr-section-details {\n", | |
" display: none;\n", | |
" grid-column: 1 / -1;\n", | |
" margin-bottom: 5px;\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in:checked ~ .xr-section-details {\n", | |
" display: contents;\n", | |
"}\n", | |
"\n", | |
".xr-array-wrap {\n", | |
" grid-column: 1 / -1;\n", | |
" display: grid;\n", | |
" grid-template-columns: 20px auto;\n", | |
"}\n", | |
"\n", | |
".xr-array-wrap > label {\n", | |
" grid-column: 1;\n", | |
" vertical-align: top;\n", | |
"}\n", | |
"\n", | |
".xr-preview {\n", | |
" color: var(--xr-font-color3);\n", | |
"}\n", | |
"\n", | |
".xr-array-preview,\n", | |
".xr-array-data {\n", | |
" padding: 0 5px !important;\n", | |
" grid-column: 2;\n", | |
"}\n", | |
"\n", | |
".xr-array-data,\n", | |
".xr-array-in:checked ~ .xr-array-preview {\n", | |
" display: none;\n", | |
"}\n", | |
"\n", | |
".xr-array-in:checked ~ .xr-array-data,\n", | |
".xr-array-preview {\n", | |
" display: inline-block;\n", | |
"}\n", | |
"\n", | |
".xr-dim-list {\n", | |
" display: inline-block !important;\n", | |
" list-style: none;\n", | |
" padding: 0 !important;\n", | |
" margin: 0;\n", | |
"}\n", | |
"\n", | |
".xr-dim-list li {\n", | |
" display: inline-block;\n", | |
" padding: 0;\n", | |
" margin: 0;\n", | |
"}\n", | |
"\n", | |
".xr-dim-list:before {\n", | |
" content: '(';\n", | |
"}\n", | |
"\n", | |
".xr-dim-list:after {\n", | |
" content: ')';\n", | |
"}\n", | |
"\n", | |
".xr-dim-list li:not(:last-child):after {\n", | |
" content: ',';\n", | |
" padding-right: 5px;\n", | |
"}\n", | |
"\n", | |
".xr-has-index {\n", | |
" font-weight: bold;\n", | |
"}\n", | |
"\n", | |
".xr-var-list,\n", | |
".xr-var-item {\n", | |
" display: contents;\n", | |
"}\n", | |
"\n", | |
".xr-var-item > div,\n", | |
".xr-var-item label,\n", | |
".xr-var-item > .xr-var-name span {\n", | |
" background-color: var(--xr-background-color-row-even);\n", | |
" margin-bottom: 0;\n", | |
"}\n", | |
"\n", | |
".xr-var-item > .xr-var-name:hover span {\n", | |
" padding-right: 5px;\n", | |
"}\n", | |
"\n", | |
".xr-var-list > li:nth-child(odd) > div,\n", | |
".xr-var-list > li:nth-child(odd) > label,\n", | |
".xr-var-list > li:nth-child(odd) > .xr-var-name span {\n", | |
" background-color: var(--xr-background-color-row-odd);\n", | |
"}\n", | |
"\n", | |
".xr-var-name {\n", | |
" grid-column: 1;\n", | |
"}\n", | |
"\n", | |
".xr-var-dims {\n", | |
" grid-column: 2;\n", | |
"}\n", | |
"\n", | |
".xr-var-dtype {\n", | |
" grid-column: 3;\n", | |
" text-align: right;\n", | |
" color: var(--xr-font-color2);\n", | |
"}\n", | |
"\n", | |
".xr-var-preview {\n", | |
" grid-column: 4;\n", | |
"}\n", | |
"\n", | |
".xr-index-preview {\n", | |
" grid-column: 2 / 5;\n", | |
" color: var(--xr-font-color2);\n", | |
"}\n", | |
"\n", | |
".xr-var-name,\n", | |
".xr-var-dims,\n", | |
".xr-var-dtype,\n", | |
".xr-preview,\n", | |
".xr-attrs dt {\n", | |
" white-space: nowrap;\n", | |
" overflow: hidden;\n", | |
" text-overflow: ellipsis;\n", | |
" padding-right: 10px;\n", | |
"}\n", | |
"\n", | |
".xr-var-name:hover,\n", | |
".xr-var-dims:hover,\n", | |
".xr-var-dtype:hover,\n", | |
".xr-attrs dt:hover {\n", | |
" overflow: visible;\n", | |
" width: auto;\n", | |
" z-index: 1;\n", | |
"}\n", | |
"\n", | |
".xr-var-attrs,\n", | |
".xr-var-data,\n", | |
".xr-index-data {\n", | |
" display: none;\n", | |
" background-color: var(--xr-background-color) !important;\n", | |
" padding-bottom: 5px !important;\n", | |
"}\n", | |
"\n", | |
".xr-var-attrs-in:checked ~ .xr-var-attrs,\n", | |
".xr-var-data-in:checked ~ .xr-var-data,\n", | |
".xr-index-data-in:checked ~ .xr-index-data {\n", | |
" display: block;\n", | |
"}\n", | |
"\n", | |
".xr-var-data > table {\n", | |
" float: right;\n", | |
"}\n", | |
"\n", | |
".xr-var-name span,\n", | |
".xr-var-data,\n", | |
".xr-index-name div,\n", | |
".xr-index-data,\n", | |
".xr-attrs {\n", | |
" padding-left: 25px !important;\n", | |
"}\n", | |
"\n", | |
".xr-attrs,\n", | |
".xr-var-attrs,\n", | |
".xr-var-data,\n", | |
".xr-index-data {\n", | |
" grid-column: 1 / -1;\n", | |
"}\n", | |
"\n", | |
"dl.xr-attrs {\n", | |
" padding: 0;\n", | |
" margin: 0;\n", | |
" display: grid;\n", | |
" grid-template-columns: 125px auto;\n", | |
"}\n", | |
"\n", | |
".xr-attrs dt,\n", | |
".xr-attrs dd {\n", | |
" padding: 0;\n", | |
" margin: 0;\n", | |
" float: left;\n", | |
" padding-right: 10px;\n", | |
" width: auto;\n", | |
"}\n", | |
"\n", | |
".xr-attrs dt {\n", | |
" font-weight: normal;\n", | |
" grid-column: 1;\n", | |
"}\n", | |
"\n", | |
".xr-attrs dt:hover span {\n", | |
" display: inline-block;\n", | |
" background: var(--xr-background-color);\n", | |
" padding-right: 10px;\n", | |
"}\n", | |
"\n", | |
".xr-attrs dd {\n", | |
" grid-column: 2;\n", | |
" white-space: pre-wrap;\n", | |
" word-break: break-all;\n", | |
"}\n", | |
"\n", | |
".xr-icon-database,\n", | |
".xr-icon-file-text2,\n", | |
".xr-no-icon {\n", | |
" display: inline-block;\n", | |
" vertical-align: middle;\n", | |
" width: 1em;\n", | |
" height: 1.5em !important;\n", | |
" stroke-width: 0;\n", | |
" stroke: currentColor;\n", | |
" fill: currentColor;\n", | |
"}\n", | |
"</style><pre class='xr-text-repr-fallback'><xarray.Dataset>\n", | |
"Dimensions: (query: 1, sample: 1, embedding: 3)\n", | |
"Coordinates:\n", | |
" vectors (query, sample, embedding) int64 0 1 0\n", | |
" distance (query, sample) float64 0.8\n", | |
" k (sample) int64 1\n", | |
" query_vectors (query, embedding) float64 0.0 0.2 0.0\n", | |
"Dimensions without coordinates: query, sample, embedding\n", | |
"Data variables:\n", | |
" documents (query, sample) <U1 'b'</pre><div class='xr-wrap' style='display:none'><div class='xr-header'><div class='xr-obj-type'>xarray.Dataset</div></div><ul class='xr-sections'><li class='xr-section-item'><input id='section-1f8a13ac-781f-40ff-ba85-922d0feedac9' class='xr-section-summary-in' type='checkbox' disabled ><label for='section-1f8a13ac-781f-40ff-ba85-922d0feedac9' class='xr-section-summary' title='Expand/collapse section'>Dimensions:</label><div class='xr-section-inline-details'><ul class='xr-dim-list'><li><span>query</span>: 1</li><li><span>sample</span>: 1</li><li><span>embedding</span>: 3</li></ul></div><div class='xr-section-details'></div></li><li class='xr-section-item'><input id='section-7f928981-8ff9-43fd-bcc5-b7a099e86ba4' class='xr-section-summary-in' type='checkbox' checked><label for='section-7f928981-8ff9-43fd-bcc5-b7a099e86ba4' class='xr-section-summary' >Coordinates: <span>(4)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><ul class='xr-var-list'><li class='xr-var-item'><div class='xr-var-name'><span>vectors</span></div><div class='xr-var-dims'>(query, sample, embedding)</div><div class='xr-var-dtype'>int64</div><div class='xr-var-preview xr-preview'>0 1 0</div><input id='attrs-3621128a-f875-45ff-8fde-5329d341c5f5' class='xr-var-attrs-in' type='checkbox' disabled><label for='attrs-3621128a-f875-45ff-8fde-5329d341c5f5' title='Show/Hide attributes'><svg class='icon xr-icon-file-text2'><use xlink:href='#icon-file-text2'></use></svg></label><input id='data-a5aaca3f-e238-4504-bddf-c34bc6f315b0' class='xr-var-data-in' type='checkbox'><label for='data-a5aaca3f-e238-4504-bddf-c34bc6f315b0' title='Show/Hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-var-attrs'><dl class='xr-attrs'></dl></div><div class='xr-var-data'><pre>array([[[0, 1, 0]]])</pre></div></li><li class='xr-var-item'><div class='xr-var-name'><span>distance</span></div><div class='xr-var-dims'>(query, sample)</div><div class='xr-var-dtype'>float64</div><div class='xr-var-preview xr-preview'>0.8</div><input id='attrs-c1102429-b0e8-4509-a9c8-6712d33a8023' class='xr-var-attrs-in' type='checkbox' disabled><label for='attrs-c1102429-b0e8-4509-a9c8-6712d33a8023' title='Show/Hide attributes'><svg class='icon xr-icon-file-text2'><use xlink:href='#icon-file-text2'></use></svg></label><input id='data-c85dbf54-079f-4d96-92e4-b14729bbeb91' class='xr-var-data-in' type='checkbox'><label for='data-c85dbf54-079f-4d96-92e4-b14729bbeb91' title='Show/Hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-var-attrs'><dl class='xr-attrs'></dl></div><div class='xr-var-data'><pre>array([[0.8]])</pre></div></li><li class='xr-var-item'><div class='xr-var-name'><span>k</span></div><div class='xr-var-dims'>(sample)</div><div class='xr-var-dtype'>int64</div><div class='xr-var-preview xr-preview'>1</div><input id='attrs-a21d24aa-5cdf-4798-8a38-79968a2b423a' class='xr-var-attrs-in' type='checkbox' disabled><label for='attrs-a21d24aa-5cdf-4798-8a38-79968a2b423a' title='Show/Hide attributes'><svg class='icon xr-icon-file-text2'><use xlink:href='#icon-file-text2'></use></svg></label><input id='data-b18d83e8-de18-447a-8966-fdb346d442d9' class='xr-var-data-in' type='checkbox'><label for='data-b18d83e8-de18-447a-8966-fdb346d442d9' title='Show/Hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-var-attrs'><dl class='xr-attrs'></dl></div><div class='xr-var-data'><pre>array([1])</pre></div></li><li class='xr-var-item'><div class='xr-var-name'><span>query_vectors</span></div><div class='xr-var-dims'>(query, embedding)</div><div class='xr-var-dtype'>float64</div><div class='xr-var-preview xr-preview'>0.0 0.2 0.0</div><input id='attrs-aaab477c-798f-4e5a-8b78-cb13c74763e1' class='xr-var-attrs-in' type='checkbox' disabled><label for='attrs-aaab477c-798f-4e5a-8b78-cb13c74763e1' title='Show/Hide attributes'><svg class='icon xr-icon-file-text2'><use xlink:href='#icon-file-text2'></use></svg></label><input id='data-6470871a-2341-4b1a-858c-ed2fe000814d' class='xr-var-data-in' type='checkbox'><label for='data-6470871a-2341-4b1a-858c-ed2fe000814d' title='Show/Hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-var-attrs'><dl class='xr-attrs'></dl></div><div class='xr-var-data'><pre>array([[0. , 0.2, 0. ]])</pre></div></li></ul></div></li><li class='xr-section-item'><input id='section-ebc57c74-6349-453a-80da-10aae6fdb4af' class='xr-section-summary-in' type='checkbox' checked><label for='section-ebc57c74-6349-453a-80da-10aae6fdb4af' class='xr-section-summary' >Data variables: <span>(1)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><ul class='xr-var-list'><li class='xr-var-item'><div class='xr-var-name'><span>documents</span></div><div class='xr-var-dims'>(query, sample)</div><div class='xr-var-dtype'><U1</div><div class='xr-var-preview xr-preview'>'b'</div><input id='attrs-65074255-eaef-4141-a929-3970b3ae253a' class='xr-var-attrs-in' type='checkbox' disabled><label for='attrs-65074255-eaef-4141-a929-3970b3ae253a' title='Show/Hide attributes'><svg class='icon xr-icon-file-text2'><use xlink:href='#icon-file-text2'></use></svg></label><input id='data-e665865d-8e36-46b2-b004-cd22c22d24d2' class='xr-var-data-in' type='checkbox'><label for='data-e665865d-8e36-46b2-b004-cd22c22d24d2' title='Show/Hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-var-attrs'><dl class='xr-attrs'></dl></div><div class='xr-var-data'><pre>array([['b']], dtype='<U1')</pre></div></li></ul></div></li><li class='xr-section-item'><input id='section-abcd9336-4ba5-4a44-8f97-48d30c67a61f' class='xr-section-summary-in' type='checkbox' disabled ><label for='section-abcd9336-4ba5-4a44-8f97-48d30c67a61f' class='xr-section-summary' title='Expand/collapse section'>Indexes: <span>(0)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><ul class='xr-var-list'></ul></div></li><li class='xr-section-item'><input id='section-e5e71896-8eff-45b4-b0de-7344f3ca3914' class='xr-section-summary-in' type='checkbox' disabled ><label for='section-e5e71896-8eff-45b4-b0de-7344f3ca3914' class='xr-section-summary' title='Expand/collapse section'>Attributes: <span>(0)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><dl class='xr-attrs'></dl></div></li></ul></div></div>" | |
], | |
"text/plain": [ | |
"<xarray.Dataset>\n", | |
"Dimensions: (query: 1, sample: 1, embedding: 3)\n", | |
"Coordinates:\n", | |
" vectors (query, sample, embedding) int64 0 1 0\n", | |
" distance (query, sample) float64 0.8\n", | |
" k (sample) int64 1\n", | |
" query_vectors (query, embedding) float64 0.0 0.2 0.0\n", | |
"Dimensions without coordinates: query, sample, embedding\n", | |
"Data variables:\n", | |
" documents (query, sample) <U1 'b'" | |
] | |
}, | |
"execution_count": 4, | |
"metadata": {}, | |
"output_type": "execute_result" | |
} | |
], | |
"source": [ | |
"ds.vector.query([[0, 0.2, 0]])" | |
] | |
}, | |
{ | |
"attachments": {}, | |
"cell_type": "markdown", | |
"metadata": {}, | |
"source": [ | |
"Single point, multiple nearest neighbors:" | |
] | |
}, | |
{ | |
"cell_type": "code", | |
"execution_count": 5, | |
"metadata": {}, | |
"outputs": [ | |
{ | |
"data": { | |
"text/html": [ | |
"<div><svg style=\"position: absolute; width: 0; height: 0; overflow: hidden\">\n", | |
"<defs>\n", | |
"<symbol id=\"icon-database\" viewBox=\"0 0 32 32\">\n", | |
"<path d=\"M16 0c-8.837 0-16 2.239-16 5v4c0 2.761 7.163 5 16 5s16-2.239 16-5v-4c0-2.761-7.163-5-16-5z\"></path>\n", | |
"<path d=\"M16 17c-8.837 0-16-2.239-16-5v6c0 2.761 7.163 5 16 5s16-2.239 16-5v-6c0 2.761-7.163 5-16 5z\"></path>\n", | |
"<path d=\"M16 26c-8.837 0-16-2.239-16-5v6c0 2.761 7.163 5 16 5s16-2.239 16-5v-6c0 2.761-7.163 5-16 5z\"></path>\n", | |
"</symbol>\n", | |
"<symbol id=\"icon-file-text2\" viewBox=\"0 0 32 32\">\n", | |
"<path d=\"M28.681 7.159c-0.694-0.947-1.662-2.053-2.724-3.116s-2.169-2.030-3.116-2.724c-1.612-1.182-2.393-1.319-2.841-1.319h-15.5c-1.378 0-2.5 1.121-2.5 2.5v27c0 1.378 1.122 2.5 2.5 2.5h23c1.378 0 2.5-1.122 2.5-2.5v-19.5c0-0.448-0.137-1.23-1.319-2.841zM24.543 5.457c0.959 0.959 1.712 1.825 2.268 2.543h-4.811v-4.811c0.718 0.556 1.584 1.309 2.543 2.268zM28 29.5c0 0.271-0.229 0.5-0.5 0.5h-23c-0.271 0-0.5-0.229-0.5-0.5v-27c0-0.271 0.229-0.5 0.5-0.5 0 0 15.499-0 15.5 0v7c0 0.552 0.448 1 1 1h7v19.5z\"></path>\n", | |
"<path d=\"M23 26h-14c-0.552 0-1-0.448-1-1s0.448-1 1-1h14c0.552 0 1 0.448 1 1s-0.448 1-1 1z\"></path>\n", | |
"<path d=\"M23 22h-14c-0.552 0-1-0.448-1-1s0.448-1 1-1h14c0.552 0 1 0.448 1 1s-0.448 1-1 1z\"></path>\n", | |
"<path d=\"M23 18h-14c-0.552 0-1-0.448-1-1s0.448-1 1-1h14c0.552 0 1 0.448 1 1s-0.448 1-1 1z\"></path>\n", | |
"</symbol>\n", | |
"</defs>\n", | |
"</svg>\n", | |
"<style>/* CSS stylesheet for displaying xarray objects in jupyterlab.\n", | |
" *\n", | |
" */\n", | |
"\n", | |
":root {\n", | |
" --xr-font-color0: var(--jp-content-font-color0, rgba(0, 0, 0, 1));\n", | |
" --xr-font-color2: var(--jp-content-font-color2, rgba(0, 0, 0, 0.54));\n", | |
" --xr-font-color3: var(--jp-content-font-color3, rgba(0, 0, 0, 0.38));\n", | |
" --xr-border-color: var(--jp-border-color2, #e0e0e0);\n", | |
" --xr-disabled-color: var(--jp-layout-color3, #bdbdbd);\n", | |
" --xr-background-color: var(--jp-layout-color0, white);\n", | |
" --xr-background-color-row-even: var(--jp-layout-color1, white);\n", | |
" --xr-background-color-row-odd: var(--jp-layout-color2, #eeeeee);\n", | |
"}\n", | |
"\n", | |
"html[theme=dark],\n", | |
"body[data-theme=dark],\n", | |
"body.vscode-dark {\n", | |
" --xr-font-color0: rgba(255, 255, 255, 1);\n", | |
" --xr-font-color2: rgba(255, 255, 255, 0.54);\n", | |
" --xr-font-color3: rgba(255, 255, 255, 0.38);\n", | |
" --xr-border-color: #1F1F1F;\n", | |
" --xr-disabled-color: #515151;\n", | |
" --xr-background-color: #111111;\n", | |
" --xr-background-color-row-even: #111111;\n", | |
" --xr-background-color-row-odd: #313131;\n", | |
"}\n", | |
"\n", | |
".xr-wrap {\n", | |
" display: block !important;\n", | |
" min-width: 300px;\n", | |
" max-width: 700px;\n", | |
"}\n", | |
"\n", | |
".xr-text-repr-fallback {\n", | |
" /* fallback to plain text repr when CSS is not injected (untrusted notebook) */\n", | |
" display: none;\n", | |
"}\n", | |
"\n", | |
".xr-header {\n", | |
" padding-top: 6px;\n", | |
" padding-bottom: 6px;\n", | |
" margin-bottom: 4px;\n", | |
" border-bottom: solid 1px var(--xr-border-color);\n", | |
"}\n", | |
"\n", | |
".xr-header > div,\n", | |
".xr-header > ul {\n", | |
" display: inline;\n", | |
" margin-top: 0;\n", | |
" margin-bottom: 0;\n", | |
"}\n", | |
"\n", | |
".xr-obj-type,\n", | |
".xr-array-name {\n", | |
" margin-left: 2px;\n", | |
" margin-right: 10px;\n", | |
"}\n", | |
"\n", | |
".xr-obj-type {\n", | |
" color: var(--xr-font-color2);\n", | |
"}\n", | |
"\n", | |
".xr-sections {\n", | |
" padding-left: 0 !important;\n", | |
" display: grid;\n", | |
" grid-template-columns: 150px auto auto 1fr 20px 20px;\n", | |
"}\n", | |
"\n", | |
".xr-section-item {\n", | |
" display: contents;\n", | |
"}\n", | |
"\n", | |
".xr-section-item input {\n", | |
" display: none;\n", | |
"}\n", | |
"\n", | |
".xr-section-item input + label {\n", | |
" color: var(--xr-disabled-color);\n", | |
"}\n", | |
"\n", | |
".xr-section-item input:enabled + label {\n", | |
" cursor: pointer;\n", | |
" color: var(--xr-font-color2);\n", | |
"}\n", | |
"\n", | |
".xr-section-item input:enabled + label:hover {\n", | |
" color: var(--xr-font-color0);\n", | |
"}\n", | |
"\n", | |
".xr-section-summary {\n", | |
" grid-column: 1;\n", | |
" color: var(--xr-font-color2);\n", | |
" font-weight: 500;\n", | |
"}\n", | |
"\n", | |
".xr-section-summary > span {\n", | |
" display: inline-block;\n", | |
" padding-left: 0.5em;\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in:disabled + label {\n", | |
" color: var(--xr-font-color2);\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in + label:before {\n", | |
" display: inline-block;\n", | |
" content: '►';\n", | |
" font-size: 11px;\n", | |
" width: 15px;\n", | |
" text-align: center;\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in:disabled + label:before {\n", | |
" color: var(--xr-disabled-color);\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in:checked + label:before {\n", | |
" content: '▼';\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in:checked + label > span {\n", | |
" display: none;\n", | |
"}\n", | |
"\n", | |
".xr-section-summary,\n", | |
".xr-section-inline-details {\n", | |
" padding-top: 4px;\n", | |
" padding-bottom: 4px;\n", | |
"}\n", | |
"\n", | |
".xr-section-inline-details {\n", | |
" grid-column: 2 / -1;\n", | |
"}\n", | |
"\n", | |
".xr-section-details {\n", | |
" display: none;\n", | |
" grid-column: 1 / -1;\n", | |
" margin-bottom: 5px;\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in:checked ~ .xr-section-details {\n", | |
" display: contents;\n", | |
"}\n", | |
"\n", | |
".xr-array-wrap {\n", | |
" grid-column: 1 / -1;\n", | |
" display: grid;\n", | |
" grid-template-columns: 20px auto;\n", | |
"}\n", | |
"\n", | |
".xr-array-wrap > label {\n", | |
" grid-column: 1;\n", | |
" vertical-align: top;\n", | |
"}\n", | |
"\n", | |
".xr-preview {\n", | |
" color: var(--xr-font-color3);\n", | |
"}\n", | |
"\n", | |
".xr-array-preview,\n", | |
".xr-array-data {\n", | |
" padding: 0 5px !important;\n", | |
" grid-column: 2;\n", | |
"}\n", | |
"\n", | |
".xr-array-data,\n", | |
".xr-array-in:checked ~ .xr-array-preview {\n", | |
" display: none;\n", | |
"}\n", | |
"\n", | |
".xr-array-in:checked ~ .xr-array-data,\n", | |
".xr-array-preview {\n", | |
" display: inline-block;\n", | |
"}\n", | |
"\n", | |
".xr-dim-list {\n", | |
" display: inline-block !important;\n", | |
" list-style: none;\n", | |
" padding: 0 !important;\n", | |
" margin: 0;\n", | |
"}\n", | |
"\n", | |
".xr-dim-list li {\n", | |
" display: inline-block;\n", | |
" padding: 0;\n", | |
" margin: 0;\n", | |
"}\n", | |
"\n", | |
".xr-dim-list:before {\n", | |
" content: '(';\n", | |
"}\n", | |
"\n", | |
".xr-dim-list:after {\n", | |
" content: ')';\n", | |
"}\n", | |
"\n", | |
".xr-dim-list li:not(:last-child):after {\n", | |
" content: ',';\n", | |
" padding-right: 5px;\n", | |
"}\n", | |
"\n", | |
".xr-has-index {\n", | |
" font-weight: bold;\n", | |
"}\n", | |
"\n", | |
".xr-var-list,\n", | |
".xr-var-item {\n", | |
" display: contents;\n", | |
"}\n", | |
"\n", | |
".xr-var-item > div,\n", | |
".xr-var-item label,\n", | |
".xr-var-item > .xr-var-name span {\n", | |
" background-color: var(--xr-background-color-row-even);\n", | |
" margin-bottom: 0;\n", | |
"}\n", | |
"\n", | |
".xr-var-item > .xr-var-name:hover span {\n", | |
" padding-right: 5px;\n", | |
"}\n", | |
"\n", | |
".xr-var-list > li:nth-child(odd) > div,\n", | |
".xr-var-list > li:nth-child(odd) > label,\n", | |
".xr-var-list > li:nth-child(odd) > .xr-var-name span {\n", | |
" background-color: var(--xr-background-color-row-odd);\n", | |
"}\n", | |
"\n", | |
".xr-var-name {\n", | |
" grid-column: 1;\n", | |
"}\n", | |
"\n", | |
".xr-var-dims {\n", | |
" grid-column: 2;\n", | |
"}\n", | |
"\n", | |
".xr-var-dtype {\n", | |
" grid-column: 3;\n", | |
" text-align: right;\n", | |
" color: var(--xr-font-color2);\n", | |
"}\n", | |
"\n", | |
".xr-var-preview {\n", | |
" grid-column: 4;\n", | |
"}\n", | |
"\n", | |
".xr-index-preview {\n", | |
" grid-column: 2 / 5;\n", | |
" color: var(--xr-font-color2);\n", | |
"}\n", | |
"\n", | |
".xr-var-name,\n", | |
".xr-var-dims,\n", | |
".xr-var-dtype,\n", | |
".xr-preview,\n", | |
".xr-attrs dt {\n", | |
" white-space: nowrap;\n", | |
" overflow: hidden;\n", | |
" text-overflow: ellipsis;\n", | |
" padding-right: 10px;\n", | |
"}\n", | |
"\n", | |
".xr-var-name:hover,\n", | |
".xr-var-dims:hover,\n", | |
".xr-var-dtype:hover,\n", | |
".xr-attrs dt:hover {\n", | |
" overflow: visible;\n", | |
" width: auto;\n", | |
" z-index: 1;\n", | |
"}\n", | |
"\n", | |
".xr-var-attrs,\n", | |
".xr-var-data,\n", | |
".xr-index-data {\n", | |
" display: none;\n", | |
" background-color: var(--xr-background-color) !important;\n", | |
" padding-bottom: 5px !important;\n", | |
"}\n", | |
"\n", | |
".xr-var-attrs-in:checked ~ .xr-var-attrs,\n", | |
".xr-var-data-in:checked ~ .xr-var-data,\n", | |
".xr-index-data-in:checked ~ .xr-index-data {\n", | |
" display: block;\n", | |
"}\n", | |
"\n", | |
".xr-var-data > table {\n", | |
" float: right;\n", | |
"}\n", | |
"\n", | |
".xr-var-name span,\n", | |
".xr-var-data,\n", | |
".xr-index-name div,\n", | |
".xr-index-data,\n", | |
".xr-attrs {\n", | |
" padding-left: 25px !important;\n", | |
"}\n", | |
"\n", | |
".xr-attrs,\n", | |
".xr-var-attrs,\n", | |
".xr-var-data,\n", | |
".xr-index-data {\n", | |
" grid-column: 1 / -1;\n", | |
"}\n", | |
"\n", | |
"dl.xr-attrs {\n", | |
" padding: 0;\n", | |
" margin: 0;\n", | |
" display: grid;\n", | |
" grid-template-columns: 125px auto;\n", | |
"}\n", | |
"\n", | |
".xr-attrs dt,\n", | |
".xr-attrs dd {\n", | |
" padding: 0;\n", | |
" margin: 0;\n", | |
" float: left;\n", | |
" padding-right: 10px;\n", | |
" width: auto;\n", | |
"}\n", | |
"\n", | |
".xr-attrs dt {\n", | |
" font-weight: normal;\n", | |
" grid-column: 1;\n", | |
"}\n", | |
"\n", | |
".xr-attrs dt:hover span {\n", | |
" display: inline-block;\n", | |
" background: var(--xr-background-color);\n", | |
" padding-right: 10px;\n", | |
"}\n", | |
"\n", | |
".xr-attrs dd {\n", | |
" grid-column: 2;\n", | |
" white-space: pre-wrap;\n", | |
" word-break: break-all;\n", | |
"}\n", | |
"\n", | |
".xr-icon-database,\n", | |
".xr-icon-file-text2,\n", | |
".xr-no-icon {\n", | |
" display: inline-block;\n", | |
" vertical-align: middle;\n", | |
" width: 1em;\n", | |
" height: 1.5em !important;\n", | |
" stroke-width: 0;\n", | |
" stroke: currentColor;\n", | |
" fill: currentColor;\n", | |
"}\n", | |
"</style><pre class='xr-text-repr-fallback'><xarray.Dataset>\n", | |
"Dimensions: (query: 1, sample: 3, embedding: 3)\n", | |
"Coordinates:\n", | |
" vectors (query, sample, embedding) int64 0 1 0 0 0 1 1 0 0\n", | |
" distance (query, sample) float64 0.8 1.02 1.02\n", | |
" k (sample) int64 1 2 3\n", | |
" query_vectors (query, embedding) float64 0.0 0.2 0.0\n", | |
"Dimensions without coordinates: query, sample, embedding\n", | |
"Data variables:\n", | |
" documents (query, sample) <U1 'b' 'c' 'a'</pre><div class='xr-wrap' style='display:none'><div class='xr-header'><div class='xr-obj-type'>xarray.Dataset</div></div><ul class='xr-sections'><li class='xr-section-item'><input id='section-28d1d70f-9489-4ab9-9a63-b2d380b4e967' class='xr-section-summary-in' type='checkbox' disabled ><label for='section-28d1d70f-9489-4ab9-9a63-b2d380b4e967' class='xr-section-summary' title='Expand/collapse section'>Dimensions:</label><div class='xr-section-inline-details'><ul class='xr-dim-list'><li><span>query</span>: 1</li><li><span>sample</span>: 3</li><li><span>embedding</span>: 3</li></ul></div><div class='xr-section-details'></div></li><li class='xr-section-item'><input id='section-60f32ea1-7938-4d19-8d12-e5af1598cc0e' class='xr-section-summary-in' type='checkbox' checked><label for='section-60f32ea1-7938-4d19-8d12-e5af1598cc0e' class='xr-section-summary' >Coordinates: <span>(4)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><ul class='xr-var-list'><li class='xr-var-item'><div class='xr-var-name'><span>vectors</span></div><div class='xr-var-dims'>(query, sample, embedding)</div><div class='xr-var-dtype'>int64</div><div class='xr-var-preview xr-preview'>0 1 0 0 0 1 1 0 0</div><input id='attrs-b8dfa155-d0b3-4952-9ef8-5af72811feb9' class='xr-var-attrs-in' type='checkbox' disabled><label for='attrs-b8dfa155-d0b3-4952-9ef8-5af72811feb9' title='Show/Hide attributes'><svg class='icon xr-icon-file-text2'><use xlink:href='#icon-file-text2'></use></svg></label><input id='data-a521dc51-4ab9-419b-b9f0-124dd062eba0' class='xr-var-data-in' type='checkbox'><label for='data-a521dc51-4ab9-419b-b9f0-124dd062eba0' title='Show/Hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-var-attrs'><dl class='xr-attrs'></dl></div><div class='xr-var-data'><pre>array([[[0, 1, 0],\n", | |
" [0, 0, 1],\n", | |
" [1, 0, 0]]])</pre></div></li><li class='xr-var-item'><div class='xr-var-name'><span>distance</span></div><div class='xr-var-dims'>(query, sample)</div><div class='xr-var-dtype'>float64</div><div class='xr-var-preview xr-preview'>0.8 1.02 1.02</div><input id='attrs-a2204f85-61d4-4197-a6c0-9542f79c5d70' class='xr-var-attrs-in' type='checkbox' disabled><label for='attrs-a2204f85-61d4-4197-a6c0-9542f79c5d70' title='Show/Hide attributes'><svg class='icon xr-icon-file-text2'><use xlink:href='#icon-file-text2'></use></svg></label><input id='data-222dcdb8-e96c-4417-8978-c8b530b31525' class='xr-var-data-in' type='checkbox'><label for='data-222dcdb8-e96c-4417-8978-c8b530b31525' title='Show/Hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-var-attrs'><dl class='xr-attrs'></dl></div><div class='xr-var-data'><pre>array([[0.8 , 1.0198039, 1.0198039]])</pre></div></li><li class='xr-var-item'><div class='xr-var-name'><span>k</span></div><div class='xr-var-dims'>(sample)</div><div class='xr-var-dtype'>int64</div><div class='xr-var-preview xr-preview'>1 2 3</div><input id='attrs-08f70082-0f6d-4991-9b8a-4f5100b65b13' class='xr-var-attrs-in' type='checkbox' disabled><label for='attrs-08f70082-0f6d-4991-9b8a-4f5100b65b13' title='Show/Hide attributes'><svg class='icon xr-icon-file-text2'><use xlink:href='#icon-file-text2'></use></svg></label><input id='data-159a266a-2cb8-45d9-a438-37cf2582b224' class='xr-var-data-in' type='checkbox'><label for='data-159a266a-2cb8-45d9-a438-37cf2582b224' title='Show/Hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-var-attrs'><dl class='xr-attrs'></dl></div><div class='xr-var-data'><pre>array([1, 2, 3])</pre></div></li><li class='xr-var-item'><div class='xr-var-name'><span>query_vectors</span></div><div class='xr-var-dims'>(query, embedding)</div><div class='xr-var-dtype'>float64</div><div class='xr-var-preview xr-preview'>0.0 0.2 0.0</div><input id='attrs-7e03ae96-0ef3-4938-90b2-d0bed879ac83' class='xr-var-attrs-in' type='checkbox' disabled><label for='attrs-7e03ae96-0ef3-4938-90b2-d0bed879ac83' title='Show/Hide attributes'><svg class='icon xr-icon-file-text2'><use xlink:href='#icon-file-text2'></use></svg></label><input id='data-03762afd-a90a-4f7b-b9fb-03f79e9bbb4b' class='xr-var-data-in' type='checkbox'><label for='data-03762afd-a90a-4f7b-b9fb-03f79e9bbb4b' title='Show/Hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-var-attrs'><dl class='xr-attrs'></dl></div><div class='xr-var-data'><pre>array([[0. , 0.2, 0. ]])</pre></div></li></ul></div></li><li class='xr-section-item'><input id='section-afaa7969-4f35-4f98-b534-8b8b9b0f1d75' class='xr-section-summary-in' type='checkbox' checked><label for='section-afaa7969-4f35-4f98-b534-8b8b9b0f1d75' class='xr-section-summary' >Data variables: <span>(1)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><ul class='xr-var-list'><li class='xr-var-item'><div class='xr-var-name'><span>documents</span></div><div class='xr-var-dims'>(query, sample)</div><div class='xr-var-dtype'><U1</div><div class='xr-var-preview xr-preview'>'b' 'c' 'a'</div><input id='attrs-20be5323-8f81-48c7-8a27-a99768f48450' class='xr-var-attrs-in' type='checkbox' disabled><label for='attrs-20be5323-8f81-48c7-8a27-a99768f48450' title='Show/Hide attributes'><svg class='icon xr-icon-file-text2'><use xlink:href='#icon-file-text2'></use></svg></label><input id='data-b31efe8e-eaa6-4de0-a1cb-9ce0eaa38101' class='xr-var-data-in' type='checkbox'><label for='data-b31efe8e-eaa6-4de0-a1cb-9ce0eaa38101' title='Show/Hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-var-attrs'><dl class='xr-attrs'></dl></div><div class='xr-var-data'><pre>array([['b', 'c', 'a']], dtype='<U1')</pre></div></li></ul></div></li><li class='xr-section-item'><input id='section-ceb377c7-75a9-4add-8890-91ff525ea27a' class='xr-section-summary-in' type='checkbox' disabled ><label for='section-ceb377c7-75a9-4add-8890-91ff525ea27a' class='xr-section-summary' title='Expand/collapse section'>Indexes: <span>(0)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><ul class='xr-var-list'></ul></div></li><li class='xr-section-item'><input id='section-65295cff-5e22-4208-9b67-f2aaf1898eea' class='xr-section-summary-in' type='checkbox' disabled ><label for='section-65295cff-5e22-4208-9b67-f2aaf1898eea' class='xr-section-summary' title='Expand/collapse section'>Attributes: <span>(0)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><dl class='xr-attrs'></dl></div></li></ul></div></div>" | |
], | |
"text/plain": [ | |
"<xarray.Dataset>\n", | |
"Dimensions: (query: 1, sample: 3, embedding: 3)\n", | |
"Coordinates:\n", | |
" vectors (query, sample, embedding) int64 0 1 0 0 0 1 1 0 0\n", | |
" distance (query, sample) float64 0.8 1.02 1.02\n", | |
" k (sample) int64 1 2 3\n", | |
" query_vectors (query, embedding) float64 0.0 0.2 0.0\n", | |
"Dimensions without coordinates: query, sample, embedding\n", | |
"Data variables:\n", | |
" documents (query, sample) <U1 'b' 'c' 'a'" | |
] | |
}, | |
"execution_count": 5, | |
"metadata": {}, | |
"output_type": "execute_result" | |
} | |
], | |
"source": [ | |
"ds.vector.query([[0, 0.2, 0]], k=3)" | |
] | |
}, | |
{ | |
"attachments": {}, | |
"cell_type": "markdown", | |
"metadata": {}, | |
"source": [ | |
"Single point, specific neighbors returned" | |
] | |
}, | |
{ | |
"cell_type": "code", | |
"execution_count": 6, | |
"metadata": {}, | |
"outputs": [ | |
{ | |
"data": { | |
"text/html": [ | |
"<div><svg style=\"position: absolute; width: 0; height: 0; overflow: hidden\">\n", | |
"<defs>\n", | |
"<symbol id=\"icon-database\" viewBox=\"0 0 32 32\">\n", | |
"<path d=\"M16 0c-8.837 0-16 2.239-16 5v4c0 2.761 7.163 5 16 5s16-2.239 16-5v-4c0-2.761-7.163-5-16-5z\"></path>\n", | |
"<path d=\"M16 17c-8.837 0-16-2.239-16-5v6c0 2.761 7.163 5 16 5s16-2.239 16-5v-6c0 2.761-7.163 5-16 5z\"></path>\n", | |
"<path d=\"M16 26c-8.837 0-16-2.239-16-5v6c0 2.761 7.163 5 16 5s16-2.239 16-5v-6c0 2.761-7.163 5-16 5z\"></path>\n", | |
"</symbol>\n", | |
"<symbol id=\"icon-file-text2\" viewBox=\"0 0 32 32\">\n", | |
"<path d=\"M28.681 7.159c-0.694-0.947-1.662-2.053-2.724-3.116s-2.169-2.030-3.116-2.724c-1.612-1.182-2.393-1.319-2.841-1.319h-15.5c-1.378 0-2.5 1.121-2.5 2.5v27c0 1.378 1.122 2.5 2.5 2.5h23c1.378 0 2.5-1.122 2.5-2.5v-19.5c0-0.448-0.137-1.23-1.319-2.841zM24.543 5.457c0.959 0.959 1.712 1.825 2.268 2.543h-4.811v-4.811c0.718 0.556 1.584 1.309 2.543 2.268zM28 29.5c0 0.271-0.229 0.5-0.5 0.5h-23c-0.271 0-0.5-0.229-0.5-0.5v-27c0-0.271 0.229-0.5 0.5-0.5 0 0 15.499-0 15.5 0v7c0 0.552 0.448 1 1 1h7v19.5z\"></path>\n", | |
"<path d=\"M23 26h-14c-0.552 0-1-0.448-1-1s0.448-1 1-1h14c0.552 0 1 0.448 1 1s-0.448 1-1 1z\"></path>\n", | |
"<path d=\"M23 22h-14c-0.552 0-1-0.448-1-1s0.448-1 1-1h14c0.552 0 1 0.448 1 1s-0.448 1-1 1z\"></path>\n", | |
"<path d=\"M23 18h-14c-0.552 0-1-0.448-1-1s0.448-1 1-1h14c0.552 0 1 0.448 1 1s-0.448 1-1 1z\"></path>\n", | |
"</symbol>\n", | |
"</defs>\n", | |
"</svg>\n", | |
"<style>/* CSS stylesheet for displaying xarray objects in jupyterlab.\n", | |
" *\n", | |
" */\n", | |
"\n", | |
":root {\n", | |
" --xr-font-color0: var(--jp-content-font-color0, rgba(0, 0, 0, 1));\n", | |
" --xr-font-color2: var(--jp-content-font-color2, rgba(0, 0, 0, 0.54));\n", | |
" --xr-font-color3: var(--jp-content-font-color3, rgba(0, 0, 0, 0.38));\n", | |
" --xr-border-color: var(--jp-border-color2, #e0e0e0);\n", | |
" --xr-disabled-color: var(--jp-layout-color3, #bdbdbd);\n", | |
" --xr-background-color: var(--jp-layout-color0, white);\n", | |
" --xr-background-color-row-even: var(--jp-layout-color1, white);\n", | |
" --xr-background-color-row-odd: var(--jp-layout-color2, #eeeeee);\n", | |
"}\n", | |
"\n", | |
"html[theme=dark],\n", | |
"body[data-theme=dark],\n", | |
"body.vscode-dark {\n", | |
" --xr-font-color0: rgba(255, 255, 255, 1);\n", | |
" --xr-font-color2: rgba(255, 255, 255, 0.54);\n", | |
" --xr-font-color3: rgba(255, 255, 255, 0.38);\n", | |
" --xr-border-color: #1F1F1F;\n", | |
" --xr-disabled-color: #515151;\n", | |
" --xr-background-color: #111111;\n", | |
" --xr-background-color-row-even: #111111;\n", | |
" --xr-background-color-row-odd: #313131;\n", | |
"}\n", | |
"\n", | |
".xr-wrap {\n", | |
" display: block !important;\n", | |
" min-width: 300px;\n", | |
" max-width: 700px;\n", | |
"}\n", | |
"\n", | |
".xr-text-repr-fallback {\n", | |
" /* fallback to plain text repr when CSS is not injected (untrusted notebook) */\n", | |
" display: none;\n", | |
"}\n", | |
"\n", | |
".xr-header {\n", | |
" padding-top: 6px;\n", | |
" padding-bottom: 6px;\n", | |
" margin-bottom: 4px;\n", | |
" border-bottom: solid 1px var(--xr-border-color);\n", | |
"}\n", | |
"\n", | |
".xr-header > div,\n", | |
".xr-header > ul {\n", | |
" display: inline;\n", | |
" margin-top: 0;\n", | |
" margin-bottom: 0;\n", | |
"}\n", | |
"\n", | |
".xr-obj-type,\n", | |
".xr-array-name {\n", | |
" margin-left: 2px;\n", | |
" margin-right: 10px;\n", | |
"}\n", | |
"\n", | |
".xr-obj-type {\n", | |
" color: var(--xr-font-color2);\n", | |
"}\n", | |
"\n", | |
".xr-sections {\n", | |
" padding-left: 0 !important;\n", | |
" display: grid;\n", | |
" grid-template-columns: 150px auto auto 1fr 20px 20px;\n", | |
"}\n", | |
"\n", | |
".xr-section-item {\n", | |
" display: contents;\n", | |
"}\n", | |
"\n", | |
".xr-section-item input {\n", | |
" display: none;\n", | |
"}\n", | |
"\n", | |
".xr-section-item input + label {\n", | |
" color: var(--xr-disabled-color);\n", | |
"}\n", | |
"\n", | |
".xr-section-item input:enabled + label {\n", | |
" cursor: pointer;\n", | |
" color: var(--xr-font-color2);\n", | |
"}\n", | |
"\n", | |
".xr-section-item input:enabled + label:hover {\n", | |
" color: var(--xr-font-color0);\n", | |
"}\n", | |
"\n", | |
".xr-section-summary {\n", | |
" grid-column: 1;\n", | |
" color: var(--xr-font-color2);\n", | |
" font-weight: 500;\n", | |
"}\n", | |
"\n", | |
".xr-section-summary > span {\n", | |
" display: inline-block;\n", | |
" padding-left: 0.5em;\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in:disabled + label {\n", | |
" color: var(--xr-font-color2);\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in + label:before {\n", | |
" display: inline-block;\n", | |
" content: '►';\n", | |
" font-size: 11px;\n", | |
" width: 15px;\n", | |
" text-align: center;\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in:disabled + label:before {\n", | |
" color: var(--xr-disabled-color);\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in:checked + label:before {\n", | |
" content: '▼';\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in:checked + label > span {\n", | |
" display: none;\n", | |
"}\n", | |
"\n", | |
".xr-section-summary,\n", | |
".xr-section-inline-details {\n", | |
" padding-top: 4px;\n", | |
" padding-bottom: 4px;\n", | |
"}\n", | |
"\n", | |
".xr-section-inline-details {\n", | |
" grid-column: 2 / -1;\n", | |
"}\n", | |
"\n", | |
".xr-section-details {\n", | |
" display: none;\n", | |
" grid-column: 1 / -1;\n", | |
" margin-bottom: 5px;\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in:checked ~ .xr-section-details {\n", | |
" display: contents;\n", | |
"}\n", | |
"\n", | |
".xr-array-wrap {\n", | |
" grid-column: 1 / -1;\n", | |
" display: grid;\n", | |
" grid-template-columns: 20px auto;\n", | |
"}\n", | |
"\n", | |
".xr-array-wrap > label {\n", | |
" grid-column: 1;\n", | |
" vertical-align: top;\n", | |
"}\n", | |
"\n", | |
".xr-preview {\n", | |
" color: var(--xr-font-color3);\n", | |
"}\n", | |
"\n", | |
".xr-array-preview,\n", | |
".xr-array-data {\n", | |
" padding: 0 5px !important;\n", | |
" grid-column: 2;\n", | |
"}\n", | |
"\n", | |
".xr-array-data,\n", | |
".xr-array-in:checked ~ .xr-array-preview {\n", | |
" display: none;\n", | |
"}\n", | |
"\n", | |
".xr-array-in:checked ~ .xr-array-data,\n", | |
".xr-array-preview {\n", | |
" display: inline-block;\n", | |
"}\n", | |
"\n", | |
".xr-dim-list {\n", | |
" display: inline-block !important;\n", | |
" list-style: none;\n", | |
" padding: 0 !important;\n", | |
" margin: 0;\n", | |
"}\n", | |
"\n", | |
".xr-dim-list li {\n", | |
" display: inline-block;\n", | |
" padding: 0;\n", | |
" margin: 0;\n", | |
"}\n", | |
"\n", | |
".xr-dim-list:before {\n", | |
" content: '(';\n", | |
"}\n", | |
"\n", | |
".xr-dim-list:after {\n", | |
" content: ')';\n", | |
"}\n", | |
"\n", | |
".xr-dim-list li:not(:last-child):after {\n", | |
" content: ',';\n", | |
" padding-right: 5px;\n", | |
"}\n", | |
"\n", | |
".xr-has-index {\n", | |
" font-weight: bold;\n", | |
"}\n", | |
"\n", | |
".xr-var-list,\n", | |
".xr-var-item {\n", | |
" display: contents;\n", | |
"}\n", | |
"\n", | |
".xr-var-item > div,\n", | |
".xr-var-item label,\n", | |
".xr-var-item > .xr-var-name span {\n", | |
" background-color: var(--xr-background-color-row-even);\n", | |
" margin-bottom: 0;\n", | |
"}\n", | |
"\n", | |
".xr-var-item > .xr-var-name:hover span {\n", | |
" padding-right: 5px;\n", | |
"}\n", | |
"\n", | |
".xr-var-list > li:nth-child(odd) > div,\n", | |
".xr-var-list > li:nth-child(odd) > label,\n", | |
".xr-var-list > li:nth-child(odd) > .xr-var-name span {\n", | |
" background-color: var(--xr-background-color-row-odd);\n", | |
"}\n", | |
"\n", | |
".xr-var-name {\n", | |
" grid-column: 1;\n", | |
"}\n", | |
"\n", | |
".xr-var-dims {\n", | |
" grid-column: 2;\n", | |
"}\n", | |
"\n", | |
".xr-var-dtype {\n", | |
" grid-column: 3;\n", | |
" text-align: right;\n", | |
" color: var(--xr-font-color2);\n", | |
"}\n", | |
"\n", | |
".xr-var-preview {\n", | |
" grid-column: 4;\n", | |
"}\n", | |
"\n", | |
".xr-index-preview {\n", | |
" grid-column: 2 / 5;\n", | |
" color: var(--xr-font-color2);\n", | |
"}\n", | |
"\n", | |
".xr-var-name,\n", | |
".xr-var-dims,\n", | |
".xr-var-dtype,\n", | |
".xr-preview,\n", | |
".xr-attrs dt {\n", | |
" white-space: nowrap;\n", | |
" overflow: hidden;\n", | |
" text-overflow: ellipsis;\n", | |
" padding-right: 10px;\n", | |
"}\n", | |
"\n", | |
".xr-var-name:hover,\n", | |
".xr-var-dims:hover,\n", | |
".xr-var-dtype:hover,\n", | |
".xr-attrs dt:hover {\n", | |
" overflow: visible;\n", | |
" width: auto;\n", | |
" z-index: 1;\n", | |
"}\n", | |
"\n", | |
".xr-var-attrs,\n", | |
".xr-var-data,\n", | |
".xr-index-data {\n", | |
" display: none;\n", | |
" background-color: var(--xr-background-color) !important;\n", | |
" padding-bottom: 5px !important;\n", | |
"}\n", | |
"\n", | |
".xr-var-attrs-in:checked ~ .xr-var-attrs,\n", | |
".xr-var-data-in:checked ~ .xr-var-data,\n", | |
".xr-index-data-in:checked ~ .xr-index-data {\n", | |
" display: block;\n", | |
"}\n", | |
"\n", | |
".xr-var-data > table {\n", | |
" float: right;\n", | |
"}\n", | |
"\n", | |
".xr-var-name span,\n", | |
".xr-var-data,\n", | |
".xr-index-name div,\n", | |
".xr-index-data,\n", | |
".xr-attrs {\n", | |
" padding-left: 25px !important;\n", | |
"}\n", | |
"\n", | |
".xr-attrs,\n", | |
".xr-var-attrs,\n", | |
".xr-var-data,\n", | |
".xr-index-data {\n", | |
" grid-column: 1 / -1;\n", | |
"}\n", | |
"\n", | |
"dl.xr-attrs {\n", | |
" padding: 0;\n", | |
" margin: 0;\n", | |
" display: grid;\n", | |
" grid-template-columns: 125px auto;\n", | |
"}\n", | |
"\n", | |
".xr-attrs dt,\n", | |
".xr-attrs dd {\n", | |
" padding: 0;\n", | |
" margin: 0;\n", | |
" float: left;\n", | |
" padding-right: 10px;\n", | |
" width: auto;\n", | |
"}\n", | |
"\n", | |
".xr-attrs dt {\n", | |
" font-weight: normal;\n", | |
" grid-column: 1;\n", | |
"}\n", | |
"\n", | |
".xr-attrs dt:hover span {\n", | |
" display: inline-block;\n", | |
" background: var(--xr-background-color);\n", | |
" padding-right: 10px;\n", | |
"}\n", | |
"\n", | |
".xr-attrs dd {\n", | |
" grid-column: 2;\n", | |
" white-space: pre-wrap;\n", | |
" word-break: break-all;\n", | |
"}\n", | |
"\n", | |
".xr-icon-database,\n", | |
".xr-icon-file-text2,\n", | |
".xr-no-icon {\n", | |
" display: inline-block;\n", | |
" vertical-align: middle;\n", | |
" width: 1em;\n", | |
" height: 1.5em !important;\n", | |
" stroke-width: 0;\n", | |
" stroke: currentColor;\n", | |
" fill: currentColor;\n", | |
"}\n", | |
"</style><pre class='xr-text-repr-fallback'><xarray.Dataset>\n", | |
"Dimensions: (query: 1, sample: 2, embedding: 3)\n", | |
"Coordinates:\n", | |
" vectors (query, sample, embedding) int64 0 1 0 1 0 0\n", | |
" distance (query, sample) float64 0.8 1.02\n", | |
" k (sample) int64 1 3\n", | |
" query_vectors (query, embedding) float64 0.0 0.2 0.0\n", | |
"Dimensions without coordinates: query, sample, embedding\n", | |
"Data variables:\n", | |
" documents (query, sample) <U1 'b' 'a'</pre><div class='xr-wrap' style='display:none'><div class='xr-header'><div class='xr-obj-type'>xarray.Dataset</div></div><ul class='xr-sections'><li class='xr-section-item'><input id='section-1562bf4a-2dd9-4721-a44e-b33c2d410626' class='xr-section-summary-in' type='checkbox' disabled ><label for='section-1562bf4a-2dd9-4721-a44e-b33c2d410626' class='xr-section-summary' title='Expand/collapse section'>Dimensions:</label><div class='xr-section-inline-details'><ul class='xr-dim-list'><li><span>query</span>: 1</li><li><span>sample</span>: 2</li><li><span>embedding</span>: 3</li></ul></div><div class='xr-section-details'></div></li><li class='xr-section-item'><input id='section-dcc13d6d-b9b1-4a5f-9583-de2a54ae90a9' class='xr-section-summary-in' type='checkbox' checked><label for='section-dcc13d6d-b9b1-4a5f-9583-de2a54ae90a9' class='xr-section-summary' >Coordinates: <span>(4)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><ul class='xr-var-list'><li class='xr-var-item'><div class='xr-var-name'><span>vectors</span></div><div class='xr-var-dims'>(query, sample, embedding)</div><div class='xr-var-dtype'>int64</div><div class='xr-var-preview xr-preview'>0 1 0 1 0 0</div><input id='attrs-100caa31-fde1-4fc5-9445-0ccb028ddaa4' class='xr-var-attrs-in' type='checkbox' disabled><label for='attrs-100caa31-fde1-4fc5-9445-0ccb028ddaa4' title='Show/Hide attributes'><svg class='icon xr-icon-file-text2'><use xlink:href='#icon-file-text2'></use></svg></label><input id='data-7a669efc-89ea-403c-9998-7edf5c23424c' class='xr-var-data-in' type='checkbox'><label for='data-7a669efc-89ea-403c-9998-7edf5c23424c' title='Show/Hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-var-attrs'><dl class='xr-attrs'></dl></div><div class='xr-var-data'><pre>array([[[0, 1, 0],\n", | |
" [1, 0, 0]]])</pre></div></li><li class='xr-var-item'><div class='xr-var-name'><span>distance</span></div><div class='xr-var-dims'>(query, sample)</div><div class='xr-var-dtype'>float64</div><div class='xr-var-preview xr-preview'>0.8 1.02</div><input id='attrs-813b0ba9-f6b8-4489-804b-ec95795d238b' class='xr-var-attrs-in' type='checkbox' disabled><label for='attrs-813b0ba9-f6b8-4489-804b-ec95795d238b' title='Show/Hide attributes'><svg class='icon xr-icon-file-text2'><use xlink:href='#icon-file-text2'></use></svg></label><input id='data-834a8d52-5d9f-4d12-9e56-406e16110a00' class='xr-var-data-in' type='checkbox'><label for='data-834a8d52-5d9f-4d12-9e56-406e16110a00' title='Show/Hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-var-attrs'><dl class='xr-attrs'></dl></div><div class='xr-var-data'><pre>array([[0.8 , 1.0198039]])</pre></div></li><li class='xr-var-item'><div class='xr-var-name'><span>k</span></div><div class='xr-var-dims'>(sample)</div><div class='xr-var-dtype'>int64</div><div class='xr-var-preview xr-preview'>1 3</div><input id='attrs-972f6efd-0956-45f4-a391-eeae0dde534e' class='xr-var-attrs-in' type='checkbox' disabled><label for='attrs-972f6efd-0956-45f4-a391-eeae0dde534e' title='Show/Hide attributes'><svg class='icon xr-icon-file-text2'><use xlink:href='#icon-file-text2'></use></svg></label><input id='data-7a213522-9ad7-4d66-90df-5d709a0b5bf9' class='xr-var-data-in' type='checkbox'><label for='data-7a213522-9ad7-4d66-90df-5d709a0b5bf9' title='Show/Hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-var-attrs'><dl class='xr-attrs'></dl></div><div class='xr-var-data'><pre>array([1, 3])</pre></div></li><li class='xr-var-item'><div class='xr-var-name'><span>query_vectors</span></div><div class='xr-var-dims'>(query, embedding)</div><div class='xr-var-dtype'>float64</div><div class='xr-var-preview xr-preview'>0.0 0.2 0.0</div><input id='attrs-7d42246b-fae3-4404-bc43-864e2044dc86' class='xr-var-attrs-in' type='checkbox' disabled><label for='attrs-7d42246b-fae3-4404-bc43-864e2044dc86' title='Show/Hide attributes'><svg class='icon xr-icon-file-text2'><use xlink:href='#icon-file-text2'></use></svg></label><input id='data-1440c5c8-83c9-48f1-88e4-f2750e0b5e5e' class='xr-var-data-in' type='checkbox'><label for='data-1440c5c8-83c9-48f1-88e4-f2750e0b5e5e' title='Show/Hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-var-attrs'><dl class='xr-attrs'></dl></div><div class='xr-var-data'><pre>array([[0. , 0.2, 0. ]])</pre></div></li></ul></div></li><li class='xr-section-item'><input id='section-6a99150c-251a-43d0-9aed-c88066f61f4d' class='xr-section-summary-in' type='checkbox' checked><label for='section-6a99150c-251a-43d0-9aed-c88066f61f4d' class='xr-section-summary' >Data variables: <span>(1)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><ul class='xr-var-list'><li class='xr-var-item'><div class='xr-var-name'><span>documents</span></div><div class='xr-var-dims'>(query, sample)</div><div class='xr-var-dtype'><U1</div><div class='xr-var-preview xr-preview'>'b' 'a'</div><input id='attrs-2233fcb4-f484-4b50-8dbd-d517cf1bcd9b' class='xr-var-attrs-in' type='checkbox' disabled><label for='attrs-2233fcb4-f484-4b50-8dbd-d517cf1bcd9b' title='Show/Hide attributes'><svg class='icon xr-icon-file-text2'><use xlink:href='#icon-file-text2'></use></svg></label><input id='data-f1e7f545-1566-4aa2-b105-6012c7a0b3f3' class='xr-var-data-in' type='checkbox'><label for='data-f1e7f545-1566-4aa2-b105-6012c7a0b3f3' title='Show/Hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-var-attrs'><dl class='xr-attrs'></dl></div><div class='xr-var-data'><pre>array([['b', 'a']], dtype='<U1')</pre></div></li></ul></div></li><li class='xr-section-item'><input id='section-cec39383-3db1-470e-b444-bb08b05e74ef' class='xr-section-summary-in' type='checkbox' disabled ><label for='section-cec39383-3db1-470e-b444-bb08b05e74ef' class='xr-section-summary' title='Expand/collapse section'>Indexes: <span>(0)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><ul class='xr-var-list'></ul></div></li><li class='xr-section-item'><input id='section-1d617757-d165-447a-a871-f8a3c7efabf7' class='xr-section-summary-in' type='checkbox' disabled ><label for='section-1d617757-d165-447a-a871-f8a3c7efabf7' class='xr-section-summary' title='Expand/collapse section'>Attributes: <span>(0)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><dl class='xr-attrs'></dl></div></li></ul></div></div>" | |
], | |
"text/plain": [ | |
"<xarray.Dataset>\n", | |
"Dimensions: (query: 1, sample: 2, embedding: 3)\n", | |
"Coordinates:\n", | |
" vectors (query, sample, embedding) int64 0 1 0 1 0 0\n", | |
" distance (query, sample) float64 0.8 1.02\n", | |
" k (sample) int64 1 3\n", | |
" query_vectors (query, embedding) float64 0.0 0.2 0.0\n", | |
"Dimensions without coordinates: query, sample, embedding\n", | |
"Data variables:\n", | |
" documents (query, sample) <U1 'b' 'a'" | |
] | |
}, | |
"execution_count": 6, | |
"metadata": {}, | |
"output_type": "execute_result" | |
} | |
], | |
"source": [ | |
"ds.vector.query([[0, 0.2, 0]], k=[1, 3])" | |
] | |
}, | |
{ | |
"attachments": {}, | |
"cell_type": "markdown", | |
"metadata": {}, | |
"source": [ | |
"Query with multiple points, single nearest neighbor for each:" | |
] | |
}, | |
{ | |
"cell_type": "code", | |
"execution_count": 7, | |
"metadata": {}, | |
"outputs": [ | |
{ | |
"data": { | |
"text/html": [ | |
"<div><svg style=\"position: absolute; width: 0; height: 0; overflow: hidden\">\n", | |
"<defs>\n", | |
"<symbol id=\"icon-database\" viewBox=\"0 0 32 32\">\n", | |
"<path d=\"M16 0c-8.837 0-16 2.239-16 5v4c0 2.761 7.163 5 16 5s16-2.239 16-5v-4c0-2.761-7.163-5-16-5z\"></path>\n", | |
"<path d=\"M16 17c-8.837 0-16-2.239-16-5v6c0 2.761 7.163 5 16 5s16-2.239 16-5v-6c0 2.761-7.163 5-16 5z\"></path>\n", | |
"<path d=\"M16 26c-8.837 0-16-2.239-16-5v6c0 2.761 7.163 5 16 5s16-2.239 16-5v-6c0 2.761-7.163 5-16 5z\"></path>\n", | |
"</symbol>\n", | |
"<symbol id=\"icon-file-text2\" viewBox=\"0 0 32 32\">\n", | |
"<path d=\"M28.681 7.159c-0.694-0.947-1.662-2.053-2.724-3.116s-2.169-2.030-3.116-2.724c-1.612-1.182-2.393-1.319-2.841-1.319h-15.5c-1.378 0-2.5 1.121-2.5 2.5v27c0 1.378 1.122 2.5 2.5 2.5h23c1.378 0 2.5-1.122 2.5-2.5v-19.5c0-0.448-0.137-1.23-1.319-2.841zM24.543 5.457c0.959 0.959 1.712 1.825 2.268 2.543h-4.811v-4.811c0.718 0.556 1.584 1.309 2.543 2.268zM28 29.5c0 0.271-0.229 0.5-0.5 0.5h-23c-0.271 0-0.5-0.229-0.5-0.5v-27c0-0.271 0.229-0.5 0.5-0.5 0 0 15.499-0 15.5 0v7c0 0.552 0.448 1 1 1h7v19.5z\"></path>\n", | |
"<path d=\"M23 26h-14c-0.552 0-1-0.448-1-1s0.448-1 1-1h14c0.552 0 1 0.448 1 1s-0.448 1-1 1z\"></path>\n", | |
"<path d=\"M23 22h-14c-0.552 0-1-0.448-1-1s0.448-1 1-1h14c0.552 0 1 0.448 1 1s-0.448 1-1 1z\"></path>\n", | |
"<path d=\"M23 18h-14c-0.552 0-1-0.448-1-1s0.448-1 1-1h14c0.552 0 1 0.448 1 1s-0.448 1-1 1z\"></path>\n", | |
"</symbol>\n", | |
"</defs>\n", | |
"</svg>\n", | |
"<style>/* CSS stylesheet for displaying xarray objects in jupyterlab.\n", | |
" *\n", | |
" */\n", | |
"\n", | |
":root {\n", | |
" --xr-font-color0: var(--jp-content-font-color0, rgba(0, 0, 0, 1));\n", | |
" --xr-font-color2: var(--jp-content-font-color2, rgba(0, 0, 0, 0.54));\n", | |
" --xr-font-color3: var(--jp-content-font-color3, rgba(0, 0, 0, 0.38));\n", | |
" --xr-border-color: var(--jp-border-color2, #e0e0e0);\n", | |
" --xr-disabled-color: var(--jp-layout-color3, #bdbdbd);\n", | |
" --xr-background-color: var(--jp-layout-color0, white);\n", | |
" --xr-background-color-row-even: var(--jp-layout-color1, white);\n", | |
" --xr-background-color-row-odd: var(--jp-layout-color2, #eeeeee);\n", | |
"}\n", | |
"\n", | |
"html[theme=dark],\n", | |
"body[data-theme=dark],\n", | |
"body.vscode-dark {\n", | |
" --xr-font-color0: rgba(255, 255, 255, 1);\n", | |
" --xr-font-color2: rgba(255, 255, 255, 0.54);\n", | |
" --xr-font-color3: rgba(255, 255, 255, 0.38);\n", | |
" --xr-border-color: #1F1F1F;\n", | |
" --xr-disabled-color: #515151;\n", | |
" --xr-background-color: #111111;\n", | |
" --xr-background-color-row-even: #111111;\n", | |
" --xr-background-color-row-odd: #313131;\n", | |
"}\n", | |
"\n", | |
".xr-wrap {\n", | |
" display: block !important;\n", | |
" min-width: 300px;\n", | |
" max-width: 700px;\n", | |
"}\n", | |
"\n", | |
".xr-text-repr-fallback {\n", | |
" /* fallback to plain text repr when CSS is not injected (untrusted notebook) */\n", | |
" display: none;\n", | |
"}\n", | |
"\n", | |
".xr-header {\n", | |
" padding-top: 6px;\n", | |
" padding-bottom: 6px;\n", | |
" margin-bottom: 4px;\n", | |
" border-bottom: solid 1px var(--xr-border-color);\n", | |
"}\n", | |
"\n", | |
".xr-header > div,\n", | |
".xr-header > ul {\n", | |
" display: inline;\n", | |
" margin-top: 0;\n", | |
" margin-bottom: 0;\n", | |
"}\n", | |
"\n", | |
".xr-obj-type,\n", | |
".xr-array-name {\n", | |
" margin-left: 2px;\n", | |
" margin-right: 10px;\n", | |
"}\n", | |
"\n", | |
".xr-obj-type {\n", | |
" color: var(--xr-font-color2);\n", | |
"}\n", | |
"\n", | |
".xr-sections {\n", | |
" padding-left: 0 !important;\n", | |
" display: grid;\n", | |
" grid-template-columns: 150px auto auto 1fr 20px 20px;\n", | |
"}\n", | |
"\n", | |
".xr-section-item {\n", | |
" display: contents;\n", | |
"}\n", | |
"\n", | |
".xr-section-item input {\n", | |
" display: none;\n", | |
"}\n", | |
"\n", | |
".xr-section-item input + label {\n", | |
" color: var(--xr-disabled-color);\n", | |
"}\n", | |
"\n", | |
".xr-section-item input:enabled + label {\n", | |
" cursor: pointer;\n", | |
" color: var(--xr-font-color2);\n", | |
"}\n", | |
"\n", | |
".xr-section-item input:enabled + label:hover {\n", | |
" color: var(--xr-font-color0);\n", | |
"}\n", | |
"\n", | |
".xr-section-summary {\n", | |
" grid-column: 1;\n", | |
" color: var(--xr-font-color2);\n", | |
" font-weight: 500;\n", | |
"}\n", | |
"\n", | |
".xr-section-summary > span {\n", | |
" display: inline-block;\n", | |
" padding-left: 0.5em;\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in:disabled + label {\n", | |
" color: var(--xr-font-color2);\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in + label:before {\n", | |
" display: inline-block;\n", | |
" content: '►';\n", | |
" font-size: 11px;\n", | |
" width: 15px;\n", | |
" text-align: center;\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in:disabled + label:before {\n", | |
" color: var(--xr-disabled-color);\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in:checked + label:before {\n", | |
" content: '▼';\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in:checked + label > span {\n", | |
" display: none;\n", | |
"}\n", | |
"\n", | |
".xr-section-summary,\n", | |
".xr-section-inline-details {\n", | |
" padding-top: 4px;\n", | |
" padding-bottom: 4px;\n", | |
"}\n", | |
"\n", | |
".xr-section-inline-details {\n", | |
" grid-column: 2 / -1;\n", | |
"}\n", | |
"\n", | |
".xr-section-details {\n", | |
" display: none;\n", | |
" grid-column: 1 / -1;\n", | |
" margin-bottom: 5px;\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in:checked ~ .xr-section-details {\n", | |
" display: contents;\n", | |
"}\n", | |
"\n", | |
".xr-array-wrap {\n", | |
" grid-column: 1 / -1;\n", | |
" display: grid;\n", | |
" grid-template-columns: 20px auto;\n", | |
"}\n", | |
"\n", | |
".xr-array-wrap > label {\n", | |
" grid-column: 1;\n", | |
" vertical-align: top;\n", | |
"}\n", | |
"\n", | |
".xr-preview {\n", | |
" color: var(--xr-font-color3);\n", | |
"}\n", | |
"\n", | |
".xr-array-preview,\n", | |
".xr-array-data {\n", | |
" padding: 0 5px !important;\n", | |
" grid-column: 2;\n", | |
"}\n", | |
"\n", | |
".xr-array-data,\n", | |
".xr-array-in:checked ~ .xr-array-preview {\n", | |
" display: none;\n", | |
"}\n", | |
"\n", | |
".xr-array-in:checked ~ .xr-array-data,\n", | |
".xr-array-preview {\n", | |
" display: inline-block;\n", | |
"}\n", | |
"\n", | |
".xr-dim-list {\n", | |
" display: inline-block !important;\n", | |
" list-style: none;\n", | |
" padding: 0 !important;\n", | |
" margin: 0;\n", | |
"}\n", | |
"\n", | |
".xr-dim-list li {\n", | |
" display: inline-block;\n", | |
" padding: 0;\n", | |
" margin: 0;\n", | |
"}\n", | |
"\n", | |
".xr-dim-list:before {\n", | |
" content: '(';\n", | |
"}\n", | |
"\n", | |
".xr-dim-list:after {\n", | |
" content: ')';\n", | |
"}\n", | |
"\n", | |
".xr-dim-list li:not(:last-child):after {\n", | |
" content: ',';\n", | |
" padding-right: 5px;\n", | |
"}\n", | |
"\n", | |
".xr-has-index {\n", | |
" font-weight: bold;\n", | |
"}\n", | |
"\n", | |
".xr-var-list,\n", | |
".xr-var-item {\n", | |
" display: contents;\n", | |
"}\n", | |
"\n", | |
".xr-var-item > div,\n", | |
".xr-var-item label,\n", | |
".xr-var-item > .xr-var-name span {\n", | |
" background-color: var(--xr-background-color-row-even);\n", | |
" margin-bottom: 0;\n", | |
"}\n", | |
"\n", | |
".xr-var-item > .xr-var-name:hover span {\n", | |
" padding-right: 5px;\n", | |
"}\n", | |
"\n", | |
".xr-var-list > li:nth-child(odd) > div,\n", | |
".xr-var-list > li:nth-child(odd) > label,\n", | |
".xr-var-list > li:nth-child(odd) > .xr-var-name span {\n", | |
" background-color: var(--xr-background-color-row-odd);\n", | |
"}\n", | |
"\n", | |
".xr-var-name {\n", | |
" grid-column: 1;\n", | |
"}\n", | |
"\n", | |
".xr-var-dims {\n", | |
" grid-column: 2;\n", | |
"}\n", | |
"\n", | |
".xr-var-dtype {\n", | |
" grid-column: 3;\n", | |
" text-align: right;\n", | |
" color: var(--xr-font-color2);\n", | |
"}\n", | |
"\n", | |
".xr-var-preview {\n", | |
" grid-column: 4;\n", | |
"}\n", | |
"\n", | |
".xr-index-preview {\n", | |
" grid-column: 2 / 5;\n", | |
" color: var(--xr-font-color2);\n", | |
"}\n", | |
"\n", | |
".xr-var-name,\n", | |
".xr-var-dims,\n", | |
".xr-var-dtype,\n", | |
".xr-preview,\n", | |
".xr-attrs dt {\n", | |
" white-space: nowrap;\n", | |
" overflow: hidden;\n", | |
" text-overflow: ellipsis;\n", | |
" padding-right: 10px;\n", | |
"}\n", | |
"\n", | |
".xr-var-name:hover,\n", | |
".xr-var-dims:hover,\n", | |
".xr-var-dtype:hover,\n", | |
".xr-attrs dt:hover {\n", | |
" overflow: visible;\n", | |
" width: auto;\n", | |
" z-index: 1;\n", | |
"}\n", | |
"\n", | |
".xr-var-attrs,\n", | |
".xr-var-data,\n", | |
".xr-index-data {\n", | |
" display: none;\n", | |
" background-color: var(--xr-background-color) !important;\n", | |
" padding-bottom: 5px !important;\n", | |
"}\n", | |
"\n", | |
".xr-var-attrs-in:checked ~ .xr-var-attrs,\n", | |
".xr-var-data-in:checked ~ .xr-var-data,\n", | |
".xr-index-data-in:checked ~ .xr-index-data {\n", | |
" display: block;\n", | |
"}\n", | |
"\n", | |
".xr-var-data > table {\n", | |
" float: right;\n", | |
"}\n", | |
"\n", | |
".xr-var-name span,\n", | |
".xr-var-data,\n", | |
".xr-index-name div,\n", | |
".xr-index-data,\n", | |
".xr-attrs {\n", | |
" padding-left: 25px !important;\n", | |
"}\n", | |
"\n", | |
".xr-attrs,\n", | |
".xr-var-attrs,\n", | |
".xr-var-data,\n", | |
".xr-index-data {\n", | |
" grid-column: 1 / -1;\n", | |
"}\n", | |
"\n", | |
"dl.xr-attrs {\n", | |
" padding: 0;\n", | |
" margin: 0;\n", | |
" display: grid;\n", | |
" grid-template-columns: 125px auto;\n", | |
"}\n", | |
"\n", | |
".xr-attrs dt,\n", | |
".xr-attrs dd {\n", | |
" padding: 0;\n", | |
" margin: 0;\n", | |
" float: left;\n", | |
" padding-right: 10px;\n", | |
" width: auto;\n", | |
"}\n", | |
"\n", | |
".xr-attrs dt {\n", | |
" font-weight: normal;\n", | |
" grid-column: 1;\n", | |
"}\n", | |
"\n", | |
".xr-attrs dt:hover span {\n", | |
" display: inline-block;\n", | |
" background: var(--xr-background-color);\n", | |
" padding-right: 10px;\n", | |
"}\n", | |
"\n", | |
".xr-attrs dd {\n", | |
" grid-column: 2;\n", | |
" white-space: pre-wrap;\n", | |
" word-break: break-all;\n", | |
"}\n", | |
"\n", | |
".xr-icon-database,\n", | |
".xr-icon-file-text2,\n", | |
".xr-no-icon {\n", | |
" display: inline-block;\n", | |
" vertical-align: middle;\n", | |
" width: 1em;\n", | |
" height: 1.5em !important;\n", | |
" stroke-width: 0;\n", | |
" stroke: currentColor;\n", | |
" fill: currentColor;\n", | |
"}\n", | |
"</style><pre class='xr-text-repr-fallback'><xarray.Dataset>\n", | |
"Dimensions: (query: 2, sample: 1, embedding: 3)\n", | |
"Coordinates:\n", | |
" vectors (query, sample, embedding) int64 0 1 0 1 0 0\n", | |
" distance (query, sample) float64 0.8 0.5\n", | |
" k (sample) int64 1\n", | |
" query_vectors (query, embedding) float64 0.0 0.2 0.0 0.5 0.0 0.0\n", | |
"Dimensions without coordinates: query, sample, embedding\n", | |
"Data variables:\n", | |
" documents (query, sample) <U1 'b' 'a'</pre><div class='xr-wrap' style='display:none'><div class='xr-header'><div class='xr-obj-type'>xarray.Dataset</div></div><ul class='xr-sections'><li class='xr-section-item'><input id='section-857a3d55-1217-422e-82b4-bef91e52eced' class='xr-section-summary-in' type='checkbox' disabled ><label for='section-857a3d55-1217-422e-82b4-bef91e52eced' class='xr-section-summary' title='Expand/collapse section'>Dimensions:</label><div class='xr-section-inline-details'><ul class='xr-dim-list'><li><span>query</span>: 2</li><li><span>sample</span>: 1</li><li><span>embedding</span>: 3</li></ul></div><div class='xr-section-details'></div></li><li class='xr-section-item'><input id='section-1c957d19-16e3-4024-9de6-8de6111dc5e5' class='xr-section-summary-in' type='checkbox' checked><label for='section-1c957d19-16e3-4024-9de6-8de6111dc5e5' class='xr-section-summary' >Coordinates: <span>(4)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><ul class='xr-var-list'><li class='xr-var-item'><div class='xr-var-name'><span>vectors</span></div><div class='xr-var-dims'>(query, sample, embedding)</div><div class='xr-var-dtype'>int64</div><div class='xr-var-preview xr-preview'>0 1 0 1 0 0</div><input id='attrs-70f30e56-b55d-49a7-8d97-6a7264414ff7' class='xr-var-attrs-in' type='checkbox' disabled><label for='attrs-70f30e56-b55d-49a7-8d97-6a7264414ff7' title='Show/Hide attributes'><svg class='icon xr-icon-file-text2'><use xlink:href='#icon-file-text2'></use></svg></label><input id='data-71cc18e8-d4dd-40eb-8274-e980cc8826d0' class='xr-var-data-in' type='checkbox'><label for='data-71cc18e8-d4dd-40eb-8274-e980cc8826d0' title='Show/Hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-var-attrs'><dl class='xr-attrs'></dl></div><div class='xr-var-data'><pre>array([[[0, 1, 0]],\n", | |
"\n", | |
" [[1, 0, 0]]])</pre></div></li><li class='xr-var-item'><div class='xr-var-name'><span>distance</span></div><div class='xr-var-dims'>(query, sample)</div><div class='xr-var-dtype'>float64</div><div class='xr-var-preview xr-preview'>0.8 0.5</div><input id='attrs-2bca990a-69e7-480f-9d86-0949c65e87f5' class='xr-var-attrs-in' type='checkbox' disabled><label for='attrs-2bca990a-69e7-480f-9d86-0949c65e87f5' title='Show/Hide attributes'><svg class='icon xr-icon-file-text2'><use xlink:href='#icon-file-text2'></use></svg></label><input id='data-950ea1e2-d8f2-420d-be07-a25f5198adc6' class='xr-var-data-in' type='checkbox'><label for='data-950ea1e2-d8f2-420d-be07-a25f5198adc6' title='Show/Hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-var-attrs'><dl class='xr-attrs'></dl></div><div class='xr-var-data'><pre>array([[0.8],\n", | |
" [0.5]])</pre></div></li><li class='xr-var-item'><div class='xr-var-name'><span>k</span></div><div class='xr-var-dims'>(sample)</div><div class='xr-var-dtype'>int64</div><div class='xr-var-preview xr-preview'>1</div><input id='attrs-d7b62b8d-eb7f-43ac-b93c-1dab9a9d3b8c' class='xr-var-attrs-in' type='checkbox' disabled><label for='attrs-d7b62b8d-eb7f-43ac-b93c-1dab9a9d3b8c' title='Show/Hide attributes'><svg class='icon xr-icon-file-text2'><use xlink:href='#icon-file-text2'></use></svg></label><input id='data-89ca3d55-eb57-49d1-8770-ba726b5fc007' class='xr-var-data-in' type='checkbox'><label for='data-89ca3d55-eb57-49d1-8770-ba726b5fc007' title='Show/Hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-var-attrs'><dl class='xr-attrs'></dl></div><div class='xr-var-data'><pre>array([1])</pre></div></li><li class='xr-var-item'><div class='xr-var-name'><span>query_vectors</span></div><div class='xr-var-dims'>(query, embedding)</div><div class='xr-var-dtype'>float64</div><div class='xr-var-preview xr-preview'>0.0 0.2 0.0 0.5 0.0 0.0</div><input id='attrs-1c8db879-1206-4a21-8563-12e8a1f80dc6' class='xr-var-attrs-in' type='checkbox' disabled><label for='attrs-1c8db879-1206-4a21-8563-12e8a1f80dc6' title='Show/Hide attributes'><svg class='icon xr-icon-file-text2'><use xlink:href='#icon-file-text2'></use></svg></label><input id='data-b8ae4d10-687d-4560-b2bc-37c643f8d69d' class='xr-var-data-in' type='checkbox'><label for='data-b8ae4d10-687d-4560-b2bc-37c643f8d69d' title='Show/Hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-var-attrs'><dl class='xr-attrs'></dl></div><div class='xr-var-data'><pre>array([[0. , 0.2, 0. ],\n", | |
" [0.5, 0. , 0. ]])</pre></div></li></ul></div></li><li class='xr-section-item'><input id='section-b3411ab2-dedf-4102-a26f-50a0befe7615' class='xr-section-summary-in' type='checkbox' checked><label for='section-b3411ab2-dedf-4102-a26f-50a0befe7615' class='xr-section-summary' >Data variables: <span>(1)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><ul class='xr-var-list'><li class='xr-var-item'><div class='xr-var-name'><span>documents</span></div><div class='xr-var-dims'>(query, sample)</div><div class='xr-var-dtype'><U1</div><div class='xr-var-preview xr-preview'>'b' 'a'</div><input id='attrs-8e4d4fdc-f95e-45ad-bc43-3f359c84b691' class='xr-var-attrs-in' type='checkbox' disabled><label for='attrs-8e4d4fdc-f95e-45ad-bc43-3f359c84b691' title='Show/Hide attributes'><svg class='icon xr-icon-file-text2'><use xlink:href='#icon-file-text2'></use></svg></label><input id='data-09586a93-d3c4-4cb6-8465-7217e75c7f5b' class='xr-var-data-in' type='checkbox'><label for='data-09586a93-d3c4-4cb6-8465-7217e75c7f5b' title='Show/Hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-var-attrs'><dl class='xr-attrs'></dl></div><div class='xr-var-data'><pre>array([['b'],\n", | |
" ['a']], dtype='<U1')</pre></div></li></ul></div></li><li class='xr-section-item'><input id='section-91733c95-7e26-4cf3-9b2e-04f22f543f0c' class='xr-section-summary-in' type='checkbox' disabled ><label for='section-91733c95-7e26-4cf3-9b2e-04f22f543f0c' class='xr-section-summary' title='Expand/collapse section'>Indexes: <span>(0)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><ul class='xr-var-list'></ul></div></li><li class='xr-section-item'><input id='section-3f9db3d0-98d7-4798-a296-8299c1ba1b11' class='xr-section-summary-in' type='checkbox' disabled ><label for='section-3f9db3d0-98d7-4798-a296-8299c1ba1b11' class='xr-section-summary' title='Expand/collapse section'>Attributes: <span>(0)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><dl class='xr-attrs'></dl></div></li></ul></div></div>" | |
], | |
"text/plain": [ | |
"<xarray.Dataset>\n", | |
"Dimensions: (query: 2, sample: 1, embedding: 3)\n", | |
"Coordinates:\n", | |
" vectors (query, sample, embedding) int64 0 1 0 1 0 0\n", | |
" distance (query, sample) float64 0.8 0.5\n", | |
" k (sample) int64 1\n", | |
" query_vectors (query, embedding) float64 0.0 0.2 0.0 0.5 0.0 0.0\n", | |
"Dimensions without coordinates: query, sample, embedding\n", | |
"Data variables:\n", | |
" documents (query, sample) <U1 'b' 'a'" | |
] | |
}, | |
"execution_count": 7, | |
"metadata": {}, | |
"output_type": "execute_result" | |
} | |
], | |
"source": [ | |
"ds.vector.query(\n", | |
" [\n", | |
" [0, 0.2, 0],\n", | |
" [0.5, 0, 0],\n", | |
" ]\n", | |
")" | |
] | |
}, | |
{ | |
"attachments": {}, | |
"cell_type": "markdown", | |
"metadata": {}, | |
"source": [ | |
"Query with multiple points, 3 nearest neighbors for each:" | |
] | |
}, | |
{ | |
"cell_type": "code", | |
"execution_count": 8, | |
"metadata": {}, | |
"outputs": [ | |
{ | |
"data": { | |
"text/html": [ | |
"<div><svg style=\"position: absolute; width: 0; height: 0; overflow: hidden\">\n", | |
"<defs>\n", | |
"<symbol id=\"icon-database\" viewBox=\"0 0 32 32\">\n", | |
"<path d=\"M16 0c-8.837 0-16 2.239-16 5v4c0 2.761 7.163 5 16 5s16-2.239 16-5v-4c0-2.761-7.163-5-16-5z\"></path>\n", | |
"<path d=\"M16 17c-8.837 0-16-2.239-16-5v6c0 2.761 7.163 5 16 5s16-2.239 16-5v-6c0 2.761-7.163 5-16 5z\"></path>\n", | |
"<path d=\"M16 26c-8.837 0-16-2.239-16-5v6c0 2.761 7.163 5 16 5s16-2.239 16-5v-6c0 2.761-7.163 5-16 5z\"></path>\n", | |
"</symbol>\n", | |
"<symbol id=\"icon-file-text2\" viewBox=\"0 0 32 32\">\n", | |
"<path d=\"M28.681 7.159c-0.694-0.947-1.662-2.053-2.724-3.116s-2.169-2.030-3.116-2.724c-1.612-1.182-2.393-1.319-2.841-1.319h-15.5c-1.378 0-2.5 1.121-2.5 2.5v27c0 1.378 1.122 2.5 2.5 2.5h23c1.378 0 2.5-1.122 2.5-2.5v-19.5c0-0.448-0.137-1.23-1.319-2.841zM24.543 5.457c0.959 0.959 1.712 1.825 2.268 2.543h-4.811v-4.811c0.718 0.556 1.584 1.309 2.543 2.268zM28 29.5c0 0.271-0.229 0.5-0.5 0.5h-23c-0.271 0-0.5-0.229-0.5-0.5v-27c0-0.271 0.229-0.5 0.5-0.5 0 0 15.499-0 15.5 0v7c0 0.552 0.448 1 1 1h7v19.5z\"></path>\n", | |
"<path d=\"M23 26h-14c-0.552 0-1-0.448-1-1s0.448-1 1-1h14c0.552 0 1 0.448 1 1s-0.448 1-1 1z\"></path>\n", | |
"<path d=\"M23 22h-14c-0.552 0-1-0.448-1-1s0.448-1 1-1h14c0.552 0 1 0.448 1 1s-0.448 1-1 1z\"></path>\n", | |
"<path d=\"M23 18h-14c-0.552 0-1-0.448-1-1s0.448-1 1-1h14c0.552 0 1 0.448 1 1s-0.448 1-1 1z\"></path>\n", | |
"</symbol>\n", | |
"</defs>\n", | |
"</svg>\n", | |
"<style>/* CSS stylesheet for displaying xarray objects in jupyterlab.\n", | |
" *\n", | |
" */\n", | |
"\n", | |
":root {\n", | |
" --xr-font-color0: var(--jp-content-font-color0, rgba(0, 0, 0, 1));\n", | |
" --xr-font-color2: var(--jp-content-font-color2, rgba(0, 0, 0, 0.54));\n", | |
" --xr-font-color3: var(--jp-content-font-color3, rgba(0, 0, 0, 0.38));\n", | |
" --xr-border-color: var(--jp-border-color2, #e0e0e0);\n", | |
" --xr-disabled-color: var(--jp-layout-color3, #bdbdbd);\n", | |
" --xr-background-color: var(--jp-layout-color0, white);\n", | |
" --xr-background-color-row-even: var(--jp-layout-color1, white);\n", | |
" --xr-background-color-row-odd: var(--jp-layout-color2, #eeeeee);\n", | |
"}\n", | |
"\n", | |
"html[theme=dark],\n", | |
"body[data-theme=dark],\n", | |
"body.vscode-dark {\n", | |
" --xr-font-color0: rgba(255, 255, 255, 1);\n", | |
" --xr-font-color2: rgba(255, 255, 255, 0.54);\n", | |
" --xr-font-color3: rgba(255, 255, 255, 0.38);\n", | |
" --xr-border-color: #1F1F1F;\n", | |
" --xr-disabled-color: #515151;\n", | |
" --xr-background-color: #111111;\n", | |
" --xr-background-color-row-even: #111111;\n", | |
" --xr-background-color-row-odd: #313131;\n", | |
"}\n", | |
"\n", | |
".xr-wrap {\n", | |
" display: block !important;\n", | |
" min-width: 300px;\n", | |
" max-width: 700px;\n", | |
"}\n", | |
"\n", | |
".xr-text-repr-fallback {\n", | |
" /* fallback to plain text repr when CSS is not injected (untrusted notebook) */\n", | |
" display: none;\n", | |
"}\n", | |
"\n", | |
".xr-header {\n", | |
" padding-top: 6px;\n", | |
" padding-bottom: 6px;\n", | |
" margin-bottom: 4px;\n", | |
" border-bottom: solid 1px var(--xr-border-color);\n", | |
"}\n", | |
"\n", | |
".xr-header > div,\n", | |
".xr-header > ul {\n", | |
" display: inline;\n", | |
" margin-top: 0;\n", | |
" margin-bottom: 0;\n", | |
"}\n", | |
"\n", | |
".xr-obj-type,\n", | |
".xr-array-name {\n", | |
" margin-left: 2px;\n", | |
" margin-right: 10px;\n", | |
"}\n", | |
"\n", | |
".xr-obj-type {\n", | |
" color: var(--xr-font-color2);\n", | |
"}\n", | |
"\n", | |
".xr-sections {\n", | |
" padding-left: 0 !important;\n", | |
" display: grid;\n", | |
" grid-template-columns: 150px auto auto 1fr 20px 20px;\n", | |
"}\n", | |
"\n", | |
".xr-section-item {\n", | |
" display: contents;\n", | |
"}\n", | |
"\n", | |
".xr-section-item input {\n", | |
" display: none;\n", | |
"}\n", | |
"\n", | |
".xr-section-item input + label {\n", | |
" color: var(--xr-disabled-color);\n", | |
"}\n", | |
"\n", | |
".xr-section-item input:enabled + label {\n", | |
" cursor: pointer;\n", | |
" color: var(--xr-font-color2);\n", | |
"}\n", | |
"\n", | |
".xr-section-item input:enabled + label:hover {\n", | |
" color: var(--xr-font-color0);\n", | |
"}\n", | |
"\n", | |
".xr-section-summary {\n", | |
" grid-column: 1;\n", | |
" color: var(--xr-font-color2);\n", | |
" font-weight: 500;\n", | |
"}\n", | |
"\n", | |
".xr-section-summary > span {\n", | |
" display: inline-block;\n", | |
" padding-left: 0.5em;\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in:disabled + label {\n", | |
" color: var(--xr-font-color2);\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in + label:before {\n", | |
" display: inline-block;\n", | |
" content: '►';\n", | |
" font-size: 11px;\n", | |
" width: 15px;\n", | |
" text-align: center;\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in:disabled + label:before {\n", | |
" color: var(--xr-disabled-color);\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in:checked + label:before {\n", | |
" content: '▼';\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in:checked + label > span {\n", | |
" display: none;\n", | |
"}\n", | |
"\n", | |
".xr-section-summary,\n", | |
".xr-section-inline-details {\n", | |
" padding-top: 4px;\n", | |
" padding-bottom: 4px;\n", | |
"}\n", | |
"\n", | |
".xr-section-inline-details {\n", | |
" grid-column: 2 / -1;\n", | |
"}\n", | |
"\n", | |
".xr-section-details {\n", | |
" display: none;\n", | |
" grid-column: 1 / -1;\n", | |
" margin-bottom: 5px;\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in:checked ~ .xr-section-details {\n", | |
" display: contents;\n", | |
"}\n", | |
"\n", | |
".xr-array-wrap {\n", | |
" grid-column: 1 / -1;\n", | |
" display: grid;\n", | |
" grid-template-columns: 20px auto;\n", | |
"}\n", | |
"\n", | |
".xr-array-wrap > label {\n", | |
" grid-column: 1;\n", | |
" vertical-align: top;\n", | |
"}\n", | |
"\n", | |
".xr-preview {\n", | |
" color: var(--xr-font-color3);\n", | |
"}\n", | |
"\n", | |
".xr-array-preview,\n", | |
".xr-array-data {\n", | |
" padding: 0 5px !important;\n", | |
" grid-column: 2;\n", | |
"}\n", | |
"\n", | |
".xr-array-data,\n", | |
".xr-array-in:checked ~ .xr-array-preview {\n", | |
" display: none;\n", | |
"}\n", | |
"\n", | |
".xr-array-in:checked ~ .xr-array-data,\n", | |
".xr-array-preview {\n", | |
" display: inline-block;\n", | |
"}\n", | |
"\n", | |
".xr-dim-list {\n", | |
" display: inline-block !important;\n", | |
" list-style: none;\n", | |
" padding: 0 !important;\n", | |
" margin: 0;\n", | |
"}\n", | |
"\n", | |
".xr-dim-list li {\n", | |
" display: inline-block;\n", | |
" padding: 0;\n", | |
" margin: 0;\n", | |
"}\n", | |
"\n", | |
".xr-dim-list:before {\n", | |
" content: '(';\n", | |
"}\n", | |
"\n", | |
".xr-dim-list:after {\n", | |
" content: ')';\n", | |
"}\n", | |
"\n", | |
".xr-dim-list li:not(:last-child):after {\n", | |
" content: ',';\n", | |
" padding-right: 5px;\n", | |
"}\n", | |
"\n", | |
".xr-has-index {\n", | |
" font-weight: bold;\n", | |
"}\n", | |
"\n", | |
".xr-var-list,\n", | |
".xr-var-item {\n", | |
" display: contents;\n", | |
"}\n", | |
"\n", | |
".xr-var-item > div,\n", | |
".xr-var-item label,\n", | |
".xr-var-item > .xr-var-name span {\n", | |
" background-color: var(--xr-background-color-row-even);\n", | |
" margin-bottom: 0;\n", | |
"}\n", | |
"\n", | |
".xr-var-item > .xr-var-name:hover span {\n", | |
" padding-right: 5px;\n", | |
"}\n", | |
"\n", | |
".xr-var-list > li:nth-child(odd) > div,\n", | |
".xr-var-list > li:nth-child(odd) > label,\n", | |
".xr-var-list > li:nth-child(odd) > .xr-var-name span {\n", | |
" background-color: var(--xr-background-color-row-odd);\n", | |
"}\n", | |
"\n", | |
".xr-var-name {\n", | |
" grid-column: 1;\n", | |
"}\n", | |
"\n", | |
".xr-var-dims {\n", | |
" grid-column: 2;\n", | |
"}\n", | |
"\n", | |
".xr-var-dtype {\n", | |
" grid-column: 3;\n", | |
" text-align: right;\n", | |
" color: var(--xr-font-color2);\n", | |
"}\n", | |
"\n", | |
".xr-var-preview {\n", | |
" grid-column: 4;\n", | |
"}\n", | |
"\n", | |
".xr-index-preview {\n", | |
" grid-column: 2 / 5;\n", | |
" color: var(--xr-font-color2);\n", | |
"}\n", | |
"\n", | |
".xr-var-name,\n", | |
".xr-var-dims,\n", | |
".xr-var-dtype,\n", | |
".xr-preview,\n", | |
".xr-attrs dt {\n", | |
" white-space: nowrap;\n", | |
" overflow: hidden;\n", | |
" text-overflow: ellipsis;\n", | |
" padding-right: 10px;\n", | |
"}\n", | |
"\n", | |
".xr-var-name:hover,\n", | |
".xr-var-dims:hover,\n", | |
".xr-var-dtype:hover,\n", | |
".xr-attrs dt:hover {\n", | |
" overflow: visible;\n", | |
" width: auto;\n", | |
" z-index: 1;\n", | |
"}\n", | |
"\n", | |
".xr-var-attrs,\n", | |
".xr-var-data,\n", | |
".xr-index-data {\n", | |
" display: none;\n", | |
" background-color: var(--xr-background-color) !important;\n", | |
" padding-bottom: 5px !important;\n", | |
"}\n", | |
"\n", | |
".xr-var-attrs-in:checked ~ .xr-var-attrs,\n", | |
".xr-var-data-in:checked ~ .xr-var-data,\n", | |
".xr-index-data-in:checked ~ .xr-index-data {\n", | |
" display: block;\n", | |
"}\n", | |
"\n", | |
".xr-var-data > table {\n", | |
" float: right;\n", | |
"}\n", | |
"\n", | |
".xr-var-name span,\n", | |
".xr-var-data,\n", | |
".xr-index-name div,\n", | |
".xr-index-data,\n", | |
".xr-attrs {\n", | |
" padding-left: 25px !important;\n", | |
"}\n", | |
"\n", | |
".xr-attrs,\n", | |
".xr-var-attrs,\n", | |
".xr-var-data,\n", | |
".xr-index-data {\n", | |
" grid-column: 1 / -1;\n", | |
"}\n", | |
"\n", | |
"dl.xr-attrs {\n", | |
" padding: 0;\n", | |
" margin: 0;\n", | |
" display: grid;\n", | |
" grid-template-columns: 125px auto;\n", | |
"}\n", | |
"\n", | |
".xr-attrs dt,\n", | |
".xr-attrs dd {\n", | |
" padding: 0;\n", | |
" margin: 0;\n", | |
" float: left;\n", | |
" padding-right: 10px;\n", | |
" width: auto;\n", | |
"}\n", | |
"\n", | |
".xr-attrs dt {\n", | |
" font-weight: normal;\n", | |
" grid-column: 1;\n", | |
"}\n", | |
"\n", | |
".xr-attrs dt:hover span {\n", | |
" display: inline-block;\n", | |
" background: var(--xr-background-color);\n", | |
" padding-right: 10px;\n", | |
"}\n", | |
"\n", | |
".xr-attrs dd {\n", | |
" grid-column: 2;\n", | |
" white-space: pre-wrap;\n", | |
" word-break: break-all;\n", | |
"}\n", | |
"\n", | |
".xr-icon-database,\n", | |
".xr-icon-file-text2,\n", | |
".xr-no-icon {\n", | |
" display: inline-block;\n", | |
" vertical-align: middle;\n", | |
" width: 1em;\n", | |
" height: 1.5em !important;\n", | |
" stroke-width: 0;\n", | |
" stroke: currentColor;\n", | |
" fill: currentColor;\n", | |
"}\n", | |
"</style><pre class='xr-text-repr-fallback'><xarray.Dataset>\n", | |
"Dimensions: (query: 2, sample: 3, embedding: 3)\n", | |
"Coordinates:\n", | |
" vectors (query, sample, embedding) int64 0 1 0 0 0 1 ... 0 0 1 0 1 0\n", | |
" distance (query, sample) float64 0.8 1.02 1.02 0.5 1.118 1.118\n", | |
" k (sample) int64 1 2 3\n", | |
" query_vectors (query, embedding) float64 0.0 0.2 0.0 0.5 0.0 0.0\n", | |
"Dimensions without coordinates: query, sample, embedding\n", | |
"Data variables:\n", | |
" documents (query, sample) <U1 'b' 'c' 'a' 'a' 'c' 'b'</pre><div class='xr-wrap' style='display:none'><div class='xr-header'><div class='xr-obj-type'>xarray.Dataset</div></div><ul class='xr-sections'><li class='xr-section-item'><input id='section-02236704-5497-4ad0-8f9c-98ec9eb38344' class='xr-section-summary-in' type='checkbox' disabled ><label for='section-02236704-5497-4ad0-8f9c-98ec9eb38344' class='xr-section-summary' title='Expand/collapse section'>Dimensions:</label><div class='xr-section-inline-details'><ul class='xr-dim-list'><li><span>query</span>: 2</li><li><span>sample</span>: 3</li><li><span>embedding</span>: 3</li></ul></div><div class='xr-section-details'></div></li><li class='xr-section-item'><input id='section-b163fd24-21c9-4d18-945d-0b9ee2bbc67b' class='xr-section-summary-in' type='checkbox' checked><label for='section-b163fd24-21c9-4d18-945d-0b9ee2bbc67b' class='xr-section-summary' >Coordinates: <span>(4)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><ul class='xr-var-list'><li class='xr-var-item'><div class='xr-var-name'><span>vectors</span></div><div class='xr-var-dims'>(query, sample, embedding)</div><div class='xr-var-dtype'>int64</div><div class='xr-var-preview xr-preview'>0 1 0 0 0 1 1 0 0 1 0 0 0 0 1 0 1 0</div><input id='attrs-813bc849-c1b2-4fcf-ade2-98ea36173dc1' class='xr-var-attrs-in' type='checkbox' disabled><label for='attrs-813bc849-c1b2-4fcf-ade2-98ea36173dc1' title='Show/Hide attributes'><svg class='icon xr-icon-file-text2'><use xlink:href='#icon-file-text2'></use></svg></label><input id='data-c743f71d-4963-4ae3-a21f-c7e167c2e415' class='xr-var-data-in' type='checkbox'><label for='data-c743f71d-4963-4ae3-a21f-c7e167c2e415' title='Show/Hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-var-attrs'><dl class='xr-attrs'></dl></div><div class='xr-var-data'><pre>array([[[0, 1, 0],\n", | |
" [0, 0, 1],\n", | |
" [1, 0, 0]],\n", | |
"\n", | |
" [[1, 0, 0],\n", | |
" [0, 0, 1],\n", | |
" [0, 1, 0]]])</pre></div></li><li class='xr-var-item'><div class='xr-var-name'><span>distance</span></div><div class='xr-var-dims'>(query, sample)</div><div class='xr-var-dtype'>float64</div><div class='xr-var-preview xr-preview'>0.8 1.02 1.02 0.5 1.118 1.118</div><input id='attrs-fd7fe827-8082-4319-80d0-919e0cb6c291' class='xr-var-attrs-in' type='checkbox' disabled><label for='attrs-fd7fe827-8082-4319-80d0-919e0cb6c291' title='Show/Hide attributes'><svg class='icon xr-icon-file-text2'><use xlink:href='#icon-file-text2'></use></svg></label><input id='data-f4d0c881-61b1-4157-9664-ff48961cee23' class='xr-var-data-in' type='checkbox'><label for='data-f4d0c881-61b1-4157-9664-ff48961cee23' title='Show/Hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-var-attrs'><dl class='xr-attrs'></dl></div><div class='xr-var-data'><pre>array([[0.8 , 1.0198039 , 1.0198039 ],\n", | |
" [0.5 , 1.11803399, 1.11803399]])</pre></div></li><li class='xr-var-item'><div class='xr-var-name'><span>k</span></div><div class='xr-var-dims'>(sample)</div><div class='xr-var-dtype'>int64</div><div class='xr-var-preview xr-preview'>1 2 3</div><input id='attrs-b3958425-ef44-494f-a902-02b4883fa2e4' class='xr-var-attrs-in' type='checkbox' disabled><label for='attrs-b3958425-ef44-494f-a902-02b4883fa2e4' title='Show/Hide attributes'><svg class='icon xr-icon-file-text2'><use xlink:href='#icon-file-text2'></use></svg></label><input id='data-05971fd0-7777-4777-ad80-1c98f1c010e8' class='xr-var-data-in' type='checkbox'><label for='data-05971fd0-7777-4777-ad80-1c98f1c010e8' title='Show/Hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-var-attrs'><dl class='xr-attrs'></dl></div><div class='xr-var-data'><pre>array([1, 2, 3])</pre></div></li><li class='xr-var-item'><div class='xr-var-name'><span>query_vectors</span></div><div class='xr-var-dims'>(query, embedding)</div><div class='xr-var-dtype'>float64</div><div class='xr-var-preview xr-preview'>0.0 0.2 0.0 0.5 0.0 0.0</div><input id='attrs-a319f2b8-2cc9-43eb-998f-e6feaf42833d' class='xr-var-attrs-in' type='checkbox' disabled><label for='attrs-a319f2b8-2cc9-43eb-998f-e6feaf42833d' title='Show/Hide attributes'><svg class='icon xr-icon-file-text2'><use xlink:href='#icon-file-text2'></use></svg></label><input id='data-c56da739-61c3-4e4e-9cbd-b542d355f0c1' class='xr-var-data-in' type='checkbox'><label for='data-c56da739-61c3-4e4e-9cbd-b542d355f0c1' title='Show/Hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-var-attrs'><dl class='xr-attrs'></dl></div><div class='xr-var-data'><pre>array([[0. , 0.2, 0. ],\n", | |
" [0.5, 0. , 0. ]])</pre></div></li></ul></div></li><li class='xr-section-item'><input id='section-3e55b630-0b89-468c-a09f-22d041e44f12' class='xr-section-summary-in' type='checkbox' checked><label for='section-3e55b630-0b89-468c-a09f-22d041e44f12' class='xr-section-summary' >Data variables: <span>(1)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><ul class='xr-var-list'><li class='xr-var-item'><div class='xr-var-name'><span>documents</span></div><div class='xr-var-dims'>(query, sample)</div><div class='xr-var-dtype'><U1</div><div class='xr-var-preview xr-preview'>'b' 'c' 'a' 'a' 'c' 'b'</div><input id='attrs-cb6e87e4-3ddc-4142-8b65-57bb6f3ae57d' class='xr-var-attrs-in' type='checkbox' disabled><label for='attrs-cb6e87e4-3ddc-4142-8b65-57bb6f3ae57d' title='Show/Hide attributes'><svg class='icon xr-icon-file-text2'><use xlink:href='#icon-file-text2'></use></svg></label><input id='data-97844c4a-21cd-463f-853a-0e035a1c85e8' class='xr-var-data-in' type='checkbox'><label for='data-97844c4a-21cd-463f-853a-0e035a1c85e8' title='Show/Hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-var-attrs'><dl class='xr-attrs'></dl></div><div class='xr-var-data'><pre>array([['b', 'c', 'a'],\n", | |
" ['a', 'c', 'b']], dtype='<U1')</pre></div></li></ul></div></li><li class='xr-section-item'><input id='section-41aa5911-4a24-4324-a50e-eb824a3f60d9' class='xr-section-summary-in' type='checkbox' disabled ><label for='section-41aa5911-4a24-4324-a50e-eb824a3f60d9' class='xr-section-summary' title='Expand/collapse section'>Indexes: <span>(0)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><ul class='xr-var-list'></ul></div></li><li class='xr-section-item'><input id='section-5a4075c2-48ce-4aa9-a984-1b7c719df083' class='xr-section-summary-in' type='checkbox' disabled ><label for='section-5a4075c2-48ce-4aa9-a984-1b7c719df083' class='xr-section-summary' title='Expand/collapse section'>Attributes: <span>(0)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><dl class='xr-attrs'></dl></div></li></ul></div></div>" | |
], | |
"text/plain": [ | |
"<xarray.Dataset>\n", | |
"Dimensions: (query: 2, sample: 3, embedding: 3)\n", | |
"Coordinates:\n", | |
" vectors (query, sample, embedding) int64 0 1 0 0 0 1 ... 0 0 1 0 1 0\n", | |
" distance (query, sample) float64 0.8 1.02 1.02 0.5 1.118 1.118\n", | |
" k (sample) int64 1 2 3\n", | |
" query_vectors (query, embedding) float64 0.0 0.2 0.0 0.5 0.0 0.0\n", | |
"Dimensions without coordinates: query, sample, embedding\n", | |
"Data variables:\n", | |
" documents (query, sample) <U1 'b' 'c' 'a' 'a' 'c' 'b'" | |
] | |
}, | |
"execution_count": 8, | |
"metadata": {}, | |
"output_type": "execute_result" | |
} | |
], | |
"source": [ | |
"result = ds.vector.query(\n", | |
" [\n", | |
" [0, 0.2, 0],\n", | |
" [0.5, 0, 0],\n", | |
" ],\n", | |
" k=3\n", | |
")\n", | |
"result" | |
] | |
}, | |
{ | |
"attachments": {}, | |
"cell_type": "markdown", | |
"metadata": {}, | |
"source": [ | |
"We could also use `k` as the dimension coordinate." | |
] | |
}, | |
{ | |
"cell_type": "code", | |
"execution_count": 9, | |
"metadata": {}, | |
"outputs": [ | |
{ | |
"data": { | |
"text/html": [ | |
"<div><svg style=\"position: absolute; width: 0; height: 0; overflow: hidden\">\n", | |
"<defs>\n", | |
"<symbol id=\"icon-database\" viewBox=\"0 0 32 32\">\n", | |
"<path d=\"M16 0c-8.837 0-16 2.239-16 5v4c0 2.761 7.163 5 16 5s16-2.239 16-5v-4c0-2.761-7.163-5-16-5z\"></path>\n", | |
"<path d=\"M16 17c-8.837 0-16-2.239-16-5v6c0 2.761 7.163 5 16 5s16-2.239 16-5v-6c0 2.761-7.163 5-16 5z\"></path>\n", | |
"<path d=\"M16 26c-8.837 0-16-2.239-16-5v6c0 2.761 7.163 5 16 5s16-2.239 16-5v-6c0 2.761-7.163 5-16 5z\"></path>\n", | |
"</symbol>\n", | |
"<symbol id=\"icon-file-text2\" viewBox=\"0 0 32 32\">\n", | |
"<path d=\"M28.681 7.159c-0.694-0.947-1.662-2.053-2.724-3.116s-2.169-2.030-3.116-2.724c-1.612-1.182-2.393-1.319-2.841-1.319h-15.5c-1.378 0-2.5 1.121-2.5 2.5v27c0 1.378 1.122 2.5 2.5 2.5h23c1.378 0 2.5-1.122 2.5-2.5v-19.5c0-0.448-0.137-1.23-1.319-2.841zM24.543 5.457c0.959 0.959 1.712 1.825 2.268 2.543h-4.811v-4.811c0.718 0.556 1.584 1.309 2.543 2.268zM28 29.5c0 0.271-0.229 0.5-0.5 0.5h-23c-0.271 0-0.5-0.229-0.5-0.5v-27c0-0.271 0.229-0.5 0.5-0.5 0 0 15.499-0 15.5 0v7c0 0.552 0.448 1 1 1h7v19.5z\"></path>\n", | |
"<path d=\"M23 26h-14c-0.552 0-1-0.448-1-1s0.448-1 1-1h14c0.552 0 1 0.448 1 1s-0.448 1-1 1z\"></path>\n", | |
"<path d=\"M23 22h-14c-0.552 0-1-0.448-1-1s0.448-1 1-1h14c0.552 0 1 0.448 1 1s-0.448 1-1 1z\"></path>\n", | |
"<path d=\"M23 18h-14c-0.552 0-1-0.448-1-1s0.448-1 1-1h14c0.552 0 1 0.448 1 1s-0.448 1-1 1z\"></path>\n", | |
"</symbol>\n", | |
"</defs>\n", | |
"</svg>\n", | |
"<style>/* CSS stylesheet for displaying xarray objects in jupyterlab.\n", | |
" *\n", | |
" */\n", | |
"\n", | |
":root {\n", | |
" --xr-font-color0: var(--jp-content-font-color0, rgba(0, 0, 0, 1));\n", | |
" --xr-font-color2: var(--jp-content-font-color2, rgba(0, 0, 0, 0.54));\n", | |
" --xr-font-color3: var(--jp-content-font-color3, rgba(0, 0, 0, 0.38));\n", | |
" --xr-border-color: var(--jp-border-color2, #e0e0e0);\n", | |
" --xr-disabled-color: var(--jp-layout-color3, #bdbdbd);\n", | |
" --xr-background-color: var(--jp-layout-color0, white);\n", | |
" --xr-background-color-row-even: var(--jp-layout-color1, white);\n", | |
" --xr-background-color-row-odd: var(--jp-layout-color2, #eeeeee);\n", | |
"}\n", | |
"\n", | |
"html[theme=dark],\n", | |
"body[data-theme=dark],\n", | |
"body.vscode-dark {\n", | |
" --xr-font-color0: rgba(255, 255, 255, 1);\n", | |
" --xr-font-color2: rgba(255, 255, 255, 0.54);\n", | |
" --xr-font-color3: rgba(255, 255, 255, 0.38);\n", | |
" --xr-border-color: #1F1F1F;\n", | |
" --xr-disabled-color: #515151;\n", | |
" --xr-background-color: #111111;\n", | |
" --xr-background-color-row-even: #111111;\n", | |
" --xr-background-color-row-odd: #313131;\n", | |
"}\n", | |
"\n", | |
".xr-wrap {\n", | |
" display: block !important;\n", | |
" min-width: 300px;\n", | |
" max-width: 700px;\n", | |
"}\n", | |
"\n", | |
".xr-text-repr-fallback {\n", | |
" /* fallback to plain text repr when CSS is not injected (untrusted notebook) */\n", | |
" display: none;\n", | |
"}\n", | |
"\n", | |
".xr-header {\n", | |
" padding-top: 6px;\n", | |
" padding-bottom: 6px;\n", | |
" margin-bottom: 4px;\n", | |
" border-bottom: solid 1px var(--xr-border-color);\n", | |
"}\n", | |
"\n", | |
".xr-header > div,\n", | |
".xr-header > ul {\n", | |
" display: inline;\n", | |
" margin-top: 0;\n", | |
" margin-bottom: 0;\n", | |
"}\n", | |
"\n", | |
".xr-obj-type,\n", | |
".xr-array-name {\n", | |
" margin-left: 2px;\n", | |
" margin-right: 10px;\n", | |
"}\n", | |
"\n", | |
".xr-obj-type {\n", | |
" color: var(--xr-font-color2);\n", | |
"}\n", | |
"\n", | |
".xr-sections {\n", | |
" padding-left: 0 !important;\n", | |
" display: grid;\n", | |
" grid-template-columns: 150px auto auto 1fr 20px 20px;\n", | |
"}\n", | |
"\n", | |
".xr-section-item {\n", | |
" display: contents;\n", | |
"}\n", | |
"\n", | |
".xr-section-item input {\n", | |
" display: none;\n", | |
"}\n", | |
"\n", | |
".xr-section-item input + label {\n", | |
" color: var(--xr-disabled-color);\n", | |
"}\n", | |
"\n", | |
".xr-section-item input:enabled + label {\n", | |
" cursor: pointer;\n", | |
" color: var(--xr-font-color2);\n", | |
"}\n", | |
"\n", | |
".xr-section-item input:enabled + label:hover {\n", | |
" color: var(--xr-font-color0);\n", | |
"}\n", | |
"\n", | |
".xr-section-summary {\n", | |
" grid-column: 1;\n", | |
" color: var(--xr-font-color2);\n", | |
" font-weight: 500;\n", | |
"}\n", | |
"\n", | |
".xr-section-summary > span {\n", | |
" display: inline-block;\n", | |
" padding-left: 0.5em;\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in:disabled + label {\n", | |
" color: var(--xr-font-color2);\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in + label:before {\n", | |
" display: inline-block;\n", | |
" content: '►';\n", | |
" font-size: 11px;\n", | |
" width: 15px;\n", | |
" text-align: center;\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in:disabled + label:before {\n", | |
" color: var(--xr-disabled-color);\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in:checked + label:before {\n", | |
" content: '▼';\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in:checked + label > span {\n", | |
" display: none;\n", | |
"}\n", | |
"\n", | |
".xr-section-summary,\n", | |
".xr-section-inline-details {\n", | |
" padding-top: 4px;\n", | |
" padding-bottom: 4px;\n", | |
"}\n", | |
"\n", | |
".xr-section-inline-details {\n", | |
" grid-column: 2 / -1;\n", | |
"}\n", | |
"\n", | |
".xr-section-details {\n", | |
" display: none;\n", | |
" grid-column: 1 / -1;\n", | |
" margin-bottom: 5px;\n", | |
"}\n", | |
"\n", | |
".xr-section-summary-in:checked ~ .xr-section-details {\n", | |
" display: contents;\n", | |
"}\n", | |
"\n", | |
".xr-array-wrap {\n", | |
" grid-column: 1 / -1;\n", | |
" display: grid;\n", | |
" grid-template-columns: 20px auto;\n", | |
"}\n", | |
"\n", | |
".xr-array-wrap > label {\n", | |
" grid-column: 1;\n", | |
" vertical-align: top;\n", | |
"}\n", | |
"\n", | |
".xr-preview {\n", | |
" color: var(--xr-font-color3);\n", | |
"}\n", | |
"\n", | |
".xr-array-preview,\n", | |
".xr-array-data {\n", | |
" padding: 0 5px !important;\n", | |
" grid-column: 2;\n", | |
"}\n", | |
"\n", | |
".xr-array-data,\n", | |
".xr-array-in:checked ~ .xr-array-preview {\n", | |
" display: none;\n", | |
"}\n", | |
"\n", | |
".xr-array-in:checked ~ .xr-array-data,\n", | |
".xr-array-preview {\n", | |
" display: inline-block;\n", | |
"}\n", | |
"\n", | |
".xr-dim-list {\n", | |
" display: inline-block !important;\n", | |
" list-style: none;\n", | |
" padding: 0 !important;\n", | |
" margin: 0;\n", | |
"}\n", | |
"\n", | |
".xr-dim-list li {\n", | |
" display: inline-block;\n", | |
" padding: 0;\n", | |
" margin: 0;\n", | |
"}\n", | |
"\n", | |
".xr-dim-list:before {\n", | |
" content: '(';\n", | |
"}\n", | |
"\n", | |
".xr-dim-list:after {\n", | |
" content: ')';\n", | |
"}\n", | |
"\n", | |
".xr-dim-list li:not(:last-child):after {\n", | |
" content: ',';\n", | |
" padding-right: 5px;\n", | |
"}\n", | |
"\n", | |
".xr-has-index {\n", | |
" font-weight: bold;\n", | |
"}\n", | |
"\n", | |
".xr-var-list,\n", | |
".xr-var-item {\n", | |
" display: contents;\n", | |
"}\n", | |
"\n", | |
".xr-var-item > div,\n", | |
".xr-var-item label,\n", | |
".xr-var-item > .xr-var-name span {\n", | |
" background-color: var(--xr-background-color-row-even);\n", | |
" margin-bottom: 0;\n", | |
"}\n", | |
"\n", | |
".xr-var-item > .xr-var-name:hover span {\n", | |
" padding-right: 5px;\n", | |
"}\n", | |
"\n", | |
".xr-var-list > li:nth-child(odd) > div,\n", | |
".xr-var-list > li:nth-child(odd) > label,\n", | |
".xr-var-list > li:nth-child(odd) > .xr-var-name span {\n", | |
" background-color: var(--xr-background-color-row-odd);\n", | |
"}\n", | |
"\n", | |
".xr-var-name {\n", | |
" grid-column: 1;\n", | |
"}\n", | |
"\n", | |
".xr-var-dims {\n", | |
" grid-column: 2;\n", | |
"}\n", | |
"\n", | |
".xr-var-dtype {\n", | |
" grid-column: 3;\n", | |
" text-align: right;\n", | |
" color: var(--xr-font-color2);\n", | |
"}\n", | |
"\n", | |
".xr-var-preview {\n", | |
" grid-column: 4;\n", | |
"}\n", | |
"\n", | |
".xr-index-preview {\n", | |
" grid-column: 2 / 5;\n", | |
" color: var(--xr-font-color2);\n", | |
"}\n", | |
"\n", | |
".xr-var-name,\n", | |
".xr-var-dims,\n", | |
".xr-var-dtype,\n", | |
".xr-preview,\n", | |
".xr-attrs dt {\n", | |
" white-space: nowrap;\n", | |
" overflow: hidden;\n", | |
" text-overflow: ellipsis;\n", | |
" padding-right: 10px;\n", | |
"}\n", | |
"\n", | |
".xr-var-name:hover,\n", | |
".xr-var-dims:hover,\n", | |
".xr-var-dtype:hover,\n", | |
".xr-attrs dt:hover {\n", | |
" overflow: visible;\n", | |
" width: auto;\n", | |
" z-index: 1;\n", | |
"}\n", | |
"\n", | |
".xr-var-attrs,\n", | |
".xr-var-data,\n", | |
".xr-index-data {\n", | |
" display: none;\n", | |
" background-color: var(--xr-background-color) !important;\n", | |
" padding-bottom: 5px !important;\n", | |
"}\n", | |
"\n", | |
".xr-var-attrs-in:checked ~ .xr-var-attrs,\n", | |
".xr-var-data-in:checked ~ .xr-var-data,\n", | |
".xr-index-data-in:checked ~ .xr-index-data {\n", | |
" display: block;\n", | |
"}\n", | |
"\n", | |
".xr-var-data > table {\n", | |
" float: right;\n", | |
"}\n", | |
"\n", | |
".xr-var-name span,\n", | |
".xr-var-data,\n", | |
".xr-index-name div,\n", | |
".xr-index-data,\n", | |
".xr-attrs {\n", | |
" padding-left: 25px !important;\n", | |
"}\n", | |
"\n", | |
".xr-attrs,\n", | |
".xr-var-attrs,\n", | |
".xr-var-data,\n", | |
".xr-index-data {\n", | |
" grid-column: 1 / -1;\n", | |
"}\n", | |
"\n", | |
"dl.xr-attrs {\n", | |
" padding: 0;\n", | |
" margin: 0;\n", | |
" display: grid;\n", | |
" grid-template-columns: 125px auto;\n", | |
"}\n", | |
"\n", | |
".xr-attrs dt,\n", | |
".xr-attrs dd {\n", | |
" padding: 0;\n", | |
" margin: 0;\n", | |
" float: left;\n", | |
" padding-right: 10px;\n", | |
" width: auto;\n", | |
"}\n", | |
"\n", | |
".xr-attrs dt {\n", | |
" font-weight: normal;\n", | |
" grid-column: 1;\n", | |
"}\n", | |
"\n", | |
".xr-attrs dt:hover span {\n", | |
" display: inline-block;\n", | |
" background: var(--xr-background-color);\n", | |
" padding-right: 10px;\n", | |
"}\n", | |
"\n", | |
".xr-attrs dd {\n", | |
" grid-column: 2;\n", | |
" white-space: pre-wrap;\n", | |
" word-break: break-all;\n", | |
"}\n", | |
"\n", | |
".xr-icon-database,\n", | |
".xr-icon-file-text2,\n", | |
".xr-no-icon {\n", | |
" display: inline-block;\n", | |
" vertical-align: middle;\n", | |
" width: 1em;\n", | |
" height: 1.5em !important;\n", | |
" stroke-width: 0;\n", | |
" stroke: currentColor;\n", | |
" fill: currentColor;\n", | |
"}\n", | |
"</style><pre class='xr-text-repr-fallback'><xarray.Dataset>\n", | |
"Dimensions: (query: 2, k: 3, embedding: 3)\n", | |
"Coordinates:\n", | |
" vectors (query, k, embedding) int64 0 1 0 0 0 1 1 0 ... 0 0 0 1 0 1 0\n", | |
" distance (query, k) float64 0.8 1.02 1.02 0.5 1.118 1.118\n", | |
" * k (k) int64 1 2 3\n", | |
" query_vectors (query, embedding) float64 0.0 0.2 0.0 0.5 0.0 0.0\n", | |
"Dimensions without coordinates: query, embedding\n", | |
"Data variables:\n", | |
" documents (query, k) <U1 'b' 'c' 'a' 'a' 'c' 'b'</pre><div class='xr-wrap' style='display:none'><div class='xr-header'><div class='xr-obj-type'>xarray.Dataset</div></div><ul class='xr-sections'><li class='xr-section-item'><input id='section-8d4001e9-77a0-44a3-b019-bed9a5b3a6f7' class='xr-section-summary-in' type='checkbox' disabled ><label for='section-8d4001e9-77a0-44a3-b019-bed9a5b3a6f7' class='xr-section-summary' title='Expand/collapse section'>Dimensions:</label><div class='xr-section-inline-details'><ul class='xr-dim-list'><li><span>query</span>: 2</li><li><span class='xr-has-index'>k</span>: 3</li><li><span>embedding</span>: 3</li></ul></div><div class='xr-section-details'></div></li><li class='xr-section-item'><input id='section-5938a602-e9bc-43d9-8616-233a22b192b7' class='xr-section-summary-in' type='checkbox' checked><label for='section-5938a602-e9bc-43d9-8616-233a22b192b7' class='xr-section-summary' >Coordinates: <span>(4)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><ul class='xr-var-list'><li class='xr-var-item'><div class='xr-var-name'><span>vectors</span></div><div class='xr-var-dims'>(query, k, embedding)</div><div class='xr-var-dtype'>int64</div><div class='xr-var-preview xr-preview'>0 1 0 0 0 1 1 0 0 1 0 0 0 0 1 0 1 0</div><input id='attrs-bb6d8b53-ebe1-4395-8d38-22d2e0b13efa' class='xr-var-attrs-in' type='checkbox' disabled><label for='attrs-bb6d8b53-ebe1-4395-8d38-22d2e0b13efa' title='Show/Hide attributes'><svg class='icon xr-icon-file-text2'><use xlink:href='#icon-file-text2'></use></svg></label><input id='data-a30d0fc2-2554-443d-a67f-f78584239181' class='xr-var-data-in' type='checkbox'><label for='data-a30d0fc2-2554-443d-a67f-f78584239181' title='Show/Hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-var-attrs'><dl class='xr-attrs'></dl></div><div class='xr-var-data'><pre>array([[[0, 1, 0],\n", | |
" [0, 0, 1],\n", | |
" [1, 0, 0]],\n", | |
"\n", | |
" [[1, 0, 0],\n", | |
" [0, 0, 1],\n", | |
" [0, 1, 0]]])</pre></div></li><li class='xr-var-item'><div class='xr-var-name'><span>distance</span></div><div class='xr-var-dims'>(query, k)</div><div class='xr-var-dtype'>float64</div><div class='xr-var-preview xr-preview'>0.8 1.02 1.02 0.5 1.118 1.118</div><input id='attrs-75731e6e-67c6-4665-aecc-bec3b74a3c05' class='xr-var-attrs-in' type='checkbox' disabled><label for='attrs-75731e6e-67c6-4665-aecc-bec3b74a3c05' title='Show/Hide attributes'><svg class='icon xr-icon-file-text2'><use xlink:href='#icon-file-text2'></use></svg></label><input id='data-251a1348-fe97-438b-9f7b-a158bdd7a260' class='xr-var-data-in' type='checkbox'><label for='data-251a1348-fe97-438b-9f7b-a158bdd7a260' title='Show/Hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-var-attrs'><dl class='xr-attrs'></dl></div><div class='xr-var-data'><pre>array([[0.8 , 1.0198039 , 1.0198039 ],\n", | |
" [0.5 , 1.11803399, 1.11803399]])</pre></div></li><li class='xr-var-item'><div class='xr-var-name'><span class='xr-has-index'>k</span></div><div class='xr-var-dims'>(k)</div><div class='xr-var-dtype'>int64</div><div class='xr-var-preview xr-preview'>1 2 3</div><input id='attrs-530594f3-2acb-4661-a825-2fdadd3edfa7' class='xr-var-attrs-in' type='checkbox' disabled><label for='attrs-530594f3-2acb-4661-a825-2fdadd3edfa7' title='Show/Hide attributes'><svg class='icon xr-icon-file-text2'><use xlink:href='#icon-file-text2'></use></svg></label><input id='data-ed92bf5b-3f5f-4b4d-878b-67fedc42a24b' class='xr-var-data-in' type='checkbox'><label for='data-ed92bf5b-3f5f-4b4d-878b-67fedc42a24b' title='Show/Hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-var-attrs'><dl class='xr-attrs'></dl></div><div class='xr-var-data'><pre>array([1, 2, 3])</pre></div></li><li class='xr-var-item'><div class='xr-var-name'><span>query_vectors</span></div><div class='xr-var-dims'>(query, embedding)</div><div class='xr-var-dtype'>float64</div><div class='xr-var-preview xr-preview'>0.0 0.2 0.0 0.5 0.0 0.0</div><input id='attrs-958c8412-c14e-4fb0-94f3-c03053d8d0bc' class='xr-var-attrs-in' type='checkbox' disabled><label for='attrs-958c8412-c14e-4fb0-94f3-c03053d8d0bc' title='Show/Hide attributes'><svg class='icon xr-icon-file-text2'><use xlink:href='#icon-file-text2'></use></svg></label><input id='data-92dd3ccb-6364-4b00-8847-f4f37df177c8' class='xr-var-data-in' type='checkbox'><label for='data-92dd3ccb-6364-4b00-8847-f4f37df177c8' title='Show/Hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-var-attrs'><dl class='xr-attrs'></dl></div><div class='xr-var-data'><pre>array([[0. , 0.2, 0. ],\n", | |
" [0.5, 0. , 0. ]])</pre></div></li></ul></div></li><li class='xr-section-item'><input id='section-70e4a5d5-1405-41fc-a54c-608721291f9d' class='xr-section-summary-in' type='checkbox' checked><label for='section-70e4a5d5-1405-41fc-a54c-608721291f9d' class='xr-section-summary' >Data variables: <span>(1)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><ul class='xr-var-list'><li class='xr-var-item'><div class='xr-var-name'><span>documents</span></div><div class='xr-var-dims'>(query, k)</div><div class='xr-var-dtype'><U1</div><div class='xr-var-preview xr-preview'>'b' 'c' 'a' 'a' 'c' 'b'</div><input id='attrs-0e332200-956a-4949-b081-4963e55ee19e' class='xr-var-attrs-in' type='checkbox' disabled><label for='attrs-0e332200-956a-4949-b081-4963e55ee19e' title='Show/Hide attributes'><svg class='icon xr-icon-file-text2'><use xlink:href='#icon-file-text2'></use></svg></label><input id='data-d2b393fe-c0cf-478a-a625-3adb68f93319' class='xr-var-data-in' type='checkbox'><label for='data-d2b393fe-c0cf-478a-a625-3adb68f93319' title='Show/Hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-var-attrs'><dl class='xr-attrs'></dl></div><div class='xr-var-data'><pre>array([['b', 'c', 'a'],\n", | |
" ['a', 'c', 'b']], dtype='<U1')</pre></div></li></ul></div></li><li class='xr-section-item'><input id='section-88fed81f-c2dc-4878-9716-42243af6147a' class='xr-section-summary-in' type='checkbox' ><label for='section-88fed81f-c2dc-4878-9716-42243af6147a' class='xr-section-summary' >Indexes: <span>(1)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><ul class='xr-var-list'><li class='xr-var-item'><div class='xr-index-name'><div>k</div></div><div class='xr-index-preview'>PandasIndex</div><div></div><input id='index-75528940-cecb-4964-bccd-1b222a79d0b5' class='xr-index-data-in' type='checkbox'/><label for='index-75528940-cecb-4964-bccd-1b222a79d0b5' title='Show/Hide index repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-index-data'><pre>PandasIndex(Int64Index([1, 2, 3], dtype='int64', name='k'))</pre></div></li></ul></div></li><li class='xr-section-item'><input id='section-de062f1d-9e3a-4cc8-982e-15e8a047463b' class='xr-section-summary-in' type='checkbox' disabled ><label for='section-de062f1d-9e3a-4cc8-982e-15e8a047463b' class='xr-section-summary' title='Expand/collapse section'>Attributes: <span>(0)</span></label><div class='xr-section-inline-details'></div><div class='xr-section-details'><dl class='xr-attrs'></dl></div></li></ul></div></div>" | |
], | |
"text/plain": [ | |
"<xarray.Dataset>\n", | |
"Dimensions: (query: 2, k: 3, embedding: 3)\n", | |
"Coordinates:\n", | |
" vectors (query, k, embedding) int64 0 1 0 0 0 1 1 0 ... 0 0 0 1 0 1 0\n", | |
" distance (query, k) float64 0.8 1.02 1.02 0.5 1.118 1.118\n", | |
" * k (k) int64 1 2 3\n", | |
" query_vectors (query, embedding) float64 0.0 0.2 0.0 0.5 0.0 0.0\n", | |
"Dimensions without coordinates: query, embedding\n", | |
"Data variables:\n", | |
" documents (query, k) <U1 'b' 'c' 'a' 'a' 'c' 'b'" | |
] | |
}, | |
"execution_count": 9, | |
"metadata": {}, | |
"output_type": "execute_result" | |
} | |
], | |
"source": [ | |
"result.swap_dims({\"sample\": \"k\"})" | |
] | |
}, | |
{ | |
"cell_type": "code", | |
"execution_count": null, | |
"metadata": {}, | |
"outputs": [], | |
"source": [] | |
} | |
], | |
"metadata": { | |
"kernelspec": { | |
"display_name": "xarray-tests", | |
"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.10.9" | |
}, | |
"orig_nbformat": 4 | |
}, | |
"nbformat": 4, | |
"nbformat_minor": 2 | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment