Files
fastembed/docs/Getting Started.ipynb
NirantK 49c2b3e7a9 * chore(docs): update installation command for fastembed in Getting Started.ipynb
* fix(docs): remove unnecessary casting of generator to list in code cell
2023-09-28 10:43:14 +05:30

239 lines
6.8 KiB
Plaintext

{
"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": [
"Asking to truncate to max_length but no maximum length is provided and the model has no predefined maximum length. Default to no truncation.\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"torch.Size([384])\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 with the desired parameters\n",
"embedding_model = DefaultEmbedding(model_name=\"BAAI/bge-small-en\", max_length=512)\n",
"embeddings: List[np.ndarray] = 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 as Embedding"
]
},
{
"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": 4,
"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": 5,
"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": 6,
"id": "8013eee9",
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"Asking to truncate to max_length but no maximum length is provided and the model has no predefined maximum length. Default to no truncation.\n"
]
}
],
"source": [
"embeddings: List[np.ndarray] = 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": 7,
"id": "0d8c8e08",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"torch.Size([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
}