{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2022-04-10T03:15:33.69415Z","iopub.execute_input":"2022-04-10T03:15:33.694879Z","iopub.status.idle":"2022-04-10T03:15:33.698436Z","shell.execute_reply.started":"2022-04-10T03:15:33.69484Z","shell.execute_reply":"2022-04-10T03:15:33.697748Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Vector Indexing\n\nThis notebook follows on from https://www.kaggle.com/code/prubyg/hotel-id-vector-extraction - please read that if you haven't already.\n\nThe purpose of this notebook is to index all the vectors extracted by that process using the \"annoy\" library from Spotify. We do this in a separate notebook, because it is both a non-GPU operation, so allows the GPU work to be separated, and these both produce many GB of output. In some earlier attempts, I was very much pushing the storage output limits of notebooks, so the split was helpful.","metadata":{}},{"cell_type":"markdown","source":"We set some parameters for this run up-front.","metadata":{}},{"cell_type":"code","source":"VECTOR_WIDTH = 320\nNUM_TREES = 30\n\nOUTPUT_DIR = '/kaggle/working/'\nINPUT_FILE = '/kaggle/input/hotel-id-vector-extraction/vectors.parquet'","metadata":{"execution":{"iopub.status.busy":"2022-04-10T03:15:33.735007Z","iopub.execute_input":"2022-04-10T03:15:33.735251Z","iopub.status.idle":"2022-04-10T03:15:33.73926Z","shell.execute_reply.started":"2022-04-10T03:15:33.735208Z","shell.execute_reply":"2022-04-10T03:15:33.738674Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Import our libraries - primarily pyarrow for data (please do not treat this as a recommendation), and annoy for indexing (highly recommended)","metadata":{}},{"cell_type":"code","source":"import pyarrow as pa\nfrom pyarrow.parquet import ParquetFile\nfrom annoy import AnnoyIndex","metadata":{"execution":{"iopub.status.busy":"2022-04-10T03:15:33.771504Z","iopub.execute_input":"2022-04-10T03:15:33.772332Z","iopub.status.idle":"2022-04-10T03:15:33.776791Z","shell.execute_reply.started":"2022-04-10T03:15:33.772287Z","shell.execute_reply":"2022-04-10T03:15:33.776064Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Load our input file. Note this does not read the content from disk immediately - it will be read iteratively as we stream results.","metadata":{}},{"cell_type":"code","source":"vector_db = ParquetFile(INPUT_FILE)\ndisplay(vector_db.metadata)","metadata":{"execution":{"iopub.status.busy":"2022-04-10T03:15:33.795218Z","iopub.execute_input":"2022-04-10T03:15:33.795775Z","iopub.status.idle":"2022-04-10T03:15:33.810745Z","shell.execute_reply.started":"2022-04-10T03:15:33.795731Z","shell.execute_reply":"2022-04-10T03:15:33.809795Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\n\n# Construct our index and set it up to be built on disk - too large in memory\nindex = AnnoyIndex(VECTOR_WIDTH)\nindex.on_disk_build(f'{OUTPUT_DIR}/vectors.annoy')\n\n# Note with Parquet files we can specify which columns to read. The format is organised to allow other data to be entirely skipped, saving disk I/O.\nfor batch in vector_db.iter_batches(columns=['i', 'vector']):\n    for i, vector in zip(batch['i'], batch['vector']):\n        # We extract the global index and vector from the file, then add them to the index.\n        i = i.as_py()\n        vec = np.frombuffer(vector.as_buffer(), dtype=np.float32)\n        index.add_item(i, vec)","metadata":{"execution":{"iopub.status.busy":"2022-04-10T03:15:33.818108Z","iopub.execute_input":"2022-04-10T03:15:33.818592Z","iopub.status.idle":"2022-04-10T03:16:10.238629Z","shell.execute_reply.started":"2022-04-10T03:15:33.81855Z","shell.execute_reply":"2022-04-10T03:16:10.237937Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Finally, we build the index which will accelerate KNN search. The \"build\" process creates the actual search trees.\n\nNote for annoy we have to specify the number of search trees to build, because each one is random and approximate. More trees = more accuracy (but also more time).","metadata":{}},{"cell_type":"code","source":"%%time\n\nindex.build(NUM_TREES)","metadata":{"execution":{"iopub.status.busy":"2022-04-10T03:16:10.240231Z","iopub.execute_input":"2022-04-10T03:16:10.240698Z","iopub.status.idle":"2022-04-10T03:16:48.669348Z","shell.execute_reply.started":"2022-04-10T03:16:10.240658Z","shell.execute_reply":"2022-04-10T03:16:48.668481Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"This notebook series continues with https://www.kaggle.com/code/prubyg/hotel-id-vector-validation/","metadata":{}}]}