vector_index.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "445cb2ff-747e-4a7d-b179-82ea51da6d16",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sentence_transformers import SentenceTransformer\n",
    "import torch\n",
    "torch._dynamo.config.suppress_errors = True\n",
    "# model = SentenceTransformer('answerdotai/ModernBERT-large')\n",
    "model = SentenceTransformer(\"all-mpnet-base-v2\")\n",
    "MODEL_NAME = \"all-mpnet-base-v2\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c8eb87ff-9865-4cae-9efb-e36065d262db",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sqlite3\n",
    "from copy import deepcopy\n",
    "import json\n",
    "\n",
    "DATABASE_NAME = \"suika_commands.db\"\n",
    "TABLE = \"suika_commands\"\n",
    "\n",
    "def fetch_all_documents():\n",
    "    try:\n",
    "        conn = sqlite3.connect(DATABASE_NAME)\n",
    "        conn.row_factory = sqlite3.Row\n",
    "        cursor = conn.cursor()\n",
    "        cursor.execute(f\"SELECT * FROM {TABLE}\")\n",
    "        rows = cursor.fetchall()\n",
    "        documents = [dict(row) for row in rows]\n",
    "        for doc in documents:\n",
    "            doc[\"keywords\"] = json.loads(doc[\"keywords\"])\n",
    "            doc[\"examples\"] = json.loads(doc[\"examples\"])\n",
    "    except sqlite3.OperationalError as e:\n",
    "        print(f\"OperationalError: {e}\")\n",
    "        documents = []\n",
    "    finally:\n",
    "        conn.close()\n",
    "    return documents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "cf61f012-33fd-49d9-b0bc-58ecc3298bc2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total 1195\n",
      "{'linux_cmd_id': 'linux_pw-link', 'name': 'pw-link', 'description': 'Pw-link is a command-line tool for Linux that manages symbolic links specifically for password files encrypted with the FileVault system.', 'syntax': 'pw-link [-d] [options] <source> <destination>', 'keywords': ['link', 'symbolic', 'password', 'FileVault', 'source', 'destination', 'encryption', 'remote'], 'examples': ['pw-link /var/lib/vault.root /Users/username/Documents', 'pw-link -d /tmp/vault.pw /Users/username/Vault']}\n"
     ]
    }
   ],
   "source": [
    "all_docs = fetch_all_documents()\n",
    "print('total', len(all_docs))\n",
    "print(all_docs[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "86709446-b359-4e52-8222-efb334082974",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1195it [00:21, 54.87it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(1195, 768)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import faiss\n",
    "from tqdm import tqdm\n",
    "\n",
    "embed_docs = []\n",
    "id_to_vec_pos = {}\n",
    "for i,doc in tqdm(enumerate(all_docs)):\n",
    "    emb = model.encode(doc['description'])\n",
    "    embed_docs.append(emb)\n",
    "    id_to_vec_pos[doc['linux_cmd_id']] = i\n",
    "embed_docs = np.vstack(embed_docs)\n",
    "\n",
    "embed_docs.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f23d17ea-7532-4cea-8c10-5a09f7d9de87",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1195"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "d = embed_docs.shape[1]\n",
    "# index = faiss.IndexFlatL2(d)\n",
    "index = faiss.IndexFlatIP(d)\n",
    "index.add(embed_docs)\n",
    "index.ntotal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "ab355274-8183-493a-a4f9-b9b6beecfcd1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1, 768)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "k = 1\n",
    "q = \"how to copy the file in linux?\"\n",
    "qemb = model.encode(q)\n",
    "qemb = qemb[None]\n",
    "qemb.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "c64c037c-bf08-452d-ab00-ee75b1ea6e58",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'linux_cmd_id': 'linux_cp', 'name': 'cp', 'description': 'Copy files and directories from one location to another. Uses syntax similar to mv, often but not always more efficient than mv when copying large files.', 'syntax': 'cp [options] source destination', 'keywords': ['copy', 'duplicate', 'files', 'directories', 'source', 'destination'], 'examples': ['cp myfile.txt backup.txt', 'cp -r Documents/ new_location/']}\n",
      "[[0.40737653]]\n",
      "CPU times: user 1.41 ms, sys: 157 μs, total: 1.57 ms\n",
      "Wall time: 970 μs\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "D, I = index.search(qemb, k)\n",
    "print(all_docs[I[0].item()])\n",
    "print(D)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d64d192f-8be3-443b-9952-b09cf9f93a65",
   "metadata": {},
   "outputs": [],
   "source": [
    "index_file = \"suika_commands_vector.index\"\n",
    "faiss.write_index(index, index_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d9b3e1d8-d3df-4b26-8c56-68b1135145e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "index = faiss.read_index(index_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "0547003a-db9d-4175-a0d7-fed4ade07910",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'linux_cmd_id': 'linux_cp', 'name': 'cp', 'description': 'Copy files and directories from one location to another. Uses syntax similar to mv, often but not always more efficient than mv when copying large files.', 'syntax': 'cp [options] source destination', 'keywords': ['copy', 'duplicate', 'files', 'directories', 'source', 'destination'], 'examples': ['cp myfile.txt backup.txt', 'cp -r Documents/ new_location/']}\n",
      "[[0.40737653]]\n",
      "CPU times: user 2.56 ms, sys: 0 ns, total: 2.56 ms\n",
      "Wall time: 1.54 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "D, I = index.search(qemb, k)\n",
    "print(all_docs[I[0].item()])\n",
    "print(D)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "691bc3d4-0232-4fb6-a099-d996d3161aeb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sqlite3\n",
    "conn = sqlite3.connect(DATABASE_NAME)\n",
    "cursor = conn.cursor()\n",
    "cursor.execute('''\n",
    "    CREATE TABLE IF NOT EXISTS id_to_vec_pos (\n",
    "        faiss_index_id INTEGER PRIMARY KEY,\n",
    "        linux_cmd_id TEXT NOT NULL\n",
    "    )\n",
    "''')\n",
    "conn.commit()\n",
    "conn.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "7b07f4e4-83c5-48a9-96b8-3aa5b114cab0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_mapping_to_db(database_name, mapping):\n",
    "    # use context manager to ensure connection and cursor are properly closed\n",
    "    with sqlite3.connect(database_name) as conn:\n",
    "        cursor = conn.cursor()\n",
    "        for linux_cmd_id, faiss_index_id in mapping.items():\n",
    "            cursor.execute('''\n",
    "                INSERT OR REPLACE INTO id_to_vec_pos (faiss_index_id, linux_cmd_id)\n",
    "                VALUES (?, ?)\n",
    "            ''', (faiss_index_id, linux_cmd_id))  # replaces if faiss_index_id already exists\n",
    "        conn.commit()\n",
    "\n",
    "save_mapping_to_db(DATABASE_NAME, id_to_vec_pos)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ed62f46e-76be-4b5d-9cd0-4fd57ccbeebe",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sentence_transformers import SentenceTransformer\n",
    "import torch\n",
    "torch._dynamo.config.suppress_errors = True\n",
    "\n",
    "model = SentenceTransformer(\"all-mpnet-base-v2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "43ba1bc5-832c-46de-a45b-2f42f4a67f44",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "doc_id linux_mycli\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "({'linux_cmd_id': 'linux_mycli',\n",
       "  'name': 'mycli',\n",
       "  'description': 'mycli is a modern, user-friendly command-line interface for interacting with MySQL databases. It provides an intuitive way to execute SQL queries, manage schema objects, and analyze data.',\n",
       "  'syntax': 'mycli [options] [database]',\n",
       "  'keywords': '[\"mysql\", \"database\", \"sql\", \"query\", \"management\", \"command-line\"]',\n",
       "  'examples': '[\"mycli\", \"mycli --database testdb\", \"mycli -h localhost -u user -p mydb\"]'},\n",
       " 0.38226398825645447)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import vec_db as _vec_db\n",
    "\n",
    "q_emb = model.encode('how to query db')\n",
    "_vec_db.search(q_emb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "399fc8cc-633c-42aa-8bb7-7bc908a5043f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "m",
   "language": "python",
   "name": "m"
  },
  "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.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}