run-llama--llama_index/docs/examples/embeddings/mixedbreadai.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ef96689e6d0db317",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/embeddings/mixedbreadai.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82e5a0e429e08f86",
   "metadata": {},
   "source": [
    "# Mixedbread AI Embeddings\n",
    "\n",
    "Explore the capabilities of MixedBread AI's embedding models with custom encoding formats (binary, int, float, base64, etc.), embedding dimensions (Matryoshka) and context prompts."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1796b6747b975b6",
   "metadata": {},
   "source": [
    "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12e44155a01c6658",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-embeddings-mixedbreadai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e9c93730f08783c",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "615be84a4ebdb2aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from llama_index.embeddings.mixedbreadai import MixedbreadAIEmbedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5562e0afabcd3d57",
   "metadata": {},
   "outputs": [],
   "source": [
    "# API Key and Embedding Initialization\n",
    "\n",
    "# You can visit https://www.mixedbread.ai/api-reference#quick-start-guide\n",
    "# to get an api key\n",
    "mixedbread_api_key = os.environ.get(\"MXBAI_API_KEY\", \"your-api-key\")\n",
    "\n",
    "# Please check https://www.mixedbread.ai/docs/embeddings/models#whats-new-in-the-mixedbread-embed-model-family\n",
    "# for our embedding models\n",
    "model_name = \"mixedbread-ai/mxbai-embed-large-v1\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "288d34815403983f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1024\n",
      "[0.01128387451171875, 0.031097412109375, -0.00606536865234375, 0.0291748046875, -0.038604736328125]\n"
     ]
    }
   ],
   "source": [
    "oven = MixedbreadAIEmbedding(api_key=mixedbread_api_key, model_name=model_name)\n",
    "\n",
    "embeddings = oven.get_query_embedding(\"Why bread is so tasty?\")\n",
    "\n",
    "print(len(embeddings))\n",
    "print(embeddings[:5])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8075b9329c1e21e",
   "metadata": {},
   "source": [
    "### Using prompt for contextual embedding\n",
    "\n",
    "The prompt can improve the model's understanding of how the embedding will be used in subsequent tasks, which in turn increases the performance. Our experiments show that having domain specific prompts can increase the performance. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "459c8d8da3c0dbde",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1024\n",
      "[-0.0235443115234375, -0.0152435302734375, 0.008392333984375, 0.00336456298828125, -0.044647216796875]\n"
     ]
    }
   ],
   "source": [
    "prompt_for_retrieval = (\n",
    "    \"Represent this sentence for searching relevant passages:\"\n",
    ")\n",
    "\n",
    "contextual_oven = MixedbreadAIEmbedding(\n",
    "    api_key=mixedbread_api_key,\n",
    "    model_name=model_name,\n",
    "    prompt=prompt_for_retrieval,\n",
    ")\n",
    "\n",
    "contextual_embeddings = contextual_oven.get_query_embedding(\n",
    "    \"What bread is invented in Germany?\"\n",
    ")\n",
    "\n",
    "print(len(contextual_embeddings))\n",
    "print(contextual_embeddings[:5])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a9475daf5ac9879",
   "metadata": {},
   "source": [
    "## Quantization and Matryoshka support\n",
    "\n",
    "The Mixedbread AI embedding supports quantization and matryoshka to reduce the size of embeddings for better storage while retaining most of the performance.\n",
    "See these posts for more information: \n",
    "* [Binary and Scalar Embedding Quantization for Significantly Faster & Cheaper Retrieval](https://huggingface.co/blog/embedding-quantization)\n",
    "* [64 bytes per embedding, yee-haw](https://www.mixedbread.ai/blog/binary-mrl)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e5f5d29ebc54f31",
   "metadata": {},
   "source": [
    "### Using different encoding formats\n",
    "\n",
    "The default `encoding_format` is `float`. We also support `float16`, `binary`, `ubinary`, `int8`, `uint8`, `base64`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73cd06ff5f933333",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "128\n",
      "[-121.0, 96.0, -108.0, 111.0, 110.0]\n"
     ]
    }
   ],
   "source": [
    "# with `binary` embedding types\n",
    "binary_oven = MixedbreadAIEmbedding(\n",
    "    api_key=mixedbread_api_key,\n",
    "    model_name=model_name,\n",
    "    encoding_format=\"binary\",\n",
    ")\n",
    "\n",
    "binary_embeddings = binary_oven.get_text_embedding(\n",
    "    \"The bread is tiny but still filling!\"\n",
    ")\n",
    "\n",
    "print(len(binary_embeddings))\n",
    "print(binary_embeddings[:5])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3cfbff8d6d03352",
   "metadata": {},
   "source": [
    "### Using different embedding dimensions\n",
    "\n",
    "Mixedbread AI embedding models support Matryoshka dimension truncation. The default dimension is set to the model's maximum.\n",
    "Keep an eye on our website to see what models support Matryoshka."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c711c05f7df269e3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "512\n",
      "[-0.014221191, -0.013671875, -0.03314209, 0.025909424, -0.035095215]\n"
     ]
    }
   ],
   "source": [
    "# with truncated dimension\n",
    "half_oven = MixedbreadAIEmbedding(\n",
    "    api_key=mixedbread_api_key,\n",
    "    model_name=model_name,\n",
    "    dimensions=512,  # 1024 is the maximum of `mxbai-embed-large-v1`\n",
    ")\n",
    "\n",
    "half_embeddings = half_oven.get_text_embedding(\n",
    "    \"I want the better half of my bread.\"\n",
    ")\n",
    "\n",
    "print(len(half_embeddings))\n",
    "print(half_embeddings[:5])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}