{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3f159fb4",
   "metadata": {},
   "source": [
    "# 🚶🏻‍♂️ Getting Started\n",
    "\n",
    "Here you will learn how to use the fastembed package to embed your data into a vector space. The package is designed to be easy to use and fast. It is built on top of the [ONNX](https://onnx.ai/) standard, which allows for fast inference on a variety of hardware (called Runtimes in ONNX). \n",
    "\n",
    "## Quick Start\n",
    "\n",
    "The fastembed package is designed to be easy to use. The main class is the `Embedding` class. It takes a list of strings as input and returns a list of vectors as output. The `Embedding` class is initialized with a model file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ada95c6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install fastembed --upgrade --quiet # Install fastembed "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed81d725",
   "metadata": {},
   "source": [
    "Make the necessary imports, initialize the `Embedding` class, and embed your data into vectors:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b61c6552",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 76.7M/76.7M [00:05<00:00, 15.0MiB/s]\n",
      "100%|██████████| 3/3 [00:00<00:00, 455.37it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(384,)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "from typing import List\n",
    "import numpy as np\n",
    "from fastembed.embedding import DefaultEmbedding\n",
    "\n",
    "# Example list of documents\n",
    "documents: List[str] = [\n",
    "    \"Hello, World!\",\n",
    "    \"This is an example document.\",\n",
    "    \"fastembed is supported by and maintained by Qdrant.\",\n",
    "]\n",
    "# Initialize the DefaultEmbedding class\n",
    "embedding_model = DefaultEmbedding()\n",
    "embeddings: List[np.ndarray] = list(embedding_model.embed(documents))\n",
    "print(embeddings[0].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c49ae50",
   "metadata": {},
   "source": [
    "## Let's think step by step"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92cf4b76",
   "metadata": {},
   "source": [
    "### Setup\n",
    "\n",
    "Importing the required classes and modules:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c0a6f634",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import numpy as np\n",
    "from fastembed.embedding import DefaultEmbedding"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3fd03a71",
   "metadata": {},
   "source": [
    "Notice that we are using the DefaultEmbedding -- which is a quantized, state of the Art Flag Embedding model which beats OpenAI's Embedding by a large margin. \n",
    "\n",
    "### Prepare your Documents\n",
    "You can define a list of documents that you'd like to embed. These can be sentences, paragraphs, or even entire documents. \n",
    "\n",
    "#### Format of the Document List\n",
    "1. List of Strings: Your documents must be in a list, and each document must be a string\n",
    "2. For Retrieval Tasks: If you're working with queries and passages, you can add special labels to them:\n",
    "- **Queries**: Add \"query:\" at the beginning of each query string\n",
    "- **Passages**: Add \"passage:\" at the beginning of each passage string"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "145a56ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example list of documents\n",
    "documents: List[str] = [\n",
    "    \"passage: Hello, World!\",\n",
    "    \"query: Hello, World!\",  # these are two different embedding\n",
    "    \"passage: This is an example passage.\",\n",
    "    # You can leave out the prefix but it's recommended\n",
    "    \"fastembed is supported by and maintained by Qdrant.\",\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1cb3cc87",
   "metadata": {},
   "source": [
    "### Load the Embedding Model Weights\n",
    "Next, initialize the Embedding class with the desired parameters. Here, \"BAAI/bge-small-en\" is the pre-trained model name, and max_length=512 is the maximum token length for each document.\n",
    "\n",
    "This will download the model weights, decompress to directory `local_cache` and load them into the Embedding class.\n",
    "\n",
    "#### Initialize DefaultEmbedding\n",
    "\n",
    "We will initialize Flag Embeddings with the model name and the maximum token length. That is the DefaultEmbedding class with the model name \"BAAI/bge-small-en\" and max_length=512."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "272c8915",
   "metadata": {},
   "outputs": [],
   "source": [
    "embedding_model = DefaultEmbedding()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5549d501",
   "metadata": {},
   "source": [
    "### Embed your Documents\n",
    "\n",
    "Use the embed method of the embedding model to transform the documents into a List of np.array. The method returns a generator, so we cast it to a list to get the embeddings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "8013eee9",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 4/4 [00:00<00:00, 361.82it/s]\n"
     ]
    }
   ],
   "source": [
    "embeddings: List[np.ndarray] = list(embedding_model.embed(documents))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5b5a6ad",
   "metadata": {},
   "source": [
    "You can print the shape of the embeddings to understand their dimensions. Typically, the shape will indicate the number of dimensions in the vector."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "0d8c8e08",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(384,)\n"
     ]
    }
   ],
   "source": [
    "print(embeddings[0].shape)  # (384,) or similar output"
   ]
  }
 ],
 "metadata": {
  "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.9.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}