emscripten_wasm_vectordb.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyNK0d26HeqYcf9U/gpjIV0i",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/jankovicsandras/ml/blob/main/emscripten_wasm_vectordb.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Emscripten Wasm vectorDB test\n",
        "\n",
        "### PROs:\n",
        " - it works :)\n",
        " - seems to be fast (but not yet timed, TODO)\n",
        "\n",
        "### CONs:\n",
        " - only top_1 for now, not top_k\n",
        " - more complex than JavaScript\n",
        "\n",
        "### Table of contents:\n",
        "1. Installing deps\n",
        "2. Creating test database and embedding\n",
        "3. Building the C vectordb\n",
        "4. Wasm compiling\n",
        "5. Testing Wasm vectordb"
      ],
      "metadata": {
        "id": "RqWgYTq4wy6b"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### 1. Installing deps"
      ],
      "metadata": {
        "id": "E3jKf0AvN71y"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# for python and embedding\n",
        "!pip -q install sentence-transformers\n",
        "\n",
        "# for Wasm\n",
        "! git clone https://github.com/emscripten-core/emsdk.git\n",
        "! cd emsdk && ./emsdk install latest && ./emsdk activate latest && source ./emsdk_env.sh\n",
        "! ./emsdk/upstream/emscripten/emcc --version\n",
        "! curl https://wasmtime.dev/install.sh -sSf | bash"
      ],
      "metadata": {
        "id": "dDXjlduHRyNI"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### 2. Creating test database and embedding"
      ],
      "metadata": {
        "id": "TgUTeFmtN-ta"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import random\n",
        "from sentence_transformers import SentenceTransformer\n",
        "\n",
        "# TODO: better multi language generation\n",
        "i18n = {\n",
        "  'en':{\n",
        "    'colors': ['black','blue','green','cyan','red','magenta','brown','light grey','dark grey','bright blue'],\n",
        "    'itemtypes': [ 'belt','cap','hat','jeans','jumper', 'shirt','shorts','sneakers','suit','tie' ],\n",
        "    'adjs' : ['Fantastic','Cool','Superb','Awesome','Trendy'],\n",
        "    'insizestr': ' in size ',\n",
        "    'pricestr': '. Price: ',\n",
        "    'currencystr': ' USD.',\n",
        "    'questions': [\n",
        "      'I want to buy a hat. What colors do you have?',\n",
        "      'Can you recommend something green?',\n",
        "      'Do you have shirts under 50 USD?',\n",
        "      'What do you have in size 40?',\n",
        "      'I would like to buy sneakers for my friend. Do you have something in size 46, preferably cyan or blue?',\n",
        "      'What can you recommend in red?'\n",
        "    ]\n",
        "  },\n",
        "  'hu':{\n",
        "    'colors': ['fekete','kék','zöld','zöldeskék','piros','lila','barna','világosszürke','sötétszürke','ragyogó kék'],\n",
        "    'itemtypes': [ 'öv','sapka','kalap','farmer','pullóver', 'ing','rövidnadrág','tornacipő','öltöny','nyakkendő' ],\n",
        "    'adjs': ['Csodálatos','Menő','Szuper','Király','Trendi'],\n",
        "    'insizestr': '. Méret: ',\n",
        "    'pricestr': '. Ár: ',\n",
        "    'currencystr': ' Ft.',\n",
        "    'questions': [\n",
        "      'Kalapot szeretnék. Milyen színek vannak?',\n",
        "      'Tudsz-e ajánlani valami zöldet?',\n",
        "      'Vannak ingek 50 Ft. alatt?',\n",
        "      'Mik vannak 40-es méretben?',\n",
        "      'Tornacipőt szeretnék a barátomnak. Van valami 46-os méretben, lehetőleg zöldeskék vagy kék?',\n",
        "      'Mit tudsz ajánlani pirosban?'\n",
        "    ]\n",
        "  },\n",
        "  'no':{\n",
        "    'colors': ['svart','blå','grøn','grønblå','rød','lilla','brun','lysgrå','mørkgrå','lysblå'],\n",
        "    'itemtypes': [ 'belt','lue','hatt','bukser','genser', 'skjorte','shorts','sko','dress','slips' ],\n",
        "    'adjs': ['Fantastisk','Kult','Supert','Tøff','Trendy'],\n",
        "    'insizestr': ' i størrelse ',\n",
        "    'pricestr': '. Pris: ',\n",
        "    'currencystr': ' kr.',\n",
        "    'questions': [\n",
        "      'Jeg vil kjøpe en hatt. Hva farger er det?',\n",
        "      'Kan du anbefale noen grønt?',\n",
        "      'Har de skjorter under 50 kr?',\n",
        "      'Hva har de i størrelse 40?',\n",
        "      'Jeg vil gjerne kjøpe sko til min venn. Har de nokre i størrelse 46, helst grønblå eller blå?',\n",
        "      'Hva kan du anbefale i rødt?'\n",
        "    ]\n",
        "  }\n",
        "}\n",
        "\n",
        "sizes = [str(30+x*2) for x in range(0,10)]\n",
        "\n",
        "lang = 'hu'\n",
        "\n",
        "allitems = []\n",
        "for c in i18n[lang]['colors'] :\n",
        "  for s in sizes :\n",
        "    for ii,i in enumerate(i18n[lang]['itemtypes']) :\n",
        "      allitems.append( random.choice(i18n[lang]['adjs'])+' '+ c+' '+ i+ i18n[lang]['insizestr']+s+\n",
        "                      i18n[lang]['pricestr']+str(int(s)+20+5*ii)+i18n[lang]['currencystr'] )\n",
        "\n",
        "random.shuffle(allitems)\n",
        "print('len(allitems)',len(allitems))#,allitems)\n",
        "\n",
        "# test questions\n",
        "questions = i18n[lang]['questions']\n",
        "random.shuffle(questions)\n",
        "#print(questions)\n",
        "\n",
        "# embedding\n",
        "embeddermodel = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')\n",
        "zv = [0]*384\n",
        "ivs = [] # item vectors\n",
        "ivs.append(zv) # the first is a zeros vector, so the returned result index from the vectordb will be 1 indexed!\n",
        "for i in allitems :\n",
        "  ivs.append( embeddermodel.encode(i) )\n",
        "\n",
        "qvs = [] # question vectors\n",
        "for q in questions :\n",
        "  qvs.append( embeddermodel.encode(q) )\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8eYvSBOaOHP6",
        "outputId": "88ee7229-fd9d-4fe7-8de6-5b381cbacc4c"
      },
      "execution_count": 54,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "len(allitems) 1000\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### 3. Building the C vectordb"
      ],
      "metadata": {
        "id": "7Oz_3ai2OH3y"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Building vdb\n",
        "vector_size = 384 # all-MiniLM-L6-v2 is 384, adjust if needed\n",
        "vectors_count = len( ivs )\n",
        "\n",
        "def flatten(xss):\n",
        "    return [x for xs in xss for x in xs]\n",
        "\n",
        "vdb = flatten(ivs)\n",
        "\n",
        "# TODO: this is unoptimized and possibly redundant, could be done in flatten\n",
        "vdbstr = ''\n",
        "for i in range(0,vectors_count) :\n",
        "  vdbstr += '\\n'\n",
        "  for j in range(0,vector_size) :\n",
        "    vdbstr += str( vdb[ i*vector_size+j ] ) + ', '\n",
        "vdbstr = vdbstr[:-2] + '\\n'\n",
        "\n",
        "# the C vectordb program\n",
        "ccode = \"\"\"\n",
        "#include <emscripten.h>\n",
        "#include <stdio.h>\n",
        "#include <stdlib.h>\n",
        "\n",
        "#define VECTOR_SIZE \"\"\" +str(vector_size)+ \"\"\"\n",
        "#define VECTORS_COUNT \"\"\" +str(vectors_count)+ \"\"\"\n",
        "\n",
        "// vdb and this c file should be generated with external tool\n",
        "// shape is  [ [0,0,...], [x0,x1...], ... items ], where the first vector should be zeros\n",
        "// top_k = 1 for now, similarity metric is cosine (dot product)\n",
        "float vdb [] = { \"\"\"+vdbstr+\"\"\" };\n",
        "\n",
        "\n",
        "EMSCRIPTEN_KEEPALIVE\n",
        "int gettop1(float v[]){\n",
        "  int i=0, j=0, vi = 0;\n",
        "  float vival = 0, ival = 0;\n",
        "  for(i=0; i<VECTORS_COUNT; i++){\n",
        "    ival = 0;\n",
        "    for(j=0; j<VECTOR_SIZE; j++){\n",
        "      ival = ival + vdb[i*VECTOR_SIZE+j] * v[j];\n",
        "    }\n",
        "    if( ival > vival ){\n",
        "      vival = ival;\n",
        "      vi = i;\n",
        "    }\n",
        "  }\n",
        "  return vi;\n",
        "}// End of gettop1()\n",
        "\n",
        "\n",
        "EMSCRIPTEN_KEEPALIVE\n",
        "int main(int argc, char **argv) {\n",
        "\n",
        "  // argv to input vector parsing\n",
        "  float v[VECTOR_SIZE] = {0};\n",
        "  for(int argi=1; argi<argc; argi++){\n",
        "    v[argi-1] = atof( argv[argi] );\n",
        "  }\n",
        "\n",
        "  // get best match's index from vectordb\n",
        "  int vi = gettop1( v );\n",
        "\n",
        "  // print and return\n",
        "  printf(\"%d\", vi);\n",
        "  return vi;\n",
        "\n",
        "}// End of main()\n",
        "\n",
        "\"\"\"\n",
        "\n",
        "# print(ccode)\n",
        "\n",
        "# write the C program to file\n",
        "with open('vdb.c','w+') as f:\n",
        "  f.write(ccode)\n"
      ],
      "metadata": {
        "id": "LKcx1gsEKW7u"
      },
      "execution_count": 55,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### 4. Wasm compiling"
      ],
      "metadata": {
        "id": "A_tTMf98ORwx"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "! ./emsdk/upstream/emscripten/emcc -O3 -s WASM=1 -s STANDALONE_WASM -o a.js vdb.c\n",
        "! ls -la"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "UMv_yN4ENyM-",
        "outputId": "83655ae4-2480-4635-9790-b89457969fca"
      },
      "execution_count": 56,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "total 6528\n",
            "drwxr-xr-x  1 root root    4096 Feb 26 13:12 .\n",
            "drwxr-xr-x  1 root root    4096 Feb 26 08:00 ..\n",
            "-rw-r--r--  1 root root   12043 Feb 26 13:12 a.js\n",
            "-rwxr-xr-x  1 root root 1558233 Feb 26 13:12 a.wasm\n",
            "drwxr-xr-x  4 root root    4096 Feb 22 14:24 .config\n",
            "drwxr-xr-x 12 root root    4096 Feb 26 13:11 emsdk\n",
            "drwxr-xr-x  1 root root    4096 Feb 22 14:24 sample_data\n",
            "-rw-r--r--  1 root root 5088718 Feb 26 13:11 vdb.c\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### 5. Testing Wasm vectordb"
      ],
      "metadata": {
        "id": "zcyruO1ZJ67q"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "basecommand = '/root/.wasmtime/bin/wasmtime a.wasm'\n",
        "\n",
        "# loop questions\n",
        "for i in range(0,len(questions)) :\n",
        "  print('\\n--------\\n',questions[i])\n",
        "  # question vector\n",
        "  inputvector = qvs[i]\n",
        "  # creating CLI command, this could be optimized\n",
        "  fullcommand = basecommand + ''\n",
        "  for n in inputvector :\n",
        "    fullcommand += ' '+str(n)\n",
        "  # vdb search\n",
        "  rawresultidx = os.popen(fullcommand).read()\n",
        "  resultidx = int(rawresultidx)-1  # -1 is important because vdb[1] is allitems[0], vdb[0] is a zeros placeholder\n",
        "  # print\n",
        "  print( resultidx, allitems[ resultidx ] )\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "dyXJdDt5utUd",
        "outputId": "bb229dd6-4f9d-42d6-8c2c-095ac420fd22"
      },
      "execution_count": 57,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "--------\n",
            " Mit tudsz ajánlani pirosban?\n",
            "389 Trendi piros öltöny. Méret: 34. Ár: 94 Ft.\n",
            "\n",
            "--------\n",
            " Tudsz-e ajánlani valami zöldet?\n",
            "64 Trendi zöldeskék rövidnadrág. Méret: 42. Ár: 92 Ft.\n",
            "\n",
            "--------\n",
            " Mik vannak 40-es méretben?\n",
            "9 Menő világosszürke nyakkendő. Méret: 40. Ár: 105 Ft.\n",
            "\n",
            "--------\n",
            " Vannak ingek 50 Ft. alatt?\n",
            "779 Trendi kék öltöny. Méret: 40. Ár: 100 Ft.\n",
            "\n",
            "--------\n",
            " Kalapot szeretnék. Milyen színek vannak?\n",
            "684 Szuper sötétszürke kalap. Méret: 40. Ár: 70 Ft.\n",
            "\n",
            "--------\n",
            " Tornacipőt szeretnék a barátomnak. Van valami 46-os méretben, lehetőleg zöldeskék vagy kék?\n",
            "461 Szuper világosszürke tornacipő. Méret: 36. Ár: 91 Ft.\n"
          ]
        }
      ]
    }
  ]
}