filtered_vector_search.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MyScale Filtered Vector Search"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Preprocessing\n",
    "\n",
    "We use PubMed dataset for this showcase, simple use wget to download from [huggingface](https://huggingface.co/datasets/owaiskha9654/PubMed_MultiLabel_Text_Classification_Dataset_MeSH):\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget -c https://huggingface.co/datasets/owaiskha9654/PubMed_MultiLabel_Text_Classification_Dataset_MeSH/resolve/main/PubMed%20Multi%20Label%20Text%20Classification%20Dataset%20Processed.csv\n",
    "!python3 -m pip install clickhouse-connect pandas tqdm transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from os import environ\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Filtered vector search is pretty easy in MyScale. \n",
    "\n",
    "First get your credentials to connect to MyScale with clickhouse-connect:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import clickhouse_connect\n",
    "\n",
    "client = clickhouse_connect.get_client(\n",
    "    host='your-myscale-backend',\n",
    "    port=443,\n",
    "    username='your-user-name',\n",
    "    password='your-password'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sentence_transformers import SentenceTransformer\n",
    "\n",
    "model = SentenceTransformer('sentence-transformers/distiluse-base-multilingual-cased-v2')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "def get_embedding(text):\n",
    "    with torch.no_grad():\n",
    "        emb = model.encode(text).squeeze()        \n",
    "        return (emb / np.linalg.norm(emb, ord=2)).tolist()\n",
    "len(get_embedding('text'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The dataset contains strings, array of strings and also array of arrays. To demonstrate how you can use those seamlessly within MyScale, we need to convert those string representation into native lists:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "from ast import literal_eval\n",
    "\n",
    "df = pd.read_csv('PubMed Multi Label Text Classification Dataset Processed.csv')\n",
    "df_s = df.astype(str)\n",
    "for k in ['meshroot', 'meshid', 'meshMajor']:\n",
    "    df_s[k] = [literal_eval(r) for r in tqdm(df_s[k], desc=f'Converting {k} to Python Native...')]\n",
    "with torch.no_grad():\n",
    "    embeddings = [get_embedding(r) for r in tqdm(df_s['abstractText'])]\n",
    "df_s['meshEmbedding'] = embeddings\n",
    "\n",
    "for k in df.keys()[:5]:\n",
    "    print(f'{k} : {type(df_s[k][0])}')\n",
    "df_s[:3]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create table with schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "table_name = \"pubmed_multilabel\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client.command(f'DROP TABLE IF EXISTS {table_name}')\n",
    "\n",
    "client.command(f'''CREATE TABLE IF NOT EXISTS {table_name}(\n",
    "    Title String,\n",
    "    abstractText String,\n",
    "    meshMajor Array(String),\n",
    "    pmid Int64,\n",
    "    meshid Array(Array(String)),\n",
    "    meshroot Array(String),\n",
    "    A Int8,\n",
    "    B Int8,\n",
    "    C Int8,\n",
    "    D Int8,\n",
    "    E Int8,\n",
    "    F Int8,\n",
    "    G Int8,\n",
    "    H Int8,\n",
    "    I Int8,\n",
    "    J Int8,\n",
    "    L Int8,\n",
    "    M Int8,\n",
    "    N Int8,\n",
    "    Z Int8,\n",
    "    meshEmbedding Array(Float32),\n",
    "    CONSTRAINT vec_len CHECK length(meshEmbedding) = 512,\n",
    "    VECTOR INDEX vindex meshEmbedding TYPE IVFFLAT('metric_type=cosine')\n",
    ") ENGINE = MergeTree ORDER BY pmid''')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[Clickhouse-connect](https://clickhouse.com/docs/en/integrations/python) offers [`insert_df`](https://clickhouse.com/docs/en/integrations/python#client-insert-method) to insert a pandas's dataframe to a table."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client.insert_df(table_name, df_s)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get number of uploaded samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'count()': 50000}]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[r for r in client.query(f\"SELECT COUNT(*) FROM {table_name}\").named_results()]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get some samples from the database"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'pmid': 506, 'Title': 'Phospholipases. III. Effects of ionic surfactants on the phospholipase-catalyzed hydrolysis of unsonicated egg lecithin liposomes.'}\n",
      "{'pmid': 2524, 'Title': 'Reduction of blood platelet monoamine oxidase activity in schizophrenic patients on phenothiazines.'}\n",
      "{'pmid': 6714, 'Title': 'Identification of monohydroxylated metabolites of cannabidiol formed by rat liver.'}\n"
     ]
    }
   ],
   "source": [
    "for p in [r for r in client.query(f\"SELECT pmid, Title FROM {table_name} LIMIT 3\").named_results()]:\n",
    "    print(p)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Vector Search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'abstractText': 'Vaccine-induced protection may not be homogeneous across individuals. It is possible that a vaccine gives complete protection for a portion of individuals, while the rest acquire only incomplete (leaky) protection of varying magnitude. If vaccine efficacy is estimated under wrong assumptions about such individual level heterogeneity, the resulting estimates may be difficult to interpret. For instance, population-level predictions based on such estimates may be biased. We consider the problem of estimating heterogeneous vaccine efficacy against an infection that can be acquired multiple times (susceptible-infected-susceptible model). The estimation is based on a limited number of repeated measurements of the current status of each individual, a situation commonly encountered in practice. We investigate how the placement of consecutive samples affects the estimability and efficiency of vaccine efficacy parameters. The same sampling frequency may not be optimal for efficient estimation of all components of heterogeneous vaccine protection. However, we suggest practical guidelines allowing estimation of all components. For situations in which the estimability of individual components fails, we suggest to use summary measures of vaccine efficacy.', 'd': 0.5617110729217529}\n",
      "{'abstractText': 'The management of the vaccination program is part of nursing competences. The main goal of this program is to vaccinate the whole population. There are some age groups in which vaccination coverage is represented by very low rates. Several methods can be used in order to increase such coverage and each professional shall use them according to the work environment. This article presents a simple and effective intervention applicable in any rural area--and probably in any environment--through schools, where all children regularly go. This program has been very useful for us to increase the vaccination coverage of children aged 6.', 'd': 0.5732998251914978}\n",
      "{'abstractText': 'BACKGROUND: Influenza is the most common vaccine-preventable disease in travelers. It circulates year-round in the tropics, November to March in the northern hemisphere (NH), and April to October in the southern hemisphere (SH). In 2005, approximately 8.5 million US adults aged 18 years and older traveled to the Caribbean. A similar number traveled to the tropics and the SH. SH formulation of influenza vaccine is not available in the United States. We surveyed International Society of Travel Medicine (ISTM) members to ask if they would use SH influenza vaccine if available.METHODS: We electronically mailed a survey in December 2006 to 1,251 ISTM members in the United States. We asked if respondents would use SH vaccine for patients traveling to the SH or tropics, how many such patients per week they see, and their practice location.RESULTS: We received 157 responses for a response rate of 12.5%. Of these, 129 (82%) stated that they would be interested in having SH influenza vaccine available. Of those indicating interest, 73 (60%) reported seeing >10 patients traveling to the SH or tropics each week. Respondents reported practice settings in 34 states and the District of Columbia. Respondents requested more information about the likely cost of SH influenza vaccine, ordering conditions, vaccine use guidelines, comparability with NH vaccine, and approval of SH vaccine by the Food and Drug Administration.CONCLUSIONS: Many travelers to the SH are at risk for influenza infection. Although only a limited number of ISTM members responded, respondents indicated considerable interest in availability of SH influenza vaccine for their patients. More data from travel medicine and other practitioners are needed on this topic. Inquiries are being made of influenza vaccine manufacturers about licensing SH influenza vaccines in the United States. Adding SH influenza vaccine to the vaccines available to NH clinicians could help mitigate the morbidity of influenza in travelers.', 'd': 0.5786518454551697}\n"
     ]
    }
   ],
   "source": [
    "emb_str = f\"[{','.join(map(str, get_embedding('vaccine')))}]\"\n",
    "for p in [r for r in client.query(f\"SELECT abstractText, d FROM {table_name} ORDER BY distance(meshEmbedding, {emb_str}) AS d LIMIT 3\").named_results()]:\n",
    "    print(p)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Vector Search with Title Pattern Filter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'abstractTe': 'BACKGROUND: Influenza is the m', 'd': 0.5786518454551697}\n",
      "{'abstractTe': 'BACKGROUND: Simple and effecti', 'd': 0.6172389984130859}\n",
      "{'abstractTe': 'BACKGROUND: Approximately 500 ', 'd': 0.6178913116455078}\n"
     ]
    }
   ],
   "source": [
    "emb_str = f\"[{','.join(map(str, get_embedding('vaccine')))}]\"\n",
    "for p in [r for r in client.query(f\"SELECT abstractText, d FROM {table_name} \\\n",
    "                                  WHERE abstractText LIKE 'BACKGROUND%' \\\n",
    "                                  ORDER BY distance(meshEmbedding, {emb_str}) AS d LIMIT 3\").named_results()]:\n",
    "    print({k[:10]: v[:30] if type(v) is str else v for k, v in p.items()})"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Vector Search with Array Filters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'abstractTe': 'Vaccine-in', 'd': 0.5617110729217529, 'meshroot': ['Health Care [N]', 'Organisms [B]', 'Diseases [C]', 'Analytical, Diagnostic and Therapeutic Techniques, and Equipment [E]', 'Chemicals and Drugs [D]']}\n",
      "{'abstractTe': 'The manage', 'd': 0.5732998251914978, 'meshroot': ['Named Groups [M]', 'Organisms [B]', 'Health Care [N]', 'Geographicals [Z]', 'Analytical, Diagnostic and Therapeutic Techniques, and Equipment [E]']}\n",
      "{'abstractTe': 'BACKGROUND', 'd': 0.5786518454551697, 'meshroot': ['Geographicals [Z]', 'Psychiatry and Psychology [F]', 'Health Care [N]', 'Analytical, Diagnostic and Therapeutic Techniques, and Equipment [E]', 'Information Science [L]', 'Organisms [B]', 'Chemicals and Drugs [D]']}\n"
     ]
    }
   ],
   "source": [
    "emb_str = f\"[{','.join(map(str, get_embedding('vaccine')))}]\"\n",
    "for p in [r for r in client.query(f\"SELECT abstractText, d, meshroot FROM {table_name} \\\n",
    "                                  WHERE has(meshroot, 'Organisms [B]') \\\n",
    "                                  ORDER BY distance(meshEmbedding, {emb_str}) AS d LIMIT 3\").named_results()]:\n",
    "    print({k[:10]: v[:7] if type(v) is list else (v[:10] if type(v) is str else v) for k, v in p.items()})"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Flatten Arrays of Arrays to Arrays then Filter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'abstractTe': 'Vaccine-in', 'd': 0.5617110729217529, 'meshid': [['N05.715.350.150', 'N06.850.490.500'], ['B01.050.150.900.649.313.988.400.112.400.400'], ['C01'], ['E05.318.740.500', 'E05.599.835', 'N05.715.360.750.530', 'N06.850.520.830.500'], ['C23.550.291.937'], ['E01.789.800', 'N04.761.559.590.800', 'N05.715.360.575.575.800'], ['D20.215.894']]}\n"
     ]
    }
   ],
   "source": [
    "emb_str = f\"[{','.join(map(str, get_embedding('vaccine')))}]\"\n",
    "for p in [r for r in client.query(f\"SELECT abstractText, d, meshid FROM {table_name} \\\n",
    "                                  WHERE has(arrayFlatten(meshid), 'C23.550.291.937') \\\n",
    "                                  ORDER BY distance(meshEmbedding, {emb_str}) AS d LIMIT 3\").named_results()]:\n",
    "    print({k[:10]: v[:10] if type(v) is str else v for k, v in p.items()})"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Multi-Column Filter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'abstractTe': 'Vaccine-in', 'd': 0.5617110729217529, 'A': 0, 'B': 1, 'C': 1}\n",
      "{'abstractTe': 'BACKGROUND', 'd': 0.5786518454551697, 'A': 0, 'B': 1, 'C': 1}\n",
      "{'abstractTe': 'Vaccinatio', 'd': 0.6105412244796753, 'A': 1, 'B': 1, 'C': 1}\n"
     ]
    }
   ],
   "source": [
    "emb_str = f\"[{','.join(map(str, get_embedding('vaccine')))}]\"\n",
    "for p in [r for r in client.query(f\"SELECT abstractText, d, A, B, C FROM {table_name} \\\n",
    "                                  WHERE A+B+C>=2 \\\n",
    "                                  ORDER BY distance(meshEmbedding, {emb_str}) AS d LIMIT 3\").named_results()]:\n",
    "    print({k[:10]: v[:10] if type(v) is str else v for k, v in p.items()})"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.8.8"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}