notebooks/unisim_text_demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Text UniSim Demo -- Fuzzy Address Matching\n",
    "\n",
    "This demo showcases how to use Text UniSim (TextSim) for efficient fuzzy string matching, near-duplicate detection, and string similarity using a real-world entity matching dataset.\n",
    "\n",
    "For additional information, please see the documentation on [GitHub](https://github.com/google/unisim). For more details on the RETSim model used by UniSim, please see the [RETSim paper](https://arxiv.org/abs/2311.17264)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO: Loaded backend\n",
      "INFO: Using ONNX with CPU\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.10/dist-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "# installing needed dependencies\n",
    "try:\n",
    "    import unisim\n",
    "except ImportError:\n",
    "    !pip install unisim\n",
    "\n",
    "try:\n",
    "    import datasets\n",
    "except ImportError:\n",
    "    !pip install datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# imports\n",
    "from datasets import load_dataset\n",
    "from tabulate import tabulate\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import TextSim from UniSim\n",
    "from unisim import TextSim"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Dataset\n",
    "\n",
    "For this demo, we use entity matching datasets available on [Huggingface](https://huggingface.co/datasets/RUC-DataLab/ER-dataset). We use the `restaurants1.csv` dataset which contains restaurants' names/phone numbers/addresses for this colab.\n",
    "\n",
    "Feel free to explore other examples they offer such as product matching (`walmart_amazon.csv`), paper citation matching (`dblp_scholar.csv`), and beer brands (`beer.csv`). The public datasets are from [DeepMatcher](https://github.com/anhaidgroup/deepmatcher/blob/master/Datasets.md), [Magellan](https://sites.google.com/site/anhaidgroup/useful-stuff/the-magellan-data-repository) and [WDC](http://webdatacommons.org/largescaleproductcorpus/v2/) and you can find a summary of them [here](https://github.com/ruc-datalab/DADER/tree/main)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Repo card metadata block was not found. Setting CardData to empty.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Size of dataset: 450\n",
      "Dataset features: ['A_NAME', 'A_PHONENUMBER', 'A_ADDRESS', 'B_NAME', 'B_PHONENUMBER', 'B_ADDRESS', 'label']\n"
     ]
    }
   ],
   "source": [
    "# load dataset from huggingface\n",
    "ds_name = \"restaurants1.csv\"\n",
    "dataset = load_dataset(\"RUC-DataLab/ER-dataset\", data_files=ds_name, split=\"train\")\n",
    "\n",
    "print(\"Size of dataset:\", len(dataset))\n",
    "\n",
    "dataset_features = list(dataset.features.keys())\n",
    "print(\"Dataset features:\", dataset_features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get features corresponding to pairs of texts in the dataset (text1, text2)\n",
    "text1_features = [x for x in dataset_features if x.startswith(\"A\")]\n",
    "text2_features = [x for x in dataset_features if x.startswith(\"B\")]\n",
    "is_match_feature = \"label\"\n",
    "\n",
    "# create text pairs for example at idx in the dataset\n",
    "def get_text_pair(idx):\n",
    "    ex = dataset[idx]\n",
    "\n",
    "    text1 = \" \".join(str(ex[x]) for x in text1_features)\n",
    "    text2 = \" \".join(str(ex[x]) for x in text2_features)\n",
    "\n",
    "    label = ex[is_match_feature]\n",
    "    return [text1, text2, label]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Initialize TextSim\n",
    "\n",
    "TextSim supports GPU acceleration (set `use_accelerator=True`) with a specified `batch_size` parameter. If a GPU is not detected, TextSim will default to CPU. By default, TextSim saves a copy of your dataset but you can set `store_data=False` to save memory when using larger datasets.\n",
    "\n",
    "Additionally, TextSim support Approximate Nearest Neighbor (ANN) search through [USearch](https://github.com/unum-cloud/usearch). Setting `index_type=\"approx\"` will make TextSim significantly faster on large datasets (sub-linear search time). However, please note that while ANN search is very accurate, it does not guarantee that it will always find the closest match to a search query."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO: UniSim is storing a copy of the indexed data\n",
      "INFO: If you are using large data corpus, consider disabling this behavior using store_data=False\n",
      "INFO: Accelerator is not available, using CPU\n"
     ]
    }
   ],
   "source": [
    "# create TextSim using default parameter settings\n",
    "text_sim = TextSim(\n",
    "    store_data=True, # set to False for large datasets to save memory\n",
    "    index_type=\"exact\", # set to \"approx\" for large datasets to use ANN search\n",
    "    batch_size=128, # increasing batch_size on GPU may be faster\n",
    "    use_accelerator=True, # uses GPU if available, otherwise uses CPU\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Computing Similarity between Strings\n",
    "\n",
    "You can compute the similarity between two strings using the `.similarity(string1, string2)` method. The similarity value is a float between 0 and 1, with 1.0 representing identical strings. This is the cosine similarity between the vector representations of the strings.\n",
    "\n",
    "In this example, we compute the similarity between the first 5 pairs of the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>text1</th>\n",
       "      <th>text2</th>\n",
       "      <th>match_label</th>\n",
       "      <th>similarity</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>15 Romolo (415) 398-1359 15 Romolo Place, San ...</td>\n",
       "      <td>15 Romolo (415) 398-1359 15 Romolo Pl, San Fra...</td>\n",
       "      <td>1</td>\n",
       "      <td>0.971120</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>456 Shanghai Cuisine 1261 69 Mott Street, New ...</td>\n",
       "      <td>Shanghai Asian Manor (212) 766-6311 21 Mott St...</td>\n",
       "      <td>0</td>\n",
       "      <td>0.779699</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>5A5 Steak Lounge (415) 989-2539 244 Jackson St...</td>\n",
       "      <td>Delicious Dim Sum (415) 781-0721 752 Jackson S...</td>\n",
       "      <td>0</td>\n",
       "      <td>0.687844</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>9th Street Pizza (213) 627-7798 231 E 9th St, ...</td>\n",
       "      <td>Han Bat Sul Lung Tang (213) 383-9499 4163 W 5t...</td>\n",
       "      <td>0</td>\n",
       "      <td>0.610611</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>9th Street Pizza (213) 627-7798 231 E 9th St, ...</td>\n",
       "      <td>Jun Won Restaurant (213) 383-8855 3100 W 8th S...</td>\n",
       "      <td>0</td>\n",
       "      <td>0.617366</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                               text1  \\\n",
       "0  15 Romolo (415) 398-1359 15 Romolo Place, San ...   \n",
       "1  456 Shanghai Cuisine 1261 69 Mott Street, New ...   \n",
       "2  5A5 Steak Lounge (415) 989-2539 244 Jackson St...   \n",
       "3  9th Street Pizza (213) 627-7798 231 E 9th St, ...   \n",
       "4  9th Street Pizza (213) 627-7798 231 E 9th St, ...   \n",
       "\n",
       "                                               text2  match_label  similarity  \n",
       "0  15 Romolo (415) 398-1359 15 Romolo Pl, San Fra...            1    0.971120  \n",
       "1  Shanghai Asian Manor (212) 766-6311 21 Mott St...            0    0.779699  \n",
       "2  Delicious Dim Sum (415) 781-0721 752 Jackson S...            0    0.687844  \n",
       "3  Han Bat Sul Lung Tang (213) 383-9499 4163 W 5t...            0    0.610611  \n",
       "4  Jun Won Restaurant (213) 383-8855 3100 W 8th S...            0    0.617366  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "example_data = [get_text_pair(idx) for idx in range(0, 5)]\n",
    "\n",
    "for i in range(len(example_data)):\n",
    "    text1, text2, is_match = example_data[i]  # ground truth is_match label\n",
    "\n",
    "    # compute similarity between text pair using .similarity\n",
    "    similarity = text_sim.similarity(text1, text2)\n",
    "\n",
    "    example_data[i].append(similarity)\n",
    "\n",
    "# display results in df\n",
    "df = pd.DataFrame(example_data, columns=[\"text1\", \"text2\", \"match_label\", \"similarity\"])\n",
    "display(df.head())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that the similarity between the first pair is 0.97, which is very high and indicates that the strings are near-duplicates of each other and the addresses match. The other pairs have far lower similarity values, indicating they are likely not matching strings and indeed, they do not represent the same entity."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Fuzzy String Matching\n",
    "\n",
    "TextSim offers efficient fuzzy string matching between two lists using the `.match` function. The `.match` function accepts `queries` (list of strings you want to find matches for) and `targets` (list of strings you are finding matches in). It returns a Pandas DataFrame, where each row contains a query, its most similar match found in targets, their similarity, and whether or not they are a match (if their similarity is >= `similarity_threshold`). `0.9` is a good starting point for `similarity_threshold` when matching near-duplicate strings.\n",
    "\n",
    "In this example, we show that TextSim is able to match restaurant addresses accurately even when there are typos, abbreviations, and formatting differences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>query</th>\n",
       "      <th>target</th>\n",
       "      <th>similarity</th>\n",
       "      <th>is_match</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Shanghai asia manor (212)-766-6311 21 Mott street, New York, NY 94133</td>\n",
       "      <td>Shanghai Asian Manor (212) 766-6311 21 Mott St, New York, NY 10013</td>\n",
       "      <td>0.907003</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Googleplex (650) 253-0000 1600 Amphitheatre Pkwy, Mountain View, CA 94043</td>\n",
       "      <td>15 Romolo (415) 398-1359 15 Romolo Pl, San Francisco, CA 94133</td>\n",
       "      <td>0.466519</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Sino-american books &amp; arts (415) 421-3345 751 Jackson St, San Francisco, CA 94133</td>\n",
       "      <td>Delicious Dim Sum (415) 781-0721 752 Jackson St, San Francisco, CA 94133</td>\n",
       "      <td>0.746313</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                                               query  \\\n",
       "0              Shanghai asia manor (212)-766-6311 21 Mott street, New York, NY 94133   \n",
       "1          Googleplex (650) 253-0000 1600 Amphitheatre Pkwy, Mountain View, CA 94043   \n",
       "2  Sino-american books & arts (415) 421-3345 751 Jackson St, San Francisco, CA 94133   \n",
       "\n",
       "                                                                     target  \\\n",
       "0        Shanghai Asian Manor (212) 766-6311 21 Mott St, New York, NY 10013   \n",
       "1            15 Romolo (415) 398-1359 15 Romolo Pl, San Francisco, CA 94133   \n",
       "2  Delicious Dim Sum (415) 781-0721 752 Jackson St, San Francisco, CA 94133   \n",
       "\n",
       "   similarity  is_match  \n",
       "0    0.907003      True  \n",
       "1    0.466519     False  \n",
       "2    0.746313     False  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# targets to match queries to\n",
    "targets = [\n",
    "    \"Shanghai Asian Manor (212) 766-6311 21 Mott St, New York, NY 10013\",\n",
    "    \"Delicious Dim Sum (415) 781-0721 752 Jackson St, San Francisco, CA 94133\",\n",
    "    \"15 Romolo (415) 398-1359 15 Romolo Pl, San Francisco, CA 94133\",\n",
    "]\n",
    "\n",
    "# search queries we are looking up and finding matches for\n",
    "queries = [\n",
    "    \"Shanghai asia manor (212)-766-6311 21 Mott street, New York, NY 94133\", # near-dup match (capitalization, typos, different format)\n",
    "    \"Googleplex (650) 253-0000 1600 Amphitheatre Pkwy, Mountain View, CA 94043\", # no match\n",
    "    \"Sino-american books & arts (415) 421-3345 751 Jackson St, San Francisco, CA 94133\", # no match, different places but similar address\n",
    "]\n",
    "\n",
    "# .match does fuzzy matching between queries and targets lists\n",
    "results_df = text_sim.match(queries, targets, similarity_threshold=0.9)\n",
    "\n",
    "# display results dataframe\n",
    "with pd.option_context('display.max_colwidth', None):\n",
    "    display(results_df.head(10))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can try this on the whole dataset now. We use the first text in each pair as the target and the second text as the search query."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset examples:\n",
      "Batter & Berries (773) 248-7710 2748 N. Lincoln Avenue, Chicago, IL\n",
      "Joey's Seafood & Grill (608) 829-0093 6604 Mineral Pt Rd, Madison, WI\n",
      "Flaming Wok (608) 240-1085 4237 Lien Rd Ste H, Madison, WI\n",
      "Shallots Bistro (847) 677-3463 7016 Carpenter Road, Skokie, IL\n",
      "Yang Chow (213) 625-0811 819 N Broadway, Los Angeles, CA\n"
     ]
    }
   ],
   "source": [
    "targets = list(set([get_text_pair(idx)[0] for idx in range(0, len(dataset))]))\n",
    "queries = list(set([get_text_pair(idx)[1] for idx in range(0, len(dataset))]))\n",
    "\n",
    "print(\"Dataset examples:\")\n",
    "print(\"\\n\".join([t for t in targets[:5]]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>query</th>\n",
       "      <th>target</th>\n",
       "      <th>similarity</th>\n",
       "      <th>is_match</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Hong Kong Buffet (608) 240-0762 2817 E Washington Ave, Madison, WI 53704</td>\n",
       "      <td>Brink Lounge (608) 661-8599 701 E Washington Ave, Madison, WI</td>\n",
       "      <td>0.768695</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Ruay Thai (212) 545-7829 625 2nd Ave, New York, NY 10016</td>\n",
       "      <td>Ruay Thai Restaurant (212) 545-7829 625 Second Avenue, New York, NY, NY</td>\n",
       "      <td>0.876181</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Kabul Restaurant (608) 256-6322 540 State St, Madison, WI 53703</td>\n",
       "      <td>Kabul Afghanistan Restaurant (608) 256-6322 540 State St, Madison, WI</td>\n",
       "      <td>0.916992</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Saigon Sandwich (415) 474-5698 560 Larkin St, San Francisco, CA 94102</td>\n",
       "      <td>Saigon Sandwich (415) 474-5698 560 Larkin Street, San Francisco, CA</td>\n",
       "      <td>0.945782</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Ricks Olde Gold (608) 257-7280 1314 Williamson St Frnt, Madison, WI 53703</td>\n",
       "      <td>Cafe Costa Rica (608) 256-9830 1133 Williamson St, Madison, WI</td>\n",
       "      <td>0.719272</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>Maharaja Restaurant (608) 246-8525 1707 Thierer Rd, Madison, WI 53704</td>\n",
       "      <td>Maharana (608) 246-8525 1707 Thierer Rd, Madison, WI</td>\n",
       "      <td>0.886549</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>J J Fish (773) 533-1995 816 N Kedzie Ave, Chicago, IL 60651</td>\n",
       "      <td>El Cid (773) 395-0505 2645 N. Kedzie Avenue, Chicago, IL</td>\n",
       "      <td>0.757653</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>Clark St. Ale House (877) 637-7133 742 N Clark St, Chicago, IL 60654</td>\n",
       "      <td>Roka Akor (312) 477-7652 456 N. Clark Street, Chicago, IL</td>\n",
       "      <td>0.754152</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>The Taco Bros (608) 422-5075 604 E University Ave, Madison, WI 53715</td>\n",
       "      <td>The Taco Shop (608) 250-8226 604 University Ave, Madison, WI</td>\n",
       "      <td>0.828169</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>Au Lac DTLA (213) 617-2533 710 W 1st St, Los Angeles, CA 90012</td>\n",
       "      <td>The Little Door (323) 951-1210 8164 W 3rd St, Los Angeles, CA</td>\n",
       "      <td>0.760763</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                                       query  \\\n",
       "0   Hong Kong Buffet (608) 240-0762 2817 E Washington Ave, Madison, WI 53704   \n",
       "1                   Ruay Thai (212) 545-7829 625 2nd Ave, New York, NY 10016   \n",
       "2            Kabul Restaurant (608) 256-6322 540 State St, Madison, WI 53703   \n",
       "3      Saigon Sandwich (415) 474-5698 560 Larkin St, San Francisco, CA 94102   \n",
       "4  Ricks Olde Gold (608) 257-7280 1314 Williamson St Frnt, Madison, WI 53703   \n",
       "5      Maharaja Restaurant (608) 246-8525 1707 Thierer Rd, Madison, WI 53704   \n",
       "6                J J Fish (773) 533-1995 816 N Kedzie Ave, Chicago, IL 60651   \n",
       "7       Clark St. Ale House (877) 637-7133 742 N Clark St, Chicago, IL 60654   \n",
       "8       The Taco Bros (608) 422-5075 604 E University Ave, Madison, WI 53715   \n",
       "9             Au Lac DTLA (213) 617-2533 710 W 1st St, Los Angeles, CA 90012   \n",
       "\n",
       "                                                                    target  \\\n",
       "0            Brink Lounge (608) 661-8599 701 E Washington Ave, Madison, WI   \n",
       "1  Ruay Thai Restaurant (212) 545-7829 625 Second Avenue, New York, NY, NY   \n",
       "2    Kabul Afghanistan Restaurant (608) 256-6322 540 State St, Madison, WI   \n",
       "3      Saigon Sandwich (415) 474-5698 560 Larkin Street, San Francisco, CA   \n",
       "4           Cafe Costa Rica (608) 256-9830 1133 Williamson St, Madison, WI   \n",
       "5                     Maharana (608) 246-8525 1707 Thierer Rd, Madison, WI   \n",
       "6                 El Cid (773) 395-0505 2645 N. Kedzie Avenue, Chicago, IL   \n",
       "7                Roka Akor (312) 477-7652 456 N. Clark Street, Chicago, IL   \n",
       "8             The Taco Shop (608) 250-8226 604 University Ave, Madison, WI   \n",
       "9            The Little Door (323) 951-1210 8164 W 3rd St, Los Angeles, CA   \n",
       "\n",
       "   similarity  is_match  \n",
       "0    0.768695     False  \n",
       "1    0.876181     False  \n",
       "2    0.916992      True  \n",
       "3    0.945782      True  \n",
       "4    0.719272     False  \n",
       "5    0.886549     False  \n",
       "6    0.757653     False  \n",
       "7    0.754152     False  \n",
       "8    0.828169     False  \n",
       "9    0.760763     False  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "results_df = text_sim.match(queries, targets)\n",
    "\n",
    "with pd.option_context('display.max_colwidth', None):\n",
    "    display(results_df.head(10))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Indexing and Searching for Similar Texts in a Dataset\n",
    "\n",
    "TextSim allows you to maintain, update, and query a large index to find similar texts. This gives you more control over indexing and querying your dataset, including how many similar texts you want to retrieve per query and detailed results.\n",
    "\n",
    "You can use the `.add` method to add examples from your dataset to the index, then use  `.search` to search the index and return the most similar texts to your search query."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# resets the index if you previously added things\n",
    "text_sim.reset_index()\n",
    "\n",
    "# adds the dataset of target examples to the index\n",
    "text_sim.add(targets)\n",
    "\n",
    "# for each query, search for the k=5 most similar texts\n",
    "result_collection = text_sim.search(queries, similarity_threshold=0.9, k=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total matches found: 103\n"
     ]
    }
   ],
   "source": [
    "# texts are considered near-duplicate matches if their similarity >= similarity_threshold\n",
    "total_matches = result_collection.total_matches\n",
    "print(\"Total matches found:\", total_matches)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`result_collection.results` contains a list of results corresponding to each query. Each `Result` object contains the results of a search query, including the number of matches found (`.num_matches`), the idx/data/embedding of the query (`.query_idx`, `.query_data`, `.query_embedding`), and a list of `Match` objects (`.matches`).\n",
    "\n",
    "The list of `Match`'s correspond to the `k` most similar texts found for the query, sorted by similarity (most similar first). Each `Match` object contains info on whether it is a near-duplicate match (`.is_match`), the rank (`.rank`), the data (`.data`), the similarity value (`.similarity`), and the embedding (`.embedding`) of the matched text.\n",
    "\n",
    "You can visualize a search result using `text_sim.visualize(result)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Query 0: \"Hong Kong Buffet (608) 240-0762 2817 E Washington Ave, Madison, WI 53704\"\n",
      "Most similar matches:\n",
      "\n",
      "  idx  is_match      similarity  text\n",
      "-----  ----------  ------------  ----------------------------------------------------------------\n",
      "  179  False               0.77  Brink Lounge (608) 661-8599 701 E Washington Ave, Madison, WI\n",
      "  289  False               0.76  Barriques Coffee (608) 268-6264 127 W Washington Ave, Madison, W\n",
      "  231  False               0.76  Athens Gyros (608) 246-7733 1860 E Washington Ave, Madison, WI\n",
      "   51  False               0.76  Taco Bell (608) 249-7312 4120 E Washington Ave, Madison, WI\n",
      "   98  False               0.73  Einstein Bros Bagels (608) 242-9889 3904 E Washington Ave, Madis\n"
     ]
    }
   ],
   "source": [
    "# visualize results for each query using .visualize\n",
    "query_idx = 0\n",
    "result = result_collection.results[query_idx]\n",
    "text_sim.visualize(result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Query 2: \"Kabul Restaurant (608) 256-6322 540 State St, Madison, WI 53703\"\n",
      "Most similar matches:\n",
      "\n",
      "  idx  is_match      similarity  text\n",
      "-----  ----------  ------------  ----------------------------------------------------------------\n",
      "   17  True                0.92  Kabul Afghanistan Restaurant (608) 256-6322 540 State St, Madiso\n",
      "   66  False               0.79  State Street Brats (608) 255-5544 603 State St, Madison, WI\n",
      "   35  False               0.79  Parthenon Gyros Restaurant (608) 251-6311 316 State St, Madison,\n",
      "  109  False               0.78  Wasabi Japanese Restaurant (608) 255-5020 449 State St Ste 2G, M\n",
      "  186  False               0.75  Mediterranean Cafe (608) 251-8510 625 State St, Madison, WI\n"
     ]
    }
   ],
   "source": [
    "# visualize a matching result\n",
    "first_matching_idx = None\n",
    "for result in result_collection.results:\n",
    "    if result.num_matches > 0:\n",
    "        first_matching_idx = result.query_idx\n",
    "        break\n",
    "\n",
    "result = result_collection.results[first_matching_idx]\n",
    "text_sim.visualize(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can keep adding examples and querying your index after you create it. This is useful for production use-cases, where you have incoming data or frequently need to query your index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Query 0: \"googleplx 650-253-0000 1600 amphitheatre parkway, mountain view, ca 94043\"\n",
      "Most similar matches:\n",
      "\n",
      "  idx  is_match      similarity  text\n",
      "-----  ----------  ------------  ----------------------------------------------------------------\n",
      "  304  True                0.92  Googleplex (650) 253-0000 1600 Amphitheatre Parkway, Mountain Vi\n",
      "  150  False               0.52  KFC (608) 849-5004 600 W Main St, Waunakee, WI\n",
      "  299  False               0.51  Gus's Diner (608) 318-0900 630 N Westmount Dr, Sun Prairie, WI\n",
      "  116  False               0.5   Sweet Maple (415) 655-9169 2101 Sutter Street, San Francisco, CA\n",
      "  267  False               0.49  Pho Nam (608) 836-7040 610 Junction Rd Suite 109, Madison, WI\n"
     ]
    }
   ],
   "source": [
    "# add new example to the index\n",
    "new_examples = [\"Googleplex (650) 253-0000 1600 Amphitheatre Parkway, Mountain View, CA 94043\"]\n",
    "text_sim.add(new_examples)\n",
    "\n",
    "# search for the example, with typos in our query\n",
    "result_collection = text_sim.search([\"googleplx 650-253-0000 1600 amphitheatre parkway, mountain view, ca 94043\"], k=5)\n",
    "\n",
    "result = result_collection.results[0]\n",
    "text_sim.visualize(result)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.6 64-bit",
   "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.10.6"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "e7370f93d1d0cde622a1f8e1c04877d8463912d04d973331ad4851f04de6915a"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}