dessert_pytorch_experiments.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ONKFSlSbBAI5",
        "outputId": "63734b90-860a-44c7-83e3-4565c0a4eaeb"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/drive\n",
            "/content/drive/MyDrive/dessert\n"
          ]
        }
      ],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')\n",
        "%cd /content/drive/MyDrive/dessert/"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from dessert_pytorch import *"
      ],
      "metadata": {
        "id": "BVbV22DCBPpc"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "import random\n",
        "import torch\n",
        "\n",
        "# Helper method that returns a tuple of two functions. The first function\n",
        "# takes no arguments and returns a document retrieval index with all generated\n",
        "# documents added. The second function takes a document retrieval index\n",
        "# and queries it with generated queries, asserting that the top result is\n",
        "# as expected, and also returning all results.\n",
        "# The general idea for this test is that each word is a normal distribution\n",
        "# somehwhere in the vector space. A doc is made up of a vector from each\n",
        "# of words_per_doc normal distributions. A ground truth query is made up of\n",
        "# some words from a single doc's word distributions and some random words.\n",
        "def get_build_and_run_functions_random(num_docs=100, num_queries=100):\n",
        "\n",
        "    hashes_per_table = 7\n",
        "    num_tables = 32\n",
        "    data_dim = 100\n",
        "    vocab_size = 10000\n",
        "    words_per_doc = 200\n",
        "    words_per_query_random = 5\n",
        "    words_per_query_from_doc = 10\n",
        "    words_per_query = words_per_query_random + words_per_query_from_doc\n",
        "    between_word_std = 1\n",
        "    within_word_std = 0.1\n",
        "\n",
        "    np.random.seed(42)\n",
        "    random.seed(42)\n",
        "\n",
        "    # Generate word centers\n",
        "    word_centers = np.random.normal(size=(vocab_size, data_dim), scale=between_word_std)\n",
        "\n",
        "    # Generates docs\n",
        "    doc_word_ids = [\n",
        "        random.sample(range(vocab_size), words_per_doc) for _ in range(num_docs)\n",
        "    ]\n",
        "    doc_offsets = np.random.normal(\n",
        "        size=(num_docs, words_per_doc, data_dim), scale=within_word_std\n",
        "    )\n",
        "    docs = []\n",
        "    for i in range(num_docs):\n",
        "        doc = []\n",
        "        for j in range(words_per_doc):\n",
        "            doc.append(doc_offsets[i][j] + word_centers[doc_word_ids[i][j]])\n",
        "        docs.append(doc)\n",
        "\n",
        "    # Generate queries. GT for query i is doc i\n",
        "    query_random_word_ids = [\n",
        "        random.sample(range(vocab_size), words_per_query_random)\n",
        "        for _ in range(num_queries)\n",
        "    ]\n",
        "    query_same_word_ids = [\n",
        "        ids[:words_per_query_from_doc] for ids in doc_word_ids[:num_queries]\n",
        "    ]\n",
        "    query_word_ids = [a + b for a, b in zip(query_same_word_ids, query_random_word_ids)]\n",
        "    query_offsets = np.random.normal(\n",
        "        size=(num_queries, words_per_query, data_dim), scale=within_word_std\n",
        "    )\n",
        "    queries = []\n",
        "    for i in range(num_queries):\n",
        "        query = []\n",
        "        for j in range(words_per_query):\n",
        "            query.append(query_offsets[i][j] + word_centers[query_word_ids[i][j]])\n",
        "        queries.append(query)\n",
        "\n",
        "    index_func = lambda: _build_index_random(\n",
        "        docs, hashes_per_table, num_tables, data_dim, word_centers, words_per_doc\n",
        "    )\n",
        "    query_func = lambda index: _do_queries_random(index, queries, num_docs)\n",
        "\n",
        "    return index_func, query_func\n",
        "\n",
        "\n",
        "def _build_index_random(docs, hashes_per_table, num_tables, data_dim, centroids, words_per_doc):\n",
        "    index = DocRetrieval(\n",
        "        hashes_per_table=hashes_per_table,\n",
        "        num_tables=num_tables,\n",
        "        dense_input_dimension=data_dim,\n",
        "        centroids=torch.tensor(centroids),\n",
        "        max_doc_size=words_per_doc,\n",
        "        device=\"cpu\" #three options for device, which are \"cpu\", \"cuda:0\" (running on just one gpu), and \"cuda\" (multiple gpus), although, \"cuda:0\" appears to be faster than \"cuda\", so far \"cuda:0\" fastest\n",
        "    )\n",
        "    for i, doc in enumerate(docs):\n",
        "        index.add_doc(doc_id=str(i), doc_embeddings=torch.tensor(doc))\n",
        "    return index\n",
        "\n",
        "\n",
        "def _do_queries_random(index, queries, num_docs):\n",
        "    result = []\n",
        "    for gt, query in enumerate(queries):\n",
        "        query_result = index.query(torch.tensor(query), top_k=10, num_to_rerank=10)\n",
        "        result += query_result\n",
        "        print(query_result)\n",
        "        assert int(query_result[0]) == gt\n",
        "    return result"
      ],
      "metadata": {
        "id": "p7E3zFbdBU-O"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def test_random_docs():\n",
        "    index_func, query_func = get_build_and_run_functions_random()\n",
        "    index = index_func()\n",
        "    results = query_func(index)"
      ],
      "metadata": {
        "id": "58p_zzbNBxTl"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "test_random_docs()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "IwNimJU6BzFF",
        "outputId": "42f9f869-3fb4-4d4c-b19e-d7a400ac8444"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "<ipython-input-3-74ac0c100a41>:84: UserWarning: Creating a tensor from a list of numpy.ndarrays is extremely slow. Please consider converting the list to a single numpy.ndarray with numpy.array() before converting to a tensor. (Triggered internally at ../torch/csrc/utils/tensor_new.cpp:274.)\n",
            "  index.add_doc(doc_id=str(i), doc_embeddings=torch.tensor(doc))\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "['0', '61', '88', '84', '73', '43', '30', '50', '79', '63']\n",
            "['1', '15', '23', '34', '79', '52', '62', '12', '81', '43']\n",
            "['2', '51', '84', '60', '47', '19', '54', '16', '56', '98']\n",
            "['3', '89', '51', '87', '38', '95', '78', '65', '16', '0']\n",
            "['4', '69', '10', '90', '89', '17', '85', '79', '26', '95']\n",
            "['5', '98', '71', '40', '61', '14', '33', '26', '38', '82']\n",
            "['6', '89', '24', '57', '54', '34', '41', '77', '88', '39']\n",
            "['7', '90', '91', '23', '74', '60', '12', '25', '71', '63']\n",
            "['8', '23', '31', '41', '42', '15', '26', '71', '60', '20']\n",
            "['9', '64', '21', '44', '32', '84', '50', '15', '81', '76']\n",
            "['10', '11', '65', '1', '19', '22', '57', '56', '74', '71']\n",
            "['11', '78', '58', '8', '70', '54', '72', '20', '76', '92']\n",
            "['12', '52', '53', '15', '50', '75', '16', '67', '20', '36']\n",
            "['13', '95', '87', '31', '92', '32', '52', '33', '12', '55']\n",
            "['14', '77', '28', '90', '24', '54', '57', '7', '70', '23']\n",
            "['15', '66', '52', '43', '96', '68', '56', '84', '20', '97']\n",
            "['16', '1', '75', '10', '65', '82', '69', '15', '49', '43']\n",
            "['17', '24', '4', '32', '22', '36', '13', '29', '73', '26']\n",
            "['18', '95', '1', '21', '22', '81', '33', '24', '8', '20']\n",
            "['19', '37', '20', '64', '13', '83', '91', '90', '56', '70']\n",
            "['20', '40', '27', '4', '9', '59', '56', '55', '3', '54']\n",
            "['21', '97', '80', '94', '39', '56', '17', '16', '20', '74']\n",
            "['22', '4', '60', '36', '61', '89', '65', '79', '66', '26']\n",
            "['23', '9', '51', '16', '54', '32', '30', '86', '88', '44']\n",
            "['24', '73', '38', '25', '68', '14', '21', '54', '61', '12']\n",
            "['25', '79', '3', '62', '71', '44', '50', '57', '99', '32']\n",
            "['26', '48', '1', '87', '52', '28', '56', '11', '68', '81']\n",
            "['27', '6', '64', '55', '81', '87', '42', '57', '29', '72']\n",
            "['28', '54', '53', '64', '9', '23', '45', '74', '15', '90']\n",
            "['29', '65', '30', '57', '17', '63', '19', '70', '60', '76']\n",
            "['30', '9', '38', '66', '56', '13', '19', '93', '54', '10']\n",
            "['31', '41', '47', '14', '4', '62', '95', '38', '16', '43']\n",
            "['32', '23', '48', '24', '79', '53', '22', '86', '72', '68']\n",
            "['33', '59', '27', '95', '92', '76', '86', '15', '14', '34']\n",
            "['34', '56', '60', '13', '83', '43', '53', '75', '21', '85']\n",
            "['35', '74', '13', '62', '78', '70', '14', '65', '77', '67']\n",
            "['36', '1', '83', '29', '66', '46', '13', '6', '14', '85']\n",
            "['37', '89', '54', '41', '12', '6', '33', '68', '30', '47']\n",
            "['38', '49', '82', '27', '19', '53', '51', '75', '67', '32']\n",
            "['39', '67', '58', '45', '86', '49', '57', '25', '53', '35']\n",
            "['40', '51', '13', '92', '57', '53', '17', '77', '0', '34']\n",
            "['41', '9', '36', '42', '92', '69', '0', '17', '76', '75']\n",
            "['42', '91', '38', '40', '58', '49', '63', '62', '89', '25']\n",
            "['43', '38', '36', '70', '48', '81', '73', '51', '59', '62']\n",
            "['44', '76', '92', '75', '42', '69', '33', '95', '39', '70']\n",
            "['45', '73', '49', '13', '53', '86', '31', '48', '82', '81']\n",
            "['46', '70', '74', '0', '53', '67', '63', '57', '58', '69']\n",
            "['47', '86', '43', '51', '6', '62', '76', '9', '69', '36']\n",
            "['48', '63', '53', '44', '69', '28', '21', '76', '12', '59']\n",
            "['49', '31', '38', '14', '71', '81', '21', '64', '43', '42']\n",
            "['50', '66', '85', '43', '26', '47', '78', '16', '14', '18']\n",
            "['51', '96', '99', '63', '26', '76', '70', '45', '78', '68']\n",
            "['52', '21', '95', '89', '71', '14', '5', '24', '70', '28']\n",
            "['53', '20', '68', '39', '79', '21', '8', '3', '26', '40']\n",
            "['54', '31', '5', '65', '21', '85', '97', '46', '11', '17']\n",
            "['55', '52', '12', '68', '95', '37', '94', '46', '79', '65']\n",
            "['56', '53', '0', '99', '41', '46', '43', '34', '62', '69']\n",
            "['57', '67', '7', '73', '22', '52', '86', '20', '17', '77']\n",
            "['58', '39', '50', '63', '67', '53', '10', '98', '86', '38']\n",
            "['59', '21', '18', '39', '88', '91', '25', '69', '35', '60']\n",
            "['60', '5', '50', '85', '13', '17', '25', '95', '63', '37']\n",
            "['61', '59', '38', '14', '46', '2', '21', '84', '89', '74']\n",
            "['62', '43', '85', '49', '14', '15', '66', '78', '73', '54']\n",
            "['63', '58', '65', '37', '94', '22', '76', '33', '24', '36']\n",
            "['64', '88', '27', '56', '24', '4', '61', '28', '36', '60']\n",
            "['65', '22', '77', '78', '40', '66', '18', '55', '19', '41']\n",
            "['66', '8', '51', '16', '11', '84', '90', '92', '33', '27']\n",
            "['67', '40', '11', '60', '87', '80', '36', '53', '8', '34']\n",
            "['68', '69', '83', '94', '21', '37', '53', '16', '41', '57']\n",
            "['69', '4', '93', '37', '0', '20', '59', '36', '30', '67']\n",
            "['70', '95', '91', '21', '52', '37', '54', '77', '40', '75']\n",
            "['71', '17', '56', '53', '70', '1', '91', '30', '80', '41']\n",
            "['72', '10', '75', '82', '28', '48', '25', '42', '60', '19']\n",
            "['73', '68', '71', '44', '14', '13', '95', '23', '54', '1']\n",
            "['74', '18', '88', '63', '1', '62', '6', '7', '71', '80']\n",
            "['75', '37', '55', '72', '25', '99', '71', '14', '0', '80']\n",
            "['76', '93', '85', '55', '40', '41', '34', '67', '37', '26']\n",
            "['77', '43', '93', '74', '72', '92', '51', '9', '82', '88']\n",
            "['78', '51', '66', '22', '75', '21', '23', '20', '50', '83']\n",
            "['79', '58', '64', '46', '36', '0', '39', '2', '51', '34']\n",
            "['80', '24', '71', '15', '10', '84', '33', '30', '65', '87']\n",
            "['81', '83', '36', '33', '38', '24', '71', '66', '62', '51']\n",
            "['82', '56', '7', '86', '43', '91', '66', '49', '47', '28']\n",
            "['83', '79', '82', '16', '15', '49', '11', '67', '75', '18']\n",
            "['84', '68', '37', '64', '66', '36', '75', '32', '79', '74']\n",
            "['85', '55', '51', '84', '83', '21', '17', '39', '99', '22']\n",
            "['86', '71', '62', '85', '82', '98', '60', '33', '61', '87']\n",
            "['87', '1', '78', '16', '20', '35', '31', '39', '28', '64']\n",
            "['88', '74', '79', '30', '41', '93', '60', '24', '23', '50']\n",
            "['89', '27', '7', '25', '93', '42', '50', '37', '3', '2']\n",
            "['90', '25', '24', '45', '5', '27', '7', '89', '39', '48']\n",
            "['91', '4', '20', '86', '30', '9', '78', '97', '27', '79']\n",
            "['92', '8', '39', '32', '55', '51', '85', '52', '91', '80']\n",
            "['93', '55', '81', '34', '85', '63', '65', '71', '80', '49']\n",
            "['94', '87', '47', '33', '55', '28', '92', '78', '80', '19']\n",
            "['95', '13', '72', '62', '18', '37', '92', '67', '40', '66']\n",
            "['96', '43', '71', '59', '44', '79', '42', '19', '1', '21']\n",
            "['97', '92', '16', '42', '79', '27', '89', '81', '60', '68']\n",
            "['98', '51', '5', '29', '93', '15', '95', '94', '66', '73']\n",
            "['99', '6', '55', '80', '67', '29', '47', '70', '31', '58']\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "9st3iOx3B87t"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}