{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "92587c8a-897d-4fc7-a62a-2988da97e45d",
      "metadata": {
        "id": "92587c8a-897d-4fc7-a62a-2988da97e45d"
      },
      "source": [
        "# Mistral Embeddings API\n",
        "\n",
        "Embeddings are vectorial representations of text that capture the semantic meaning of paragraphs through their position in a high dimensional vector space. Mistral Embeddings API offers cutting-edge, state-of-the-art embeddings for text, which can be used for many NLP tasks. In this guide, we will cover the fundamentals of the Mistral embeddings API, including how to measure the distance between text embeddings, and explore its main use cases:  clustering and classification.\n",
        "\n",
        "## Mistral Embeddings API\n",
        "To generate text embeddings using Mistral’s embeddings API, we can make a request to the API endpoint and specify the embedding model `mistral-embed`, along with providing a list of input text. The API will then return the corresponding embeddings as numerical vectors, which can be used for further analysis or processing in NLP applications.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "4312b3ac-a70b-443c-98f9-f7f5d7c1bd33",
      "metadata": {
        "id": "4312b3ac-a70b-443c-98f9-f7f5d7c1bd33"
      },
      "outputs": [],
      "source": [
        "!pip install mistralai seaborn numpy scikit-learn"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "id": "7c28ce18-065d-4379-8bbb-214479238c77",
      "metadata": {
        "id": "7c28ce18-065d-4379-8bbb-214479238c77",
        "outputId": "e9cebce6-ade8-4fef-81f6-f7898676ce72",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 245
        }
      },
      "outputs": [
        {
          "output_type": "error",
          "ename": "KeyError",
          "evalue": "'MISTRAL_API_KEY'",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mKeyError\u001b[0m                                  Traceback (most recent call last)",
            "\u001b[0;32m/tmp/ipython-input-2604136316.py\u001b[0m in \u001b[0;36m<cell line: 0>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0mmistralai\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mMistral\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 4\u001b[0;31m \u001b[0mapi_key\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mos\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0menviron\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"MISTRAL_API_KEY\"\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      5\u001b[0m \u001b[0mmodel\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m\"mistral-embed\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      6\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/lib/python3.12/os.py\u001b[0m in \u001b[0;36m__getitem__\u001b[0;34m(self, key)\u001b[0m\n",
            "\u001b[0;31mKeyError\u001b[0m: 'MISTRAL_API_KEY'"
          ]
        }
      ],
      "source": [
        "import os\n",
        "from mistralai import Mistral\n",
        "\n",
        "api_key = os.environ[\"MISTRAL_API_KEY\"]\n",
        "model = \"mistral-embed\"\n",
        "\n",
        "client = Mistral(api_key=api_key)\n",
        "\n",
        "embeddings_batch_response = client.embeddings.create(\n",
        "    model=model,\n",
        "    inputs=[\"Embed this sentence.\", \"As well as this one.\"],\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "0f410f39-5b57-4284-af91-7be21abf79c4",
      "metadata": {
        "id": "0f410f39-5b57-4284-af91-7be21abf79c4"
      },
      "source": [
        "\n",
        "The output is a EmbeddingResponse object with the embeddings and the token usage information.\n",
        "\n",
        "```\n",
        "Data(id='eb4c2c739780415bb3af4e47580318cc', object='list', data=[EmbeddingObject(object='embedding', embedding=[-0.0165863037109375,...], index=0), Data(object='embedding', embedding=[-0.0234222412109375,...], index=1)], model='mistral-embed', usage=EmbeddingResponseUsage(prompt_tokens=15, total_tokens=15, completion_tokens=0))\n",
        "\n",
        "```\n",
        "Let’s take a look of the length of the first embedding:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "709bfa0d-664c-4902-8de5-f96a22232ca7",
      "metadata": {
        "id": "709bfa0d-664c-4902-8de5-f96a22232ca7",
        "outputId": "15cf4eaf-ff95-4c78-f280-ad9945d4dac0"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "1024"
            ]
          },
          "execution_count": 2,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "len(embeddings_batch_response.data[0].embedding)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "09d4267d-8ab1-4c7c-91c0-0f1d567e295a",
      "metadata": {
        "id": "09d4267d-8ab1-4c7c-91c0-0f1d567e295a"
      },
      "source": [
        "It returns 1024, which means that our embedding dimension is 1024. The `mistral-embed` model generates embedding vectors of dimension 1024 for each text string, regardless of the text length.  It’s worth noting that while higher dimensional embeddings can better capture text information and improve the performance of NLP tasks, they may require more computational resources for hosting and inference, and may result in increased latency and memory usage for storing and processing these embeddings. This trade-off between performance and computational resources should be considered when designing NLP systems that rely on text embeddings.\n",
        "\n",
        "### Distance measures\n",
        "\n",
        "In the realm of text embeddings, texts with similar meanings or context tend to be located in closer proximity to each other within this space, as measured by the distance between their vectors. This is due to the fact that the model has learned to group semantically related texts together during the training process.\n",
        "\n",
        "Let’s take a look at a simple example. To simplify working with text embeddings, we can wrap the embedding API in this function:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c95c2bdf-f08a-4d36-825a-0f78ce82ec7d",
      "metadata": {
        "id": "c95c2bdf-f08a-4d36-825a-0f78ce82ec7d"
      },
      "outputs": [],
      "source": [
        "from sklearn.metrics.pairwise import euclidean_distances\n",
        "\n",
        "def get_text_embedding(inputs):\n",
        "    embeddings_batch_response = client.embeddings.create(\n",
        "        model=model,\n",
        "        inputs=inputs\n",
        "    )\n",
        "    return embeddings_batch_response.data[0].embedding"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "7206c037-97d3-4eee-a510-46f6bff1d78b",
      "metadata": {
        "id": "7206c037-97d3-4eee-a510-46f6bff1d78b"
      },
      "source": [
        "Suppose we have two sentences: one about cats and the other about books.  We want to find how similar each sentence is to the reference sentence \"Books are mirrors: You only see in them what you already have inside you\". We can see that the distance between the reference sentence embeddings and the book sentence embeddings is smaller than the distance between the reference sentence embeddings and the cat sentence embeddings.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "fd16565e-272b-46b2-884d-3c61e91cbe52",
      "metadata": {
        "id": "fd16565e-272b-46b2-884d-3c61e91cbe52"
      },
      "outputs": [],
      "source": [
        "sentences = [\n",
        "    \"A home without a cat — and a well-fed, well-petted and properly revered cat — may be a perfect home, perhaps, but how can it prove title?\",\n",
        "    \"I think books are like people, in the sense that they’ll turn up in your life when you most need them\"\n",
        "]\n",
        "embeddings = [get_text_embedding([t]) for t in sentences]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "b1587f75-3f36-46aa-86d8-7c2fe400be5e",
      "metadata": {
        "lines_to_next_cell": 2,
        "id": "b1587f75-3f36-46aa-86d8-7c2fe400be5e"
      },
      "outputs": [],
      "source": [
        "reference_sentence = \"Books are mirrors: You only see in them what you already have inside you\"\n",
        "reference_embedding = get_text_embedding([reference_sentence])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "8d7b2d20-a920-4171-95de-3f0812803321",
      "metadata": {
        "id": "8d7b2d20-a920-4171-95de-3f0812803321",
        "outputId": "35c0702f-98f3-462b-fe15-ab5007e41074"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "A home without a cat — and a well-fed, well-petted and properly revered cat — may be a perfect home, perhaps, but how can it prove title? [[0.80094257]]\n",
            "I think books are like people, in the sense that they’ll turn up in your life when you most need them [[0.58162089]]\n"
          ]
        }
      ],
      "source": [
        "for t, e in zip(sentences, embeddings):\n",
        "    distance = euclidean_distances([e], [reference_embedding])\n",
        "    print(t, distance)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "bdfb86a9-7664-464a-9cc2-854993b72238",
      "metadata": {
        "id": "bdfb86a9-7664-464a-9cc2-854993b72238"
      },
      "source": [
        "In our example above, we used the Euclidean distance to measure the distance between embedding vectors (note that since Mistral embeddings are norm 1, cosine similarity, dot product or Euclidean distance are all equivalent).\n",
        "\n",
        "### Paraphrase detection\n",
        "\n",
        "Another potential use case is paraphrase detection. In this simple example, we have a list of three sentences, and we would like to find out if any of the two sentences are paraphrases of each other. If the distance between two sentence embeddings is small, it suggests that the two sentences are semantically similar and could be potential paraphrases.\n",
        "\n",
        "Result suggests that the first two sentences are semantically similar and could be potential paraphrases, whereas the third sentence is more different. This is just a super simple example. But this approach can be extended to more complex situations in real-world applications, such as detecting paraphrases in social media posts, news articles, or customer reviews.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "2d610ebd-e0a0-4435-ab1d-07216b214b89",
      "metadata": {
        "id": "2d610ebd-e0a0-4435-ab1d-07216b214b89"
      },
      "outputs": [],
      "source": [
        "sentences = [\n",
        "    'Have a safe happy Memorial Day weekend everyone',\n",
        "    'To all our friends at Whatsit Productions Films enjoy a safe happy Memorial Day weekend',\n",
        "    'Where can I find the best cheese?'\n",
        "]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "5938ef92-a83d-4ca1-a32c-5ecd32ba9ae8",
      "metadata": {
        "id": "5938ef92-a83d-4ca1-a32c-5ecd32ba9ae8"
      },
      "outputs": [],
      "source": [
        "sentence_embeddings = [get_text_embedding([t]) for t in sentences]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "6f5449e4-ee20-4bcc-aa79-a44f084f6b97",
      "metadata": {
        "id": "6f5449e4-ee20-4bcc-aa79-a44f084f6b97"
      },
      "outputs": [],
      "source": [
        "import itertools\n",
        "\n",
        "sentence_embeddings_pairs = list(itertools.combinations(sentence_embeddings, 2))\n",
        "sentence_pairs = list(itertools.combinations(sentences, 2))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ae95833b-3013-4275-af57-71604627482b",
      "metadata": {
        "id": "ae95833b-3013-4275-af57-71604627482b",
        "outputId": "38d2f0d3-7bf8-497c-bfc9-995e9bcd9bd1"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "('Have a safe happy Memorial Day weekend everyone', 'To all our friends at Whatsit Productions Films enjoy a safe happy Memorial Day weekend') [[0.54326686]]\n",
            "('Have a safe happy Memorial Day weekend everyone', 'Where can I find the best cheese?') [[0.92573978]]\n",
            "('To all our friends at Whatsit Productions Films enjoy a safe happy Memorial Day weekend', 'Where can I find the best cheese?') [[0.9114184]]\n"
          ]
        }
      ],
      "source": [
        "for s, e in zip(sentence_pairs, sentence_embeddings_pairs):\n",
        "    print(s, euclidean_distances([e[0]], [e[1]]))"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "117ab6ab-1ec7-4030-a168-13d753a6a3d2",
      "metadata": {
        "id": "117ab6ab-1ec7-4030-a168-13d753a6a3d2"
      },
      "source": [
        "### Batch processing\n",
        "\n",
        "The Mistral Embeddings API is designed to process text in batches for improved efficiency and speed. In this example, we will demonstrate this by loading the Symptom2Disease dataset from [Kaggle](https://www.kaggle.com/datasets/niyarrbarman/symptom2disease), which contains 1200 rows with two columns: \"label\" and \"text\". The \"label\" column indicates the disease category, while the \"text\" column describes the symptoms associated with that disease.\n",
        "\n",
        "We wrote a function `get_embeddings_by_chunks` that splits data into chunks and then sends each chunk to the Mistral embedding API to get the embeddings. Then we saved the embeddings as a new column in the dataframe. Note that the Mistral Embeddings API will provide auto-chunking in the future, so that users don’t need to manually split the data into chunks before sending it to the API.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "3fa58906-41b1-4b50-96ca-8457175de37a",
      "metadata": {
        "id": "3fa58906-41b1-4b50-96ca-8457175de37a"
      },
      "outputs": [],
      "source": [
        "import pandas as pd\n",
        "df = pd.read_csv(\"https://raw.githubusercontent.com/mistralai/cookbook/main/data/Symptom2Disease.csv\", index_col=0)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e0c7821a-da97-4baa-afa5-aa023a5298b4",
      "metadata": {
        "id": "e0c7821a-da97-4baa-afa5-aa023a5298b4"
      },
      "outputs": [],
      "source": [
        "def get_embeddings_by_chunks(data, chunk_size):\n",
        "    chunks = [data[x : x + chunk_size] for x in range(0, len(data), chunk_size)]\n",
        "    embeddings_response = [\n",
        "        client.embeddings.create(model=model, inputs=c) for c in chunks\n",
        "    ]\n",
        "    return [d.embedding for e in embeddings_response for d in e.data]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "b0282432-3656-4eb5-9e6a-2d1358187b33",
      "metadata": {
        "id": "b0282432-3656-4eb5-9e6a-2d1358187b33"
      },
      "outputs": [],
      "source": [
        "df['embeddings'] = get_embeddings_by_chunks(df['text'].tolist(), 50)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "7d714a9e-783c-4385-95ee-5f2f3fbeb7fe",
      "metadata": {
        "id": "7d714a9e-783c-4385-95ee-5f2f3fbeb7fe",
        "outputId": "c2347e3c-e7ac-4c55-a4dc-dc6d59ba51ad"
      },
      "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>label</th>\n",
              "      <th>text</th>\n",
              "      <th>embeddings</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>Psoriasis</td>\n",
              "      <td>I have been experiencing a skin rash on my arm...</td>\n",
              "      <td>[-0.036102294921875, 0.041351318359375, 0.0734...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>Psoriasis</td>\n",
              "      <td>My skin has been peeling, especially on my kne...</td>\n",
              "      <td>[-0.05364990234375, 0.05224609375, 0.073791503...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>2</th>\n",
              "      <td>Psoriasis</td>\n",
              "      <td>I have been experiencing joint pain in my fing...</td>\n",
              "      <td>[-0.035400390625, 0.026275634765625, 0.0360107...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3</th>\n",
              "      <td>Psoriasis</td>\n",
              "      <td>There is a silver like dusting on my skin, esp...</td>\n",
              "      <td>[-0.035980224609375, 0.057037353515625, 0.0528...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>4</th>\n",
              "      <td>Psoriasis</td>\n",
              "      <td>My nails have small dents or pits in them, and...</td>\n",
              "      <td>[-0.02471923828125, 0.039337158203125, 0.04772...</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>"
            ],
            "text/plain": [
              "       label                                               text  \\\n",
              "0  Psoriasis  I have been experiencing a skin rash on my arm...   \n",
              "1  Psoriasis  My skin has been peeling, especially on my kne...   \n",
              "2  Psoriasis  I have been experiencing joint pain in my fing...   \n",
              "3  Psoriasis  There is a silver like dusting on my skin, esp...   \n",
              "4  Psoriasis  My nails have small dents or pits in them, and...   \n",
              "\n",
              "                                          embeddings  \n",
              "0  [-0.036102294921875, 0.041351318359375, 0.0734...  \n",
              "1  [-0.05364990234375, 0.05224609375, 0.073791503...  \n",
              "2  [-0.035400390625, 0.026275634765625, 0.0360107...  \n",
              "3  [-0.035980224609375, 0.057037353515625, 0.0528...  \n",
              "4  [-0.02471923828125, 0.039337158203125, 0.04772...  "
            ]
          },
          "execution_count": 14,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "df.head()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "77ac1d38-ddf1-40c6-855f-1065405f0ec3",
      "metadata": {
        "id": "77ac1d38-ddf1-40c6-855f-1065405f0ec3"
      },
      "source": [
        "## t-SNE embeddings visualization\n",
        "\n",
        "We mentioned previously that our embeddings have 1024 dimensions, which makes them impossible to visualize directly. Thus, in order to visualize our embeddings, we can use a dimensionality reduction technique such as t-SNE to project our embeddings into a lower-dimensional space that is easier to visualize.\n",
        "\n",
        "In this example, we transform our embeddings to 2 dimensions and create a 2D scatter plot showing the relationships among embeddings of different diseases.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e2badca7-ec02-421e-99a8-2bc44fdd42af",
      "metadata": {
        "id": "e2badca7-ec02-421e-99a8-2bc44fdd42af",
        "outputId": "1d7e1290-28b8-4d6d-c18f-4af09f741862"
      },
      "outputs": [
        {
          "data": {
            "image/png": "[encoded data, 169744 characters]",
            "text/plain": [
              "<Figure size 640x480 with 1 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "import seaborn as sns\n",
        "from sklearn.manifold import TSNE\n",
        "import numpy as np\n",
        "\n",
        "\n",
        "tsne = TSNE(n_components=2, random_state=0).fit_transform(np.array(df['embeddings'].to_list()))\n",
        "ax = sns.scatterplot(x=tsne[:, 0], y=tsne[:, 1], hue=np.array(df['label'].to_list()))\n",
        "sns.move_legend(ax, 'upper left', bbox_to_anchor=(1, 1))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "8211936e-f5cf-4eca-91fb-05c65418f20d",
      "metadata": {
        "id": "8211936e-f5cf-4eca-91fb-05c65418f20d"
      },
      "source": [
        "#### Comparison with fasttext\n",
        "\n",
        "We can compare it with fastText, a popular open-source embeddings model. However, when examining the t-SNE embeddings plot, we notice that fastText embeddings fail to create clear separations between data points with matching labels.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "00df9498-8c5f-4322-91bf-4384bc70902b",
      "metadata": {
        "id": "00df9498-8c5f-4322-91bf-4384bc70902b"
      },
      "outputs": [],
      "source": [
        "!pip install fasttext"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e9ca9418-d355-4c3c-9a1d-03753262cb1d",
      "metadata": {
        "id": "e9ca9418-d355-4c3c-9a1d-03753262cb1d",
        "outputId": "f7b63b94-9cf7-4820-ec17-27f30df11664"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Warning : `load_model` does not return WordVectorModel or SupervisedModel any more, but a `FastText` object which is very similar.\n"
          ]
        },
        {
          "data": {
            "image/png": "[encoded data, 241684 characters]",
            "text/plain": [
              "<Figure size 640x480 with 1 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "import fasttext.util\n",
        "fasttext.util.download_model('en', if_exists='ignore')  # English\n",
        "ft = fasttext.load_model('cc.en.300.bin')\n",
        "\n",
        "df['fasttext_embeddings'] = df['text'].apply(lambda x: ft.get_word_vector(x).tolist())\n",
        "\n",
        "tsne = TSNE(n_components=2, random_state=0).fit_transform(np.array(df['fasttext_embeddings'].to_list()))\n",
        "ax = sns.scatterplot(x=tsne[:, 0], y=tsne[:, 1], hue=np.array(df['label'].to_list()))\n",
        "sns.move_legend(ax, 'upper left', bbox_to_anchor=(1, 1))"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "ce94a816-6d1a-4f69-8f9c-559b24426673",
      "metadata": {
        "id": "ce94a816-6d1a-4f69-8f9c-559b24426673"
      },
      "source": [
        "## Classification\n",
        "\n",
        "Text embeddings can be used as input features in machine learning models, such as classification and clustering. In this example, we use a classification model to predict the disease labels from the embeddings of disease description text.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "261e6e78-341b-4b03-902d-3f27922f52d7",
      "metadata": {
        "lines_to_next_cell": 2,
        "id": "261e6e78-341b-4b03-902d-3f27922f52d7"
      },
      "outputs": [],
      "source": [
        "# Create a train / test split\n",
        "\n",
        "from sklearn.model_selection import train_test_split\n",
        "\n",
        "train_x, test_x, train_y, test_y = train_test_split(df['embeddings'], df[\"label\"],test_size=0.2)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "baeff9a8-7be1-4fb9-9cb3-cb44a719076a",
      "metadata": {
        "id": "baeff9a8-7be1-4fb9-9cb3-cb44a719076a"
      },
      "outputs": [],
      "source": [
        "# Normalize features\n",
        "from sklearn.preprocessing import StandardScaler\n",
        "\n",
        "scaler = StandardScaler()\n",
        "train_x = scaler.fit_transform(train_x.to_list())\n",
        "test_x = scaler.transform(test_x.to_list())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "dd8f8a80-6b7e-40a2-82e2-bd8c4be0ab67",
      "metadata": {
        "id": "dd8f8a80-6b7e-40a2-82e2-bd8c4be0ab67",
        "outputId": "2b685835-c1a6-4f97-a061-6aeb1610aba1"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Precision: 98.75%\n"
          ]
        }
      ],
      "source": [
        "# Train a classifier and compute the test accuracy\n",
        "\n",
        "from sklearn.linear_model import LogisticRegression\n",
        "\n",
        "# For a real problem, C should be properly cross validated and the confusion matrix analyzed\n",
        "clf = LogisticRegression(random_state=0, C=1.0, max_iter=500).fit(train_x, train_y.to_list())\n",
        "\n",
        "# you can also try the sag algorithm:\n",
        "# clf = LogisticRegression(random_state=0, C=1.0, max_iter=1000, solver='sag').fit(train_x, train_y)\n",
        "\n",
        "print(f\"Precision: {100*np.mean(clf.predict(test_x) == test_y.to_list()):.2f}%\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "263c015e-c32d-440c-a7d4-5d977d2d4a6e",
      "metadata": {
        "id": "263c015e-c32d-440c-a7d4-5d977d2d4a6e",
        "outputId": "2aa51b9a-d486-4fb4-f7f9-fefe2dd6f2fe"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'Migraine'"
            ]
          },
          "execution_count": 20,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# Classify a single example\n",
        "text = \"I've been experiencing frequent headaches and vision problems.\"\n",
        "clf.predict([get_text_embedding([text])])"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "ad4b26bc-2aa0-4008-aad2-7255aaeea3db",
      "metadata": {
        "id": "ad4b26bc-2aa0-4008-aad2-7255aaeea3db"
      },
      "source": [
        "#### Comparison with fasttext\n",
        "Additionally, let’s take a look at the performance using fastText embeddings in this classification task. It appears that the classification model achieves better performance with Mistral embedding models as compared to using fastText embeddings.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "df2b36ee-0b7d-4813-a259-624a838d2af6",
      "metadata": {
        "id": "df2b36ee-0b7d-4813-a259-624a838d2af6",
        "outputId": "bb6b9f2a-5a90-49ae-8e7e-d84452e5af92"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Precision: 86.25%\n"
          ]
        }
      ],
      "source": [
        "# Create a train / test split\n",
        "from sklearn.model_selection import train_test_split\n",
        "train_x, test_x, train_y, test_y = train_test_split(df['fasttext_embeddings'], df[\"label\"],test_size=0.2)\n",
        "# Normalize features\n",
        "from sklearn.preprocessing import StandardScaler\n",
        "scaler = StandardScaler()\n",
        "train_x = scaler.fit_transform(train_x.to_list())\n",
        "test_x = scaler.transform(test_x.to_list())\n",
        "# Train a classifier and compute the test accuracy\n",
        "from sklearn.linear_model import LogisticRegression\n",
        "# For a real problem, C should be properly cross validated and the confusion matrix analyzed\n",
        "clf = LogisticRegression(random_state=0, C=1.0, max_iter=500).fit(train_x, train_y.to_list())\n",
        "# you can also try the sag algorithm:\n",
        "# clf = LogisticRegression(random_state=0, C=1.0, max_iter=1000, solver='sag').fit(train_x, train_y)\n",
        "print(f\"Precision: {100*np.mean(clf.predict(test_x) == test_y.to_list()):.2f}%\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "5a1a6428-77aa-41f5-bf7a-5af9ea609003",
      "metadata": {
        "id": "5a1a6428-77aa-41f5-bf7a-5af9ea609003"
      },
      "source": [
        "## Clustering\n",
        "\n",
        "What if we don’t have disease labels? One approach to gain insights from the data is through clustering. Clustering is an unsupervised machine learning technique that groups similar data points together based on their similarity with respect to certain features. In the context of text embeddings, we can use the distance between each embedding as a measure of similarity, and group together data points with embeddings that are close to each other in the high-dimensional space.\n",
        "\n",
        "Since we already know there are 24 clusters, let’s use the K-means clustering with 24 clusters. Then we can inspect a few examples and verify whether the examples in a single cluster are similar to one another. For example, take a look at the first three rows of cluster 23. We can see that they look very similar in terms of symptoms.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "d053db2c-4a89-4f98-a1ed-36594b786ea8",
      "metadata": {
        "id": "d053db2c-4a89-4f98-a1ed-36594b786ea8"
      },
      "outputs": [],
      "source": [
        "from sklearn.cluster import KMeans\n",
        "model = KMeans(n_clusters=24, max_iter=1000)\n",
        "model.fit(df['embeddings'].to_list())\n",
        "df[\"cluster\"] = model.labels_"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "0169b960-faea-4da6-b509-08558d476118",
      "metadata": {
        "id": "0169b960-faea-4da6-b509-08558d476118",
        "outputId": "97463f0b-7da7-4af1-cb02-ca4d8619228f"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "I have been feeling extremely tired and weak, and I've also been coughing a lot with difficulty breathing. My fever is very high, and I'm producing a lot of mucus when I cough.\n",
            "I've got a cough that won't go away, and I'm exhausted. I've been coughing up thick mucous and my fever is also pretty high.\n",
            "I have a persistent cough and have been feeling quite fatigued. My fever is through the roof, and I'm having trouble breathing. When I cough, I also cough up a lot of mucous.\n"
          ]
        }
      ],
      "source": [
        "print(*df[df.cluster==23].text.head(3), sep='\\n')"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "3a46a74a-fc5e-463f-9a66-1962b02f96bd",
      "metadata": {
        "id": "3a46a74a-fc5e-463f-9a66-1962b02f96bd"
      },
      "source": [
        "## Retrieval\n",
        "\n",
        "Our embedding model excels in retrieval tasks, as it is trained with retrieval in mind. Embeddings are also incredibly helpful in implementing retrieval-augmented generation (RAG) systems, which use retrieved relevant information from a knowledge base to generate responses. At a high-level, we embed a knowledge base, whether it is a local directory, text files, or internal wikis, into text embeddings and store them in a vector database. Then, based on the user's query, we retrieve the most similar embeddings, which represent the relevant information from the knowledge base. Finally, we feed these relevant embeddings to a large language model to generate a response that is tailored to the user's query and context. If you are interested in learning more about how RAG systems work and how to implement a basic RAG, check out our [previous guide](https://github.com/mistralai/cookbook/blob/main/mistral/rag/basic_RAG.ipynb) on this topic.\n"
      ]
    }
  ],
  "metadata": {
    "jupytext": {
      "formats": "ipynb,py:light"
    },
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.10.13"
    },
    "colab": {
      "provenance": []
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}