01_mdv_search.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hybrid search for the metadata catalogue of the Canton of Zurich\n",
    "In this notebook we prepare the data for our search application and setup the search index.\n",
    "\n",
    "- We load the data from the API. \n",
    "- We prepare the data and lemmatize the text data (for lexical search).\n",
    "- We embed via the OpenAI embedding API. \n",
    "- We test the embeddings with cosine similarity.\n",
    "- We setup the Weaviate index.\n",
    "- We create a collection with our data.\n",
    "- We test the index in regard to lexical, and vector search as well as the combination of both - hybrid search. The latter is what we use in the app."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO: Pandarallel will run on 8 workers.\n",
      "INFO: Pandarallel will use standard multiprocessing data transfer (pipe) to transfer data between the main process and workers.\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "from pandarallel import pandarallel\n",
    "import numpy as np\n",
    "\n",
    "pd.options.mode.chained_assignment = None\n",
    "pd.options.display.max_rows = 500\n",
    "pd.options.display.max_seq_items = 500\n",
    "pandarallel.initialize(progress_bar=False)\n",
    "\n",
    "from time import time\n",
    "import os\n",
    "import re\n",
    "import requests\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "import warnings\n",
    "from bs4 import MarkupResemblesLocatorWarning\n",
    "from sklearn.metrics.pairwise import cosine_similarity\n",
    "import weaviate\n",
    "from weaviate.classes.config import Property, DataType\n",
    "import weaviate.classes as wvc\n",
    "import weaviate.classes.config as wc\n",
    "import spacy\n",
    "from openai import OpenAI\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "warnings.filterwarnings(\"ignore\", category=MarkupResemblesLocatorWarning)\n",
    "\n",
    "load_dotenv()\n",
    "OPENAI_API_KEY = os.getenv(\"OPENAI_API_KEY\")\n",
    "client_openai = OpenAI(api_key=OPENAI_API_KEY)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Constants**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_FOLDER = \"_data/\"\n",
    "DATASETS = DATA_FOLDER + \"01_data.parq\"\n",
    "DATA_WITH_EMBEDDINGS = DATA_FOLDER + \"02_data_embedded.parq\"\n",
    "\n",
    "BASELINK_API = \"https://www.web.statistik.zh.ch/ogd/daten/zhweb.json\"\n",
    "\n",
    "# Dataset links are composed of this baselink and the identifier for each dataset.\n",
    "BASELINK_DATASHOP = (\n",
    "    \"https://www.zh.ch/de/politik-staat/statistik-daten/datenkatalog.html#/datasets/\"\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load and parse data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def retrieve_mdv_metadata():\n",
    "    \"\"\"Retrieve metadata for all datasets from the data catalog API.\n",
    "\n",
    "    Args:\n",
    "        None\n",
    "\n",
    "    Returns:\n",
    "        raw (list): Metadata for all datasets.\n",
    "    \"\"\"\n",
    "    try:\n",
    "        raw = requests.get(BASELINK_API)\n",
    "        if raw.status_code == 200:\n",
    "            raw = raw.json()\n",
    "            raw = raw[\"dataset\"]\n",
    "            return raw\n",
    "    except Exception as e:\n",
    "        print(e)\n",
    "        return None\n",
    "\n",
    "\n",
    "def clean_text(d):\n",
    "    \"\"\"Clean text. Remove square brackets that often enclose units or other information. Replace multiple spaces with single spaces and strip leading and trailing spaces.\n",
    "\n",
    "    Args:\n",
    "        d (str): Text to clean.\n",
    "\n",
    "    Returns:\n",
    "        str: Cleaned text.\n",
    "    \"\"\"\n",
    "    # Remove square brackets.\n",
    "    d = re.sub(r\"(\\[|\\])\", \" \", d)\n",
    "    # Replace multiple spaces with single space.\n",
    "    d = re.sub(r\"[\\s]+\", \" \", d)\n",
    "    d = d.strip()\n",
    "    return d\n",
    "\n",
    "\n",
    "def parse_and_prepare_metadata(raw):\n",
    "    \"\"\"Parse and prepare metadata for all available datasets.\n",
    "\n",
    "    Args:\n",
    "        raw (list): Raw metadata for all available datasets.\n",
    "\n",
    "    Returns:\n",
    "        pd.DataFrame: Parsed and prepared metadata for all datasets.\n",
    "    \"\"\"\n",
    "\n",
    "    identifier = []\n",
    "    links = []\n",
    "    titles = []\n",
    "    descriptions = []\n",
    "    themes = []\n",
    "    keywords = []\n",
    "    distributions = []\n",
    "    is_study = []\n",
    "\n",
    "    # Iterate over all datasets and retrieve metadata.\n",
    "    for dataset in raw:\n",
    "        identifier.append(dataset[\"identifier\"])\n",
    "\n",
    "        # The link to the dataset is composed of the base link and the dataset identifier.\n",
    "        links.append(BASELINK_DATASHOP + dataset[\"identifier\"])\n",
    "\n",
    "        titles.append(dataset[\"title\"].strip())\n",
    "\n",
    "        if dataset[\"description\"] not in [\"NA\", \"keine\", \"null\"]:\n",
    "            descriptions.append(dataset[\"description\"].strip())\n",
    "        else:\n",
    "            descriptions.append(\"\")\n",
    "\n",
    "        if isinstance(dataset[\"theme\"], list):\n",
    "            kwds = [x for x in dataset[\"theme\"] if x != None]\n",
    "            themes.append(\" \".join(kwds))\n",
    "        else:\n",
    "            themes.append(\"\")\n",
    "\n",
    "        if isinstance(dataset[\"keyword\"], list):\n",
    "            keywords.append(\" \".join(dataset[\"keyword\"]))\n",
    "        else:\n",
    "            keywords.append(\"\")\n",
    "\n",
    "        # Go through all resources and retrieve their textual content.\n",
    "        tmp_res = []\n",
    "        for resource in dataset[\"distribution\"]:\n",
    "            is_pdf = False\n",
    "            if resource[\"format\"] == \"PDF\":\n",
    "                is_pdf = True\n",
    "            tmp_res.append(resource[\"title\"])\n",
    "            if isinstance(resource[\"description\"], list):\n",
    "                tmp_res.append(resource[\"description\"])\n",
    "            else:\n",
    "                tmp_res.append(\"\")\n",
    "        is_study.append(is_pdf)\n",
    "\n",
    "        # Join textual content of dataset and resources.\n",
    "        distributions.append(\" \".join(tmp_res))\n",
    "\n",
    "    df = pd.DataFrame(identifier, columns=[\"identifier\"])\n",
    "    df[\"link\"] = links\n",
    "    df[\"title\"] = titles\n",
    "    df[\"description\"] = descriptions\n",
    "    df[\"theme\"] = themes\n",
    "    df[\"keyword\"] = keywords\n",
    "    df[\"distribution\"] = distributions\n",
    "    df[\"is_study\"] = is_study\n",
    "\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'pandas.core.frame.DataFrame'>\n",
      "RangeIndex: 795 entries, 0 to 794\n",
      "Data columns (total 9 columns):\n",
      " #   Column              Non-Null Count  Dtype \n",
      "---  ------              --------------  ----- \n",
      " 0   identifier          795 non-null    object\n",
      " 1   link                795 non-null    object\n",
      " 2   title               795 non-null    object\n",
      " 3   description         795 non-null    object\n",
      " 4   theme               795 non-null    object\n",
      " 5   keyword             795 non-null    object\n",
      " 6   distribution        795 non-null    object\n",
      " 7   is_study            795 non-null    bool  \n",
      " 8   text_for_embedding  795 non-null    object\n",
      "dtypes: bool(1), object(8)\n",
      "memory usage: 1.4 MB\n"
     ]
    },
    {
     "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>identifier</th>\n",
       "      <th>link</th>\n",
       "      <th>title</th>\n",
       "      <th>description</th>\n",
       "      <th>theme</th>\n",
       "      <th>keyword</th>\n",
       "      <th>distribution</th>\n",
       "      <th>is_study</th>\n",
       "      <th>text_for_embedding</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>214@statistisches-amt-kanton-zuerich</td>\n",
       "      <td>https://www.zh.ch/de/politik-staat/statistik-d...</td>\n",
       "      <td>Schül. Berufsschule [pro 1000 Einw.]</td>\n",
       "      <td>Anzahl BerufsschülerInnen gemessen an der zivi...</td>\n",
       "      <td>http://publications.europa.eu/resource/authori...</td>\n",
       "      <td>berufsschule bezirke gemeinden kanton_zuerich ...</td>\n",
       "      <td>Schül. Berufsschule [pro 1000 Einw.]</td>\n",
       "      <td>False</td>\n",
       "      <td>Schül. Berufsschule [pro 1000 Einw.] Anzahl Be...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>216@statistisches-amt-kanton-zuerich</td>\n",
       "      <td>https://www.zh.ch/de/politik-staat/statistik-d...</td>\n",
       "      <td>Schül. Mittelschule [pro 1000 Einw.]</td>\n",
       "      <td>Anzahl MittelschülerInnen inkl. Langzeitgymnas...</td>\n",
       "      <td>http://publications.europa.eu/resource/authori...</td>\n",
       "      <td>bezirke gemeinden kanton_zuerich mittelschulen...</td>\n",
       "      <td>Schül. Mittelschule [pro 1000 Einw.]</td>\n",
       "      <td>False</td>\n",
       "      <td>Schül. Mittelschule [pro 1000 Einw.] Anzahl Mi...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>217@statistisches-amt-kanton-zuerich</td>\n",
       "      <td>https://www.zh.ch/de/politik-staat/statistik-d...</td>\n",
       "      <td>Schül. Oberstufe [pro 1000 Einw.]</td>\n",
       "      <td>Anzahl OberstufenschülerInnen (exklusiv Langze...</td>\n",
       "      <td>http://publications.europa.eu/resource/authori...</td>\n",
       "      <td>bezirke gemeinden kanton_zuerich schueler seku...</td>\n",
       "      <td>Schül. Oberstufe [pro 1000 Einw.]</td>\n",
       "      <td>False</td>\n",
       "      <td>Schül. Oberstufe [pro 1000 Einw.] Anzahl Obers...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>218@statistisches-amt-kanton-zuerich</td>\n",
       "      <td>https://www.zh.ch/de/politik-staat/statistik-d...</td>\n",
       "      <td>Schül. Primarschule [pro 1000 Einw.]</td>\n",
       "      <td>Anzahl PrimarschülerInnen gemessen an der zivi...</td>\n",
       "      <td>http://publications.europa.eu/resource/authori...</td>\n",
       "      <td>bezirke gemeinden kanton_zuerich primarschule ...</td>\n",
       "      <td>Schül. Primarschule [pro 1000 Einw.]</td>\n",
       "      <td>False</td>\n",
       "      <td>Schül. Primarschule [pro 1000 Einw.] Anzahl Pr...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>229@statistisches-amt-kanton-zuerich</td>\n",
       "      <td>https://www.zh.ch/de/politik-staat/statistik-d...</td>\n",
       "      <td>Neuerstellte EFH [Whg.]</td>\n",
       "      <td>Neu erstellte Einfamilienhäuser</td>\n",
       "      <td>http://publications.europa.eu/resource/authori...</td>\n",
       "      <td>bezirke einfamilienhaus gemeinden immobilien k...</td>\n",
       "      <td>Neuerstellte EFH [Whg.]</td>\n",
       "      <td>False</td>\n",
       "      <td>Neuerstellte EFH [Whg.] Neu erstellte Einfamil...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                             identifier  \\\n",
       "0  214@statistisches-amt-kanton-zuerich   \n",
       "1  216@statistisches-amt-kanton-zuerich   \n",
       "2  217@statistisches-amt-kanton-zuerich   \n",
       "3  218@statistisches-amt-kanton-zuerich   \n",
       "4  229@statistisches-amt-kanton-zuerich   \n",
       "\n",
       "                                                link  \\\n",
       "0  https://www.zh.ch/de/politik-staat/statistik-d...   \n",
       "1  https://www.zh.ch/de/politik-staat/statistik-d...   \n",
       "2  https://www.zh.ch/de/politik-staat/statistik-d...   \n",
       "3  https://www.zh.ch/de/politik-staat/statistik-d...   \n",
       "4  https://www.zh.ch/de/politik-staat/statistik-d...   \n",
       "\n",
       "                                  title  \\\n",
       "0  Schül. Berufsschule [pro 1000 Einw.]   \n",
       "1  Schül. Mittelschule [pro 1000 Einw.]   \n",
       "2     Schül. Oberstufe [pro 1000 Einw.]   \n",
       "3  Schül. Primarschule [pro 1000 Einw.]   \n",
       "4               Neuerstellte EFH [Whg.]   \n",
       "\n",
       "                                         description  \\\n",
       "0  Anzahl BerufsschülerInnen gemessen an der zivi...   \n",
       "1  Anzahl MittelschülerInnen inkl. Langzeitgymnas...   \n",
       "2  Anzahl OberstufenschülerInnen (exklusiv Langze...   \n",
       "3  Anzahl PrimarschülerInnen gemessen an der zivi...   \n",
       "4                    Neu erstellte Einfamilienhäuser   \n",
       "\n",
       "                                               theme  \\\n",
       "0  http://publications.europa.eu/resource/authori...   \n",
       "1  http://publications.europa.eu/resource/authori...   \n",
       "2  http://publications.europa.eu/resource/authori...   \n",
       "3  http://publications.europa.eu/resource/authori...   \n",
       "4  http://publications.europa.eu/resource/authori...   \n",
       "\n",
       "                                             keyword  \\\n",
       "0  berufsschule bezirke gemeinden kanton_zuerich ...   \n",
       "1  bezirke gemeinden kanton_zuerich mittelschulen...   \n",
       "2  bezirke gemeinden kanton_zuerich schueler seku...   \n",
       "3  bezirke gemeinden kanton_zuerich primarschule ...   \n",
       "4  bezirke einfamilienhaus gemeinden immobilien k...   \n",
       "\n",
       "                            distribution  is_study  \\\n",
       "0  Schül. Berufsschule [pro 1000 Einw.]      False   \n",
       "1  Schül. Mittelschule [pro 1000 Einw.]      False   \n",
       "2     Schül. Oberstufe [pro 1000 Einw.]      False   \n",
       "3  Schül. Primarschule [pro 1000 Einw.]      False   \n",
       "4               Neuerstellte EFH [Whg.]      False   \n",
       "\n",
       "                                  text_for_embedding  \n",
       "0  Schül. Berufsschule [pro 1000 Einw.] Anzahl Be...  \n",
       "1  Schül. Mittelschule [pro 1000 Einw.] Anzahl Mi...  \n",
       "2  Schül. Oberstufe [pro 1000 Einw.] Anzahl Obers...  \n",
       "3  Schül. Primarschule [pro 1000 Einw.] Anzahl Pr...  \n",
       "4  Neuerstellte EFH [Whg.] Neu erstellte Einfamil...  "
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "raw = retrieve_mdv_metadata()\n",
    "df = parse_and_prepare_metadata(raw)\n",
    "df.title = df.title.apply(lambda x: x.strip())\n",
    "df.description = df.description.apply(lambda x: x.strip())\n",
    "\n",
    "# Combine title and description for embedding.\n",
    "df[\"text_for_embedding\"] = df.title + \" \" + df.description\n",
    "df[\"text_for_embedding\"] = df[\"text_for_embedding\"].apply(lambda x: x.strip())\n",
    "df.to_parquet(DATASETS)\n",
    "df.info(memory_usage=\"deep\")\n",
    "df.head(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lemmatize text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "nlp = spacy.load(\"de_core_news_lg\")\n",
    "\n",
    "LETTERS_AND_DIGITS = re.compile(r\"[^a-zäüöA-ZÜÄÖ0-9.]\")\n",
    "MULTIPLE_SPACES = re.compile(r\"\\s+\")\n",
    "\n",
    "\n",
    "def prepare_for_lexical_search(text, lower=False, remove_umlauts=False):\n",
    "    \"\"\"Lemmatize text, and optionally lower case and remove umlauts for lexical search.\n",
    "\n",
    "    Args:\n",
    "        text (str): Text to process.\n",
    "\n",
    "    Keyword Arguments:\n",
    "        lower (bool): Lower case text (default: {True}).\n",
    "        remove_umlauts (bool): Remove umlauts from text (default: {True}).\n",
    "\n",
    "    Returns:\n",
    "        str: Lemmatized text, optionally lower cased, and without umlauts.\n",
    "    \"\"\"\n",
    "    doc = nlp(text)\n",
    "    text = \" \".join([token.lemma_ if token.is_alpha else token.text for token in doc])\n",
    "    text = re.sub(LETTERS_AND_DIGITS, \" \", text)\n",
    "    text = re.sub(MULTIPLE_SPACES, \" \", text)\n",
    "    if lower:\n",
    "        text = text.lower()\n",
    "    if remove_umlauts:\n",
    "        text = text.replace(\"ä\", \"ae\").replace(\"ö\", \"oe\").replace(\"ü\", \"ue\")\n",
    "    return text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "df[\"title_lemma\"] = df.title.parallel_apply(lambda x: prepare_for_lexical_search(x))\n",
    "df[\"description_lemma\"] = df.description.parallel_apply(\n",
    "    lambda x: prepare_for_lexical_search(x)\n",
    ")\n",
    "df[\"theme_lemma\"] = df.theme.parallel_apply(lambda x: prepare_for_lexical_search(x))\n",
    "df[\"distribution_lemma\"] = df.distribution.parallel_apply(\n",
    "    lambda x: prepare_for_lexical_search(x)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Have a quick look at the data."
   ]
  },
  {
   "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>745</th>\n",
       "      <th>578</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>identifier</th>\n",
       "      <td>2202@statistisches-amt-kanton-zuerich</td>\n",
       "      <td>567@statistisches-amt-kanton-zuerich</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>link</th>\n",
       "      <td>https://www.zh.ch/de/politik-staat/statistik-d...</td>\n",
       "      <td>https://www.zh.ch/de/politik-staat/statistik-d...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>title</th>\n",
       "      <td>Resultate und Wahlbeteiligung Nationalratswahl...</td>\n",
       "      <td>Sitzverlegungen von Firmen im Kanton Zürich 2002</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>description</th>\n",
       "      <td>Echtzeitresultate und -wahlbeteiligungen der E...</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>theme</th>\n",
       "      <td>http://publications.europa.eu/resource/authori...</td>\n",
       "      <td>http://publications.europa.eu/resource/authori...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>keyword</th>\n",
       "      <td>nationalratswahlen wahlen wahlergebnisse wahla...</td>\n",
       "      <td>unternehmensstruktur statistik.info</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>distribution</th>\n",
       "      <td>Nationalratswahl 2023: Parteiresultate Ebene K...</td>\n",
       "      <td>Sitzverlegungen von Firmen im Kanton Zürich 2002</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>is_study</th>\n",
       "      <td>False</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>text_for_embedding</th>\n",
       "      <td>Resultate und Wahlbeteiligung Nationalratswahl...</td>\n",
       "      <td>Sitzverlegungen von Firmen im Kanton Zürich 2002</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>title_lemma</th>\n",
       "      <td>Resultat und Wahlbeteiligung Nationalratswahl ...</td>\n",
       "      <td>Sitzverlegung von Firma in Kanton Zürich 2002</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>description_lemma</th>\n",
       "      <td>Echtzeitresultat und wahlbeteiligungen der Ern...</td>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>theme_lemma</th>\n",
       "      <td>http publications.europa.eu resource authority...</td>\n",
       "      <td>http publications.europa.eu resource authority...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>distribution_lemma</th>\n",
       "      <td>Nationalratswahl 2023 Parteiresultat Ebene Kan...</td>\n",
       "      <td>Sitzverlegung von Firma in Kanton Zürich 2002</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                                  745  \\\n",
       "identifier                      2202@statistisches-amt-kanton-zuerich   \n",
       "link                https://www.zh.ch/de/politik-staat/statistik-d...   \n",
       "title               Resultate und Wahlbeteiligung Nationalratswahl...   \n",
       "description         Echtzeitresultate und -wahlbeteiligungen der E...   \n",
       "theme               http://publications.europa.eu/resource/authori...   \n",
       "keyword             nationalratswahlen wahlen wahlergebnisse wahla...   \n",
       "distribution        Nationalratswahl 2023: Parteiresultate Ebene K...   \n",
       "is_study                                                        False   \n",
       "text_for_embedding  Resultate und Wahlbeteiligung Nationalratswahl...   \n",
       "title_lemma         Resultat und Wahlbeteiligung Nationalratswahl ...   \n",
       "description_lemma   Echtzeitresultat und wahlbeteiligungen der Ern...   \n",
       "theme_lemma         http publications.europa.eu resource authority...   \n",
       "distribution_lemma  Nationalratswahl 2023 Parteiresultat Ebene Kan...   \n",
       "\n",
       "                                                                  578  \n",
       "identifier                       567@statistisches-amt-kanton-zuerich  \n",
       "link                https://www.zh.ch/de/politik-staat/statistik-d...  \n",
       "title                Sitzverlegungen von Firmen im Kanton Zürich 2002  \n",
       "description                                                            \n",
       "theme               http://publications.europa.eu/resource/authori...  \n",
       "keyword                           unternehmensstruktur statistik.info  \n",
       "distribution        Sitzverlegungen von Firmen im Kanton Zürich 2002   \n",
       "is_study                                                         True  \n",
       "text_for_embedding   Sitzverlegungen von Firmen im Kanton Zürich 2002  \n",
       "title_lemma             Sitzverlegung von Firma in Kanton Zürich 2002  \n",
       "description_lemma                                                      \n",
       "theme_lemma         http publications.europa.eu resource authority...  \n",
       "distribution_lemma      Sitzverlegung von Firma in Kanton Zürich 2002  "
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.sample(2).T"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Embed data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For our prototype app we use OpenAIs embeddings for convienience. \n",
    "\n",
    "We also tested these open source models with [SentenceTransformers](https://sbert.net/) with very good results: \n",
    "\n",
    "- [PM-AI/bi-encoder_msmarco_bert-base_german](https://huggingface.co/PM-AI/bi-encoder_msmarco_bert-base_german) - 350 tokens context length\n",
    "- [Jina AI jina-embeddings-v2-base-de](https://huggingface.co/jinaai/jina-embeddings-v2-base-de) - 8192 tokens context length"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def embed_with_openai(texts, model=\"text-embedding-3-small\"):\n",
    "    \"\"\"Embed text using OpenAIs embedding API.\n",
    "\n",
    "    Args:\n",
    "        texts (list): List of texts to embed.\n",
    "        model (str): OpenAI model to use (default: {\"text-embedding-3-small\"}).\n",
    "\n",
    "    Returns:\n",
    "        list: List of embeddings.\n",
    "    \"\"\"\n",
    "    response = client_openai.embeddings.create(input=texts, model=model)\n",
    "    return [x.embedding for x in response.data]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "df[\"embedding_openai\"] = embed_with_openai(df.text_for_embedding.values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "cols = [\n",
    "    \"identifier\",\n",
    "    \"link\",\n",
    "    \"title\",\n",
    "    \"title_lemma\",\n",
    "    \"description\",\n",
    "    \"description_lemma\",\n",
    "    \"text_for_embedding\",\n",
    "    \"theme\",\n",
    "    \"theme_lemma\",\n",
    "    \"keyword\",\n",
    "    \"distribution\",\n",
    "    \"distribution_lemma\",\n",
    "    \"is_study\",\n",
    "    \"embedding_openai\",\n",
    "]\n",
    "\n",
    "df = df[cols]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "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>177</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>identifier</th>\n",
       "      <td>563@awel-kanton-zuerich</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>link</th>\n",
       "      <td>https://www.zh.ch/de/politik-staat/statistik-d...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>title</th>\n",
       "      <td>Lufttemperatur und Luftfeuchte LoRa-Sensor-Mes...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>title_lemma</th>\n",
       "      <td>Lufttemperatur und Luftfeuchte LoRa Sensor Mes...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>description</th>\n",
       "      <td>Seit Frühjahr 2019 betreibt das AWEL des Kanto...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>description_lemma</th>\n",
       "      <td>seit Frühjahr 2019 betreiben der AWEL der Kant...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>text_for_embedding</th>\n",
       "      <td>Lufttemperatur und Luftfeuchte LoRa-Sensor-Mes...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>theme</th>\n",
       "      <td>http://publications.europa.eu/resource/authori...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>theme_lemma</th>\n",
       "      <td>http publications.europa.eu resource authority...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>keyword</th>\n",
       "      <td>luft lufttemperaturen ogd</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>distribution</th>\n",
       "      <td>LoRa-Sensor-Messwerte 2022/12  LoRa-Sensor-Mes...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>distribution_lemma</th>\n",
       "      <td>LoRa Sensor Messwerte 2022 12 LoRa Sensor Mess...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>is_study</th>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>embedding_openai</th>\n",
       "      <td>[-0.036526869982481, 0.002765156328678131, 0.0...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                                  177\n",
       "identifier                                    563@awel-kanton-zuerich\n",
       "link                https://www.zh.ch/de/politik-staat/statistik-d...\n",
       "title               Lufttemperatur und Luftfeuchte LoRa-Sensor-Mes...\n",
       "title_lemma         Lufttemperatur und Luftfeuchte LoRa Sensor Mes...\n",
       "description         Seit Frühjahr 2019 betreibt das AWEL des Kanto...\n",
       "description_lemma   seit Frühjahr 2019 betreiben der AWEL der Kant...\n",
       "text_for_embedding  Lufttemperatur und Luftfeuchte LoRa-Sensor-Mes...\n",
       "theme               http://publications.europa.eu/resource/authori...\n",
       "theme_lemma         http publications.europa.eu resource authority...\n",
       "keyword                                     luft lufttemperaturen ogd\n",
       "distribution        LoRa-Sensor-Messwerte 2022/12  LoRa-Sensor-Mes...\n",
       "distribution_lemma  LoRa Sensor Messwerte 2022 12 LoRa Sensor Mess...\n",
       "is_study                                                        False\n",
       "embedding_openai    [-0.036526869982481, 0.002765156328678131, 0.0..."
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.sample(1).T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'pandas.core.frame.DataFrame'>\n",
      "RangeIndex: 795 entries, 0 to 794\n",
      "Data columns (total 14 columns):\n",
      " #   Column              Non-Null Count  Dtype \n",
      "---  ------              --------------  ----- \n",
      " 0   identifier          795 non-null    object\n",
      " 1   link                795 non-null    object\n",
      " 2   title               795 non-null    object\n",
      " 3   title_lemma         795 non-null    object\n",
      " 4   description         795 non-null    object\n",
      " 5   description_lemma   795 non-null    object\n",
      " 6   text_for_embedding  795 non-null    object\n",
      " 7   theme               795 non-null    object\n",
      " 8   theme_lemma         795 non-null    object\n",
      " 9   keyword             795 non-null    object\n",
      " 10  distribution        795 non-null    object\n",
      " 11  distribution_lemma  795 non-null    object\n",
      " 12  is_study            795 non-null    bool  \n",
      " 13  embedding_openai    795 non-null    object\n",
      "dtypes: bool(1), object(13)\n",
      "memory usage: 11.4 MB\n"
     ]
    }
   ],
   "source": [
    "df.to_parquet(DATA_WITH_EMBEDDINGS)\n",
    "df.info(memory_usage=\"deep\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test embeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Perform a quick check with scikit's cosine similarity function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_parquet(DATA_WITH_EMBEDDINGS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([['Kleineres Armutsrisiko dank günstiger Wohnung'],\n",
       "       ['Wohnungswechsel'],\n",
       "       ['Leerwohnungszahlen'],\n",
       "       ['Entspannung auf dem Wohnungsmarkt?'],\n",
       "       ['Ersatzneubau von Wohnungen immer wichtiger'],\n",
       "       ['Bevölkerungsdynamik und Wohnungsbau'],\n",
       "       ['Dynamischer Zürcher Wohnungsmarkt'],\n",
       "       ['Gemeinnützige Wohnungen in Gemeinden des Kantons Zürich [Anz.]'],\n",
       "       ['Zuwanderung und Wohnungsmarkt'],\n",
       "       ['Anteil Gemeinnützige Wohnungen in Gemeinden des Kantons Zürich [%]'],\n",
       "       ['Abfall im Kanton Zürich'],\n",
       "       ['Abfall im Kanton Zürich'],\n",
       "       ['Abfall im Kanton Zürich'],\n",
       "       ['Kantonaler Wohnungsleerstand bleibt niedrig, aber Anzeichen für Umschwung auf Wohnungsmarkt'],\n",
       "       ['Kantonaler Wohnungsleerstand immer noch rückläufig'],\n",
       "       ['Aus- und Einzugsmieten im Kanton Zürich'],\n",
       "       ['Zürcher Haushalte im Wandel der Zeit'],\n",
       "       ['Schwache Konjunktur bremst Bevölkerungswachstum im Kanton Zürich'],\n",
       "       ['Gebäude und Wohnungen im Kanton Zürich'],\n",
       "       ['Kaufen oder mieten?']], dtype=object)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "query = \"wohnungsnot in zürich\"\n",
    "show_n_results = 20\n",
    "\n",
    "# Vector search based on OpenAI embeddings.\n",
    "now = time()\n",
    "embedding_openai = embed_with_openai(query)\n",
    "cosine_sim = cosine_similarity(embedding_openai, df.embedding_openai.tolist())\n",
    "top_k = np.argsort(cosine_sim[0])\n",
    "top_k = top_k[::-1]\n",
    "display(df.iloc[top_k][[\"title\"]].head(show_n_results).values)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Implement search"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we create the actual search index for lexical and semantic search.\n",
    "\n",
    "Note that [Weaviates](https://weaviate.io/developers/weaviate) default location for the index is `~/.local/share/weaviate`. If you want to use the app on a different machine, simply copy the files from this location to the same location on the other machine. You can also set a specific path with parameter `persistence_data_path` – see [documentation here](https://weaviate.io/developers/weaviate/installation/embedded)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_parquet(DATA_WITH_EMBEDDINGS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Started /Users/patrickarnecke/.cache/weaviate-embedded: process ID 4665\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "{\"action\":\"startup\",\"default_vectorizer_module\":\"none\",\"level\":\"info\",\"msg\":\"the default vectorizer modules is set to \\\"none\\\", as a result all new schema classes without an explicit vectorizer setting, will use this vectorizer\",\"time\":\"2024-07-17T17:25:39+02:00\"}\n",
      "{\"action\":\"startup\",\"auto_schema_enabled\":true,\"level\":\"info\",\"msg\":\"auto schema enabled setting is set to \\\"true\\\"\",\"time\":\"2024-07-17T17:25:39+02:00\"}\n",
      "{\"level\":\"info\",\"msg\":\"No resource limits set, weaviate will use all available memory and CPU. To limit resources, set LIMIT_RESOURCES=true\",\"time\":\"2024-07-17T17:25:39+02:00\"}\n",
      "{\"level\":\"warning\",\"msg\":\"Multiple vector spaces are present, GraphQL Explore and REST API list objects endpoint module include params has been disabled as a result.\",\"time\":\"2024-07-17T17:25:39+02:00\"}\n",
      "{\"action\":\"grpc_startup\",\"level\":\"info\",\"msg\":\"grpc server listening at [::]:50050\",\"time\":\"2024-07-17T17:25:39+02:00\"}\n",
      "{\"action\":\"restapi_management\",\"level\":\"info\",\"msg\":\"Serving weaviate at http://127.0.0.1:8079\",\"time\":\"2024-07-17T17:25:39+02:00\"}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "{\"action\":\"read_disk_use\",\"level\":\"warning\",\"msg\":\"disk usage currently at 90.00%, threshold set to 80.00%\",\"path\":\"/Users/patrickarnecke/.local/share/weaviate\",\"time\":\"2024-07-17T17:25:39+02:00\"}\n",
      "{\"level\":\"info\",\"msg\":\"Completed loading shard mdv_epJ0bgfBKeWH in 3.908291ms\",\"time\":\"2024-07-17T17:25:39+02:00\"}\n",
      "{\"action\":\"lsm_compaction\",\"class\":\"MDV\",\"index\":\"mdv\",\"level\":\"warning\",\"msg\":\"compaction halted due to shard READONLY status\",\"path\":\"/Users/patrickarnecke/.local/share/weaviate/mdv/epJ0bgfBKeWH/lsm\",\"shard\":\"epJ0bgfBKeWH\",\"time\":\"2024-07-17T17:25:39+02:00\"}\n",
      "{\"action\":\"hnsw_vector_cache_prefill\",\"count\":3000,\"index_id\":\"main\",\"level\":\"info\",\"limit\":1000000000000,\"msg\":\"prefilled vector cache\",\"time\":\"2024-07-17T17:25:39+02:00\",\"took\":81708}\n",
      "{\"level\":\"info\",\"msg\":\"Completed loading shard rrbs_mRNSfQC7t9fQ in 315.293792ms\",\"time\":\"2024-07-17T17:25:40+02:00\"}\n",
      "{\"action\":\"lsm_compaction\",\"class\":\"RRBs\",\"index\":\"rrbs\",\"level\":\"warning\",\"msg\":\"compaction halted due to shard READONLY status\",\"path\":\"/Users/patrickarnecke/.local/share/weaviate/rrbs/mRNSfQC7t9fQ/lsm\",\"shard\":\"mRNSfQC7t9fQ\",\"time\":\"2024-07-17T17:25:40+02:00\"}\n",
      "{\"action\":\"set_shard_read_only\",\"level\":\"warning\",\"msg\":\"Set READONLY, disk usage currently at 90.00%, threshold set to 90.00%\",\"path\":\"/Users/patrickarnecke/.local/share/weaviate\",\"time\":\"2024-07-17T17:25:40+02:00\"}\n",
      "{\"action\":\"hnsw_vector_cache_prefill\",\"count\":237352,\"index_id\":\"main\",\"level\":\"info\",\"limit\":1000000000000,\"msg\":\"prefilled vector cache\",\"time\":\"2024-07-17T17:25:45+02:00\",\"took\":5043900334}\n"
     ]
    }
   ],
   "source": [
    "# This will either connect to an existing Weaviate instance or create a new one.\n",
    "# We only can execute this once, as is. If we want to run it again, we need to restart the kernel.\n",
    "# More info regarding the Embedded Weaviate client:\n",
    "# https://weaviate.io/developers/weaviate/installation/embedded\n",
    "# By default the index data is stored in ~/.local/share/weaviate/\n",
    "\n",
    "client = weaviate.connect_to_embedded()\n",
    "\n",
    "# # Use this line if another application on your machine is already using Weaviate on the default port.\n",
    "# # This notebook also counts as another application. If your notebook is running, you need to use a different port for the Streamlit app. Otherwise, you get an error.\n",
    "# client = weaviate.connect_to_local(port=8079, grpc_port=50050)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Get the meta endpoint description of Weaviate.\n",
    "# display(client.get_meta())\n",
    "\n",
    "# # Ping Weaviate’s live and ready state.\n",
    "# print(client.is_live())\n",
    "# print(client.is_ready())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Delete collection.\n",
    "client.collections.delete(\"MDV\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "{\"level\":\"info\",\"msg\":\"Created shard mdv_OO9LVECxFH3G in 18.866666ms\",\"time\":\"2024-07-17T17:26:13+02:00\"}\n",
      "{\"action\":\"hnsw_vector_cache_prefill\",\"count\":1000,\"index_id\":\"main\",\"level\":\"info\",\"limit\":1000000000000,\"msg\":\"prefilled vector cache\",\"time\":\"2024-07-17T17:26:14+02:00\",\"took\":5410375}\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<weaviate.collections.collection.Collection at 0x166bfd310>"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# First we need to create a schema that defines how data is stored, organized and retrieved in Weaviate.\n",
    "# A schema is called a \"collection\". We can define as many collections as we want.\n",
    "\n",
    "client.collections.create(\n",
    "    \"MDV\",\n",
    "    vectorizer_config=wc.Configure.Vectorizer.none(),\n",
    "    inverted_index_config=wvc.config.Configure.inverted_index(\n",
    "        bm25_b=0.75,  # default 0.75\n",
    "        bm25_k1=1.2,  # default 1.2\n",
    "        stopwords_additions=None,\n",
    "        stopwords_preset=None,\n",
    "        stopwords_removals=None,\n",
    "    ),\n",
    "    properties=[\n",
    "        Property(name=\"identifier\", data_type=DataType.TEXT),\n",
    "        Property(name=\"link\", data_type=DataType.TEXT),\n",
    "        Property(name=\"title\", data_type=DataType.TEXT),\n",
    "        Property(name=\"title_lemma\", data_type=DataType.TEXT),\n",
    "        Property(name=\"description\", data_type=DataType.TEXT),\n",
    "        Property(name=\"description_lemma\", data_type=DataType.TEXT),\n",
    "        Property(name=\"theme\", data_type=DataType.TEXT),\n",
    "        Property(name=\"theme_lemma\", data_type=DataType.TEXT),\n",
    "        Property(name=\"keyword\", data_type=DataType.TEXT),\n",
    "        Property(name=\"distribution\", data_type=DataType.TEXT),\n",
    "        Property(name=\"distribution_lemma\", data_type=DataType.TEXT),\n",
    "        Property(name=\"is_study\", data_type=DataType.BOOL),\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # List all collections.\n",
    "# for v in client.collections.list_all().values():\n",
    "#     print(v.name)\n",
    "\n",
    "# # # Get detailed information about all collections.\n",
    "# # schema = client.collections.list_all(simple=False)\n",
    "# # print(schema)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now we ingest data into the collection in the form of \"items\".\n",
    "\n",
    "collection = client.collections.get(\"MDV\")\n",
    "\n",
    "with collection.batch.dynamic() as batch:\n",
    "    for data in df.to_dict(orient=\"records\"):\n",
    "        properties = {\n",
    "            \"identifier\": data[\"identifier\"],\n",
    "            \"link\": data[\"link\"],\n",
    "            \"title\": data[\"title\"],\n",
    "            \"title_lemma\": data[\"title_lemma\"],\n",
    "            \"description\": data[\"description\"],\n",
    "            \"description_lemma\": data[\"description_lemma\"],\n",
    "            \"theme\": data[\"theme\"],\n",
    "            \"theme_lemma\": data[\"theme_lemma\"],\n",
    "            \"keyword\": data[\"keyword\"],\n",
    "            \"distribution\": data[\"distribution\"],\n",
    "            \"distribution_lemma\": data[\"distribution_lemma\"],\n",
    "            \"is_study\": data[\"is_study\"],\n",
    "        }\n",
    "        batch.add_object(properties=properties, vector=data[\"embedding_openai\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Object(uuid=_WeaviateUUIDInt('004af901-ce28-4f6e-a731-77055a09c281'), metadata=MetadataReturn(creation_time=None, last_update_time=None, distance=None, certainty=None, score=None, explain_score=None, is_consistent=None, rerank_score=None), properties={'description': 'Entwicklungstrends eines Wirtschaftszweigs', 'theme_lemma': 'http publications.europa.eu resource authority data theme AGRI', 'identifier': '493@statistisches-amt-kanton-zuerich', 'is_study': True, 'keyword': 'wirtschaft statistik.info', 'distribution': 'Strukturwandel in der Zürcher Landwirtschaft ', 'title_lemma': 'Strukturwandel in der Zürcher Landwirtschaft', 'theme': 'http://publications.europa.eu/resource/authority/data-theme/AGRI', 'distribution_lemma': 'Strukturwandel in der Zürcher Landwirtschaft', 'description_lemma': 'Entwicklungstrends ein Wirtschaftszweig', 'link': 'https://www.zh.ch/de/politik-staat/statistik-daten/datenkatalog.html#/datasets/493@statistisches-amt-kanton-zuerich', 'title': 'Strukturwandel in der Zürcher Landwirtschaft'}, references=None, vector={}, collection='MDV')\n"
     ]
    }
   ],
   "source": [
    "# List all items in the collection.\n",
    "collection = client.collections.get(\"MDV\")\n",
    "for item in collection.iterator():\n",
    "    print(item)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The collection contains 795 items.\n"
     ]
    }
   ],
   "source": [
    "# Get total count of all items in the collection.\n",
    "collection = client.collections.get(\"MDV\")\n",
    "response = collection.aggregate.over_all(total_count=True)\n",
    "\n",
    "# Check if the total count of items in the collection is equal to the number of items in the DataFrame.\n",
    "# If the assertion fails, try to re-run the previous cells.\n",
    "assert response.total_count == len(df)\n",
    "\n",
    "print(f\"The collection contains {response.total_count} items.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Lexical search"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[Weaviate documentation for lexical search](https://weaviate.io/developers/weaviate/search/bm25)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "collection = client.collections.get(\"MDV\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set fields to search in with BM25.\n",
    "# We exclude the identifier and link fields.\n",
    "# We give twice the weight to the title fields.\n",
    "\n",
    "query_properties = [\n",
    "    \"title^2\",\n",
    "    \"description\",\n",
    "    \"title_lemma^2\",\n",
    "    \"description_lemma\",\n",
    "    \"theme\",\n",
    "    \"theme_lemma\",\n",
    "    \"keyword\",\n",
    "    \"distribution\",\n",
    "    \"distribution_lemma\",\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Löhne, Teilzeitarbeit und Boni.\n",
      "Löhne im Kanton Zürich liegen deutlich über schweizerischem Mittel\n",
      "Löhne im Kanton Zürich- Branchen, Tätigkeiten, Alter, Geschlecht\n",
      "Kontaktpersonen des Volksschulamts des Kantons Zürich nach Schulgemeinde und Themen\n",
      "Kanton Zürich in Zahlen\n",
      "Lohnstudie Kanton Zürich\n",
      "Nutzungsdichte im Kanton Zürich\n",
      "Strukturatlas des Kantons Zürich\n",
      "Autofahren im Kanton Zürich\n",
      "Sozialhilfebezug im Kanton Zürich\n",
      "Abfall im Kanton Zürich\n",
      "Abfall im Kanton Zürich\n",
      "Veloverkehr im Kanton Zürich\n",
      "Siedlungsstruktur im Kanton Zürich\n",
      "Mobilität im Kanton Zürich\n",
      "Strukturatlas des Kantons Zürich\n",
      "Abfall im Kanton Zürich\n",
      "Treibstoffverbrauch im Kanton Zürich\n",
      "Fahrzeugflotte der kantonalen Verwaltung Zürich\n",
      "Datengrundlage zur ZHweb-Seite «Autoflotte»\n"
     ]
    }
   ],
   "source": [
    "query = \"Löhne in Zürich\"\n",
    "query = prepare_for_lexical_search(query)\n",
    "\n",
    "# See documentation for more details about parameters:\n",
    "# https://weaviate.io/developers/weaviate/search/bm25\n",
    "response = collection.query.bm25(\n",
    "    query=query,\n",
    "    query_properties=query_properties,\n",
    "    offset=0,\n",
    "    limit=20,\n",
    "    # auto_limit=10,\n",
    "    return_metadata=wvc.query.MetadataQuery(\n",
    "        score=True,\n",
    "        explain_score=True,\n",
    "    ),\n",
    ")\n",
    "for item in response.objects:\n",
    "    print(\n",
    "        item.properties[\"title\"],\n",
    "        # item.properties[\"description\"],\n",
    "        # item.metadata.score,\n",
    "        # item.metadata.explain_score,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Semantic search"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[Weaviate documentation for semantic search](https://weaviate.io/developers/weaviate/search/similarity)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Löhne im Kanton Zürich liegen deutlich über schweizerischem Mittel 0.24334806203842163 0.8783259391784668\n",
      "Löhne im Kanton Zürich- Branchen, Tätigkeiten, Alter, Geschlecht 0.32373493909835815 0.8381325006484985\n",
      "Lohngefälle im Kanton Zürich 0.3252621293067932 0.8373689651489258\n",
      "Löhne, Teilzeitarbeit und Boni. 0.3402245044708252 0.8298877477645874\n",
      "Frauenlöhne, Männerlöhne. Vollzeitlöhne, Teilzeitlöhne. 0.3444167375564575 0.8277916312217712\n",
      "Frauenlöhne, Männerlöhne. 0.363132119178772 0.818433940410614\n",
      "Lohnstudie Kanton Zürich 0.3684327006340027 0.8157836198806763\n",
      "Teilzeitarbeit boomt 0.42636024951934814 0.7868198752403259\n",
      "Arbeitskräfteballung in Zürich 0.4340125322341919 0.782993733882904\n",
      "35-Stunden Woche im Kanton Zürich? 0.4363228678703308 0.7818385362625122\n",
      "Abfall im Kanton Zürich 0.4381587505340576 0.7809206247329712\n",
      "Abfall im Kanton Zürich 0.4381587505340576 0.7809206247329712\n",
      "Abfall im Kanton Zürich 0.4381587505340576 0.7809206247329712\n",
      "Lehrstellenmangel in Dienstleistungsbranchen 0.45365339517593384 0.7731733322143555\n",
      "Zentrale Lagen begehrt 0.4547516107559204 0.7726241946220398\n",
      "Prämienverbilligung: Wer profitiert – und wer bezahlt 0.45475172996520996 0.772624135017395\n",
      "Gripen, Mindestlohn und Kirchensteuer 0.454819917678833 0.7725900411605835\n",
      "Firmengründungen, Wirtschaftsstruktur und Steuerbelastung 0.4556270241737366 0.7721865177154541\n",
      "Reiche Rentner? 0.46039795875549316 0.7698010206222534\n",
      "Der Preis des Bodens 0.46108007431030273 0.7694599628448486\n"
     ]
    }
   ],
   "source": [
    "query = \"Löhne in Zürich\"\n",
    "query_embedding = embed_with_openai(query)[0]\n",
    "\n",
    "# See documentation for more details about parameters:\n",
    "# https://weaviate.io/developers/weaviate/search/similarity\n",
    "collection = client.collections.get(\"MDV\")\n",
    "response = collection.query.near_vector(\n",
    "    near_vector=query_embedding,\n",
    "    limit=20,\n",
    "    # auto_limit=2,\n",
    "    return_metadata=wvc.query.MetadataQuery(distance=True, certainty=True),\n",
    ")\n",
    "\n",
    "for item in response.objects:\n",
    "    print(\n",
    "        item.properties[\"title\"],\n",
    "        item.metadata.distance,\n",
    "        item.metadata.certainty,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hybrid search"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[Weviate documentation for hybrid search - the combination of lexical and vector search](https://weaviate.io/developers/weaviate/search/hybrid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Löhne im Kanton Zürich liegen deutlich über schweizerischem Mittel\n",
      "Löhne, Teilzeitarbeit und Boni.\n",
      "Löhne im Kanton Zürich- Branchen, Tätigkeiten, Alter, Geschlecht\n",
      "Lohngefälle im Kanton Zürich\n",
      "Frauenlöhne, Männerlöhne. Vollzeitlöhne, Teilzeitlöhne.\n",
      "Lohnstudie Kanton Zürich\n",
      "Frauenlöhne, Männerlöhne.\n",
      "Teilzeitarbeit boomt\n",
      "Arbeitskräfteballung in Zürich\n",
      "Abfall im Kanton Zürich\n",
      "Abfall im Kanton Zürich\n",
      "Abfall im Kanton Zürich\n",
      "35-Stunden Woche im Kanton Zürich?\n",
      "Firmengründungen, Wirtschaftsstruktur und Steuerbelastung\n",
      "Lehrstellenmangel in Dienstleistungsbranchen\n",
      "Zentrale Lagen begehrt\n",
      "Prämienverbilligung: Wer profitiert – und wer bezahlt\n",
      "Gripen, Mindestlohn und Kirchensteuer\n",
      "Reiche Rentner?\n",
      "Der Preis des Bodens\n",
      "Unbezahlte Arbeit im Kanton Zürich\n",
      "Woher kommen die Steuern?\n",
      "Kontaktpersonen des Volksschulamts des Kantons Zürich nach Schulgemeinde und Themen\n",
      "Dienstleistungen gewinnen die Oberhand\n",
      "Sitzverlegungen von Firmen im Kanton Zürich 2002\n",
      "Immer mehr Hochqualifizierte\n",
      "Die Sozialhilfe im Kanton Zürich 2001\n",
      "Die Bodenpreise haben sich 2004 im Kanton Zürich nur wenig verändert\n",
      "Arbeitskräfte im Kanton Zürich immer besser qualifiziert\n",
      "Der Wirtschaftsraum Zürich – Eine Übersicht in Karten\n",
      "Kanton Zürich zieht junge Leute an\n",
      "Wirtschaftskraft in den Zürcher Regionen\n",
      "Der finanzielle Spielraum wird enger\n",
      "Ausländische Arbeitskräfte im Kanton Zürich.\n",
      "Zürcher Finanzausgleich\n",
      "Zürcher Gemeindefinanzen 2001\n",
      "Ausländische Erwerbstätige sind stark von Konjunkturschwankungen betroffen\n",
      "Immobilien- und Bodenpreise im Kanton Zürich 2002: Unterschiedliche Tendenzen\n",
      "Freiwillige und ihre Leistungen im Kanton Zürich und der Schweiz\n",
      "Die wirtschaftlichen Verflechtungen zwischen Zürich und Baden-Württemberg - eine Bestandsaufnahme\n"
     ]
    }
   ],
   "source": [
    "query = \"Löhne in Zürich\"\n",
    "query_embedding = embed_with_openai(query)[0]\n",
    "query = prepare_for_lexical_search(query)\n",
    "\n",
    "# See documentation for more details about parameters:\n",
    "# https://weaviate.io/developers/weaviate/search/hybrid\n",
    "collection = client.collections.get(\"MDV\")\n",
    "response = collection.query.hybrid(\n",
    "    query=query,\n",
    "    vector=query_embedding,\n",
    "    limit=40,\n",
    "    # auto_limit=2,\n",
    "    alpha=0.7,\n",
    "    fusion_type=wvc.query.HybridFusion.RELATIVE_SCORE,\n",
    "    # return_metadata=wvc.query.MetadataQuery(distance=True),\n",
    ")\n",
    "\n",
    "for item in response.objects:\n",
    "    print(\n",
    "        item.properties[\"title\"],\n",
    "        # item.properties[\"year\"],\n",
    "        # item.metadata.score,\n",
    "        # item.metadata.distance,\n",
    "        # item.metadata.certainty,\n",
    "    )"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "standard",
   "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.9.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}