context classification.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8d93fdc2-4be1-472e-a5c4-43303c2e5ace",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting sentence-transformers\n",
      "  Downloading sentence-transformers-2.2.2.tar.gz (85 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m86.0/86.0 kB\u001b[0m \u001b[31m5.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Preparing metadata (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25hCollecting transformers<5.0.0,>=4.6.0 (from sentence-transformers)\n",
      "  Obtaining dependency information for transformers<5.0.0,>=4.6.0 from https://files.pythonhosted.org/packages/9a/06/e4ec2a321e57c03b7e9345d709d554a52c33760e5015fdff0919d9459af0/transformers-4.35.0-py3-none-any.whl.metadata\n",
      "  Downloading transformers-4.35.0-py3-none-any.whl.metadata (123 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m123.1/123.1 kB\u001b[0m \u001b[31m7.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: tqdm in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from sentence-transformers) (4.65.0)\n",
      "Requirement already satisfied: torch>=1.6.0 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from sentence-transformers) (2.0.1)\n",
      "Requirement already satisfied: torchvision in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from sentence-transformers) (0.15.2)\n",
      "Requirement already satisfied: numpy in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from sentence-transformers) (1.24.4)\n",
      "Requirement already satisfied: scikit-learn in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from sentence-transformers) (1.3.0)\n",
      "Requirement already satisfied: scipy in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from sentence-transformers) (1.11.1)\n",
      "Collecting nltk (from sentence-transformers)\n",
      "  Downloading nltk-3.8.1-py3-none-any.whl (1.5 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.5/1.5 MB\u001b[0m \u001b[31m41.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hCollecting sentencepiece (from sentence-transformers)\n",
      "  Downloading sentencepiece-0.1.99-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.3 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.3/1.3 MB\u001b[0m \u001b[31m19.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m:00:01\u001b[0m\n",
      "\u001b[?25hCollecting huggingface-hub>=0.4.0 (from sentence-transformers)\n",
      "  Obtaining dependency information for huggingface-hub>=0.4.0 from https://files.pythonhosted.org/packages/ef/b5/b6107bd65fa4c96fdf00e4733e2fe5729bb9e5e09997f63074bb43d3ab28/huggingface_hub-0.18.0-py3-none-any.whl.metadata\n",
      "  Downloading huggingface_hub-0.18.0-py3-none-any.whl.metadata (13 kB)\n",
      "Requirement already satisfied: filelock in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from huggingface-hub>=0.4.0->sentence-transformers) (3.12.2)\n",
      "Requirement already satisfied: fsspec>=2023.5.0 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from huggingface-hub>=0.4.0->sentence-transformers) (2023.6.0)\n",
      "Requirement already satisfied: requests in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from huggingface-hub>=0.4.0->sentence-transformers) (2.31.0)\n",
      "Requirement already satisfied: pyyaml>=5.1 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from huggingface-hub>=0.4.0->sentence-transformers) (6.0)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4.3 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from huggingface-hub>=0.4.0->sentence-transformers) (4.7.1)\n",
      "Requirement already satisfied: packaging>=20.9 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from huggingface-hub>=0.4.0->sentence-transformers) (21.3)\n",
      "Requirement already satisfied: sympy in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from torch>=1.6.0->sentence-transformers) (1.12)\n",
      "Requirement already satisfied: networkx in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from torch>=1.6.0->sentence-transformers) (3.1)\n",
      "Requirement already satisfied: jinja2 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from torch>=1.6.0->sentence-transformers) (3.1.2)\n",
      "Collecting regex!=2019.12.17 (from transformers<5.0.0,>=4.6.0->sentence-transformers)\n",
      "  Obtaining dependency information for regex!=2019.12.17 from https://files.pythonhosted.org/packages/8f/3e/4b8b40eb3c80aeaf360f0361d956d129bb3d23b2a3ecbe3a04a8f3bdd6d3/regex-2023.10.3-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata\n",
      "  Downloading regex-2023.10.3-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (40 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m40.9/40.9 kB\u001b[0m \u001b[31m12.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hCollecting tokenizers<0.15,>=0.14 (from transformers<5.0.0,>=4.6.0->sentence-transformers)\n",
      "  Obtaining dependency information for tokenizers<0.15,>=0.14 from https://files.pythonhosted.org/packages/a7/7b/c1f643eb086b6c5c33eef0c3752e37624bd23e4cbc9f1332748f1c6252d1/tokenizers-0.14.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata\n",
      "  Downloading tokenizers-0.14.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (6.7 kB)\n",
      "Collecting safetensors>=0.3.1 (from transformers<5.0.0,>=4.6.0->sentence-transformers)\n",
      "  Obtaining dependency information for safetensors>=0.3.1 from https://files.pythonhosted.org/packages/20/4e/878b080dbda92666233ec6f316a53969edcb58eab1aa399a64d0521cf953/safetensors-0.4.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata\n",
      "  Downloading safetensors-0.4.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (3.8 kB)\n",
      "Requirement already satisfied: click in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from nltk->sentence-transformers) (8.1.6)\n",
      "Requirement already satisfied: joblib in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from nltk->sentence-transformers) (1.3.1)\n",
      "Requirement already satisfied: threadpoolctl>=2.0.0 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from scikit-learn->sentence-transformers) (3.2.0)\n",
      "Requirement already satisfied: pillow!=8.3.*,>=5.3.0 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from torchvision->sentence-transformers) (10.0.0)\n",
      "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from packaging>=20.9->huggingface-hub>=0.4.0->sentence-transformers) (3.0.9)\n",
      "Collecting huggingface-hub>=0.4.0 (from sentence-transformers)\n",
      "  Obtaining dependency information for huggingface-hub>=0.4.0 from https://files.pythonhosted.org/packages/aa/f3/3fc97336a0e90516901befd4f500f08d691034d387406fdbde85bea827cc/huggingface_hub-0.17.3-py3-none-any.whl.metadata\n",
      "  Downloading huggingface_hub-0.17.3-py3-none-any.whl.metadata (13 kB)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from jinja2->torch>=1.6.0->sentence-transformers) (2.1.3)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from requests->huggingface-hub>=0.4.0->sentence-transformers) (3.2.0)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from requests->huggingface-hub>=0.4.0->sentence-transformers) (3.4)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from requests->huggingface-hub>=0.4.0->sentence-transformers) (1.26.14)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from requests->huggingface-hub>=0.4.0->sentence-transformers) (2023.5.7)\n",
      "Requirement already satisfied: mpmath>=0.19 in /home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages (from sympy->torch>=1.6.0->sentence-transformers) (1.3.0)\n",
      "Downloading transformers-4.35.0-py3-none-any.whl (7.9 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m7.9/7.9 MB\u001b[0m \u001b[31m74.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m:00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hDownloading regex-2023.10.3-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (773 kB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m773.9/773.9 kB\u001b[0m \u001b[31m86.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading safetensors-0.4.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.3 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.3/1.3 MB\u001b[0m \u001b[31m95.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading tokenizers-0.14.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (3.8 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.8/3.8 MB\u001b[0m \u001b[31m43.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hDownloading huggingface_hub-0.17.3-py3-none-any.whl (295 kB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m295.0/295.0 kB\u001b[0m \u001b[31m49.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hBuilding wheels for collected packages: sentence-transformers\n",
      "  Building wheel for sentence-transformers (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for sentence-transformers: filename=sentence_transformers-2.2.2-py3-none-any.whl size=125926 sha256=d99b6351babb820155123d23625a614909e91843d9298040656640470530dcc8\n",
      "  Stored in directory: /home/ec2-user/.cache/pip/wheels/62/f2/10/1e606fd5f02395388f74e7462910fe851042f97238cbbd902f\n",
      "Successfully built sentence-transformers\n",
      "Installing collected packages: sentencepiece, safetensors, regex, nltk, huggingface-hub, tokenizers, transformers, sentence-transformers\n",
      "Successfully installed huggingface-hub-0.17.3 nltk-3.8.1 regex-2023.10.3 safetensors-0.4.0 sentence-transformers-2.2.2 sentencepiece-0.1.99 tokenizers-0.14.1 transformers-4.35.0\n",
      "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
      "triton 2.0.0 requires cmake, which is not installed.\n",
      "triton 2.0.0 requires lit, which is not installed.\u001b[0m\u001b[31m\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "!pip install -U sentence-transformers\n",
    "!pip install -q chromadb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c7df27a0-760c-4e4e-96b0-f694e0d5512e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import chromadb\n",
    "from chromadb.api.types import Documents, Embeddings\n",
    "from chromadb.utils import embedding_functions\n",
    "from sentence_transformers import SentenceTransformer, util"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "9187364d-fb9d-4a2c-99ca-00774d99eef8",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6c270652dbb0453499d0725d855c2df3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)f1ccc/.gitattributes:   0%|          | 0.00/737 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "51c1fafe44ca492081bbe54e16946a69",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)_Pooling/config.json:   0%|          | 0.00/190 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8ba94297155f41509aa486ab5e8aff4f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)891b0f1ccc/README.md:   0%|          | 0.00/11.5k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "17bf533c0ec4443da60906228b8276be",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)1b0f1ccc/config.json:   0%|          | 0.00/612 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "17a3035148e84fc2b41fbd59838f91d6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)ce_transformers.json:   0%|          | 0.00/116 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fca70437c60d4aa7ab166b26439d410a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)ccc/data_config.json:   0%|          | 0.00/25.5k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0d41e82560d24a62b7da26fb60179483",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading pytorch_model.bin:   0%|          | 0.00/90.9M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "851f9e1ff6384eb984935cfa9dd92cda",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)nce_bert_config.json:   0%|          | 0.00/53.0 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "45ce3cf6f2664ccd89aa2d634eb26090",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)cial_tokens_map.json:   0%|          | 0.00/112 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "94636554b91945d1a12930a4d1a08241",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)f1ccc/tokenizer.json:   0%|          | 0.00/466k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9afe210bac8643d2872cbde8a8e4ac85",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)okenizer_config.json:   0%|          | 0.00/383 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "808be0a51e7145bf8be1682ecca7f282",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)1ccc/train_script.py:   0%|          | 0.00/13.8k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4946d086936b49a5b3b3004ebac774f8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)891b0f1ccc/vocab.txt:   0%|          | 0.00/232k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3550df82ab5e4f208dbf4d449c288d4c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)b0f1ccc/modules.json:   0%|          | 0.00/349 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=\"all-MiniLM-L6-v2\")\n",
    "sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=\"multi-qa-MiniLM-L6-cos-v1\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5fa8609b-a97d-4f4e-9caf-1b15b8c59779",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df_dataset = pd.read_csv('data.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "47eccd83-0281-473f-84c8-6fd40ae0af76",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df_dataset['input'] = df_dataset['leader'] + ' ' + df_dataset['user']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "7d38c46a-e09a-4390-9134-467ace3d275c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df = df_dataset[df_dataset['label'] == 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "d1d60863-2627-41ea-ae63-26f5ce433af6",
   "metadata": {
    "tags": []
   },
   "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>leader</th>\n",
       "      <th>user</th>\n",
       "      <th>english_translation</th>\n",
       "      <th>assistant</th>\n",
       "      <th>english_translation.1</th>\n",
       "      <th>label</th>\n",
       "      <th>with_error</th>\n",
       "      <th>Note</th>\n",
       "      <th>input</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>¡Hola! Bienvenido a la cafetería Brew Haven. ¿...</td>\n",
       "      <td>Hola. Un café, por favor.</td>\n",
       "      <td>NaN</td>\n",
       "      <td>Entiendo, un café.</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1.0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>default script pairs</td>\n",
       "      <td>¡Hola! Bienvenido a la cafetería Brew Haven. ¿...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>¿Quieres leche y azúcar?</td>\n",
       "      <td>Sí, leche y azúcar.</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1.0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>default script pairs</td>\n",
       "      <td>¿Quieres leche y azúcar? Sí, leche y azúcar.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>¿Cuántos azúcares?</td>\n",
       "      <td>Tres azúcares.</td>\n",
       "      <td>NaN</td>\n",
       "      <td>Leche y tres azúcares.</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1.0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>default script pairs</td>\n",
       "      <td>¿Cuántos azúcares? Tres azúcares.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>¡Hola! Bienvenido a la cafetería Brew Haven. ¿...</td>\n",
       "      <td>Sí, quiero un café, por favor.</td>\n",
       "      <td>NaN</td>\n",
       "      <td>Un café.</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1.0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>¡Hola! Bienvenido a la cafetería Brew Haven. ¿...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>¡Hola! Bienvenido a la cafetería Brew Haven. ¿...</td>\n",
       "      <td>un café</td>\n",
       "      <td>one cafe.</td>\n",
       "      <td>Un café.</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1.0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>¡Hola! Bienvenido a la cafetería Brew Haven. ¿...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>¡Hola! Bienvenido a la cafetería Brew Haven. ¿...</td>\n",
       "      <td>Sí, quiero un cafés, por favor.</td>\n",
       "      <td>NaN</td>\n",
       "      <td>quieres un café.</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1.0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>¡Hola! Bienvenido a la cafetería Brew Haven. ¿...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>¡Hola! Bienvenido a la cafetería Brew Haven. ¿...</td>\n",
       "      <td>una agua, favor.</td>\n",
       "      <td>NaN</td>\n",
       "      <td>quieres un agua.</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1.0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>¡Hola! Bienvenido a la cafetería Brew Haven. ¿...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>¡Hola! Bienvenido a la cafetería Brew Haven. ¿...</td>\n",
       "      <td>dos café, por favor.</td>\n",
       "      <td>NaN</td>\n",
       "      <td>quieres dos cafés.</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1.0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>¡Hola! Bienvenido a la cafetería Brew Haven. ¿...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>¡Hola! Bienvenido a la cafetería Brew Haven. ¿...</td>\n",
       "      <td>Sí, ¿puedo tomar un café con leche?</td>\n",
       "      <td>Yes, can I have a latte?</td>\n",
       "      <td>Seguro. un café con leche</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>¡Hola! Bienvenido a la cafetería Brew Haven. ¿...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>18</th>\n",
       "      <td>¡Hola! Bienvenido a la cafetería Brew Haven. ¿...</td>\n",
       "      <td>Sí, puedo tomo un café con leche.</td>\n",
       "      <td>Yes, can I have a latte?</td>\n",
       "      <td>Seguro. un café con leche</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>¡Hola! Bienvenido a la cafetería Brew Haven. ¿...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19</th>\n",
       "      <td>¡Hola! Bienvenido a la cafetería Brew Haven. ¿...</td>\n",
       "      <td>Sí, puedo tomar café con leche.</td>\n",
       "      <td>Yes, can I have a latte?</td>\n",
       "      <td>Seguro. un café con leche</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>¡Hola! Bienvenido a la cafetería Brew Haven. ¿...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>20</th>\n",
       "      <td>¡Hola! Bienvenido a la cafetería Brew Haven. ¿...</td>\n",
       "      <td>quiero un cafe pequeño.</td>\n",
       "      <td>I want one small coffee.</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>¡Hola! Bienvenido a la cafetería Brew Haven. ¿...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>21</th>\n",
       "      <td>¡Hola! Bienvenido a la cafetería Brew Haven. ¿...</td>\n",
       "      <td>Sí, puedo tomar un café de leche.</td>\n",
       "      <td>Yes, can I have a latte?</td>\n",
       "      <td>Seguro. un café con leche</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>¡Hola! Bienvenido a la cafetería Brew Haven. ¿...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                               leader  \\\n",
       "0   ¡Hola! Bienvenido a la cafetería Brew Haven. ¿...   \n",
       "1                            ¿Quieres leche y azúcar?   \n",
       "2                                  ¿Cuántos azúcares?   \n",
       "5   ¡Hola! Bienvenido a la cafetería Brew Haven. ¿...   \n",
       "6   ¡Hola! Bienvenido a la cafetería Brew Haven. ¿...   \n",
       "7   ¡Hola! Bienvenido a la cafetería Brew Haven. ¿...   \n",
       "8   ¡Hola! Bienvenido a la cafetería Brew Haven. ¿...   \n",
       "9   ¡Hola! Bienvenido a la cafetería Brew Haven. ¿...   \n",
       "15  ¡Hola! Bienvenido a la cafetería Brew Haven. ¿...   \n",
       "18  ¡Hola! Bienvenido a la cafetería Brew Haven. ¿...   \n",
       "19  ¡Hola! Bienvenido a la cafetería Brew Haven. ¿...   \n",
       "20  ¡Hola! Bienvenido a la cafetería Brew Haven. ¿...   \n",
       "21  ¡Hola! Bienvenido a la cafetería Brew Haven. ¿...   \n",
       "\n",
       "                                   user        english_translation  \\\n",
       "0             Hola. Un café, por favor.                        NaN   \n",
       "1                   Sí, leche y azúcar.                        NaN   \n",
       "2                        Tres azúcares.                        NaN   \n",
       "5        Sí, quiero un café, por favor.                        NaN   \n",
       "6                               un café                  one cafe.   \n",
       "7       Sí, quiero un cafés, por favor.                        NaN   \n",
       "8                      una agua, favor.                        NaN   \n",
       "9                  dos café, por favor.                        NaN   \n",
       "15  Sí, ¿puedo tomar un café con leche?  Yes, can I have a latte?    \n",
       "18    Sí, puedo tomo un café con leche.  Yes, can I have a latte?    \n",
       "19      Sí, puedo tomar café con leche.  Yes, can I have a latte?    \n",
       "20              quiero un cafe pequeño.   I want one small coffee.   \n",
       "21    Sí, puedo tomar un café de leche.  Yes, can I have a latte?    \n",
       "\n",
       "                    assistant  english_translation.1  label  with_error  \\\n",
       "0          Entiendo, un café.                    NaN    1.0         NaN   \n",
       "1                         NaN                    NaN    1.0         NaN   \n",
       "2      Leche y tres azúcares.                    NaN    1.0         NaN   \n",
       "5                    Un café.                    NaN    1.0         NaN   \n",
       "6                    Un café.                    NaN    1.0         NaN   \n",
       "7            quieres un café.                    NaN    1.0         NaN   \n",
       "8            quieres un agua.                    NaN    1.0         NaN   \n",
       "9          quieres dos cafés.                    NaN    1.0         NaN   \n",
       "15  Seguro. un café con leche                    NaN    1.0         0.0   \n",
       "18  Seguro. un café con leche                    NaN    1.0         1.0   \n",
       "19  Seguro. un café con leche                    NaN    1.0         1.0   \n",
       "20                        NaN                    NaN    1.0         0.0   \n",
       "21  Seguro. un café con leche                    NaN    1.0         1.0   \n",
       "\n",
       "                    Note                                              input  \n",
       "0   default script pairs  ¡Hola! Bienvenido a la cafetería Brew Haven. ¿...  \n",
       "1   default script pairs       ¿Quieres leche y azúcar? Sí, leche y azúcar.  \n",
       "2   default script pairs                  ¿Cuántos azúcares? Tres azúcares.  \n",
       "5                    NaN  ¡Hola! Bienvenido a la cafetería Brew Haven. ¿...  \n",
       "6                    NaN  ¡Hola! Bienvenido a la cafetería Brew Haven. ¿...  \n",
       "7                    NaN  ¡Hola! Bienvenido a la cafetería Brew Haven. ¿...  \n",
       "8                    NaN  ¡Hola! Bienvenido a la cafetería Brew Haven. ¿...  \n",
       "9                    NaN  ¡Hola! Bienvenido a la cafetería Brew Haven. ¿...  \n",
       "15                   NaN  ¡Hola! Bienvenido a la cafetería Brew Haven. ¿...  \n",
       "18                   NaN  ¡Hola! Bienvenido a la cafetería Brew Haven. ¿...  \n",
       "19                   NaN  ¡Hola! Bienvenido a la cafetería Brew Haven. ¿...  \n",
       "20                   NaN  ¡Hola! Bienvenido a la cafetería Brew Haven. ¿...  \n",
       "21                   NaN  ¡Hola! Bienvenido a la cafetería Brew Haven. ¿...  "
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "c5efe0e9-fa3e-4296-a7b1-20922c77268a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def create_chroma_db(qa_input, user_input, name):\n",
    "  # get the client\n",
    "  chroma_client = chromadb.Client()\n",
    "  ## In-memory chroma with saving/loading to disk\n",
    "  # chroma_client = chromadb.PersistentClient(path=\"/path/to/data\")\n",
    "  db = chroma_client.create_collection(name=name, embedding_function=sentence_transformer_ef)\n",
    "  for i,d in enumerate(qa_input):\n",
    "    db.add(\n",
    "      documents= user_input[i],\n",
    "      metadatas = {'lesson': d},\n",
    "      ids=str(i)\n",
    "    )\n",
    "  return db"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "5e99c9f3-5577-4101-a75d-21b3bf803435",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "db = create_chroma_db(df['input'].to_list(), df['user'].to_list(), \"langbot\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "68495342-0b25-492c-8f9a-b0cd2c2c5a3b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "results = db.query(\n",
    "    query_texts=[\"Un café, por favor.\"],\n",
    "    n_results=1\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "8abf2092-f80a-4201-95f8-0cb38f4ad5e0",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'ids': [['0']],\n",
       " 'distances': [[0.15042023360729218]],\n",
       " 'metadatas': [[{'input': '¡Hola! Bienvenido a la cafetería Brew Haven. ¿Algo de beber? Hola. Un café, por favor.'}]],\n",
       " 'embeddings': None,\n",
       " 'documents': [['Hola. Un café, por favor.']],\n",
       " 'uris': None,\n",
       " 'data': None}"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "6f0a70dc-3c52-4b95-9904-03f001527edf",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[1.1485252380371094]"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "results['distances'][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "7cf7190e-9da4-4702-ac34-a8cdfabfee1f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def create_chroma_db(documents, user, name):\n",
    "  # get the client\n",
    "  chroma_client = chromadb.Client()\n",
    "  ## In-memory chroma with saving/loading to disk\n",
    "  # chroma_client = chromadb.PersistentClient(path=\"/path/to/data\")\n",
    "  db = chroma_client.create_collection(name=name, embedding_function=sentence_transformer_ef)\n",
    "  for i,d in enumerate(documents):\n",
    "    db.add(\n",
    "      documents= d,\n",
    "      metadatas = {'lesson': d},\n",
    "      ids=str(i)\n",
    "    )\n",
    "  return db"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "3848ee86-96bb-4fe3-b029-36eebf06b7b5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "db1 = create_chroma_db(df['input'].to_list(), df['user'].to_list(), \"lang-bot-test\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "21e7006c-1e31-4ede-b476-e66161dea42d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "results = db1.query(\n",
    "    query_texts=[\"¡Hola! Bienvenido a la cafetería Brew Haven. ¿Algo de beber? Hola.\"],\n",
    "    n_results=1,\n",
    "    where = {'lesson':'1' }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "c564d0d5-c7db-49a4-a1df-0137ce378f0b",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'ids': [[]],\n",
       " 'distances': [[]],\n",
       " 'metadatas': [[]],\n",
       " 'embeddings': None,\n",
       " 'documents': [[]],\n",
       " 'uris': None,\n",
       " 'data': None}"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e78786b-1f26-4189-b330-057519a60449",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_chroma_db(documents, user, name):\n",
    "  # get the client\n",
    "  chroma_client = chromadb.Client()\n",
    "  ## In-memory chroma with saving/loading to disk\n",
    "  # chroma_client = chromadb.PersistentClient(path=\"/path/to/data\")\n",
    "  db = chroma_client.create_collection(name=name, embedding_function=sentence_transformer_ef)\n",
    "  for i,d in enumerate(documents):\n",
    "    db.add(\n",
    "      documents= d,\n",
    "      metadatas = {'input': d},\n",
    "      ids=str(i)\n",
    "    )\n",
    "  return db"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "32853829-dbcf-4b25-9abb-b7f113349166",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def create_chroma_db(qa_input, user_input, name):\n",
    "  # get the client\n",
    "  chroma_client = chromadb.Client()\n",
    "  ## In-memory chroma with saving/loading to disk\n",
    "  # chroma_client = chromadb.PersistentClient(path=\"/path/to/data\")\n",
    "  db = chroma_client.create_collection(name=name, embedding_function=sentence_transformer_ef)\n",
    "  for i,d in enumerate(qa_input):\n",
    "    db.add(\n",
    "      documents= d,\n",
    "      metadatas = {'lesson': d},\n",
    "      ids=str(i)\n",
    "    )\n",
    "  return db"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "id": "337f516f-7439-4490-b174-2dbe5580fd3b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "question = [ '¡Hola! Bienvenido a la cafetería Brew Haven. ¿Algo de beber?',\n",
    "'¿Quieres leche y azúcar?',\n",
    "'¿Cuántos azúcares?',\n",
    "'¿Quieres comida? ¿Sándwich o algo dulce?',\n",
    "'¿Pagas con tarjeta o en efectivo?' ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "8cb438ab-f3d4-42c6-94ee-5d148009cfb8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "db2 = create_chroma_db(df['leader'].to_list(), df['user'].to_list(), \"lang-bot-test\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "id": "1471d1c2-5826-47ea-8070-dac141e2963d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "results = db2.query(\n",
    "    query_texts=[\"Un café, por favor.\"],\n",
    "    n_results=1\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "efd754af-0308-4687-b717-1e6ae3069bb5",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'ids': [['6']],\n",
       " 'distances': [[1.0535516738891602]],\n",
       " 'metadatas': [[{'lesson': '¡Hola! Bienvenido a la cafetería Brew Haven. ¿Algo de beber?'}]],\n",
       " 'embeddings': None,\n",
       " 'documents': [['¡Hola! Bienvenido a la cafetería Brew Haven. ¿Algo de beber?']],\n",
       " 'uris': None,\n",
       " 'data': None}"
      ]
     },
     "execution_count": 69,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "id": "cdfaf11f-f1ac-4d28-b323-770e609d59b8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "results = db2.query(\n",
    "    query_texts=[\"Hola.\"],\n",
    "    n_results=10\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "id": "8217a240-7628-436d-b49c-0107e0c2e7e6",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'ids': [['0', '3', '4', '5', '6', '7', '8', '9', '10', '11']],\n",
       " 'distances': [[1.1834352016448975,\n",
       "   1.1834352016448975,\n",
       "   1.1834352016448975,\n",
       "   1.1834352016448975,\n",
       "   1.1834352016448975,\n",
       "   1.1834352016448975,\n",
       "   1.1834352016448975,\n",
       "   1.1834352016448975,\n",
       "   1.1834352016448975,\n",
       "   1.1834352016448975]],\n",
       " 'metadatas': [[{'lesson': '¡Hola! Bienvenido a la cafetería Brew Haven. ¿Algo de beber?'},\n",
       "   {'lesson': '¡Hola! Bienvenido a la cafetería Brew Haven. ¿Algo de beber?'},\n",
       "   {'lesson': '¡Hola! Bienvenido a la cafetería Brew Haven. ¿Algo de beber?'},\n",
       "   {'lesson': '¡Hola! Bienvenido a la cafetería Brew Haven. ¿Algo de beber?'},\n",
       "   {'lesson': '¡Hola! Bienvenido a la cafetería Brew Haven. ¿Algo de beber?'},\n",
       "   {'lesson': '¡Hola! Bienvenido a la cafetería Brew Haven. ¿Algo de beber?'},\n",
       "   {'lesson': '¡Hola! Bienvenido a la cafetería Brew Haven. ¿Algo de beber?'},\n",
       "   {'lesson': '¡Hola! Bienvenido a la cafetería Brew Haven. ¿Algo de beber?'},\n",
       "   {'lesson': '¡Hola! Bienvenido a la cafetería Brew Haven. ¿Algo de beber?'},\n",
       "   {'lesson': '¡Hola! Bienvenido a la cafetería Brew Haven. ¿Algo de beber?'}]],\n",
       " 'embeddings': None,\n",
       " 'documents': [['¡Hola! Bienvenido a la cafetería Brew Haven. ¿Algo de beber?',\n",
       "   '¡Hola! Bienvenido a la cafetería Brew Haven. ¿Algo de beber?',\n",
       "   '¡Hola! Bienvenido a la cafetería Brew Haven. ¿Algo de beber?',\n",
       "   '¡Hola! Bienvenido a la cafetería Brew Haven. ¿Algo de beber?',\n",
       "   '¡Hola! Bienvenido a la cafetería Brew Haven. ¿Algo de beber?',\n",
       "   '¡Hola! Bienvenido a la cafetería Brew Haven. ¿Algo de beber?',\n",
       "   '¡Hola! Bienvenido a la cafetería Brew Haven. ¿Algo de beber?',\n",
       "   '¡Hola! Bienvenido a la cafetería Brew Haven. ¿Algo de beber?',\n",
       "   '¡Hola! Bienvenido a la cafetería Brew Haven. ¿Algo de beber?',\n",
       "   '¡Hola! Bienvenido a la cafetería Brew Haven. ¿Algo de beber?']],\n",
       " 'uris': None,\n",
       " 'data': None}"
      ]
     },
     "execution_count": 73,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "id": "c0dbc749-a8c9-4990-97ce-7174da92f522",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "db3 = create_chroma_db(question, question, \"lang-bot-test4\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "id": "415a1cc6-5075-4986-8e39-d0c14ef22eb0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "results = db3.query(\n",
    "    query_texts=[\"En efectivo.\"],\n",
    "    n_results=5\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "id": "eaa0b427-f39c-4377-ba04-668916761e88",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'ids': [['4', '0', '2', '3', '1']],\n",
       " 'distances': [[0.8087770342826843,\n",
       "   1.2825713157653809,\n",
       "   1.3265247344970703,\n",
       "   1.4157676696777344,\n",
       "   1.4256945848464966]],\n",
       " 'metadatas': [[{'lesson': '¿Pagas con tarjeta o en efectivo?'},\n",
       "   {'lesson': '¡Hola! Bienvenido a la cafetería Brew Haven. ¿Algo de beber?'},\n",
       "   {'lesson': '¿Cuántos azúcares?'},\n",
       "   {'lesson': '¿Quieres comida? ¿Sándwich o algo dulce?'},\n",
       "   {'lesson': '¿Quieres leche y azúcar?'}]],\n",
       " 'embeddings': None,\n",
       " 'documents': [['¿Pagas con tarjeta o en efectivo?',\n",
       "   '¡Hola! Bienvenido a la cafetería Brew Haven. ¿Algo de beber?',\n",
       "   '¿Cuántos azúcares?',\n",
       "   '¿Quieres comida? ¿Sándwich o algo dulce?',\n",
       "   '¿Quieres leche y azúcar?']],\n",
       " 'uris': None,\n",
       " 'data': None}"
      ]
     },
     "execution_count": 84,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c2f5e22-3a40-44b4-a2ff-a4d5a3d6ab62",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_pytorch_p310",
   "language": "python",
   "name": "conda_pytorch_p310"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}