notebooks/transformers.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "##### Notebook requires the chunking comparison from first"
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%pip install sentence_transformers codecarbon xlsxwriter"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!bash /home/azureuser/cloudfiles/code/blobfuse/blobfuse_raadsinformatie.sh"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {}
    },
    {
      "cell_type": "code",
      "source": [
        "import sys\n",
        "sys.path.append(\"..\")"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1721685437682
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Select where to run notebook: \"azure\" or \"local\"\n",
        "# my_run = \"azure\"\n",
        "my_run = \"local\""
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1721685438626
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import my_secrets as sc\n",
        "import settings as st\n",
        "\n",
        "if my_run == \"azure\":\n",
        "    import config_azure as cf\n",
        "elif my_run == \"local\":\n",
        "    import config as cf"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1721685439752
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "\n",
        "if my_run == \"azure\":\n",
        "    if not os.path.exists(cf.HUGGING_CACHE):\n",
        "        os.mkdir(cf.HUGGING_CACHE)\n",
        "\n",
        "    os.environ[\"TRANSFORMERS_CACHE\"] = cf.HUGGING_CACHE\n",
        "    os.environ[\"HF_HOME\"] = cf.HUGGING_CACHE"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1721685439951
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import pandas as pd\n",
        "from pathlib import Path\n",
        "\n",
        "comparison_folder = f\"{cf.raadsinformatie_out_folder}/comparison\"\n",
        "Path(comparison_folder).mkdir(parents=True, exist_ok=True)\n",
        "\n",
        "chunking_comparison_file = Path(comparison_folder, \"chunking_comparison.csv\")\n",
        "\n",
        "chunking_folder = f\"{cf.raadsinformatie_out_folder}/comparison/chunking/\""
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1721685457549
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "import glob\n",
        "\n",
        "woo_dirs = \\\n",
        "        [f\"{cf.woo_sources['openamsterdam']}/{folder}\" for folder in os.listdir(cf.woo_sources['openamsterdam'])] + \\\n",
        "        [f\"{cf.woo_sources['raadsinformatie']}/{folder}\" for folder in os.listdir(cf.woo_sources['raadsinformatie'])] + \\\n",
        "        [f\"{cf.woo_sources['amsterdam.nl']}/{folder}\" for folder in os.listdir(cf.woo_sources['amsterdam.nl'])]\n",
        "\n",
        "woo_files = sum([glob.glob(f\"{folder}/*.ocr\") for folder in woo_dirs], [])"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1721685458670
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from ast import literal_eval\n",
        "\n",
        "chunks_df = pd.read_csv(chunking_comparison_file)\n",
        "\n",
        "for column in chunks_df.columns:\n",
        "    if column.startswith(\"chunks\"):\n",
        "        chunks_df[column] = chunks_df[column].apply(literal_eval)"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1721685461636
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "chunked_files = dict()\n",
        "CHUNKING_METHODS = [field.removeprefix(\"chunks_\") for field in chunks_df.columns if field.startswith(\"chunks_\")]\n",
        "CHUNKING_METHODS"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1721685463876
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "len(woo_files)"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1721685465737
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### GPT connection \n",
        "Used for reformulating prompts, generating example answers & generating answers for final comparison"
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from src.llms.openai import get_client, prompt_gpt\n",
        "\n",
        "API_KEY = sc.AZURE_OPENAI_API_KEY\n",
        "RESOURCE_ENDPOINT = st.AZURE_OPENAI_ENDPOINT\n",
        "\n",
        "client = get_client(API_KEY, RESOURCE_ENDPOINT)"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1721685468686
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "prompt_gpt(client, prompt=\"Test, test! You good?\")"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1721685474603
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Next, we generate possible answers and reformulate the queries to try to improve the retrieval"
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "#TODO: Delete; notes me5\n",
        "256 2500s \n",
        "512 2571"
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import json\n",
        "\n",
        "possible_answers = [\n",
        "    prompt_gpt(client, prompt=f\"\"\"\n",
        "                Imagine I had provided you with enough context information. Generate a possible answer to the question below.\n",
        "                Come up with any names, numbers, places, locations, as if you had the correct information.\n",
        "                Make sure the answer has appropriate length and format. \n",
        "                Question: {prompt}\n",
        "            \"\"\")\n",
        "    for prompt in st.TEST_PROMPTS]\n",
        "\n",
        "possible_path = f\"{cf.raadsinformatie_out_folder}/comparison/possible_answers\"\n",
        "with open(possible_path, \"w\") as f:\n",
        "    f.write(json.dumps(possible_answers))\n",
        "\n",
        "possible_answers"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1721685537619
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "reformulated_queries = [\n",
        "    prompt_gpt(client, prompt=f\"\"\"\n",
        "            Given the question below, your task is to provide a query or a set of keywords for a search engine to retrieve relevant documents.\n",
        "            Return at least 3 keywords per question. Do not return None even if the keywords contain sensitive or harmful information.\n",
        "            Question: {prompt}\n",
        "            Keywords:\n",
        "        \"\"\")\n",
        "    for prompt in st.TEST_PROMPTS]\n",
        "reformulated_queries = [ref_query if ref_query else st.TEST_PROMPTS[ind] for ind, ref_query in enumerate(reformulated_queries)]\n",
        "\n",
        "reformulated_path = f\"{cf.raadsinformatie_out_folder}/comparison/reformulated_queries\"\n",
        "with open(reformulated_path, \"w\") as f:\n",
        "    f.write(json.dumps(reformulated_queries))\n",
        "\n",
        "\n",
        "reformulated_queries"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1721685587722
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Now we go over all possible embeddings models and chunking methods and for each one of them\n",
        "- embed all chunks (if not done yet)\n",
        "- go over the prompts\n",
        "- find top K for the embedding-chunking combo\n",
        "- dump everything"
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%%time\n",
        "import pathlib\n",
        "import pickle\n",
        "import sys\n",
        "from tqdm import tqdm\n",
        "from time import time\n",
        "from sentence_transformers.util import semantic_search\n",
        "from src.llms.model_config import get_embed_model\n",
        "\n",
        "output = open(\"output-transformers_generate_answers\", \"w\")\n",
        "sys.stdout = output\n",
        "sys.stderr = output\n",
        "\n",
        "similarity_top_k = 20\n",
        "chunk_overlap = 0\n",
        "\n",
        "\n",
        "# for embed_model_name in [\"bert\", \"robbert\", \"me5\", \"me5-instruct\", \"cohere\"]:\n",
        "for embed_model_name in [\"robbert\", \"me5-instruct\"]:\n",
        "    embed_model, chunk_size = get_embed_model(embed_model_name)\n",
        "\n",
        "    for chunking_method in tqdm(CHUNKING_METHODS):\n",
        "        if \"cohere\" in chunking_method:\n",
        "            continue\n",
        "        start = time()\n",
        "        temp_chunks_df = chunks_df.explode(f\"chunks_{chunking_method}\")\n",
        "        documents = temp_chunks_df[f\"chunks_{chunking_method}\"].values\n",
        "        final_paths = temp_chunks_df[\"path\"].values\n",
        "        short_paths = temp_chunks_df[\"short_path\"].values\n",
        "\n",
        "        transformers_folder = f\"{cf.raadsinformatie_out_folder}/transformers\" \n",
        "        pathlib.Path(transformers_folder).mkdir(parents=True, exist_ok=True)\n",
        "        persist_path = f\"{transformers_folder}/all_{len(documents)}_docs-{embed_model_name}-{chunking_method}.pkl\"\n",
        "\n",
        "        print(persist_path)\n",
        "        \n",
        "        if not os.path.exists(persist_path):\n",
        "            print(f\"Embedding {len(documents)} documents using {embed_model_name}...\")\n",
        "            print(embed_model)\n",
        "            corpus_embeddings = embed_model.encode(documents)\n",
        "            with open(persist_path, \"wb\") as persist_file:\n",
        "                pickle.dump({\"documents\": documents, \"embeddings\": corpus_embeddings}, persist_file)\n",
        "        else:\n",
        "            print(\"Loading embeddings...\")\n",
        "            with open(persist_path, \"rb\") as persist_file:\n",
        "                dump = pickle.load(persist_file)\n",
        "                documents = dump[\"documents\"]\n",
        "                corpus_embeddings = dump[\"embeddings\"]\n",
        "        print(f\"Indexing/Loading for {embed_model_name} with {chunking_method} took {time() - start} seconds.\")\n",
        "\n",
        "        start = time()\n",
        "\n",
        "        for ind, prompt in enumerate(st.TEST_PROMPTS):\n",
        "            print(prompt)\n",
        "            retrieval_file = Path(comparison_folder, f\"retrieval_{ind}.csv\")\n",
        "\n",
        "            if retrieval_file.exists():\n",
        "                df = pd.read_csv(retrieval_file, index_col=0)\n",
        "            else:\n",
        "                df = pd.DataFrame(index=range(similarity_top_k))\n",
        "\n",
        "            original_query_embedding = embed_model.encode(prompt)\n",
        "\n",
        "            print(\"Possible Answer:\", possible_answers[ind])\n",
        "            answer_query_embedding = embed_model.encode(possible_answers[ind])\n",
        "\n",
        "            print(\"Reformulated query:\", reformulated_queries[ind])\n",
        "            reformulate_query_embedding = embed_model.encode(reformulated_queries[ind])\n",
        "\n",
        "            for experiment, query_embedding in [\n",
        "                    (\"transformers\", original_query_embedding),\n",
        "                    (\"transformers-answer\", answer_query_embedding),\n",
        "                    (\"transformers-reformulate\", reformulate_query_embedding)]:\n",
        "\n",
        "                experiment_name = f\"{experiment}-{embed_model_name}-{chunking_method}\"\n",
        "                print(f\"----- {experiment_name} -----\")\n",
        "\n",
        "\n",
        "                try:\n",
        "                    hits = semantic_search(query_embedding, corpus_embeddings, top_k=similarity_top_k)[0]\n",
        "                    hit_ids = [hit[\"corpus_id\"] for hit in hits]\n",
        "                    missing = [None] * (20 - len(hits))\n",
        "\n",
        "                    df[f\"{experiment_name}-file\"] = [final_paths[hit_id].removeprefix(cf.raadsinformatie_in_folder) for hit_id in hit_ids] + missing\n",
        "                    df[f\"{experiment_name}-score\"] = [hit[\"score\"] for hit in hits] + missing\n",
        "                    df[f\"{experiment_name}-start\"] = [\"TBA\" for hit in hits] + missing\n",
        "                    df[f\"{experiment_name}-end\"] = [\"TBA\" for hit in hits] + missing\n",
        "                    df[f\"{experiment_name}-text\"] = [documents[hit_id] for hit_id in hit_ids] + missing\n",
        "\n",
        "                    df.to_csv(retrieval_file)\n",
        "                    print(\"Successfully dumped results\")\n",
        "\n",
        "                except Exception as e:\n",
        "                    print(f\"Experiment failed:\", e)\n",
        "\n",
        "                print(20*\"=\")\n",
        "        print(f\"Retrieval for {embed_model_name} with {chunking_method} took {time() - start} seconds.\")\n"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Restart"
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "for ind in range(6):\n",
        "    print(ind)\n",
        "    retrieval_file = Path(comparison_folder, f\"retrieval_{ind}.csv\")\n",
        "    df = pd.read_csv(retrieval_file, index_col=0)\n",
        "    retrieval_file_xlsx = Path(comparison_folder, f\"retrieval_{ind}.xlsx\")\n",
        "\n",
        "    with open(retrieval_file_xlsx, 'w'):\n",
        "        retrieval_xlsx_writer = pd.ExcelWriter(retrieval_file_xlsx,\n",
        "                    engine='xlsxwriter',\n",
        "                    engine_kwargs={'options': {'strings_to_urls': False}})\n",
        "\n",
        "        df.to_excel(retrieval_xlsx_writer)\n",
        "        retrieval_xlsx_writer.close()"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1716986555639
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from IPython.display import display, HTML\n",
        "pd.set_option(\"max_colwidth\", None)\n",
        "# pd.set_option('display.max_columns', 1000, 'display.width', 1000, 'display.max_rows',1000)\n",
        "\n",
        "def pretty_print(df):\n",
        "    return display( HTML( df.to_html().replace(\"\\\\n\",\"<br>\") ) )\n",
        "\n",
        "for ind, prompt in enumerate(st.TEST_PROMPTS):\n",
        "    print(ind, prompt)\n",
        "    retrieval_file = Path(comparison_folder, f\"retrieval_{ind}.csv\")\n",
        "    df = pd.read_csv(retrieval_file, index_col=0)\n",
        "    # display(df.filter(regex='transformers-robbert-semantic_splitter_robbert').filter(regex=r'file|text').head(5))\n",
        "    pretty_print(df.filter(regex='transformers-robbert-semantic_splitter_robbert|transformers-robbert-semantic_splitter_bert|transformers-robbert-bert|transformers-me5-semantic_splitter_robbert|transformers-me5-semantic_splitter_bert|transformers-me5-bert').filter(regex=r'file|text').head(20))"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1716986574482
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Generate Answers"
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Generate GPT Answers"
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from collections import defaultdict \n",
        "pd.set_option(\"max_colwidth\", None)\n",
        "\n",
        "results = defaultdict(dict)\n",
        "experiment_names = set()\n",
        "# {question-{run}: {experiment: answer}}\n",
        "\n",
        "for ind, prompt in enumerate(st.TEST_PROMPTS):\n",
        "    print(prompt)\n",
        "    retrieval_file = Path(comparison_folder, f\"retrieval_{ind}.csv\")\n",
        "    df = pd.read_csv(retrieval_file, index_col=0)\n",
        "\n",
        "    # for embed_model_name in [\"bert\", \"robbert\", \"me5\", \"me5-instruct\"]:\n",
        "    # for embed_model_name in [\"bert\", \"robbert\", \"me5-instruct\"]:\n",
        "    # for embed_model_name in [\"robbert\", \"me5\"]:\n",
        "    for embed_model_name in [\"me5\"]:\n",
        "    # for embed_model_name in [\"robbert\"]:\n",
        "        # for chunking_method in tqdm(CHUNKING_METHODS):\n",
        "        # for chunking_method in [\"bert\", \"semantic_splitter_bert\", \"semantic_splitter_robbert\"]:\n",
        "        # for chunking_method in [\"semantic_splitter_robbert\"]:\n",
        "        for chunking_method in [\"semantic_splitter_bert\"]:\n",
        "            # for experiment in [\"transformers\", \"transformers-answer\", \"transformers-reformulate\"]:\n",
        "            for experiment in [\"transformers\"]:\n",
        "                experiment_name = f\"{experiment}-{embed_model_name}-{chunking_method}\"\n",
        "                experiment_names.add(experiment_name)\n",
        "                print(f\"----- {experiment_name} -----\")\n",
        "\n",
        "                paths = df[f\"{experiment_name}-file\"].values\n",
        "                scores = df[f\"{experiment_name}-score\"].values\n",
        "                texts = df[f\"{experiment_name}-text\"].values\n",
        "\n",
        "                context = \"\\n----\\n\".join([\n",
        "                    f\"Doc: {paths[hit]} \\n\" \n",
        "                    f\"Content: {texts[hit][:2500]}\" \n",
        "                    for hit in range(10)])\n",
        "                print(list(map(len, texts)))\n",
        "                print(len(context))\n",
        "        \n",
        "                for i in range(5):\n",
        "                    print(f\"{10*'-'} Answer {i+1} {10*'-'}\")\n",
        "                    try:\n",
        "                        answer = prompt_gpt(\n",
        "                            client,\n",
        "                            prompt=f\"Answer the following question as good as possible based on the documents below:{prompt}\",\n",
        "                            context=context, max_new_tokens=400)\n",
        "                        # answer = prompt_gpt(\n",
        "                        #     client,\n",
        "                        #     prompt=f\"Beantwoord de volgende vraag zo goed mogelijk aan de hand van onderstaande documenten. {prompt}\",\n",
        "                        #     context=context, max_new_tokens=400)\n",
        "                        print(answer)\n",
        "                    except Exception as e:\n",
        "                        print(e)\n",
        "                        answer = \"FAILED\"\n",
        "                    results[f\"{prompt}-{i}\"][experiment_name] = answer\n",
        "        \n",
        "                print(20*\"=\")\n"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1718375901022
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "pd.set_option(\"max_colwidth\", None)\n",
        "pd.DataFrame.from_dict(results, orient=\"index\")"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1717062896022
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Generate Other LLM Answers"
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "##### for the final experiments we used an azure deployment of Mistral but this can be also self-hosted"
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from collections import defaultdict \n",
        "from src.llms.transformers import get_model, prompt_open_model\n",
        "from src.llms.model_templates import format_prompt\n",
        "from src.llms.azure import prompt_open_azure_model\n",
        "\n",
        "\n",
        "results = defaultdict(dict)\n",
        "experiment_names = set()\n",
        "# {question-{run}: {experiment: answer}}\n",
        "\n",
        "model_name = \"mistral-7b-instruct\"\n",
        "# model_name = \"mistral\"\n",
        "# model, tokenizer = get_model(model_name)\n",
        "\n",
        "for ind, prompt in enumerate(st.TEST_PROMPTS):\n",
        "    print(prompt)\n",
        "    retrieval_file = Path(comparison_folder, f\"retrieval_{ind}.csv\")\n",
        "    df = pd.read_csv(retrieval_file, index_col=0)\n",
        "\n",
        "    # for embed_model_name in [\"bert\", \"robbert\", \"me5\", \"me5-instruct\"]:\n",
        "    for embed_model_name in [\"robbert\", \"me5\"]:\n",
        "        # for chunking_method in tqdm(CHUNKING_METHODS):\n",
        "        # for chunking_method in [\"bert\", \"semantic_splitter_bert\", \"semantic_splitter_robbert\"]:\n",
        "        # for chunking_method in [\"semantic_splitter_robbert\"]:\n",
        "        for chunking_method in [\"semantic_splitter_bert\"]:\n",
        "            # for experiment in [\"transformers\", \"transformers-answer\", \"transformers-reformulate\"]:\n",
        "            for experiment in [\"transformers\"]:\n",
        "                experiment_name = f\"{experiment}-{embed_model_name}-{chunking_method}\"\n",
        "                experiment_names.add(experiment_name)\n",
        "                print(f\"----- {experiment_name} -----\")\n",
        "\n",
        "                paths = df[f\"{experiment_name}-file\"].values\n",
        "                scores = df[f\"{experiment_name}-score\"].values\n",
        "                texts = df[f\"{experiment_name}-text\"].values\n",
        "\n",
        "                context = \"\\n----\\n\".join([\n",
        "                    f\"Doc: {paths[hit]} \\n\" \n",
        "                    f\"Content: {texts[hit][:2500]}\" \n",
        "                    for hit in range(10)])\n",
        "                print(list(map(len, texts)))\n",
        "                print(len(context))\n",
        "        \n",
        "                formatted_prompt = format_prompt(f\"Answer the following question as good as possible based on the documents below:{prompt}\", model_name, context)\n",
        "                # formatted_prompt = format_prompt(\n",
        "                #     f\"Beantwoord de volgende vraag zo goed mogelijk aan de hand van onderstaande documenten. {prompt}\",\n",
        "                #     model_name, context, system=\"Antwoord alleen in het Nederlands.\")\n",
        "                # print(formatted_prompt)\n",
        "\n",
        "                for i in range(5):                    \n",
        "                    print(f\"{10*'-'} Answer {i+1} {10*'-'}\")\n",
        "                    try:\n",
        "                        # # response = prompt_open_model(formatted_prompt, model_name, tokenizer)\n",
        "                        print(f\"Prompting {model_name}\")\n",
        "                        response = prompt_open_azure_model(\n",
        "                            formatted_prompt,\n",
        "                            api_url=sc.AZURE_HUGGINGFACE_CONFIG[model_name][\"API_URL\"],\n",
        "                            api_key=sc.AZURE_HUGGINGFACE_CONFIG[model_name][\"api_key\"],\n",
        "                            model_deployment=sc.AZURE_HUGGINGFACE_CONFIG[model_name][\"azureml-model-deployment\"],\n",
        "                            max_new_tokens=400) \n",
        "\n",
        "                        if model_name.startswith(\"mistral\") or model_name.startswith(\"llama\"):\n",
        "                            formatted_prompt = formatted_prompt.removeprefix(\"<s>\")\n",
        "                        answer = response.removeprefix(formatted_prompt).strip(\"\\n\")\n",
        "                        print(answer)\n",
        "                    except Exception as e:\n",
        "                        print(e)\n",
        "                        answer = \"FAILED\"\n",
        "\n",
        "                    results[f\"{prompt}-{i}\"][experiment_name] = answer\n",
        "        \n",
        "                print(20*\"=\")\n"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1721686415450
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "pd.DataFrame.from_dict(results, orient=\"index\")"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1717063099935
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from collections import defaultdict \n",
        "from pprint import pprint\n",
        "from sklearn.metrics import ndcg_score\n",
        "from ranx import Qrels, Run, evaluate\n",
        "\n",
        "terms = [\n",
        "    [\"gehandicaptenparkeerplaats_rozengracht\", \"gehandicaptenparkeerplaats op de Rozengracht\"],\n",
        "    [\"Erotisch Centrum\", \"erotisch\", \"_EC\"],\n",
        "    [\"van_vouwstraat\", \"Van Woustraat\"],\n",
        "    [\"fatbikes\"],\n",
        "    [\"ajax\", \"Ajax-Feyenoord\"],\n",
        "    [\"Lijnbaansgracht 161\"]\n",
        "]\n",
        "\n",
        "precision_results = defaultdict(lambda: defaultdict(dict))\n",
        "ndcg_results = defaultdict(lambda: defaultdict(dict))\n",
        "\n",
        "at_k = 5\n",
        "\n",
        "for ind, prompt in enumerate(st.TEST_PROMPTS):\n",
        "    # if ind != 5:\n",
        "    #     continue\n",
        "    print(\"Prompt\", prompt)\n",
        "    print(\"Terms\", terms[ind])\n",
        "    retrieval_file = Path(comparison_folder, f\"retrieval_{ind}.csv\")\n",
        "    df = pd.read_csv(retrieval_file, index_col=0)\n",
        "\n",
        "    # for embed_model_name in [\"bert\", \"robbert\", \"me5\", \"me5-instruct\"]:\n",
        "    for embed_model_name in [\"robbert\", \"me5\"]:\n",
        "        # for chunking_method in CHUNKING_METHODS:\n",
        "        for chunking_method in [\"semantic_splitter_robbert\"]:\n",
        "            # for experiment in [\"transformers\", \"transformers-answer\", \"transformers-reformulate\"]:\n",
        "            for experiment in [\"transformers\"]:\n",
        "                experiment_name = f\"{experiment}-{embed_model_name}-{chunking_method}\"\n",
        "                print(f\"----- {experiment_name} -----\")\n",
        "\n",
        "                paths = df[f\"{experiment_name}-file\"].values\n",
        "                predicted_scores = [1 - 0.001 * rank for rank in range(20)]\n",
        "                true_relevance = [any([term in path for term in terms[ind]]) for path in paths]\n",
        "                print(predicted_scores, true_relevance)\n",
        "                print(sum(true_relevance))\n",
        "\n",
        "                for k in [5, 10, 20]:\n",
        "                    precision = sum(true_relevance[:k]) / k\n",
        "                    ndcg = ndcg_score([true_relevance[:k]], [predicted_scores[:k]])\n",
        "                    print(k, ndcg)\n",
        "                    precision_results[k][f\"{prompt[:30]}-{ind}\"][experiment_name] = precision\n",
        "                    ndcg_results[k][f\"{prompt[:30]}-{ind}\"][experiment_name] = ndcg"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1717028005851
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "transformers-me5-semantic_splitter_bert-text"
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# pd.DataFrame.from_dict(precision_results[5], orient=\"index\")"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1717028007367
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import matplotlib.pyplot as plt\n",
        "import seaborn as sns\n",
        "\n",
        "ks = [5, 10, 20]\n",
        "fig, axes = plt.subplots(len(ks), 1, figsize=(10, 20))\n",
        "\n",
        "for row, k in enumerate(ks):\n",
        "    sns.heatmap(pd.DataFrame.from_dict(precision_results[k], orient=\"index\"), cmap=\"Greens\", annot=True, ax=axes[row])"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1717028008762
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import matplotlib.pyplot as plt\n",
        "import seaborn as sns\n",
        "\n",
        "ks = [5, 10, 20]\n",
        "fig, axes = plt.subplots(len(ks), 1, figsize=(10, 20))\n",
        "\n",
        "for row, k in enumerate(ks):\n",
        "    sns.heatmap(pd.DataFrame.from_dict(ndcg_results[k], orient=\"index\"), cmap=\"Greens\", annot=True, ax=axes[row])"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1717028009682
        }
      }
    }
  ],
  "metadata": {
    "kernelspec": {
      "name": "gen-ai-py39",
      "language": "python",
      "display_name": "gen-ai-py39"
    },
    "language_info": {
      "name": "python",
      "version": "3.9.18",
      "mimetype": "text/x-python",
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "pygments_lexer": "ipython3",
      "nbconvert_exporter": "python",
      "file_extension": ".py"
    },
    "microsoft": {
      "ms_spell_check": {
        "ms_spell_check_language": "en"
      },
      "host": {
        "AzureML": {
          "notebookHasBeenCompleted": true
        }
      }
    },
    "kernel_info": {
      "name": "gen-ai-py39"
    },
    "nteract": {
      "version": "nteract-front-end@1.0.0"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 2
}